ReSFlow: Rectifying Categorical Flows on Statistical Manifolds for One-Step Generation

1 Introduction

• Discrete flow models on statistical manifolds can transport priors to target distributions, but their aggregate trajectories may curve away from geodesics and require costly multi-step integration.The passages identify this computational burden as a bottleneck for scalable categorical generation.Source: §1 · pp. 1–2
• On the Fisher–Rao probability simplex, Euclidean straightening may leave the simplex or diverge from its intrinsic shortest paths, so efficient generation requires geodesic alignment.A geodesic is the length-minimizing curve that follows the manifold’s Riemannian structure.Source: §1 · p. 2
• ReSFlow maps categorical distributions to the positive orthant of a hypersphere and iteratively recouples and retrains flows there, progressively aligning trajectories with geodesics.The paper reports that one-step ReSFlow matches or outperforms state-of-the-art methods using hundreds of sampling steps across vision, NLP, and bioinformatics benchmarks.Source: §1 · p. 2

2 Preliminaries

• The square-root map is a diffeomorphism from the probability simplex to the positive orthant of the unit hypersphere, making the Fisher–Rao metric isometric to the sphere’s round metric.The notation distinguishes categorical distributions from distributions over categorical distributions and their source and target roles.Source: §2 · p. 2

2.1 Statistical Flow Matching

• Statistical Flow Matching uses the square-root map to place categorical distributions on the hypersphere, where Fisher–Rao geometry supports geodesic paths and flow matching.The formulation keeps paths on the simplex and avoids quantization issues associated with Euclidean treatments.Source: §2.1 · pp. 2–3

2.2 Rectified Flow

• Rectified Flow learns velocities along linear interpolations and iteratively retrains on endpoint pairs induced by its simulated flow, progressively straightening trajectories toward one-step integration.The passages describe this procedure in Euclidean space and note that extending it to the hypersphere requires replacing straight lines with geodesics.Source: §2.2 · p. 3

3.1 Why Euclidean Reflow Fails on the Statistical Manifold

• Euclidean reflow can leave Fisher–Rao geodesics intersecting even when its straightened linear paths do not, so the method reflows directly on the sphere.Source: §3.1 · p. 4

3.2 Geodesic Interpolation on the Hypersphere

• On the hypersphere, geodesic interpolation stays on the manifold at constant speed, and its geodesic integral curves enable exact single-step Euler simulation.Source: §3.2 · p. 5

3.3 Rectified Categorical Flow via Spherical Reflow

• The method trains a tangent-space velocity field on geodesic interpolants, then repeatedly recouples simulated endpoints and retrains to straighten trajectories.Source: §3.3 · p. 5
• After reflow converges, the initial velocity and spherical exponential map generate a target sample in one step.Source: §3.3 · p. 5

3.4 Theoretical Guarantees of Spherical Reflow

• Spherical reflow preserves marginals, makes spherical transport cost non-increasing, and guarantees progressive trajectory straightening.Source: §3.4 · p. 6

4.1 Experimental setup

• The setup separates geometry from rectification with a 2 × 2 comparison, tests ReDi only on Text8, and evaluates sample quality at NFE=1 and with ODE simulation.Likelihood-based metrics use ODE integration; architectures and hyperparameters are in Appendix E.Source: §4.1 · pp. 6–7

4.2 Toy Example: Swiss Roll on the Simplex

• ReSFlow records NLL −2.25 versus −1.88 for LinearFM, −1.87 for R-LFM, and −2.38 for SFM, and samples in 0.0361s versus SFM's 0.0679s.It is competitive with SFM on NLL and has the fastest ODE sampling time among the methods.Source: Table 1 · p. 7, §4.2 · p. 7
• For one-step Swiss-roll samples, ReSFlow most faithfully captures the target structure; Euclidean reflow improves on LinearFM but captures it only partially.Non-rectified LinearFM and SFM fail to recover the structure in one step.Source: §4.2 · p. 7, Figure 2 · p. 8

4.3 Binarized MNIST

• On binarized MNIST, ReSFlow achieves the best one-step FID (5.973) and generates 1,000 samples in 0.452 seconds, versus accurate-ODE SFM’s 112.81 seconds and FID 4.748.At one function evaluation, ReSFlow has mean FID 5.973 ± 0.442; SFM’s 5.922 ± 0.262 is at 34 evaluations.Source: Table 2 · p. 8, §4.3 · p. 8

4.4 Text8

• ReSFlow achieves the best Text8 BPC at 1.261, versus 2.306 for ReDi, 1.830 for SFM, and 2.154 for LinearFM.R-LFM has BPC 3.267.Source: Table 3 · p. 9, §4.4 · p. 9

4.5 Promoter Sequence Design

• ReSFlow's one-step SP-MSE is 0.028718 versus 0.028743 under ODE, while sampling time falls from 67.94 to 0.409 seconds per 1,000 sequences.Under ODE sampling, ReSFlow is nearly twice as fast as LinearFM: 67.94 versus 119.04 seconds per 1,000 sequences.Source: §4.5 · pp. 9–10, Table 4 · p. 10

5 Conclusion

• ReSFlow performs flow matching on the probability simplex under spherical geometry, establishes a theoretical straightening-rate guarantee, and outperforms existing baselines across diverse discrete-generation tasks.The conclusion identifies single-step generation on other spherical data as future work.Source: §5 · p. 10

A Statistical Manifold of Categorical Distributions and the Spherical Projection

• The Fisher metric becomes ill-defined near one-hot targets on the simplex boundary, while spherical dynamics resolve the singularity and enable stable, exact flow matching to them.The square-root diffeomorphism maps the simplex to the positive orthant of the unit sphere while preserving intrinsic geometry up to a constant factor.Source: A Statistical Manifold of Categorical… · p. 13

B Related Work

• SFM provides a statistical-manifold foundation, whereas ReDi relies on pretrained token embeddings and uses Euclidean paths that are not simplex geodesics.The related-work discussion says ReDi's practical advantages over standard discrete flow matching remain less clear-cut than rectified flow's advantages in continuous settings.Source: B Related Work · p. 14

C.1 Tangent Space Projection

• Tangent-space projection removes the normal component of each predicted velocity so the flow remains on the sphere.The velocity at each point must lie in that point's tangent space.Source: §C.1 · p. 14

C.2 Algorithm Procedure

• Each reflow iteration trains on the current coupling, computes spherical geodesic interpolants and tangent velocities, and minimizes a projected velocity-matching loss.The algorithm samples time uniformly on [0, 1] and draws endpoint pairs from the current coupling.Source: §C.2 · p. 14
• After training converges, the learned ODE produces a new endpoint coupling, and inference applies a single Euler step through the exponential map.The output is mapped back to the probability simplex.Source: §C.2 · p. 14

D Spherical Rectified-Flow Proofs

• The proof overview states that spherical rectification has a straightening effect with an O(1/Nreflow) rate.The analysis replaces Euclidean interpolation with spherical geodesic interpolation.Source: D Spherical Rectified-Flow Proofs · pp. 14–15

D.1 Geometric and Probabilistic Preliminaries

• On the sphere, the shorter great-circle arc is a minimizing geodesic with constant speed, and path length and energy bound endpoint geodesic distance.The appendix also establishes conditional-tangent-vector identities used in the proofs.Source: §D.1 · pp. 15–16

D.2 Marginal Preservation

• For rectifiable processes, the spherical flow induced by the expected velocity field preserves the process's time marginals.The proof establishes that the original process and its rectified flow satisfy the same continuity equation.Source: §D.2 · p. 17

D.3 Spherical Transport-Cost Contraction

• Spherical rectification contracts costs given by convex nondecreasing functions of geodesic distance for spherically rectifiable couplings.This is the manifold analogue of controlling convex costs of Euclidean displacement.Source: §D.3 · p. 18

D.4 The Straightening Effect

• A spherically rectifiable coupling is straight exactly when rectification leaves it fixed, its flow matches spherical geodesic interpolation, and its spherical velocity variance is zero.The theorem also includes unchanged squared spherical transport cost among these equivalent conditions.Source: §D.4 · p. 19
• When squared spherical transport cost is unchanged, the velocity variance is zero, so the original process follows the same characteristic equation as its rectified flow.Uniqueness under spherical rectifiability then implies that the process and rectified flow coincide.Source: §D.4 · pp. 19–20
• Under the stated rectifiability assumptions, the reflow straightening bound yields an O(1/Nreflow) rate.The proof obtains the rate by summing the per-step rectification bound.Source: §D.4 · p. 20

E.1 Swiss Roll on the Simplex

• Swiss-roll samples form a fixed 1,000-point empirical target on the 2-simplex, with no separate validation or test split.Its likelihoods are relative comparisons on that same empirical distribution, not generalization metrics.Source: §E.1 · p. 21

E.2 Binarized MNIST

• Binarized MNIST uses 28 × 28 images represented as 784 binary categorical variables, with a standard 50,000/10,000/10,000 train/validation/test split.The compared methods share a convolutional backbone; ReSFlow uses spherical tangent reflow and is fine-tuned from the prior checkpoint.Source: §E.2 · p. 22

E.3 Text8

• Text8 uses character-level windows of 256 symbols from a 27-symbol vocabulary, with the standard split of 90M characters for training, 5M for validation, and 5M for testing.The reflow model is initialized from the previous checkpoint and fine-tuned; ReSFlow uses spherical parameterization, while R-LFM uses Euclidean reflow parameterization.Source: §E.3 · pp. 22–23

E.4 Promoter Sequence Design

• Promoter design uses 1,024 bp one-hot DNA sequences paired with aligned FANTOM CAGE signals, with chromosome-based train, validation, and test splits.ReSFlow’s single reflow stage keeps each condition fixed while generating endpoint pairs, then fine-tunes the base checkpoint.Source: §E.4 · pp. 23–24

E.5 Evaluation Metrics, NFE, and Sampling Protocols

• Swiss-roll NLL is a differential likelihood evaluated with adaptive ODE integration, not one-step sampling.It treats the data as samples from a continuous density on the simplex or its square-root spherical representation.Source: §E.5 · p. 24
• Binarized MNIST FID measures sample quality under both adaptive ODE and one-step sampling, with lower values indicating closer match to the reference distribution.The two settings separate continuous-time flow quality from one-step sampler quality.Source: §E.5 · p. 24
• Text8 BPC uses a variational integral objective rather than autoregressive token-wise conditionals and is reported only for ODE-based evaluation.The variational negative log-likelihood is normalized by sequence length and converted from nats to bits.Source: §E.5 · pp. 24–25
• Promoter SP-MSE compares predicted activity profiles for generated and real sequences under the same conditioning signal, with lower values indicating better preservation of the desired activity.It is a downstream functional metric rather than a measure only of sequence-space distributional similarity.Source: §E.5 · p. 25
• One-step generation uses NFE=1: ReSFlow evaluates velocity once at t = 0, applies the spherical exponential map, and samples or selects a category from the resulting distribution.NFE counts learned-vector-field evaluations during sampling.Source: §E.5 · p. 25
• Wall-clock sampling time covers generation of 1,000 samples for binarized MNIST and promoter design, excluding training, checkpoint loading, and metric computation.Text8 sampling time is omitted because its reported BPC comes from a variational likelihood integral rather than generated samples.Source: §E.5 · p. 25

F.1 FID versus NFE on Binarized MNIST

• At NFE=34, SFM first has a lower mean FID than the one-step ReSFlow mean of 5.973.The comparison uses a matched spherical exponential-map Euler sampler, three seeds, and 1,000 generated samples per seed.Source: Figure 3 · p. 26
• Figure 3 plots Binarized MNIST FID against NFE using a matched spherical exponential-map Euler sampler.The solid curve is the SFM mean and the dashed line marks the one-step ReSFlow mean.Source: Figure 3 · p. 26
• The plotted means are based on three seeds and 1,000 generated samples per seed.Source: Figure 3 · p. 26

F.2 Fisher–Rao Transport Cost Analysis of Reflow Couplings

• Across reflow rounds, ReSFlow substantially reduces Fisher–Rao coupling cost, while R-LFM achieves only a mild reduction and quickly plateaus.The independent coupling has the largest cost; mini-batch FR-OT reduces it, but the improvement is limited.Source: Figure 4 · p. 26, §F.2 · p. 27
• ReSFlow’s lower intrinsic transport cost is attributed to rectification in the square-root hyperspherical representation, whereas Euclidean reflow is less aligned with Fisher–Rao geometry.Source: §F.2 · p. 27

F.3 Riemannian Geodesic Straightness of Learned Trajectories

• The Riemannian geodesic-straightness metric decreases monotonically across successive reflow rounds on Binarized MNIST.A smaller value indicates a more geodesic learned ODE path and greater suitability for one-step generation via the exponential map.Source: §F.3 · p. 28

• Non-rectified methods’ reasonable ODE digit samples degrade under one-step sampling, while rectified methods better preserve digit structure and ReSFlow gives clear samples.Source: §F.4 · p. 28
• ReSFlow produces more structured Text8 character-level fragments than the non-rectified baselines, consistent with the BPC results in Table 3.Source: §F.4 · p. 30
• ReSFlow maintains stable nucleotide patterns under one-step sampling in 1024-bp promoter samples; the displayed sequences show only the first 120 bases, while metrics use full sequences.Source: Table 6 · p. 29, §F.4 · p. 30

G Limitations and Broader Impact

• Fisher–Rao geometry treats categorical states symmetrically and omits semantic priors, potentially limiting tasks that depend on hierarchical or pre-trained semantic correlations.For example, it does not encode that “cat” is conceptually closer to “dog” than to “apple.”Source: G Limitations and Broader Impact · p. 30
• ReSFlow is a foundational method rather than a deployment-specific system; its efficiency may benefit language, biology, and scientific applications, while faster, higher-quality generation could enable disinformation or other malicious content.The paper also associates its straightening property with reduced inference cost and energy footprint.Source: G Limitations and Broader Impact · p. 30

NeurIPS Paper Checklist

• The authors state that theoretical results include their assumptions and complete proofs in the appendix, with key formulations and proof intuitions presented in the main text.The checklist says theorem and lemma statements are numbered and cross-referenced.Source: NeurIPS Paper Checklist · p. 32
• The paper reports providing datasets, metrics, model configurations, training procedures, and baselines for reproducibility, with further details in the appendix and code instructions in the supplement.The submitted code archive is slated for public release upon acceptance.Source: NeurIPS Paper Checklist · pp. 32–33
• The checklist says error bars are reported where repeated runs are feasible; otherwise, sample-based evaluations use more generated samples and a common protocol, alongside reported hardware, GPU memory, and approximate run times.These disclosures are intended to support interpretation and reproduction of the experiments.Source: NeurIPS Paper Checklist · p. 34
• The authors state that the work uses standard benchmarks without personally identifiable information, human-subject experiments, or unsafe data collection, and does not release high-risk pretrained models or deployment-ready systems.They characterize the contribution as a foundational methodological framework.Source: NeurIPS Paper Checklist · pp. 34–35
• The authors answer “N/A,” stating that LLMs are not important, original, or non-standard components of the core method development.They direct readers to the NeurIPS handbook’s LLM policy for guidance on what to describe.Source: NeurIPS Paper Checklist · p. 37