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

Xiaohuan Jia; Renzhe Xu; Xiao Wang; Jiayun Wu; Shaohua Fan

cs.LGOct 9, 2026

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

Authors: Xiaohuan Jia, Renzhe Xu, Xiao Wang, Jiayun Wu, Shaohua Fan

Organizations: Inner Mongolia University · Shanghai University of Finance and Economics · Carnegie Mellon University

Abstract

Generative modeling for discrete data has been significantly advanced by learning continuous vector fields on the statistical manifolds to transport simple priors to target discrete distributions. However, because these continuous-time flows are not constrained to follow geodesics on the statistical manifold, thereby requiring multi-step numerical integration for simulation and leading to substantial computational costs during inference. In this work, drawing inspiration from the "straight and fast" properties of Rectified Flow in continuous Euclidean space, we propose the Rectified Statistical Flow (ReSFlow) model, which enables one-step categorical generation on statistical manifolds by explicitly learning geodesic trajectories that adhere to the intrinsic Riemannian geometry. By rectifying the flow on the statistical manifold, we theoretically guarantee that trajectories between prior and target distributions approximate geodesics. This rectification significantly reduces the transport cost by aligning the learned trajectories directly along geodesics on the statistical manifold, allowing the model to map pure noise to target discrete data in a single step. Through extensive experiments on diverse real-world generative tasks, ranging from image and text to biological domains, we demonstrate ReSFlow's superior one-step generation quality, which truly embodies geodesically optimal and rapid categorical generation. Our code is available at https://github.com/xiaohuanjia/ReSFlow.

Figures & tables

Figure 1: Overview of Spherical Rectified Flow. Top: Standard reflow straightens linear paths on the simplex, but their corresponding spherical geodesics still intersect, preventing one-step sampling. Bottom: We perform reflow directly on the sphere $\mathbb{S}_{+}$ . This straightens the geodesics themselves to be non-intersecting, ensuring faithful one-step generation.

Metric	LinearFM	SFM Cheng et al. (2024)	R-LFM	ReSFlow
NLL $\downarrow$	-1.88 ± 0.06	-2.38 ± 0.08	-1.87 ± 0.06	-2.25 ± 0.08
Time (s) $\downarrow$	0.1059	0.0679	0.0512	0.0361

Table 1: Swiss-roll results on the 2-simplex. We report ODE-based NLL and ODE sampling time for generating 1,000 samples. Bold indicates the best result, and ‘

*

’ indicates the second-best result.

Figure 2: Swiss-roll samples on the simplex under one-step Euler sampling.

Method	ODE (NFE = adaptive)		Euler one-step (NFE = 1)
Method	FID $\downarrow$	Time / 1k samples (s) $\downarrow$	FID $\downarrow$	Time / 1k samples (s) $\downarrow$
LinearFM Cheng et al. (2024)	5.187 ± 0.129	87.13 ± 0.44	25.317 ± 0.517	0.406 ± 0.178
SFM Cheng et al. (2024)	4.748 ± 0.092	112.81 ± 8.20	35.640 ± 0.487	0.719 ± 0.175
R-LFM	5.681 ± 0.252	61.19 ± 0.28	6.261 ± 0.379	0.583 ± 0.244
ReSFlow	5.541 ± 0.252	36.07 ± 0.60	5.973 ± 0.442	0.452 ± 0.243

Table 2: Binarized MNIST results. Mean

\pm

standard deviation over three random seeds; time is for generating 1,000 samples. Bold indicates the best mean, and

*

indicates the second-best mean.

Method	BPC $\downarrow$ (adaptive ODE)
LinearFM Cheng et al. (2024)	2.15448 ± 0.00024
SFM Cheng et al. (2024)	1.82998 ± 0.00028
R-LFM	3.26687 ± 0.00012
ReDi Yoo et al. (2025)	2.30562 ± 0.01010
ReSFlow	1.26146 ± 0.00036

Table 3: Text8 character-level modeling results. ODE-based BPC, reported as mean

\pm

standard deviation over three random seeds. Bold indicates the best mean, and

*

indicates the second-best mean.

Method	ODE (NFE=adaptive)		Euler one-step (NFE=1)
Method	SP-MSE $\downarrow$	Time $\downarrow$	SP-MSE $\downarrow$	Time $\downarrow$
LinearFM Cheng et al. (2024)	0.027734 ± 0.000897	119.04 ± 16.57	0.029449 ± 0.001673	0.410 ± 0.004
SFM Cheng et al. (2024)	0.028015 ± 0.000257	82.18 ± 4.26	0.035564 ± 0.007768	0.410 ± 0.005
R-LFM	0.052146 ± 0.000781	67.88 ± 4.17	0.060504 ± 0.000351	0.409 ± 0.001
ReSFlow	0.028743 ± 0.000755	67.94 ± 0.35	0.028718 ± 0.000366	0.409 ± 0.003

Table 4: Conditional promoter design results. Mean

\pm

standard deviation over three random seeds; time is for generating 1,000 sequences of length 1024. Bold indicates the best mean, and

*

indicates the second-best mean.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Binarized MNIST FID versus NFE using the matched spherical exponential-map Euler sampler, three seeds, and 1,000 generated samples per seed. The solid curve shows the SFM mean; the dashed line marks the one-step ReSFlow mean (

5.973

). SFM first has a lower mean at NFE

=34

Figure 4: Fisher–Rao transport cost and Geodesic Straightness on BMNIST.

Figure 5: Generated samples on binarized MNIST. We compare LinearFM, SFM, R-LFM, and ReSFlow under ODE and one-step Euler sampling. Non-rectified methods degrade under one-step sampling, whereas rectified methods preserve digit structure more effectively. ReSFlow gives high-quality one-step samples by rectifying trajectories under the Fisher–Rao geometry.

Model	Sampler	Generated text
LinearFM	ODE	calls_vincenti_or_their_burnt_mochery_a_distinct_writing_artiast_is_also_known_as_the_yodioi_cascoliage_it_is_a_crutter_with_bianaple_it_is_known_as_a_while_green_hole_and_even_when_using_plhuhu_snow_and_wouldistay_is_one_of_several_ancient_nuts_of_early_w
LinearFM	Euler	e_______________________________________________________________________________________________e___e_______________________e_____
SFM	ODE	tters_improvements_a_primary_work_to_most_word_genetic_cinematics_who_also_chose_a_word_for_genetic_generation_concurrence_with_a_pt_typical_sporting_the_world_is_often_the_most_for_compatisans_works_such_as_sorcery_and_the_dark_chasen_movements_the_player
SFM	Euler	izg_i_ee_e_e__fe______________t_______g_t__e__a_ote_____e____e_________e___d__a__i______t__q__t__tebe_s______e____ee___s__d_i_oe_i__s___ee__ene___o_e__tei__eoe___e_n_____e______a_ea__a__yer_ie__ei_teein__s_n__________e____o___iu___n__eo___f__ag______ho
R-LFM	ODE	_ners_vincento_wrata_orebry_n_mother_held_at_introducing_artiary_tale_two_beek_in_the_change_casckling_and_is_spyr_operating_fiama_of_rownspart_s_symbols_large_numbers_and_even_opecause_of_his_personalny_would_state_there_in_several_day_baptists_of_eyric_w
R-LFM	Euler	ieto__uaibe_and_the_early_ver_ion_an__a_aeementary_tounterlaei_eoond_the_forst_medonies_ere_ted_from_greatly_attended_from_the_naen_so_condonati_n_tiliter__head_tittre_ont_weec_sters_tore_tad_speecres_and_which_table_to_remume_expending_that_the_leies_anni

Appendix

Table 5: Generated Text8 samples. We show representative character-level samples generated by each model under ODE and one-step Euler sampling. Underscores denote spaces in the decoded text.

Model	Sampler	First 120 bases of generated promoter sequence
LinearFM	ODE	AAGAGCTCGGCCGCGGGCTCGGGTGCCCTCCCCCCGGCCC CAAAGCCCCGGGGATTTCCAAGGCCGGCAAGGCCGGGGGG CAGGGGGGAAGAGGGGCGGCACTGGGAGGGGGGATGCGAA
LinearFM	Euler	CTCGCCATCGGGCAGGGGGCTATTTTGCATTCTAACGCCA GGCAGTAGTCGTTTGGTCGCATGCAAACTCAGATGCGGGG GATTGGGTGGCATTCGGTCTCGGCGTTAGACCGAAACCAG
SFM	ODE	AAGAGCTAAGCCGCGGGCTCAGGTGCCCTCCCCCCAGCCC CAAAGCCCCGGGGAGCTCCAACGCCACCACCCCCAGGCGG CAGACCGGAACACGGGCGCCACTGCCAGGAAGGATGCGAA
SFM	Euler	AGGGGCCAAGCCCCGGCGGCGCGTGCGGGGCCGTCAGCTT CAGCCCCCCAGCGGATCCCGTAGACGCCAGCCCCAGGCGG CACTCCGGGACGCGGGCGCCGTCCCGGAGAAGGGCGCGGA
R-LFM	ODE	AAGAACTCCGCCGCGGGCTCGGGTGCCCTCCCCCCGGCCC CAAAGCCGCGGGGAACGCCACCGCCACCACAGCCGGCCGG CAGAAGGGAACAGGGGCGGCCCTAGCAGGCGGGAGGCGAA
R-LFM	Euler	AAGAACTCGGCCGCGGGCTCAGGTGCACTACCCCCGGCCC CAAAGCCGCGGGAAACGCCACCGCCAGCAAAGCCGGGGGG CAGAAGGGAACAGGGGCGGCCCTAGGAGGGGGGAGGCGAA

Appendix

Table 6: Generated promoter sequence samples. Each generated promoter has length

1024

. For readability, we show only the first

120

bases of one selected sample for each model and sampler. Spaces are inserted every

40

bases only for line breaking and are not part of the DNA sequence. Quantitative metrics are computed using the full-length sequences.

Generative Modeling

Generative modeling learns data distributions to create new samples, from images and text to molecules, simulations, and structured records.

See all 492 papers

RW-Flow: One-Step Generation on Compact Manifolds via Wasserstein Gradient Flows
Sep 30, 2026Ualibyek Nurgulan, Seungwoo Yoo, Prin Phunyaphibarn +1Wasserstein Gradient Flows One-Step Generative Modeling
KAIST
Discrete Wasserstein Flows for One-Step Generative Modeling
Oct 1, 2026Alessandro Micheli, Andrea Zerio, Samir BhattWasserstein Gradient Flows Markov Chains
Imperial College London London, United Kingdom · Department of Computer Science, Aalborg University, Copenhagen, Denmark; Centre for Frontier AI Research (CFAR), Institute of Advanced Intelligence and Computing (IAIC), A*STAR, Singapore · University of Copenhagen Copenhagen, Denmark
Riemannian MeanFlow for One-Step Generation on Manifolds
Mar 11, 2026Zichen Zhong, Haoliang Sun, Yukun Zhao +2Generative Modeling Flow-Based Generative Modeling
School of Software, Shandong University, Jinan, China

Sep 30, 2026J/K move · Enter open · S save

RW-Flow: One-Step Generation on Compact Manifolds via Wasserstein Gradient Flows

Ualibyek Nurgulan, Seungwoo Yoo, Prin Phunyaphibarn, Minhyuk Sung

KAIST

Wasserstein Gradient Flows One-Step Generative Modeling Generative Modeling Flow-Based Generative Modeling Wasserstein Distance

Open paper AskPDF arXiv

Manifold-valued data, and consequently the distributions they induce, are prevalent across many domains, ranging from the locations of geospatial events, such as earthquakes, to biomolecular torsion angles that encode information about three-dimensional structure. While diffusion and flow-based generative models have been successfully extended to compact manifolds, sampling typically requires tens or hundreds of sequential network evaluations. We introduce RW-Flow, a theoretically grounded framework for learning one-step generative models on compact manifolds via Wasserstein gradient flows. The main challenge is identifiability: driving the velocity field to zero should guarantee that the model distribution matches the target distribution. We establish a necessary and sufficient condition for identifiability on compact, connected Riemannian manifolds. We specifically show that, for a symmetric, Lipschitz-continuous cost function, the velocity field induced by the Sinkhorn divergence is identifiable if and only if the associated Gibbs kernel is nondegenerate. This characterization provides a general principle for designing identifiable costs on compact manifolds. It also reveals that the squared geodesic distance, the natural manifold analogue of the squared Euclidean distance, does not always guarantee identifiability. Across benchmarks involving geospatial events, protein side chain torsion angles, RNA backbone torsion angles, and general manifolds discretized as triangular meshes, RW-Flow outperforms existing one-step methods in nearly all settings under fair comparison conditions.

Abstract

Generative Modeling

Explore similar work

RW-Flow: One-Step Generation on Compact Manifolds via Wasserstein Gradient Flows

Discrete Wasserstein Flows for One-Step Generative Modeling

Riemannian MeanFlow for One-Step Generation on Manifolds