Riemannian MeanFlow

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.

Generative AI can help scientists design new biological objects, such as DNA sequences and protein structures. Many current methods create these objects through a long sequence of small improvements, starting from random noise. Although this approach can produce high-quality results, it may require many repeated model evaluations, which becomes expensive when scientists need to generate and test many candidates. This paper introduces Riemannian MeanFlow, a new method for faster generation of scientific data that must follow natural geometric rules. For example, proteins are three-dimensional objects with positions and rotations, while DNA sequences have restricted possible values. Our method learns to move more directly toward a final candidate, reducing the number of steps needed at generation time. We evaluate the method on DNA sequence design and protein backbone generation. The results show that it can produce high-quality samples with fewer computation steps, and can also help guide generation toward desired scientific properties.