Riemannian MeanFlow
Dongyeop Woo, Marta Skreta, Seonghyun Park, Kirill Neklyudov, Sungsoo Ahn
Abstract
Diffusion and flow models have become the dominant paradigm for generative modeling on Riemannian manifolds, with successful applications in protein backbone generation and DNA sequence design. However, these methods require tens to hundreds of neural network evaluations at inference time, which can become a computational bottleneck in large-scale scientific sampling workflows. We introduce Riemannian MeanFlow (RMF), a framework for learning flow maps directly on manifolds, enabling high-quality generations with as few as one forward pass. We derive three equivalent characterizations of the manifold average velocity (Eulerian, Lagrangian, and semigroup identities), and analyze parameterizations and stabilization techniques to improve training on high-dimensional manifolds. In promoter DNA design and protein backbone generation settings, RMF achieves comparable sample quality to prior methods while requiring up to 10 fewer function evaluations. Finally, we show that few-step flow maps enable efficient reward-guided design through reward look-ahead, where terminal states can be predicted from intermediate steps at minimal additional 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 9ed066c6-cbbb-4f32-9f89-646e8b7d0e68Builds on28
- 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
- Torsional Diffusion for Molecular Conformer GenerationBowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay et al.NeurIPS 2022 · 413 citations
- Improved Techniques for Training Consistency ModelsYang Song, Prafulla DhariwalICLR 2024 · 383 citations
Related papers
- Riemannian MeanFlow for One-Step Generation on ManifoldsZichen Zhong, Haoliang Sun, Yukun Zhao, Yongshun Gong et al.ICML 2026
- Generalised Flow Maps for Few-Step Generative Modelling on Riemannian ManifoldsOscar Davis, Michael S. Albergo, Nicholas M. Boffi, Michael M. Bronstein et al.ICLR 2026 · 12 citations
- Training Free Guided Flow-Matching with Optimal ControlLuran Wang, Chaoran Cheng, Yizhen Liao, Yanru Qu et al.ICLR 2025
- SE(3)-Stochastic Flow Matching for Protein Backbone GenerationAvishek Joey Bose, Tara Akhound-Sadegh, Guillaume Huguet, Kilian Fatras et al.ICLR 2024 · 162 citations
- Learning Distributions on Manifolds with Free-Form FlowsPeter Sorrenson, Felix Draxler, Armand Rousselot, Sander Hummerich et al.NeurIPS 2024 · 7 citations
