Riemannian MeanFlow

2026-02-08

AI Summary

Existing manifold generative models require tens to hundreds of neural network evaluations during inference, hindering their applicability to large-scale scientific sampling. This work proposes Riemannian MeanFlow, a novel framework that directly learns flow maps on Riemannian manifolds, enabling high-quality sample generation with only a single forward pass. The approach introduces three equivalent characterizations of the manifold-averaged velocity field—Eulerian, Lagrangian, and semigroup identities—and incorporates tailored parameterization and stabilization strategies to enhance training in high-dimensional settings. Furthermore, a reward-anticipating mechanism is introduced to enable efficient guided generation. Experiments on promoter DNA design and protein backbone generation demonstrate that the method achieves sample quality comparable to state-of-the-art approaches while using less than one-tenth the number of function evaluations.

Technology Category

Machine Learning: Learning with ManifoldsSearch and Optimization: Sampling/Simulation-based SearchNatural Language Processing: Generation

Application Category

Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applicationsSocial Networks and Social Media: Generative AI / large language models and their impact on social systemsUser Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

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$\times$ 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.

Problem

• Riemannian manifolds
• generative modeling
• flow models
• computational bottleneck
• inference efficiency

Innovation

• Riemannian MeanFlow
• flow maps on manifolds
• manifold average velocity
• efficient generative modeling
• reward look-ahead