ReSFlow Enables One-Step Generation for Discrete Data via Statistical Manifolds

To ensure the generation path "goes straight," we must first define what constitutes "straight" based on the geometry of the space where the data resides.

The authors of this paper are affiliated with Inner Mongolia University, Shanghai University of Finance and Economics (SUFE), and Carnegie Mellon University (CMU). The first author, Xiaohuan Jia, is a Ph.D. student at Inner Mongolia University, and the corresponding author is Researcher Shaohua Fan from the same institution. Assistant Professor Renzhe Xu and Associate Professor Xiao Wang from SUFE, along with CMU Ph.D. student Jiayun Wu, are co-authors of the paper.

Generating a handwritten digit image or a DNA sequence using Flow Matching or Diffusion Models typically requires starting from noise and undergoing multiple sampling steps to produce the final result. Each step involves calling the model, gradually improving generation quality while increasing computational overhead.

Is it possible to complete the entire generation process with just a single model call?

For discrete data such as characters, binary pixels, and DNA bases, the challenge lies in the fact that the space of categorical probabilities has its own geometric structure. How can the generation path "go straight" within this space?

The research team proposes Rectified Statistical Flow (ReSFlow): mapping categorical probabilities to a sphere, re-pairing and rectifying generation trajectories on the spherical surface, enabling the model to generate results via a single update along a path close to the geodesic. On binarized MNIST, ReSFlow achieves generation quality comparable to Statistical Flow Matching (SFM) with 34 calls, using only one model call. In conditional promoter DNA design tasks, the error between the predicted activity spectrum of one-step sampling and full ODE sampling is minimal, while the time to generate 1,000 sequences drops from 67.94 seconds to 0.409 seconds.

This work has been accepted by NeurIPS 2026.

• Paper Title: ReSFlow: Rectifying Categorical Flows on Statistical Manifolds for One-Step Generation
• Paper Link: https://arxiv.org/abs/2610.13225
• Code Link: https://github.com/xiaohuanjia/ReSFlow

Research Background

Discrete data generation typically starts with categorical probabilities: for example, A, C, G, and T in DNA can be represented as four non-negative probabilities summing to 1. The model learns how these probabilities change over time and advances them via an Ordinary Differential Equation (ODE) to ultimately sample discrete results. SFM introduces Fisher–Rao geometry, constructing training paths using geodesics in probability space. However, achieving one-step generation still faces two main challenges.

First, although the training path is a geodesic, the actual path taken during inference may not be. Conditionally independent random pairings can cause paths to cross at the same timestamp, providing conflicting directions. The model must synthesize these directions, potentially resulting in curved learned trajectories. Relying solely on a velocity prediction at the starting point makes it difficult to accurately reach the endpoint.

Second, "straightness" in Euclidean space does not equate to "straightness" in probability space. Rectified Flow straightens generation trajectories through re-pairing, but Euclidean lines differ from Fisher–Rao geodesics. Two Euclidean interpolation paths may not intersect, yet their corresponding spherical geodesics might still cross.

ReSFlow addresses these two issues: mapping categorical probabilities to a sphere, performing re-pairing and training within a space consistent with Fisher–Rao geometry, making the actual generation trajectory closer to the geodesic, and ultimately enabling one-step sampling.

Figure 1: Difference between Euclidean rectification and spherical rectification. Linear paths in Euclidean space may not intersect, but their corresponding spherical geodesics might; ReSFlow performs rectification directly on the sphere.

Core Methodology

To truly make the generation path "straight," one must first choose the appropriate space. ReSFlow uses the square root mapping to project sample points onto a sphere. It then allows the initial flow model to find respective endpoints starting from noise. By re-training with these new pairings, the model gradually learns to advance along geodesics. Finally, utilizing the spherical exponential map, it achieves one-step generation from noise to data.

Spherical Representation and Geodesic Interpolation

ReSFlow applies the square root element-wise to categorical probability vectors, mapping them from the probability simplex to the non-negative region of the unit hypersphere:

The square root mapping ensures that Fisher–Rao geodesics correspond to great circle arcs on the sphere, with the two metrics differing only by a fixed constant. Subsequently, ReSFlow uses Spherical Linear Interpolation (SLERP) to connect the given start distribution and end distribution on the sphere:

Differentiating this spherical linear interpolation yields the conditional velocity required for training the flow model on the sphere. The model's predicted velocity is mapped to the tangent space at the current position, ensuring that the model's predictions consistently adhere to Riemannian spherical geometry.

Spherical Rectification: From Random Pairing to Model-Based Pairing

Stage 1: Train the model to move from noise to data. This stage mirrors traditional flow model training, independently sampling from the prior distribution and the data distribution, mapping them to the sphere to construct geodesics, and matching the model's predicted tangential velocity to the conditional velocity:

This stage establishes a generation flow from noise to data. However, conditionally independent pairings result in numerous crossing paths, meaning the overall trajectory learned by the model remains curved.

Stage 2: Continue training the model on trajectories with reduced crossings. After completing the initial training, ReSFlow samples starting points from the prior distribution and uses the current model to solve the spherical ODE, determining the actual generated endpoint reached by each starting point:

The starting and generated endpoints are then paired together, and the model is retrained along the spherical geodesic connecting them. Since these pairs originate from the model's own ODE trajectories, under conditions where solutions exist and are unique, a trajectory cannot continue in different directions at the same time and position. Consequently, the model reorganizes motion directions at intersections of the original conditional paths. Compared to independent random pairing, this approach significantly reduces crossings between conditional paths, making the overall learned trajectories closer to spherical geodesics.

One-Step Generation

After the aforementioned pairing adjustments and training, the trajectories learned by the model gradually approximate geodesics. During ODE integration, there is no need to repeatedly update the direction along the path. Instead, the velocity is predicted once in the tangent space at the initial noise point, and the spherical exponential map is used to directly estimate the generated endpoint along the corresponding great circle arc:

Once the spherical endpoint is obtained, the coordinates are squared element-wise to recover class probabilities:

Discrete classes are then selected or sampled based on the recovered probability distribution.

Experimental Results

The authors first used the Swiss roll synthetic distribution to verify whether the model could restore the target distribution in one step, thereby validating its effectiveness. They then applied ReSFlow to real-world tasks including image generation, text generation, and DNA sequence design. In addition to comparing generation results, the paper includes analytical experiments to verify whether rectification makes the model-generated trajectories adhere more closely to spherical geodesics. The baseline comparisons include LinearFM, SFM, and R-LFM, evaluating Euclidean versus spherical geometry and the differences before and after rectification.

Validation Experiments

On the Swiss roll synthetic distribution, the ODE sampling time for generating 1,000 samples with ReSFlow was 0.0361 seconds, the shortest among the four methods; SFM required 0.0679 seconds. ReSFlow achieved an NLL of −2.25, outperforming LinearFM and R-LFM, though slightly higher than SFM’s −2.38 (lower NLL is better). These results indicate that spherical rectification shortens sampling time while maintaining good probabilistic modeling performance.

Table 1: Swiss roll experimental results. Comparison of NLL under full ODE sampling and time taken to generate 1,000 samples.

Samples generated via one-step sampling reflect this conclusion more intuitively. The true distribution exhibits a curled band structure on the simplex. Without rectification, LinearFM and SFM fail to restore this shape, with generated points concentrated in smaller regions. In contrast, ReSFlow preserves the complete curled structure, and its distribution is closer to the real samples.

Figure 2: Swiss roll true distribution vs. one-step generated samples. From left to right: True Distribution, LinearFM, SFM, R-LFM, and ReSFlow.

Main Experiments

On the binarized MNIST dataset, the one-step FID scores for unrectified LinearFM and SFM were 25.317 and 35.640, respectively. After rectification, R-LFM dropped to 6.261, and ReSFlow further improved to 5.973.

With full ODE sampling, SFM achieved the lowest FID; however, when compressed to one step, its quality degradation was significantly larger. ReSFlow maintained an FID of 5.973 in one-step sampling, compared to 5.541 with full ODE sampling.

The paper also compares generation quality across different numbers of model calls. The generation quality of ReSFlow with a single call approaches that of SFM with 34 calls. At this point, SFM’s FID is 5.922, while ReSFlow’s is 5.973.

Table 2: Binarized MNIST experimental results. Comparison of FID and time for full ODE vs. one-step sampling; time corresponds to generating 1,000 samples.

Sample images illustrate these differences more clearly. With full ODE sampling, all four methods generate relatively clear handwritten digits. When reduced to one step, LinearFM and SFM exhibit obvious missing strokes, making some digits difficult to recognize. ReSFlow, however, preserves more complete strokes and contours, and its one-step generated samples are closer to its own full ODE results.

Figure 3: Binarized MNIST generated samples. Top row: Full ODE sampling; Bottom row: One-step sampling. Columns from left to right: LinearFM, SFM, R-LFM, and ReSFlow.

In the text generation task, ReSFlow achieved a BPC of 1.26146, lower than SFM’s 1.82998. Euclidean rectification did not help here: R-LFM’s BPC rose to 3.26687, performing worse than the unrectified LinearFM. This demonstrates that repairing does not always lead to improvement, and the geometry adopted by the paths is equally critical.

Table 3: Text8 character-level modeling results. BPC is calculated via likelihood evaluation using full ODE, reporting the mean and standard deviation over three random seeds.

For the conditional promoter design task, the generation target was switched to DNA sequences of length 1,024. ReSFlow’s SP-MSE for one-step sampling was 0.028718, and for full ODE sampling it was 0.028743, indicating that the prediction error of the activity profile is similar for both methods.

The average time to generate 1,000 sequences dropped from 67.94 seconds to 0.409 seconds, representing an approximately 166x speedup. As a control, when SFM was compressed to a single step, the SP-MSE increased from 0.028015 to 0.035564, indicating a more pronounced loss in quality.

Table 4: Results for conditional promoter design. Time corresponds to generating 1,000 sequences of length 1,024, reporting the mean and standard deviation across three random seeds.

Figure 4: Generated samples for the conditional promoter design task. From top to bottom: LinearFM, SFM, R-LFM, and ReSFlow. For each group, the top row shows full ODE sampling, while the bottom row shows one-step sampling (Euler). Each row displays the first 120 bases of a promoter sequence with a total length of 1,024 bases.

Analysis Experiments

With improved generation results, do the model's generated trajectories indeed become "straighter"? The authors recorded transport costs and geodesic deviations during the rectification process on binarized MNIST.

Transport cost measures the distance penalty between noise starting points and generated endpoints. ReSFlow’s transport cost decreases significantly as rectification progresses, indicating that after re-pairing, the model tends to map noise samples to geometrically closer generated endpoints, thereby reducing unnecessary long-distance transport. In contrast, R-LFM shows a smaller decrease and reaches a plateau earlier. These results suggest that improving pairing in Euclidean space does not necessarily reduce transport costs in probability space with equal effectiveness.

Geodesic deviation checks whether the model "deviates from direction or changes speed" along the trajectory: at each position on the trajectory, the model’s actual predicted velocity is compared with the velocity required to reach the same generated endpoint along the geodesic within the remaining time. The closer these two values are, the more the model’s actual trajectory approximates a constant-speed geodesic. As the number of rectifications increases, the deviation gradually decreases, experimentally supporting ReSFlow’s core idea: spherical rectification allows the model’s actual generation paths to approach geodesics more closely. With fewer bends in the path and speed variations better conforming to constant-speed geodesic patterns, a single prediction at the starting point is more likely to describe the entire motion. This also explains why ReSFlow can accurately estimate endpoints—which originally require multi-step integration—through a single exponential mapping.

Figure 5: Fisher–Rao empirical transport costs and geodesic deviations under different rectification rounds. Lower values indicate smaller corresponding costs or deviations.

Conclusion

ReSFlow provides a pathway for one-step generation of discrete data: starting from the Fisher–Rao geometry of class probabilities, it performs re-pairing and rectification on the sphere to gradually align generation trajectories with geodesics, followed by one-step sampling via exponential mapping. Experiments on images and DNA sequences demonstrate that this design significantly reduces sampling computation while maintaining strong generation performance.

One-step inference still requires additional rectification training, and its effectiveness on larger-scale data and complex conditional tasks remains to be explored. However, this work offers a clear insight: to make generation paths "straight," one must first define what "straight" means according to the geometry of the space where the data resides.

© THE END

Please contact this official account for authorization before reprinting.

For submissions or media inquiries: [email protected]

Open in WeChat