Riemannian MeanFlow for One-Step Generation on Manifolds

📄 arXiv: 2603.10718v1 📥 PDF

作者: Zichen Zhong, Haoliang Sun, Yukun Zhao, Yongshun Gong, Yilong Yin

分类: cs.LG

发布日期: 2026-03-11

💡 一句话要点

提出黎曼MeanFlow，用于流形上的一步生成，提升质量-效率权衡。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱四：生成式动作 (Generative Motion)

关键词: 黎曼流形 生成模型 Flow Matching MeanFlow 一步生成

📋 核心要点

  1. 现有流形生成模型依赖概率流ODE的数值积分进行采样，计算成本高昂且效率较低。
  2. RMF通过平行传输定义平均速度场，并推导出黎曼MeanFlow恒等式，实现内在监督，避免轨迹模拟。
  3. 实验表明，RMF在球体、环面和SO(3)等流形上实现了高质量的单步生成，显著降低了采样成本。

📝 摘要（中文）

Flow Matching实现了黎曼流形上生成模型的无仿真训练，但采样通常仍依赖于数值积分概率流ODE。我们提出了黎曼MeanFlow (RMF)，将MeanFlow扩展到流形值生成，其中速度位于位置相关的切空间中。RMF通过平行传输定义了一个平均速度场，并推导出一个黎曼MeanFlow恒等式，将平均速度和瞬时速度联系起来，用于内在监督。我们在对数映射切线表示中使这个恒等式实用化，避免了轨迹模拟和繁重的几何计算。为了稳定优化，我们将RMF目标分解为两个项，并应用冲突感知的多任务学习来缓解梯度干扰。RMF还支持通过无分类器指导进行条件生成。在球体、环面和SO(3)上的实验表明，RMF具有竞争力的单步采样，并改善了质量-效率的权衡，并大幅降低了采样成本。

🔬 方法详解

问题定义：现有的基于Flow Matching的流形生成模型，虽然避免了训练过程中的仿真，但在采样阶段仍然需要通过数值积分求解概率流ODE，这导致采样过程计算量大，效率较低，限制了其在实际应用中的潜力。尤其是在高维流形上，数值积分的计算复杂度会显著增加。

核心思路：RMF的核心思路是利用MeanFlow的思想，直接学习一个平均速度场，从而实现一步生成。通过在黎曼流形上定义平均速度场，并利用平行传输将不同位置的切向量联系起来，RMF避免了对概率流ODE的数值积分，从而显著降低了采样成本。此外，RMF还推导了一个黎曼MeanFlow恒等式，用于内在监督，进一步提升了生成模型的性能。

技术框架：RMF的整体框架包括以下几个主要步骤：1) 定义黎曼流形上的概率分布；2) 通过平行传输定义平均速度场；3) 推导黎曼MeanFlow恒等式，将平均速度和瞬时速度联系起来；4) 利用对数映射将切向量表示为切空间中的坐标；5) 构建神经网络来学习平均速度场；6) 使用分解后的RMF目标和冲突感知的多任务学习进行优化；7) 通过一步生成实现采样。

关键创新：RMF最重要的技术创新点在于它将MeanFlow的思想扩展到了黎曼流形上，并推导出了黎曼MeanFlow恒等式。与传统的Flow Matching方法相比，RMF避免了对概率流ODE的数值积分，从而显著降低了采样成本。此外，RMF还利用平行传输来定义平均速度场，并使用对数映射将切向量表示为切空间中的坐标，从而简化了计算。

关键设计：RMF的关键设计包括：1) 使用平行传输来定义平均速度场；2) 推导黎曼MeanFlow恒等式，用于内在监督；3) 使用对数映射将切向量表示为切空间中的坐标；4) 将RMF目标分解为两个项，以稳定优化；5) 应用冲突感知的多任务学习来缓解梯度干扰；6) 使用无分类器指导进行条件生成。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，RMF在球体、环面和SO(3)等流形上实现了具有竞争力的单步采样，并改善了质量-效率的权衡。与传统的Flow Matching方法相比，RMF显著降低了采样成本，同时保持了较高的生成质量。例如，在SO(3)上的实验表明，RMF的采样成本降低了5倍以上。

🎯 应用场景

RMF在机器人运动规划、分子动力学模拟、医学图像分析等领域具有广泛的应用前景。例如，在机器人运动规划中，RMF可以用于生成满足特定约束的机器人运动轨迹。在分子动力学模拟中，RMF可以用于生成符合物理规律的分子构象。在医学图像分析中，RMF可以用于生成具有特定病理特征的医学图像。

📄 摘要（原文）

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.