Riemannian MeanFlow for One-Step Generation on Manifolds

Zichen Zhong; Haoliang Sun; Yukun Zhao; Yongshun Gong; Yilong Yin

cs.LGMar 11, 2026

Riemannian MeanFlow for One-Step Generation on Manifolds

Authors: Zichen Zhong, Haoliang Sun, Yukun Zhao, Yongshun Gong, Yilong Yin

Organizations: School of Software, Shandong University, Jinan, China

Abstract

Flow Matching enables simulation-free training of generative models on Riemannian manifolds, yet sampling typically still relies on numerically integrating a probability-flow ODE. We propose Riemannian MeanFlow (RMF), extending MeanFlow to manifold-valued generation where velocities lie in location-dependent tangent spaces. RMF defines an average-velocity field via parallel transport and derives a Riemannian MeanFlow identity that links average and instantaneous velocities for intrinsic supervision. We make this identity practical in a log-map tangent representation, avoiding trajectory simulation and heavy geometric computations. For stable optimization, we decompose the RMF objective into two terms and apply conflict-aware multi-task learning to mitigate gradient interference. RMF also supports conditional generation via classifier-free guidance. Experiments on spheres, tori, SO(3), and SE(3) demonstrate competitive one-step sampling with improved quality-efficiency trade-offs and substantially reduced sampling cost.

Figures & tables

Figure 1: Instantaneous velocity (A→B) vs. average velocity (P→Q) on Euclidean straight lines and manifold geodesics.

Panel of Figure 2 — Figure 2: Cosine similarity between full-parameter gradients $\nabla_{\theta}\mathcal{L}_{1}$ and $\nabla{\theta}\mathcal{L}_{2}$ across Earth categories during training (per-iteration values and running average); negative values indicate gradient conflicts.

	Volcano	Earthquake	Flood	Fire
Dateset size	827	6,120	4,875	12,809
RFM	0.351	0.309	0.272	0.377
RCT	0.155	0.053	0.086	0.080
G-PSD	0.139	0.045	0.067	0.046
G-ESD	0.184	0.178	0.178	0.196
G-LSD	0.115	0.032	$\underline{0.065}$	0.027

Table 1: MMD (

\downarrow

) at 1 NFE on the Spherical test dataset. Values are averaged over 5 runs with different random seeds. Best is bold ; second best is underlined .

Panel of Figure 3 — Figure 3: Earth dataset: generated sample distributions overlaid with ground-truth test data (red). Top-left: Volcano; top-right: Earthquake; bottom-left: Flood; bottom-right: Fire.

	General	Glycine	Proline	PrePro	RNA
Dataset size	138,208	13,283	7,634	6,910	9,478
RFM	0.45	0.27	0.52	0.47	0.68
RCT	0.01	0.04	0.05	0.06	0.11
G-PSD	0.11	0.05	0.07	0.08	0.14
G-ESD	0.29	0.13	0.44	0.26	0.45
G-LSD	0.02	0.03	0.04	0.05	0.08

Table 2: MMD (

\downarrow

) at 1 NFE on Torus test datasets (mean over 5 runs with different random seeds). Best is bold ; second best is underlined .

	Cone	Fisher	Line	Swiss Roll
Dateset size	20k	40k	40k	40k
RFM	0.311	0.154	0.095	0.346
RCT	0.096	$\underline{0.071}$	0.051	0.082
G-LSD	0.044	0.073	$\underline{0.037}$	0.032
RMF	0.057	0.076	0.038	0.080
RMF-MT	$\underline{0.049}$	0.039	0.035	$\underline{0.061}$

Table 3: MMD (

\downarrow

) at 1 NFE on the

\mathrm{SO}(3)

test set (mean over 5 runs with different random seeds). Best is bold ; second best is underlined .

Step	1	2	3	4	5	6	7
Dataset size	15.6M
RFM	3.2	23	38	59	80	82	88
G-LSD	60	75	81	88	88	87	90
RMF	65	80	82	86	88	88	90
RMF-MT	60	67	70	72	75	78	81

Table 4: Grasping success rate

\%

(

\uparrow

) on the

\mathrm{SE}(3)

dataset across different sampling steps. Best is bold ; second best is underlined . Notice: The metric is the grasping success rate, not MMD.

	Cone	Fisher	Swiss	Mixture
uncond	-	-	-	0.439
CFG	0.115	0.159	0.546	0.234

Table 5: MMD (

\downarrow

) at 1 NFE on the

\mathrm{SO}(3)

mixture test set.

MMD/epoch	3	4	8	16	32	64	128
EMF	0.193 / 50	0.154 / 32	0.108 / 60	0.115 / 51	0.150 / 50	0.159 / 132	0.357 / $>400$
RFM	0.088 / 18	0.080 / 16	0.073 / 2	0.054 / 2	0.040 / 2	0.037 / 2	0.031 / 2
RMF	0.066 / 18	0.056 / 10	0.055 / 6	0.038 / 6	0.035 / 6	0.028 / 4	0.026 / 6

Table 6: MMD (

\downarrow

) for 1 NFE on hyperspheres

\mathbb{S}^{d-1}

, with epochs to reach each score.

Panel of Figure 4 — Figure 4: Visualization of samples generated by CFG and unconditional models.

Method	NFE	MSE( $\downarrow$ )	k-mer corr. ( $\uparrow$ )
Dirichlet FM	100	0.034	N/A
Fisher FM	100	0.030	0.96
Eulerian RMF	1	0.030	0.96
Lagrangian RMF	1	0.027	0.88
Semigroup RMF	1	0.030	0.84
RMF(ours)	1	0.030	0.94

Table 7: Promoter DNA sequence generation results. Best is bold ; second best is underlined .

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

	Volcano	Earthquake	Flood	Fire
Dateset size	827	6,120	4,875	12,809
RMF	0.092	0.042	0.068	0.042
RMF-MT	$\underline{0.102}$	0.035	0.048	0.032
$\alpha$ RMF	0.106	0.051	0.061	0.043
iRMF	0.108	$\underline{0.039}$	$\underline{0.060}$	$\underline{0.037}$

Appendix

Table 8: MMD (

\downarrow

) on the spherical dataset (mean over 5 runs with different random seeds). Best is bold ; second best is underlined .

Panel of Figure 5 — Figure 5: iRMF: Cosine Similarity about $\nabla\mathcal{L}_{1}(\theta)$ and $\nabla\mathcal{L}_{2}(\theta)$ with $0\%$ r=t on Spherical Datasets.

	RMF	RMF-MT	G-LSD ( Davis et al., 2026 )
iteration/second	8.27	$\underline{6.42}$	3.60

Appendix

Table 9: Iterations per second (

\uparrow

) for different methods.

Panel of Figure 6 — Figure 6: Cosine Similarity about $\nabla\mathcal{L}_{1}(\theta)$ and $\nabla\mathcal{L}_{2}(\theta)$ with $0\%$ r=t

Panel of Figure 7 — Figure 7: Multi-step MMD comparison between RFM and RMF across all datasets.

	Volcano	Earthquake	Flood	Fire
Dataset size	827	6,120	4,875	12,809
RDM	$-6.61\pm 0.97$	$-0.40\pm 0.05$	$0.43\pm 0.07$	$-1.38\pm 0.05$
RFM	$-\textbf{7.93}\pm\textbf{1.67}$	$-0.28\pm 0.08$	$0.42\pm 0.05$	$-1.86\pm 0.11$
G-LSD	$-4.96\pm 0.68$	$-0.93\pm 0.01$	$-0.38\pm 0.33$	$-2.14\pm 0.42$
G-PSD	$-3.50\pm 0.22$	$-0.63\pm 0.13$	$-0.76\pm 0.13$	$-2.48\pm 0.71$
G-ESD	$-4.49\pm 0.20$	$-0.67\pm 0.08$	$-0.88\pm 0.38$	$-2.29\pm 0.08$

Appendix

Table 10: NLL(

\downarrow

) on the Spherical Dataset. Standard deviation is estimated over 5 runs.

	General	Glycine	Proline	PrePro	RNA
Dataset size	138,208	13,283	7,634	6,910	9,478
RDM	$1.04\pm 0.012$	$1.97\pm 0.012$	$0.12\pm 0.011$	$1.24\pm 0.004$	$-3.70\pm 0.592$
RFM	$1.01\pm 0.025$	$\textbf{1.90}\pm\textbf{0.055}$	$0.15\pm 0.027$	$1.18\pm 0.055$	$-\textbf{5.20}\pm\textbf{0.067}$
G-LSD	$0.99\pm 0.05$	$1.99\pm 0.02$	$0.24\pm 0.07$	$1.11\pm 0.02$	$-4.15\pm 0.09$
G-PSD	$\textbf{0.95}\pm\textbf{0.02}$	$1.94\pm 0.03$	$\textbf{0.08}\pm\textbf{0.04}$	$1.10\pm 0.04$	$-4.40\pm 0.13$
G-ESD	$0.99\pm 0.04$	$1.95\pm 0.01$	$0.19\pm 0.04$	$1.10\pm 0.02$	$-4.61\pm 0.07$

Appendix

Table 11: NLL(

\downarrow

) on the Torus Dataset. Standard deviation is estimated over 5 runs.

	Volcano	Earthquake	Flood	Fire
Dataset size	827	6,120	4,875	12,809
RFM	472	43	49	39
RCT	280	20	18	17
RMF-MT	140	16	17	19

Appendix

Table 12: KLD(

\downarrow

) on the Spherical Dataset.

step	100	10	5	1
Dataset size	3,938
FrameFlow	1.16	2.34	6.53	N/A
Semigroup RMF	4.14	2.72	4.62	5.28
RMF(ours)	7.74	6.36	7.13	7.46

Appendix

Table 13: scRMSD(

\downarrow

) on the Protein generation.

Data dim $S^{N}$	3	4	8	16	32	64	128
data dim $\times$ train samples	3 $\times$ 50K	4 $\times$ 50K	8 $\times$ 50K	16 $\times$ 50K	32 $\times$ 50K	64 $\times$ 50K	128 $\times$ 50K
network params	732K	733K	737K	745K	762K	795K	860K

Appendix

Table 14: Dataset scale vs. network size in Ablation Studies. “data dimension

\times

train samples” reports data dimension × number of training samples for each dataset; “network params” reports the number of RMF parameters used in that experiment.

Dataset	Fire	Flood	Earthquake	Volcano	General	PrePro	Glycine	Proline	RNA
D. $\times$ T. S.	2 $\times$ 12,809	2 $\times$ 4,875	2 $\times$ 6,120	2 $\times$ 827	2 $\times$ 138,208	2 $\times$ 6,910	2 $\times$ 13,283	2 $\times$ 7,634	7 $\times$ 9,478
N. P.	36M	22.6M	36M	46M	731K	731K	2.7M	731K	1.4M

Appendix

Table 15: Dataset scale vs. network size in other dataset. ”D.” is ”data dimension”. ”T. S.” is ”train samples”. ”N. P.” is ”network params”

Dataset	Fire	Flood	Earthquake	Volcano	General	PrePro	Glycine	Proline	RNA
model size	36M	22.6M	36M	46M	731K	731K	2.7M	731K	1.4M
base	0.032	0.048	0.035	0.102	0.04	0.05	0.03	0.04	0.07
1/16	0.030	0.051	0.031	0.141	0.128	0.095	0.032	0.139	0.185
1/128	0.031	0.053	0.049	0.229	-	0.122	0.081	-	-
1/1024	0.068	0.097	0.072	-	-	-	-	-	-

Appendix

Table 16: RMF model downsizing ablation (MMD

\downarrow

). “base” denotes the original model; 1/16, 1/128, and 1/1024 denote reduced-size variants. “–” indicates experiments stopped once test performance deteriorated substantially.

Dataset	Metric	Fire	Flood	Earthquake	Volcano	General	PrePro	Glycine	Proline	RNA
100 $\%$	MMD	0.032	0.048	0.035	0.102	0.04	0.05	0.03	0.04	0.07
	p-value	0.171	0.465	0.405	0.752	0.059	0.158	0.227	0.128	0.01
10 $\%$	MMD	0.036	0.084	0.064	0.100	0.085	0.077	0.081	0.074	0.275
	p-value	0.107	0.049	0.089	0.108	0.01	0.01	0.326	0.01	0.01
1 $\%$	MMD	0.066	0.085	0.133	0.225	0.110	0.125	0.127	0.167	0.281
	p-value	0.01	0.01	0.01	0.039	0.01	0.01	0.01	0.01	0.01

Appendix

Table 17: Low-data regime MMD two-sample evaluation. We report both MMD

\downarrow

and permutation-test p-value

\uparrow

, computed by comparing 1000 generated samples with 1000 held-out test samples using 100 random permutations. Values of 0.01 indicate the minimum reported p-value resolution of the test.

Training Data Fraction	Fire(A/B)	Flood(A/B)	Earthquake(A/B)	Volcano(A/B)
100 $\%$	0.7845/0.796	0.632/0.612	0.703/0.746	0.944/0.981
10 $\%$	0.7935/0.793	0.673/0.657	0.735/0.735	0.948/0.982
1 $\%$	0.8595/0.867	0.685/0.679	0.81/0.75	0.957/0.982

Appendix

Table 18: Low-data regime leave-one-out 1-NN memorization diagnostic (accuracy

\to

0.5). A/B denote train-generated / test-generated comparisons, respectively. Each point is classified by the label of its nearest neighbor after excluding itself. Results are reported only on sphere-type datasets; torus datasets are omitted because reliable geodesic nearest-neighbor comparison would require a separate periodic-distance implementation.

Panel of Figure 8 — Figure 8: Ramachandran plots on the Torus Dataset.

Panel of Figure 9 — Figure 9: SO(3) one-step sample figure

	Volcano	Earthquake	Flood	Fire	General	Glycine	Proline	PrePro	RNA
Dateset size	827	6,120	4,875	12,809	138,208	13,283	7,634	6.910	9,478
learning rate	5e-4	5e-4	5e-4	5e-4	5e-4	5e-4	5e-4	5e-4	5e-4
number of layers	12	10	7	10	4	10	4	4	6
hidden dimensions	2048	2048	2048	2048	512	512	512	512	512
r=t $\%$	75 $\%$	75 $\%$	75 $\%$	75 $\%$	75 $\%$	75 $\%$	75 $\%$	75 $\%$	75 $\%$
batch size	8192	4096	8192	8192	4096	2048	2048	2048	2048

Appendix

Table 19: The Spherical and Torus experimental hyperparameters are as follows.

	Cone	Fisher	Line	Swiss Roll	CFG
Dateset size	20K	40K	40K	40K	60k
learning rate	5e-4	5e-4	5e-4	5e-4	5e-4
number of layers	4	4	4	8	4
hidden dimensions	512	512	512	1024	512
r=t $\%$	10 $\%$	10 $\%$	10 $\%$	10 $\%$	$10\%$
batch size	1024	1024	1024	1024	1024

Appendix

Table 20: The SO(3) experimental hyperparameters are as follows.

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
Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes
May 5, 2026Aaron Havens, Brian Karrer, Neta ShaulSampling from Unnormalized Distributions Diffusion Sampling
FAIR at Meta · Weizmann Institute of Science
Riemannian Motion Generation: A Unified Framework for Human Motion Representation and Generation via Riemannian Flow Matching
Mar 16, 2026Fangran Miao, Jian Huang, Ting LiHuman Motion Generation Motion Representation Learning
Department of Data Science and Artificial Intelligence, The Hong Kong Polytechnic University · Department of Applied Mathematics, The Hong Kong Polytechnic University · Department of Data Science and Statistics, Southern University of Science and Technology

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

Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes

Riemannian Motion Generation: A Unified Framework for Human Motion Representation and Generation via Riemannian Flow Matching