🤖 machine learning

Riemannian MeanFlow

本文介绍了黎曼平均流（RMF），这是一个在流形上学习直接流映射的框架，能够以比传统扩散模型或流模型少高达 10 倍的函数评估次数实现高质量生成建模，同时促进高效的奖励引导设计。

原作者： Dongyeop Woo, Marta Skreta, Seonghyun Park, Kirill Neklyudov, Sungsoo Ahn

发布于 2026-05-04

📖 1 分钟阅读☕ 轻松阅读

原作者： Dongyeop Woo, Marta Skreta, Seonghyun Park, Kirill Neklyudov, Sungsoo Ahn

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在尝试教一个机器人设计新的蛋白质或 DNA 序列。这些序列不仅仅是字母或数字的列表；它们是复杂的三维结构，遵循严格的规则，就像拼图一样，每一块都必须完美地契合特定的几何形状。

在人工智能领域，“扩散”（Diffusion）和“流”（Flow）模型就像艺术家，缓慢地将一团噪声转化为完美的杰作。通常，它们通过数百个微小而谨慎的步骤来实现这一目标，并在每一步都检查工作，以确保不偏离路径。

问题：缓慢的行走
该论文指出了一个主要瓶颈：采取数百个步骤极其缓慢且昂贵。这就像试图通过迈出一百个微小而谨慎的步子穿过房间，而你本可以只需迈一两大步就能到达。对于那些需要设计数千种分子的科学家来说，这种缓慢是不可接受的。

解决方案：黎曼平均流（Riemannian MeanFlow, RMF）
作者引入了一种名为**黎曼平均流（RMF）**的新方法。这相当于教机器人“一次性看清整条路径”，而不仅仅是下一步。

以下是他们如何利用日常类比来实现这一点的：

1. “平均速度”捷径

想象你正从 A 城开车前往 B 城。

旧方法（扩散/流）： 你缓慢行驶，每秒检查一次 GPS，不断调整方向盘。这很准确，但耗时极长。
新方法（RMF）： 你不再每秒检查，而是计算出从 A 到 B 所需的平均速度和方向，一次性完成旅程。你不需要停下来检查；只需基于这单次计算直接开完全程。

该论文从数学上证明，你可以直接在分子所在的复杂几何形状（流形）上学习这种“平均速度”，而无需采取数百个微小步骤。

2. 描述旅程的三种方式

作者发现了三种描述这种“平均速度”的不同方式，它们在数学上是等价的。这就像从三个不同的角度描述一次公路旅行：

欧拉视角（Eulerian View）： 从固定点（如杆子上的摄像头）观察道路，看交通流如何经过你。
拉格朗日视角（Lagrangian View）： 坐在车内，感受行驶过程中速度的变化。
半群视角（Semigroup View）： 将旅程分为两部分（A 到 B，然后 B 到 C），并意识到从 A 到 C 的直接路径仅仅是这两部分的组合。

论文表明，利用这些视角可以让 AI 直接学习“一步到位”的跳跃。

3. “停止梯度”技巧

训练这些模型就像在杂耍时试图平衡一摞盘子。如果你试图计算每一个微小变化如何影响最终结果，数学运算会变得如此繁重，导致计算机崩溃（或内存耗尽）。
作者使用了一种巧妙的技巧，称为**“停止梯度”（stop-gradient）**。想象你在教学生。与其让他们每次犯错时都重新计算得出答案的整个历史过程，你只需告诉他们：“这是正确答案，现在只需让你的下一个猜测与之匹配。”这阻止了计算机被复杂的数学运算压垮，使训练更快、更稳定。

4. 预测目的地，而非速度

大多数模型试图预测每一时刻的速度。而 RMF 试图直接预测目的地（地图上的最终点）。

类比： 与其告诉司机“以 5 英里/小时的速度左转，然后以 10 英里/小时的速度右转”，你只需说：“一直开，直到你看到那个大红谷仓。”
这种方法对于高维数据（如长 DNA 序列或大型蛋白质）效果更好，并且允许模型使用现有的、专为这些形状设计的软件，而无需重写。

他们的成就
该论文在两项现实世界的科学任务中测试了这种方法：

设计 DNA 启动子： 这些就像基因的“启动按钮”。新方法在一步内生成了高质量的 DNA 序列，其效果与旧方法所需的100 步相当。
设计蛋白质骨架： 这些是蛋白质的骨架。新方法能够在1 到 5 步内生成有效的蛋白质结构，而之前的方法在被要求以如此快的速度完成时，要么崩溃，要么完全失败。

核心结论
该论文声称，通过理解数据的几何结构并利用这些“平均速度”捷径，他们能够以10 倍的速度生成复杂的生物结构（所需的计算机计算量减少 10 倍），且不会损失质量。这使得利用 AI 进行快速、迭代的科学设计成为可能，研究人员可以快速测试数千种想法，而无需等待数小时才能获得单个结果。

技术摘要：黎曼 MeanFlow

问题陈述

扩散模型和流模型已成为黎曼流形上生成建模的主导范式，已成功应用于蛋白质骨架生成和 DNA 序列设计等任务。然而，这些方法在推理过程中通常需要数十到数百次神经网络评估（函数评估，或 NFEs），以数值积分常微分方程（ODEs）或随机微分方程（SDEs）。这种计算负担构成了瓶颈，特别是在科学设计流程中，生成模型作为属性引导优化迭代循环内的提议机制。虽然一致性模型和流映射学习等少步生成技术已在欧几里得环境中得到探索，但在高维黎曼流形环境中，它们仍未被充分探索且往往不稳定。

方法论：黎曼 MeanFlow (RMF)

作者引入了黎曼 MeanFlow (RMF)，这是一个旨在直接在流形上学习流映射的框架，能够实现仅需一次前向传播的高质量生成。该方法的核心在于推导并利用平均速度的三种等价表征，这些表征关联了任意时间点 $s$ 和 $t$ 之间的流映射。

1. 理论基础

该框架将平均速度 $u_{s,t}$ 定义为：若沿从 $x_s$ 到 $x_t$ 的测地线在时间 $t-s$ 内保持该速度，则能正确传输该点的恒定速度。作者推导了三个恒等式，作为训练目标的基础：

欧拉恒等式 (Eulerian Identity)： 通过对定义关系关于源时间 $s$ 求导得出。它将平均速度与沿积分曲线的速度场的协变导数以及对数映射的导数联系起来。
拉格朗日恒等式 (Lagrangian Identity)： 通过对定义关系关于目标时间 $t$ 求导得出。它将平均速度与应用于瞬时速度对数映射的微分以及平均速度的偏时间导数联系起来。
半群恒等式 (Semigroup Identity)： 源自流映射的半群性质（ $\Phi_{r,t} \circ \Phi_{s,r} = \Phi_{s,t}$ ）的代数条件。它强制了从 $s$ 到 $t$ 的直接流与通过中间时间 $r$ 的流组合之间的一致性。

2. 训练目标与稳定化

作者将这些恒等式转化为回归目标。一个关键的设计选择是停止梯度算子 (stop-gradient operator) 的放置位置。与先前工作（如广义流映射）部分应用停止梯度不同，RMF 将整个回归目标应用停止梯度。这种表述：

消除了在反向传播期间评估高阶导数（例如，通过雅可比 - 向量积的梯度）的需求。
将内存使用量减少了近 2 倍。
提高了高维环境下的训练稳定性。

为了解决高维环境中的优化困难，本文提出了几种稳定化技术：

时间采样： 欧拉和半群目标使用有序时间对（ $s \le t$ ），拉格朗日目标使用无序对。
自适应损失加权： 根据回归目标的大小重新加权损失，以处理依赖导数项的高方差。
时间导数控制： 使用低频时间嵌入来限制网络输出中时间导数的大小。

3. 参数化

本文评估了三种平均速度参数化方案：

$v$ -预测： 预测切向量（速度）。
$x_t$ -预测： 直接预测终点。
$x_1$ -预测： 给定当前状态和时间间隔，预测 $t=1$ 时的轨迹终点。
作者发现， $x_1$ -预测对于高维流形最为鲁棒。它与现有的科学架构兼容（例如，输出蛋白质 SE(3) 帧的架构），并且可扩展至高达 $D=2048$ 的维度，而 $x_t$ -预测会失效， $v$ -预测则随着维度增加而性能下降。

4. 奖励引导推理

该框架支持无需重新训练模型的奖励引导生成。通过利用流映射从中间步骤预测终端状态的能力，作者提出了一种奖励前瞻 (reward look-ahead) 策略。他们不使用当前噪声状态 $x_t$ 处的奖励梯度，而是使用预测终点 $\Phi_{t,1}(x_t)$ 处的梯度。这种方法被证明在将生成引导至特定目标方面更为有效。

关键结果

作者在两项生物任务上验证了 RMF：启动子 DNA 设计（ $\mathbb{R}^{1024 \times 4}$ 中单纯形上的序列）和蛋白质骨架生成（ $SE(3)^N$ 上， $N$ 高达 128）。

效率： RMF 实现了与最先进多步基线（如 Fisher FM、FrameFlow、FrameDiff）相当的样本质量，同时所需的函数评估次数减少了高达10 倍。
- 在蛋白质骨架生成中，RMF 在 5 步时实现了 82% 的可设计性，在 1 步时实现了 35%，而基线方法在这些少步情形下急剧下降或完全崩溃。
- 在 DNA 设计中，RMF 变体仅用 1 步就达到了 Fisher FM 100 步的性能。
稳定性： 提出的稳定化技术（特别是半群目标与 $x_1$ -预测的结合）使得在高维空间中的稳定训练成为可能，而其他流映射方法（如 GFM）在这些空间中无法收敛或产生高误差样本。
奖励引导： 与朴素引导或无引导生成相比，使用 $x_1$ 前瞻策略的奖励引导推理显著提高了与目标调控谱（DNA）和二级结构组成（蛋白质）的对齐度。

意义与主张

本文声称建立了黎曼流形上快速采样的原则性理论基础和实用指南。其意义在于：

弥合差距： 它将少步生成能力从欧几里得空间扩展到复杂的高维黎曼流形，这是科学应用的关键一步。
可扩展性： 通过推导避免高阶导数的恒等式并利用 $x_1$ -预测，该方法可扩展至以前少步流形生成模型无法触及的维度（ $D \approx 2000$ ）和序列长度。
科学工作流中的效率： 推理成本的降低（高达 10 倍）使得生成模型在迭代科学设计循环（如药物发现或合成生物学中的属性引导优化）中作为提议机制更具可行性。
奖励前瞻： 它证明了流映射能够通过从中间步骤预测终端状态来实现高效的奖励引导设计，这是此前在流形上未曾展示的能力。

作者将 RMF 定位为并非替代所有生成建模，而是一个专门的框架，旨在计算成本成为限制因素的几何约束科学领域中实现高效、高质量的生成。