Riemannian MeanFlow for One-Step Generation on Manifolds
Zichen Zhong, Haoliang Sun, Yukun Zhao, Yongshun Gong, Yilong Yin
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 Mean-Flow 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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 2bed085f-3973-4c62-9301-0ff77b9b96abBuilds on28
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Gradient Surgery for Multi-Task LearningTianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine et al.NeurIPS 2020 · 2,261 citations
- Consistency ModelsYang Song, Prafulla Dhariwal, Mark Chen, Ilya SutskeverICML 2023 · 1,720 citations
- Score-Based Generative Modeling through Stochastic Differential EquationsYang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar et al.ICLR 2021 · 1,270 citations
- Mean Flows for One-step Generative ModelingZhengyang Geng, Mingyang Deng, Xingjian Bai, Zico Kolter et al.NeurIPS 2025 · 628 citations
Related papers
- Riemannian MeanFlowDongyeop Woo, Marta Skreta, Seonghyun Park, Kirill Neklyudov et al.ICML 2026 · 1 citation
- Metric Flow Matching for Smooth Interpolations on the Data ManifoldKacper Kapusniak, Peter Potaptchik, Teodora Reu, Leo Zhang et al.NeurIPS 2024 · 89 citations
- Flow Matching on General GeometriesRicky T. Q. Chen, Yaron LipmanICLR 2024 · 193 citations
- Learning Distributions on Manifolds with Free-Form FlowsPeter Sorrenson, Felix Draxler, Armand Rousselot, Sander Hummerich et al.NeurIPS 2024 · 7 citations
- Riemannian Variational Flow Matching for Material and Protein DesignOlga Zaghen, Floor Eijkelboom, Alison Pouplin, Cong Liu et al.ICLR 2026 · 10 citations
