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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Generation | Torus RNA (test) | MMD0.07 | 7 | |
| Generative Modeling | Torus PrePro 6,910 (test) | NLL1.02 | 7 | |
| One-step generation | Earth Volcano (test) | MMD0.092 | 7 | |
| One-step generation | Earth Flood (test) | MMD0.048 | 7 | |
| Generation | Torus Protein Glycine (test) | MMD0.03 | 7 | |
| Generation | Torus Protein Proline (test) | MMD0.04 | 7 | |
| Generation | Torus Protein PrePro (test) | MMD0.05 | 7 | |
| Generative Modeling | Torus General 138,208 (test) | NLL0.97 | 7 | |
| Generative Modeling | Torus RNA 9,478 (test) | NLL3.79 | 7 | |
| Generative Modeling | Volcano Spherical Dataset | NLL3.73 | 7 |