Riemannian MeanFlow for One-Step Generation on Manifolds
This addresses a computational bottleneck for researchers working with manifold-valued data in fields like robotics and molecular modeling, though it appears incremental as an extension of MeanFlow to manifolds.
The paper tackles the problem of slow sampling in generative models on Riemannian manifolds by proposing Riemannian MeanFlow, which enables one-step generation without numerical integration. Experiments show competitive sampling quality with substantially reduced computational cost.
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.