ReSFlow Enables One-Step Discrete Generation by Rectifying Flows on Statistical Manifolds

Research Background

Discrete data generation typically models categorical probabilities (e.g., DNA bases A, C, G, T as probability vectors). Models learn how these probabilities evolve over time and use ODE integration to sample discrete results. Statistical Flow Matching (SFM) introduces Fisher–Rao geometry, using geodesics in probability space as training paths, but one-step generation faces two challenges: (1) training paths are geodesics, but inference paths may curve due to crossing conditional paths from independent random pairing; (2) Euclidean straight lines differ from Fisher–Rao geodesics — even non-intersecting Euclidean paths can intersect when mapped to spherical geodesics.

Core Method: ReSFlow

ReSFlow addresses both issues by mapping categorical probabilities onto a sphere via the square-root mapping, then performing rectification directly on the spherical manifold.

Spherical Representation and Geodesic Interpolation

For a categorical probability vector p, ReSFlow applies element-wise square root to map it from the probability simplex to the non-negative region of the unit hypersphere: x = sqrt(p). This mapping makes Fisher–Rao geodesics correspond to great-circle arcs on the sphere, with the two metrics differing only by a constant factor. ReSFlow then uses spherical linear interpolation (slerp) to connect a source distribution x_0 and target distribution x_1 on the sphere: x_t = (sin((1-t)θ)/sin θ) x_0 + (sin(tθ)/sin θ) x_1 where θ = arccos(x_0^T x_1). The conditional velocity v_t is obtained by differentiating this interpolation. The model's predicted velocity is projected onto the tangent space at the current position, ensuring predictions respect Riemannian spherical geometry.

Spherical Rectification: From Random Pairing to Model Pairing

Stage 1: Train the model to move from noise to data. As in standard flow models, source and target samples are drawn independently from prior and data distributions, mapped to the sphere, and connected via geodesics. The model learns to predict the tangent velocity matching the conditional velocity.

This stage establishes a generative flow, but independent pairing causes many conditional path crossings, leading to curved overall trajectories.

Stage 2: Reduce crossings by re-pairing with model-generated endpoints. After initial training, ReSFlow samples starting points from the prior, solves the spherical ODE using the current model to obtain the actual generated endpoints. These (start, generated end) pairs become new training pairs, connected by spherical geodesics. Because these pairs come from the model's own ODE trajectories, under existence and uniqueness conditions, trajectories cannot cross at the same time and position with different directions. The model reorganizes directions at previous crossing points, significantly reducing conditional path crossings and making learned trajectories closer to spherical geodesics.

One-Step Generation

After rectification, the model's trajectories approximate geodesics. For generation, the model predicts a single velocity at the initial noise point in its tangent space, then uses the spherical exponential map to travel along the corresponding great-circle arc directly to the estimated endpoint: x_1 = cos(θ) x_0 + sin(θ) v_0 / ||v_0|| where θ = ||v_0||. The spherical coordinates are then squared element-wise to recover categorical probabilities: p = x_1^2, from which discrete categories are sampled or selected.

Experimental Results

Verification on Swiss Roll

On a synthetic Swiss roll distribution, ReSFlow generates 1,000 samples in 0.0361 seconds (fastest among LinearFM, SFM, R-LFM, ReSFlow). Its negative log-likelihood (NLL) is −2.25, better than LinearFM and R-LFM but slightly higher than SFM's −2.38. One-step samples from ReSFlow preserve the full curved structure of the true distribution, while LinearFM and SFM collapse to a smaller region.

Main Experiments

Binarized MNIST: One-step FID scores: LinearFM 25.317, SFM 35.640, R-LFM 6.261, ReSFlow 5.973. Full ODE FID: SFM 5.541, ReSFlow 5.973 (minimal degradation). ReSFlow with one model call matches the quality of SFM with 34 calls (SFM FID 5.922 vs ReSFlow 5.973). Visual samples show ReSFlow retains complete strokes and contours in one step.

Text8 character-level modeling: ReSFlow achieves bits per character (BPC) 1.26146, outperforming SFM (1.82998). Euclidean rectification harms performance: R-LFM BPC rises to 3.26687, worse than unrectified LinearFM, demonstrating that geometry choice is critical.

Conditional promoter DNA design (length 1,024): One-step SP-MSE 0.028718 vs full ODE 0.028743 (nearly identical predictive activity spectrum error). Generation time for 1,000 sequences drops from 67.94 seconds to 0.409 seconds (~166× speedup). SFM compressed to one step degrades more (SP-MSE from 0.028015 to 0.035564).

Analysis Experiments

On binarized MNIST, the authors track transport cost and geodesic deviation during rectification. Transport cost (distance between noise start and generated end) decreases markedly with ReSFlow, indicating re-pairing moves noise samples to geometrically closer targets, reducing unnecessary long-distance transport. R-LFM shows smaller reduction and earlier plateau, showing Euclidean rectification does not equally lower probability-space transport cost. Geodesic deviation (comparing model-predicted velocity to constant-speed geodesic velocity needed to reach the same endpoint) decreases with rectification steps, empirically confirming that spherical rectification makes generated paths approach geodesics. Straighter paths with more constant velocity explain why a single exponential map from the start can accurately estimate the endpoint otherwise requiring multi-step integration.

Conclusion

ReSFlow provides a pathway for one-step discrete generation: starting from Fisher–Rao geometry of categorical probabilities, it re-pairs and rectifies on the sphere to align trajectories with geodesics, then uses the exponential map for one-step sampling. Image and DNA experiments show drastic sampling compute reduction while preserving generation quality. One-step inference still requires extra rectification training; scaling to larger data and complex conditional tasks remains future work. The key insight: to make generation paths "go straight," one must first define "straight" according to the geometry of the data's space.

Paper: https://arxiv.org/abs/2610.13225 Code:

https://github.com/xiaohuanjia/ReSFlow