Riemannian MeanFlow for One-Step Generation on Manifolds

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, SO(3), and SE(3) demonstrate competitive one-step sampling with improved quality–efficiency trade-offs and substantially reduced sampling cost.

Many generative models are designed for ordinary flat spaces, such as images arranged on a rectangular grid. However, many types of real-world data do not naturally live in flat space. For example, locations on the Earth lie on a sphere and robotic grasp poses involve both position and orientation. Treating such data as if it were flat can ignore important geometric structure. This work proposes Riemannian MeanFlow, a method for fast generation of data on curved spaces. Existing manifold generative models can respect the geometry of these spaces, but they often require many iterative sampling steps. Our method instead learns an average motion over a time interval, allowing high-quality samples to be generated in one or very few steps while still respecting the underlying geometry. This approach can make generative modeling more efficient for scientific and engineering applications involving non-flat data, such as Earth event modeling and robotic grasp generation. By reducing sampling cost while preserving geometric consistency, Riemannian MeanFlow provides a practical tool for fast generation on curved domains.