Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Riemannian MeanFlow for One-Step Generation on Manifolds

Zichen Zhong, Haoliang Sun, Yukun Zhao, Yongshun Gong, Yilong Yin

TL;DR

Riemannian MeanFlow is proposed, extending MeanFlow to manifold-valued generation where velocities lie in location-dependent tangent spaces and derives a Riemannian MeanFlow identity that links average and instantaneous velocities for intrinsic supervision.

Abstract

Flow Matching enables simulation-free training of generative models on Riemannian manifolds, yet sampling typically still relies on numerically integrating a probability-flow ODE. We propose Riemannian MeanFlow (RMF), extending MeanFlow to manifold-valued generation where velocities lie in location-dependent tangent spaces. RMF defines an average-velocity field via parallel transport and derives a Riemannian MeanFlow identity that links average and instantaneous velocities for intrinsic supervision. We make this identity practical in a log-map tangent representation, avoiding trajectory simulation and heavy geometric computations. For stable optimization, we decompose the RMF objective into two terms and apply conflict-aware multi-task learning to mitigate gradient interference. RMF also supports conditional generation via classifier-free guidance. Experiments on spheres, tori, and SO(3) demonstrate competitive one-step sampling with improved quality-efficiency trade-offs and substantially reduced sampling cost.

Riemannian MeanFlow for One-Step Generation on Manifolds

TL;DR

Riemannian MeanFlow is proposed, extending MeanFlow to manifold-valued generation where velocities lie in location-dependent tangent spaces and derives a Riemannian MeanFlow identity that links average and instantaneous velocities for intrinsic supervision.

Abstract

Flow Matching enables simulation-free training of generative models on Riemannian manifolds, yet sampling typically still relies on numerically integrating a probability-flow ODE. We propose Riemannian MeanFlow (RMF), extending MeanFlow to manifold-valued generation where velocities lie in location-dependent tangent spaces. RMF defines an average-velocity field via parallel transport and derives a Riemannian MeanFlow identity that links average and instantaneous velocities for intrinsic supervision. We make this identity practical in a log-map tangent representation, avoiding trajectory simulation and heavy geometric computations. For stable optimization, we decompose the RMF objective into two terms and apply conflict-aware multi-task learning to mitigate gradient interference. RMF also supports conditional generation via classifier-free guidance. Experiments on spheres, tori, and SO(3) demonstrate competitive one-step sampling with improved quality-efficiency trade-offs and substantially reduced sampling cost.
Paper Structure (35 sections, 3 theorems, 60 equations, 8 figures, 11 tables, 1 algorithm)

This paper contains 35 sections, 3 theorems, 60 equations, 8 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Riemannian MeanFlow
  4. Average Velocity on Riemannian Manifolds
  5. A Multi-Task View of the Decomposed Objective
  6. Classifier-Free Guidance (CFG)
  7. Related Work
  8. Experiments
  9. Gradient Conflict Analysis of Loss
  10. Comparisons to Baselines
  11. Ablation Study
  12. Conclusion
  13. Riemannian MeanFlow in Local Coordinates
  14. Local coordinates and notation.
  15. Covariant time derivative along $\gamma$.
  16. ...and 20 more sections

Key Result

Proposition 3.1

Let $\gamma:[r,t]\rightarrow\mathcal{M}$ be a smooth trajectory with $x_\tau=\gamma(\tau)$ and $t>r$. Define $u(x_t,r,t)\in T_{x_t}\mathcal{M}$ by Eq. eq:rmf_avg_velocity. If $v(\cdot,\tau)$ is sufficiently smooth along $\gamma$, then

Figures (8)

  • Figure 1: Instantaneous velocity (A→B) vs. average velocity (P→Q) on Euclidean straight lines and manifold geodesics.
  • Figure 2: Cosine similarity between full-parameter gradients $\nabla_{\theta}\mathcal{L}1$ and $\nabla{\theta}\mathcal{L}_2$ across Earth categories during training (per-iteration values and running average); negative values indicate gradient conflicts.
  • Figure 3: Earth dataset: generated sample distributions overlaid with ground-truth test data (red). Top-left: Volcano; top-right: Earthquake; bottom-left: Flood; bottom-right: Fire.
  • Figure 4: Visualization of samples generated by CFG and unconditional models.
  • Figure 5: iRMF: Cosine Similarity about $\nabla \mathcal{L}_1(\theta)$ and $\nabla \mathcal{L}_2(\theta)$ with $0\%$ r=t on Spherical Datasets.
  • ...and 3 more figures

Theorems & Definitions (3)

  • Proposition 3.1: Riemannian MeanFlow Identity
  • Proposition 3.2: Decomposed RMF Loss
  • Corollary 3.3: PCGrad Guarantees for Intrinsic Manifold Losses