文章摘要
Tomakethegenerationpath"straight",wefirstneedtodefinewhat"straight"meansbasedonthegeometryofthedatas

To make the generation path "straight", we first need to define what "straight" means based on the geometry of the data space. This core insight is the foundation of a new research work accepted by NeurIPS 2026.

Generating discrete data like handwritten digits or DNA sequences typically requires multiple model calls. Flow matching or diffusion models start from noise and go through multiple sampling steps to get results. Each step calls the model, improving generation quality gradually but increasing computational overhead.

Can we complete the entire generation process with just one model call? For discrete data such as characters, binary pixels, and DNA bases, an additional challenge exists: the space of category probabilities has its own geometric structure—how should the generation path "go straight" in this space?

A research team proposed Rectified Statistical Flow (ReSFlow) to solve these problems. The method maps category probabilities to a sphere, re-pairs and rectifies the generation trajectory on the sphere, allowing the model to follow a path close to a geodesic and finish generation with a single update.

On the binarized MNIST dataset, ReSFlow achieved generation quality close to that of Statistical Flow Matching (SFM) with 34 model calls in just one call. In the conditional promoter DNA design task, the predicted activity spectrum error of one-step sampling was nearly identical to that of full ODE sampling, and the time to generate 1,000 sequences was reduced from 67.94 seconds to 0.409 seconds, an acceleration of about 166 times.

The authors are from Inner Mongolia University, Shanghai University of Finance and Economics, and Carnegie Mellon University. The first author is Jia Xiaohuan, a doctoral student at Inner Mongolia University, with corresponding author Fan Shaohua, a researcher at the same university. Co-authors include Xu Runzhe, an assistant professor at Shanghai University of Finance and Economics, Wang Xiao, an associate professor, and Jiayun Wu, a doctoral student at CMU.

Research Background

Discrete data generation usually starts with category probabilities. For example, DNA bases A, C, G, T can be represented as four non-negative values that sum to 1. Models learn how these probabilities change over time, then use ODE integration to sample discrete results. SFM introduces Fisher-Rao geometry and uses geodesics in probability space to construct training paths, but two main barriers stand in the way of one-step generation.

First, the training path is a geodesic, but the actual path during inference may not be. Independent randomly paired conditional paths may cross at the same time, leading to conflicting directions. The model needs to combine these directions, so the learned trajectory may bend, making it hard to accurately reach the endpoint with just one speed prediction at the start.

Second, "straight" in Euclidean space is not the same as "straight" in probability space. Rectified Flow straightens generation trajectories by re-pairing, but Euclidean straight lines differ from Fisher-Rao geodesics. Two Euclidean interpolation paths may not intersect, but they could cross when converted to spherical geodesics.

ReSFlow addresses both problems: it maps category probabilities to a sphere, re-pairs and trains in a space that follows Fisher-Rao geometry, making the actual generated trajectory closer to a geodesic, thus enabling one-step sampling.

Core Method

To make the generation path truly "straight", the right space must be chosen first. ReSFlow first maps sample points to the sphere via a square root function. Then, the initial flow model starts from noise and finds its respective endpoints. Using these new pairs for retraining, the model gradually learns to move along geodesics. Finally, spherical exponential mapping is used to achieve one-step generation from noise to data.

Spherical Representation and Geodesic Interpolation

ReSFlow takes the square root of each element of the category probability vector, mapping it from the probability simplex to the non-negative region of the unit hypersphere. The square root mapping converts Fisher-Rao geodesics to great arcs on the sphere, with only a constant difference between the two metrics. ReSFlow uses spherical linear interpolation (slerp) to connect the given start and end distributions on the sphere. Deriving this spherical linear interpolation gives the conditional speed needed for training the flow model on the sphere. The model's predicted speed is mapped to the tangent space of the current position, ensuring it always follows Riemannian spherical geometry.

Spherical Rectification: From Random Pairing to Model Pairing

Stage 1: Train the model to move from noise to data. This stage is similar to traditional flow model training: independent random sampling from the prior and data distributions, mapping to the sphere to construct geodesics, and making the model's predicted tangential speed match the conditional speed. This stage establishes a generation flow from noise to data, but the conditional paths from independent random pairing have many crossings, so the overall trajectory learned by the model is still bent.

Stage 2: Continue training on the reduced-crossing trajectory. After initial training, ReSFlow samples start points from the prior distribution, then uses the current model to solve the spherical ODE to get the actual generation endpoints the start points reach. Then, these start points and generation endpoints are used as new pairs, retrained along the spherical geodesic between them. Since these pairs come from the model's own ODE trajectory, under the condition that the solution exists and is unique, the trajectory cannot move in different directions at the same time and position. Thus, the model reorganizes the movement direction at the intersections of the original conditional paths. Compared to independent random pairing, this pairing method significantly reduces crossings between conditional paths, making the overall trajectory learned by the model closer to spherical geodesics.

One-Step Generation

After the pairing adjustment and training, the model's learned trajectory gradually approaches geodesics. During ODE integration, there is no need to update direction repeatedly along the way: only one speed prediction in the tangent space at the initial noise is needed, then spherical exponential mapping is used to directly estimate the generation endpoint along the corresponding great arc. After getting the spherical endpoint, square each coordinate to restore the category probability distribution, then select or sample discrete categories based on the recovered probability.

Experimental Results

The research team first used the Swiss roll synthetic distribution to verify if the model could restore the target distribution in one step. Then, ReSFlow was applied to real tasks like images, text, and DNA sequences. Besides comparing generation results, analysis experiments were added to verify if the rectification made the model's trajectory closer to spherical geodesics. Baselines include LinearFM, SFM, and R-LFM, to compare Euclidean vs spherical geometry and before/after rectification.

Verification Experiment

On the Swiss roll synthetic distribution, ReSFlow's ODE sampling time for generating 1,000 samples was 0.0361 seconds, the shortest among the four methods; SFM took 0.0679 seconds. ReSFlow's NLL was -2.25, better than LinearFM and R-LFM, but slightly higher than SFM's -2.38 (lower NLL is better). This result shows that spherical rectification reduces sampling time while maintaining good probability modeling performance. The one-step generated samples intuitively reflect this: the real distribution has a curled band structure on the simplex. Unrectified LinearFM and SFM cannot restore this shape, with generated points concentrated in a small area, while ReSFlow retains the complete curled structure, closer to real samples.

Main Experiments

On the binarized MNIST dataset, the one-step FID of unrectified LinearFM and SFM was 25.317 and 35.640 respectively; after rectification, R-LFM dropped to 6.261, and ReSFlow further reached 5.973. In full ODE sampling, SFM had the lowest FID, but after compressing to one step, its quality loss was much larger. ReSFlow maintained FID from 5.541 in full ODE to 5.973 in one-step sampling. The comparison of different model calls shows that ReSFlow's one-call generation quality is close to SFM's 34-call result, with SFM's FID at 5.922 and ReSFlow's at 5.973. The sample images make this difference clearer: in full ODE sampling, all four methods can generate clear handwritten digits; in one-step, LinearFM and SFM have obvious missing strokes, making some digits hard to recognize, while ReSFlow retains more complete strokes and contours, with one-step samples closer to their full ODE results.

In the text generation task, ReSFlow's BPC was 1.26146, lower than SFM's 1.82998. However, Euclidean rectification did not help here: R-LFM's BPC rose to 3.26687, worse than unrectified LinearFM. This shows that re-pairing does not always bring improvement, and the geometry used for the path is equally important.

In the conditional promoter DNA design task, the generated object was a 1,024-length DNA sequence. ReSFlow's one-step sampling SP-MSE was 0.028718, close to full ODE sampling's 0.028743. The average time to generate 1,000 sequences was shortened from 67.94 seconds to 0.409 seconds, about 166 times faster. As a comparison, after compressing SFM to one step, SP-MSE rose from 0.028015 to 0.035564, with more obvious quality loss.

Analysis Experiment

Did the generation results improve because the model's trajectory became more "straight"? The team recorded transport cost and geodesic deviation during rectification on binarized MNIST. Transport cost measures the distance cost between the noise start and generation endpoint. ReSFlow's transport cost decreased significantly with rectification, indicating that after re-pairing, the model tends to send noise samples to generation endpoints with smaller geometric distance, reducing unnecessary long-distance transport. In contrast, R-LFM had a smaller decrease and entered a plateau earlier. This result shows that improving pairing in Euclidean space does not necessarily reduce transport cost in probability space equally.

Geodesic deviation checks if the model deviates from direction or changes speed along the trajectory: at each position on the trajectory, compare the actual predicted speed with the speed needed to reach the same endpoint along the geodesic in remaining time. The closer they are, the closer the model's actual trajectory is to a constant-speed geodesic. As the number of rectification rounds increases, the deviation gradually decreases, experimentally supporting ReSFlow's core idea: spherical rectification makes the model's actual generated path closer to geodesics. With fewer bends in the path, speed changes also conform more to the constant-speed geodesic rule, making one prediction at the start more likely to describe the entire movement. This explains why ReSFlow can use one exponential mapping to accurately estimate endpoints that would otherwise require multiple steps of integration.

Conclusion

ReSFlow provides a path for one-step generation of discrete data: starting from the Fisher-Rao geometry of category probabilities, re-pairing and rectifying on the sphere, making the generation trajectory gradually approach geodesics, then completing one-step sampling via exponential mapping. Experiments on images and DNA sequences show that this design can greatly reduce sampling computation while maintaining good generation performance.

One-step inference still requires additional rectification training, and its effect on larger-scale data and complex conditional tasks remains to be explored. But this work provides a clear inspiration: to make the generation path "straight", we must first define what "straight" means based on the geometry of the data space.

以上内容不代表本平台立场，仅供读者参考