Riemannian MeanFlow for One-Step Generation on Manifolds

2026-03-11

AI Summary

This work addresses the high computational cost of numerical integration in probability flow ODEs on Riemannian manifolds by proposing Riemannian MeanFlow (RMF), the first extension of the MeanFlow framework to manifold-structured data. RMF defines a position-dependent average velocity field in the tangent space and leverages parallel transport to establish an intrinsic supervisory relationship between average and instantaneous velocities, enabling one-step generation without trajectory simulation. The method introduces the Riemannian MeanFlow identity, a logarithmic map-based tangent space representation, and conflict-aware multi-task learning. Evaluated on spherical, toroidal, and SO(3) manifolds, RMF substantially reduces sampling complexity while preserving high-quality generation, achieving a superior trade-off between efficiency and sample fidelity.

Technology Category

• Machine Learning: Learning with Manifolds
• Search and Optimization: Sampling/Simulation-based Search
• Intelligent Robots: Manipulation

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Systems and Infrastructure for Web, Mobile and WoT: Data management and stream processing for Web, mobile and wireless applications

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.

Problem

• Riemannian manifolds
• generative models
• sampling
• probability-flow ODE
• one-step generation

Innovation

• Riemannian MeanFlow
• Flow Matching
• one-step generation
• parallel transport
• manifold-valued generation