🤖 machine learning

Riemannian MeanFlow

This paper introduces Riemannian MeanFlow (RMF), a framework that learns direct flow mappings on manifolds to enable high-quality generative modeling with up to 10 $\times$ fewer function evaluations than traditional diffusion or flow models while simultaneously supporting efficient, reward-driven design.

Original authors: Dongyeop Woo, Marta Skreta, Seonghyun Park, Kirill Neklyudov, Sungsoo Ahn

Published 2026-05-04

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Dongyeop Woo, Marta Skreta, Seonghyun Park, Kirill Neklyudov, Sungsoo Ahn

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to design new proteins or DNA sequences. These are not mere lists of letters or numbers; they are complex 3D structures with strict rules, like a puzzle where every piece must fit perfectly into a specific geometric shape.

In the world of AI, "Diffusion" and "Flow" models are like artists slowly transforming a lump of noise into a perfect masterpiece. Normally, they do this by taking hundreds of tiny, careful steps, checking their work at each individual step to ensure they do not stray from the path.

The Problem: The Slow Walk
The article points out a major bottleneck: taking hundreds of steps is incredibly slow and expensive. It is like trying to cross a room by taking 100 tiny, cautious steps, even though you could cover the distance in one or two large strides. For scientists who need to design thousands of molecules, this slowness is a dealbreaker.

The Solution: Riemannian MeanFlow (RMF)
The authors introduce a new method called Riemannian MeanFlow (RMF). Imagine teaching the robot to "see" the "entire path" all at once, rather than just the next step.

Here is how they did it, using some everyday analogies:

1. The "Average Speed" Shortcut

Imagine you are driving from City A to City B.

The Old Way (Diffusion/Flow): You drive slowly, check your GPS every second, and constantly correct your steering wheel. This is accurate but takes forever.
The New Way (RMF): Instead of checking every second, you calculate the average speed and direction needed to get from A to B in one go. You do not need to stop and check; you simply drive the whole way based on that single calculation.

The article mathematically proves that this "average speed" can be learned directly on the complex geometric shapes (manifolds) where these molecules exist, without needing to take hundreds of tiny steps.

2. Three Ways to Describe the Journey

The authors discovered three different ways to describe this "average speed," all mathematically equivalent. They are like describing a road trip from three different perspectives:

The Eulerian View: Looking at the road from a fixed point (like a camera on a pole) and seeing how traffic flows past you.
The Lagrangian View: Sitting in the car and feeling how the speed changes as you drive.
The Semigroup View: Breaking the journey into two parts (from A to B, then from B to C) and realizing that the direct path from A to C is simply the combination of these two parts.

The article shows that using these perspectives allows the AI to learn the "one-step jump" directly.

3. The "Stop-Gradient" Trick

Training these models is like trying to keep a stack of plates balanced while juggling. If you try to calculate how every tiny change affects the final result, the math becomes so heavy that the computer crashes (or memory is exhausted).
The authors use a clever trick called "Stop-Gradient". Imagine you are teaching a student. Instead of making them recalculate the entire history of how they arrived at an answer every time they make a mistake, you simply say: "Here is the correct answer; now just adjust your next attempt accordingly." This prevents the computer from being overwhelmed by complex math and makes training faster and more stable.

4. Predicting the Destination, Not the Speed

Most models try to predict the velocity at every moment. RMF tries to predict the destination (the endpoint on the map) directly.

Analogy: Instead of telling a driver, "Turn left at 5 miles per hour, then right at 10 miles per hour," you simply say: "Drive until you see the big red barn."
This approach works much better for high-dimensional data (like long DNA sequences or large proteins) and allows the model to use existing software designed for these shapes without needing to be rewritten.

What They Achieved

The article tested this on two real scientific tasks:

Designing DNA Promoters: These are like "start buttons" for genes. The new method generated high-quality DNA sequences in one step, performing just as well as old methods that required 100 steps.
Designing Protein Scaffolds: These are the skeletons of proteins. The new method could generate valid protein structures in 1 to 5 steps, whereas previous methods collapsed or failed completely when asked to do this so quickly.

The Conclusion
The article claims that by understanding the geometry of the data and using these "average speed" shortcuts, they can generate complex biological structures 10 times faster (using 10 times fewer computer calculations) without losing quality. This makes it possible to use AI for rapid, iterative scientific designs, where researchers can test thousands of ideas quickly instead of waiting hours for a single result.

Technical Summary: Riemannian MeanFlow

Problem Statement

Diffusion and flow models have established themselves as the dominant paradigm for generative modeling on Riemannian manifolds and have been successfully applied to tasks such as generating protein backbone structures and designing DNA sequences. However, these methods typically require ten to one hundred neural network evaluations (function evaluations or NFEs) during inference to numerically integrate ordinary differential equations (ODEs) or stochastic differential equations (SDEs). This computational burden represents a bottleneck, particularly in scientific design pipelines where generative models serve as proposal mechanisms within iterative loops for property-driven optimization. While few-step generation techniques such as consistency models and flow map learning have been investigated in Euclidean settings, they remain underexplored and often unstable in high-dimensional Riemannian manifold settings.

Methodology: Riemannian MeanFlow (RMF)

The authors present Riemannian MeanFlow (RMF), a framework designed to learn flow maps directly on manifolds and enable high-quality generation with only a single forward pass. The core of the methodology involves deriving and utilizing three equivalent characterizations of the mean velocity on the manifold, which relate the flow map between arbitrary time points $s$ and $t$ .

1. Theoretical Foundations

The framework defines the mean velocity $u_{s,t}$ as the constant velocity that, if maintained along a geodesic from $x_s$ to $x_t$ over the time interval $t-s$ , would correctly transport the point. The authors derive three identities that serve as the basis for training objectives:

Euler Identity: Derived by differentiating the defining relation with respect to the source time $s$ . It relates the mean velocity to the covariant derivative of the velocity field along the integral curve and to the derivative of the logarithmic map.
Lagrange Identity: Derived by differentiating with respect to the target time $t$ . It relates the mean velocity to the differential of the logarithmic map applied to the instantaneous velocity and to the partial time derivative of the mean velocity.
Semigroup Identity: An algebraic condition derived from the semigroup property of flow maps ( $\Phi_{r,t} \circ \Phi_{s,r} = \Phi_{s,t}$ ). It enforces consistency between the direct flow from $s$ to $t$ and the composition of flows through an intermediate time point $r$ .

2. Training Objectives and Stabilization

The authors convert these identities into regression objectives. A critical design decision is the placement of the stop-gradient operator. In contrast to prior work (e.g., Generalized Flow Maps), which applies stop-gradients partially, RMF applies the stop-gradient to the entire regression objective. This formulation:

Eliminates the need to evaluate higher-order derivatives (e.g., gradients through Jacobian-vector products) during backpropagation.
Reduces memory consumption by nearly a factor of 2.
Improves training stability in high-dimensional settings.

To address optimization problems in high dimensions, the paper proposes several stabilization techniques:

Time Sampling: Ordered time pairs ( $s \le t$ ) for Euler and semigroup objectives, and unordered pairs for Lagrange objectives.
Adaptive Loss Weighting: Clipping the loss based on the magnitude of the regression objective to handle high variance in derivative-dependent terms.
Time Derivative Control: Using low-frequency time embeddings to limit the magnitude of time derivatives in the network output.

3. Parameterization

The paper evaluates three parameterization schemes for the mean velocity:

$v$ -prediction: Predicting the tangent vector (velocity).
$x_t$ -prediction: Directly predicting the endpoint.
$x_1$ -prediction: Predicting the trajectory endpoint at $t=1$ given the current state and the time interval.
The authors note that $x_1$ -prediction is the most robust for high-dimensional manifolds. It is compatible with existing scientific architectures (e.g., those outputting SE(3) frames for proteins) and scales well up to dimensions of $D=2048$ , whereas $x_t$ -prediction fails and $v$ -prediction loses quality as dimensionality increases.

4. Reward-Guided Inference

The framework supports reward-guided generation without retraining the model. By leveraging the flow map's ability to predict final states from intermediate steps, the authors propose a Reward Lookahead Strategy. Instead of using the gradient of the reward at the current noisy state $x_t$ , they utilize the gradient at the predicted endpoint $\Phi_{t,1}(x_t)$ . This approach has proven more effective at steering generations toward specific targets.

Key Results

The authors validate RMF on two biological tasks: Promoter DNA Design (sequences on a simplex in $\mathbb{R}^{1024 \times 4}$ ) and Protein Backbone Generation (on $SE(3)^N$ with $N$ up to 128).

Efficiency: RMF achieves sample quality comparable to state-of-the-art multi-step baselines (e.g., Fisher FM, FrameFlow, FrameDiff) but requires up to 10 times fewer function evaluations.
- In protein backbone generation, RMF achieves 82% designability at 5 steps and 35% at 1 step, whereas baselines drop significantly or collapse completely in these few-step regimes.
- In DNA design, RMF variants match the performance of Fisher FM with 100 steps using only 1 step.
Stability: The proposed stabilization techniques (particularly the combination of semigroup objectives and $x_1$ -prediction) enable stable training in high-dimensional spaces where other flow map methods (such as GFM) fail to converge or produce erroneous samples.
Reward Guidance: Reward-guided inference using the $x_1$ -lookahead strategy significantly improves alignment with target regulatory profiles (DNA) and secondary structure compositions (proteins) compared to naive guidance or unguided generation.

Significance and Claims

The paper claims to establish a principled theoretical foundation and practical guidelines for fast sampling on Riemannian manifolds. Its significance lies in:

Bridging the Gap: It extends few-step generation capabilities from Euclidean spaces to complex, high-dimensional Riemannian manifolds, a critical step for scientific applications.
Scalability: By deriving identities that avoid higher-order derivatives and utilizing $x_1$ -prediction, the method scales to dimensions ( $D \approx 2000$ ) and sequence lengths previously inaccessible to few-step manifold generative models.
Efficiency in Scientific Workflows: The reduction in inference costs (up to 10-fold) makes generative models more practical as proposal mechanisms in iterative scientific design loops, such as property-driven optimization for drug discovery or synthetic biology.
Reward Lookahead: It demonstrates that flow maps enable efficient reward-guided design by predicting final states from intermediate steps, a capability not previously demonstrated on manifolds.

The authors position RMF not as a replacement for all generative modeling, but as a specialized framework enabling efficient, high-quality generation in geometrically constrained scientific domains where computational costs are a limiting factor.