Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

Riemannian MeanFlow for One-Step Generation on Manifolds

About

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, and SO(3) demonstrate competitive one-step sampling with improved quality-efficiency trade-offs and substantially reduced sampling cost.

Zichen Zhong, Haoliang Sun, Yukun Zhao, Yongshun Gong, Yilong Yin• 2026

Related benchmarks

TaskDatasetResultRank
GenerationTorus RNA (test)
MMD0.07
7
Generative ModelingTorus PrePro 6,910 (test)
NLL1.02
7
One-step generationEarth Volcano (test)
MMD0.092
7
One-step generationEarth Flood (test)
MMD0.048
7
GenerationTorus Protein Glycine (test)
MMD0.03
7
GenerationTorus Protein Proline (test)
MMD0.04
7
GenerationTorus Protein PrePro (test)
MMD0.05
7
Generative ModelingTorus General 138,208 (test)
NLL0.97
7
Generative ModelingTorus RNA 9,478 (test)
NLL3.79
7
Generative ModelingVolcano Spherical Dataset
NLL3.73
7
Showing 10 of 22 rows

Other info

Follow for update