Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Riemannian MeanFlow

Dongyeop Woo, Marta Skreta, Seonghyun Park, Sungsoo Ahn, Kirill Neklyudov

TL;DR

The paper addresses the computational bottleneck of diffusion/flow models on Riemannian manifolds by introducing Riemannian MeanFlow (RMF), which learns flow maps directly on manifolds to enable one-pass or few-pass sampling. It derives three equivalent characterizations of the manifold average velocity (Eulerian, Lagrangian, semigroup) and identifies an $x_1$-prediction parameterization with a semigroup objective as the most stable, scalable approach, aided by stop-gradient and stabilization techniques. RMF demonstrates strong performance on high-dimensional scientific tasks, achieving comparable sample quality to multi-step baselines in promoter DNA design and protein backbone generation while reducing function evaluations by up to $10\times$, and enabling reward-guided look-ahead inference on manifolds. The framework is validated on domains like $\mathrm{SE}(3)^N$ protein backbones and $\Delta^{d-1}$ DNA promoters, highlighting its potential to accelerate geometry-aware design in biology and materials science through geometry-respecting, efficient sampling.

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.

Riemannian MeanFlow

TL;DR

The paper addresses the computational bottleneck of diffusion/flow models on Riemannian manifolds by introducing Riemannian MeanFlow (RMF), which learns flow maps directly on manifolds to enable one-pass or few-pass sampling. It derives three equivalent characterizations of the manifold average velocity (Eulerian, Lagrangian, semigroup) and identifies an -prediction parameterization with a semigroup objective as the most stable, scalable approach, aided by stop-gradient and stabilization techniques. RMF demonstrates strong performance on high-dimensional scientific tasks, achieving comparable sample quality to multi-step baselines in promoter DNA design and protein backbone generation while reducing function evaluations by up to , and enabling reward-guided look-ahead inference on manifolds. The framework is validated on domains like protein backbones and DNA promoters, highlighting its potential to accelerate geometry-aware design in biology and materials science through geometry-respecting, efficient sampling.

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.
Paper Structure (79 sections, 10 theorems, 119 equations, 13 figures, 11 tables)

This paper contains 79 sections, 10 theorems, 119 equations, 13 figures, 11 tables.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Riemannian MeanFlow Identities
  4. Background on Riemannian Geometry
  5. Flow Maps and Average Velocity on Manifolds
  6. Riemannian MeanFlow Identities
  7. Flow Map Learning with Riemannian MF
  8. Training Objectives
  9. Parameterization of the Flow Map
  10. Stabilizing Riemannian MF Training
  11. Reward-guided Inference with Flow Maps
  12. Experiments
  13. Ablation of Design Choices
  14. Promoter DNA Design
  15. Protein Backbone Design
  16. ...and 64 more sections

Key Result

Proposition 2.1

A vector field $u_{s,t}: \mathcal{M} \to T\mathcal{M}$ is the average velocity associated with $v_t$if and only if it satisfies for any integral curve $(x_t)_{t \in [0,1]}$ and any $s,t \in [0,1]$.

Figures (13)

  • Figure 1: Protein backbone samples from RMF (ours), FrameDiff, and FrameFlow for different inference budgets. RMF produces well-formed structures in one step, while baselines require more.
  • Figure 2: Parameterization choices: One-step generation results from models trained with different parameterizations ($x_1$-, $v$-, and $x_t$-pred) across ambient dimensions $D\in\{512,2048\}$.
  • Figure 3: Objective choices: (Left) Samples generated by models trained with different objectives using 1-step (top) or 100-step (bottom) sampling on $D=512$. (Right) Adaptive loss weighting for Eulerian RMF substantially improves sample quality.
  • Figure 4: Performance vs. NFE. RMF variants outperform Fisher FM (FFM) in $k$-mer correlation ($k=6$) and MSE. RMF achieves high accuracy at $1$ NFE, whereas FFM requires $\ge 32$ steps for comparable performance.
  • Figure 5: (\ref{['subfig:protein-method']}) RMF consistently outperforms baselines in designability across inference steps. (\ref{['subfig:protein-inference']}) Intermediate noise levels ($\eta \approx 0.25$--$0.45$) yield the best performance.
  • ...and 8 more figures

Theorems & Definitions (39)

  • Definition 2.1: Integral curve
  • Definition 2.2: Flow map
  • Definition 2.3: Average velocity
  • Proposition 2.1: Eulerian RMF
  • proof : Proof sketch
  • Proposition 2.2: Lagrangian RMF
  • proof : Proof sketch
  • Proposition 2.3: Semigroup RMF
  • Proposition 3.1: Riemannian MeanFlow objectives
  • Example 1.1: The $d$-dimensional sphere
  • ...and 29 more