English

Riemannian MeanFlow for One-Step Generation on Manifolds

Machine Learning 2026-05-21 v2

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

Keywords

Cite

@article{arxiv.2603.10718,
  title  = {Riemannian MeanFlow for One-Step Generation on Manifolds},
  author = {Zichen Zhong and Haoliang Sun and Yukun Zhao and Yongshun Gong and Yilong Yin},
  journal= {arXiv preprint arXiv:2603.10718},
  year   = {2026}
}

Comments

International Conference on Machine Learning