何凯明团队新作：均值流用于一步生成建模 原创

摘要  

我们提出了一种原则性且高效的一步生成模型框架。我们引入了平均速度的概念来表征流场，这与流匹配方法中建模的瞬时速度形成对比。我们推导出了平均速度与瞬时速度之间的明确关系，并用其指导神经网络训练。我们的方法，称为均值流模型，是自包含的，无需预训练、蒸馏或课程学习。均值流模型展示了强大的实证性能：在ImageNet 256×256上从头训练，仅用一次函数评估（1-NFE）即可达到3.43的Fréchet Inception Distance（FID），显著优于之前最先进的一步扩散/流模型。我们的研究大幅缩小了一步扩散/流模型与其多步前驱之间的差距，我们希望它能激励未来研究重新审视这些强大模型的基础。

图1：ImageNet 256×256上的一步生成  

我们的均值流（MF）模型在生成质量上显著优于之前最先进的一步扩散/流方法。此处，iCT [43]、Shortcut [13]和我们的MF均为1-NFE生成，而IMM的1步结果 [52]涉及2-NFE指导。详细数据见表2。所示图像由我们的1-NFE模型生成。

1、引言  

生成建模的目标是将先验分布转换为数据分布。流匹配 [28,2,30] 为构建将一个分布映射到另一个分布的流路径提供了一个直观且概念简单的框架。与扩散模型 [42,44,19] 密切相关，流匹配侧重于指导模型训练的速度场。自提出以来，流匹配已在现代生成建模中得到广泛应用 [11,33,35]。 

流匹配和扩散模型在生成过程中都需要迭代采样。最近的研究对少步 —— 尤其是一步前馈 —— 生成模型给予了极大关注。一致性模型 [46,43,15,31] 率先在这一方向上为沿同一路径采样的输入引入了网络输出的一致性约束。尽管取得了令人鼓舞的结果，但一致性约束是作为网络行为的属性强加的，而指导学习的潜在真实场的属性仍不明确。因此，训练可能不稳定，需要精心设计的 “离散化课程”[46,43,15] 来逐步约束时间域。 

在这项工作中，我们提出了一种名为 MeanFlow 的原理性强且有效的一步生成框架。核心思想是引入一个表示平均速度的新真实场，以区别于流匹配中通常建模的瞬时速度。平均速度定义为位移与时间间隔的比率，其中位移由瞬时速度的时间积分给出。仅从这个定义出发，我们推导出平均速度与瞬时速度之间明确的内在关系，这自然成为指导网络训练的原理性基础。 

（我们在这项研究里搞了个叫MeanFlow的靠谱框架，专门解决“一步生成图像”的问题。核心思路是换了个底层概念：以前的方法算的是“瞬时速度”（类似汽车某一秒的速度表读数），得一步步调整才能生成图像；而我们改用“平均速度”（类似一段路的平均速度），它等于“总位移（路程）÷时间”，其中总位移是把瞬时速度在时间里一点点加起来得到的。就凭这个定义，我们找出了平均速度和瞬时速度之间的数学关系，拿这个关系当“指导手册”来训练神经网络，让模型学一次就能直接生成图，不用像以前那样反复算了。）

基于这一基本概念，我们训练神经网络直接对平均速度场进行建模。我们引入一个损失函数，促使网络满足平均速度与瞬时速度之间的内在关系，无需额外的一致性启发式方法。真实目标场的存在确保了最优解原则上独立于特定网络，这在实践中可以带来更稳健和稳定的训练。我们进一步表明，我们的框架可以自然地将无分类器引导（CFG）[18] 纳入目标场，在使用引导时采样时不会产生额外成本。 

我们的 MeanFlow 模型在一步生成建模中表现出强大的实验性能。在 ImageNet 256×256 [7] 上，我们的方法使用 1-NFE 生成实现了 3.43 的 FID 分数。这一结果以 50% 到 70% 的相对优势显著优于同类中之前的最先进方法（图 1）。此外，我们的方法是一个自包含的生成模型：它完全从头开始训练，无需任何预训练、蒸馏或课程学习。我们的研究在很大程度上缩小了一步扩散 / 流模型与其多步前身之间的差距，希望能启发未来的工作重新考虑这些强大模型的基础。

2、相关工作  

扩散与流匹配  

在过去十年中，扩散模型 [42, 44, 19, 45] 发展成为生成模型的非常成功的框架。这些模型逐步向干净数据添加噪声，并训练神经网络来逆转这一过程。这一过程涉及求解随机微分方程（SDE），随后被重新表述为概率流常微分方程（ODE） [45, 22]。流匹配方法 [28, 2, 30] 通过建模定义分布之间流路径的速度场扩展了这一框架。流匹配也可以看作是连续时间归一化流 [36] 的一种形式。

少步扩散/流模型  

从实际和理论角度来看，减少采样步骤已成为重要考量。一种方法是将预训练的多步扩散模型蒸馏为少步模型，例如 [39, 14, 41] 或分数蒸馏 [32, 50, 53]。早期对少步模型训练的探索 [46] 建立在基于蒸馏的方法演变之上。同时，一致性模型 [46] 作为独立生成模型开发，不需要蒸馏。这些模型对不同时间步的网络输出施加一致性约束，鼓励它们在轨迹上产生相同的端点。已研究了多种一致性模型和训练策略 [46, 43, 15, 31, 49]。 

在最近的工作中，几种方法侧重于刻画基于两个时间相关变量的扩散 / 流相关量。在 [3] 中，流图被定义为两个时间步之间的流的积分，并开发了几种形式的匹配损失用于学习。与我们的方法所基于的平均速度相比，流图对应于位移。Shortcut 模型 [13] 除了流匹配之外，还引入了自一致性损失函数，捕捉不同离散时间间隔的流之间的关系。归纳矩匹配 [52] 对不同时间步的随机插值的自一致性进行建模。

3、背景：流匹配中的速度场  

图2：流匹配中的速度场 [28]

左图：条件流 [28]。给定的 z_t 可以由不同的 (x, \epsilon) 对产生，导致不同的条件速度 v_t。右图：边际流 [28]，通过对所有可能的条件速度进行边际化获得。边际速度场作为网络训练的底层真实场。图示中的所有速度本质上是瞬时速度。插图遵循 [12]。（灰色点：来自先验的样本；红色点：来自数据的样本。）

由于给定的 z_t 和其 v_t 可以由不同的 x 和 \epsilon 产生，流匹配本质上建模所有可能性的期望，称为边际速度 [28]（图2右图）：

给定边际速度场 v(z_t, t)，通过求解 z_t 的常微分方程生成样本：

4、均值流模型  

4.1 均值流

我们方法的核心是引入一个表示平均速度的新场，而流匹配中建模的速度表示瞬时速度。

平均速度  

我们定义平均速度为两个时间步 t 和 r 之间的位移（通过积分获得）除以时间间隔。形式上，平均速度 u 为：

图3：平均速度 u(z, r, t) 的场

最左图：瞬时速度 v 决定路径的切线方向，而平均速度 u(z, r, t)（方程（3）定义）通常与 v 不对齐。平均速度与位移对齐，位移为 (t-r) u(z, r, t)。右三子图：u(z, r, t) 的场依赖于 r 和 t，此处显示 t=0.5, 0.7, 1.0 的情况。

均值流模型的最终目标是用神经网络 u_\theta(z_t, r, t) 逼近平均速度。假设我们能够精确逼近这一量，我们可以通过单次评估 u_\theta(\epsilon, 0, 1) 来逼近整个流路径。换句话说，正如我们将通过实验证明的，这种方法更适合单步或少步生成，因为它在推理时无需显式逼近时间积分，而建模瞬时速度时需要这样做。然而，直接使用方程（3）定义的平均速度作为训练网络的真实值是不可行的，因为它需要在训练期间计算积分。我们的关键洞察是将平均速度的定义方程进行操作，构建一个适合训练的优化目标，即使只能访问瞬时速度。

均值流恒等式  

为了得到适合训练的公式，我们重写方程（3）为：

现在我们对 t 进行微分，将 r 视为独立于 t。这导致：

其中左侧的操作使用了乘积法则，右侧使用了微积分基本定理。重新排列项，我们得到身份：

我们将此方程称为“均值流身份”，它描述了 v 和 u 之间的关系。易于证明方程（6）和方程（4）是等价的（见附录B.3）。  

方程（6）的右侧为 u(z_t, r, t) 提供了“目标”形式，我们将利用它来构建训练神经网络的损失函数。为了作为合适的目标，我们必须进一步分解时间导数项，接下来讨论。

计算时间导数  

使用平均速度进行训练  

到目前为止，公式独立于任何网络参数化。我们现在引入一个模型来学习 u。形式上，我们参数化一个网络 u_\theta 并鼓励它满足均值流身份（方程（6））。具体来说，我们最小化以下目标：

项 u_tgt 作为有效的回归目标，由方程（6）驱动。此目标仅使用瞬时速度 v 作为真实信号，无需计算积分。虽然目标应涉及 u 的导数（即 ∂ u），但它们被替换为参数化的对应项（即 \partial u_\theta，∂u_θ）。在损失函数中，对目标 u_{\text{tgt}} 应用停止梯度（sg）操作，遵循常见实践 [46, 43, 15, 31, 13]：在我们的案例中，它消除了通过雅可比向量积进行“双重反向传播”的需要，从而避免了高阶优化。尽管采用了这些优化实践，如果 u_\theta 实现零损失，易于证明它将满足均值流身份（方程（6）），从而满足原始定义（方程（3））。  

方程（10）中的速度 v(z_t, t) 是流匹配中的边际速度 [28]（见图2右图）。我们遵循 [28] 将其替换为条件速度（图2左图）。因此，目标为：

采样  

使用均值流模型进行采样只需将时间积分替换为平均速度：

与先前工作的关系  

虽然与先前的一步生成模型 [46, 43, 15, 31, 49, 23, 13, 52] 相关，但我们的方法提供了一个更原则性的框架。我们方法的核心是两个底层场 v 和 u 之间的函数关系，这自然导致 u 必须满足的均值流身份（方程（6））。此身份不依赖于神经网络的引入。相比之下，先前工作通常依赖于施加在神经网络行为上的额外一致性约束。一致性模型 [46, 43, 15, 31] 专注于锚定在数据侧的路径：在我们的符号中，这对应于对任何 t 固定 r \equiv 0。因此，一致性模型仅条件于单个时间变量，与我们的不同。另一方面，Shortcut [13] 和 IMM [52] 模型条件于两个时间变量：它们引入了额外的双时间自一致性约束。相比之下，我们的方法仅由平均速度的定义驱动，用于训练的均值流身份（方程（6））从该定义自然推导，无需额外假设。

4.2 带指导的均值流

我们的方法自然支持无分类器指导（CFG）[18]。我们不是在采样时简单应用 CFG（这会使 NFE 加倍），而是将 CFG 视为底层真实场的属性。这种公式允许我们在采样时保持 1-NFE 行为，同时享受 CFG 的好处。

真实场  

我们构建一个新的真实场 v^{\text{cfg}}：

这是类条件场和类无条件场的线性组合：

\[

v(z_t, t \mid \mathbf{c}) \triangleq \mathbb{E}_{p_t(v_t \mid z_t, \mathbf{c})}[v_t] \quad \text{和} \quad v(z_t, t) \triangleq \mathbb{E}_{\mathbf{c}}[v(z_t, t \mid \mathbf{c})],

\]

其中 v_t 是条件速度 [28]（更精确地说，在此上下文中是样本条件速度）。遵循均值流的精神，我们引入与 v^{\text{cfg}} 对应的平均速度 u^{\text{cfg}}。根据均值流身份（方程（6）），u^{\text{cfg}} 满足：

\[

u^{\text{cfg}}(z_t, r, t \mid \mathbf{c}) = v^{\text{cfg}}(z_t, t \mid \mathbf{c}) - (t-r) \frac{d}{dt} u^{\text{cfg}}(z_t, r, t \mid \mathbf{c}),

\]

同样，v^{\text{cfg}} 和 u^{\text{cfg}} 是独立于神经网络的底层真实场。此处，v^{\text{cfg}} 如方程（13）定义，可以重写为：

\[

v^{\text{cfg}}(z_t, t \mid \mathbf{c}) = \omega v(z_t, t \mid \mathbf{c}) + (1-\omega) u^{\text{cfg}}(z_t, t, t),

\]

我们利用关系：v(z_t, t) = v^{\text{cfg}}(z_t, t)，以及 v^{\text{cfg}}(z_t, t) = u^{\text{cfg}}(z_t, t, t)。

带指导的训练  

通过方程（15）和方程（16），我们构建网络及其学习目标。我们直接通过函数 u_\theta^{\text{cfg}} 参数化 u^{\text{cfg}}。基于方程（15），我们得到目标：

其中 u_{\text{tgt}} = \bar{v}_t - (t-r) (\bar{v}_t \partial_z u_\theta^{\text{cfg}} + \partial_t u_\theta^{\text{cfg}})。  

此公式类似于方程（9），唯一区别在于它具有修改后的 \bar{v}_t：

\[

\bar{v}_t \triangleq \omega v_t + (1-\omega) u_\theta^{\text{cfg}}(z_t, t, t),

\]

这由方程（16）驱动：方程（16）中的项 v(z_t, t \mid \mathbf{c})，即边际速度，被替换为（样本）条件速度 v_t，遵循 [28]。如果 \omega=1，此损失函数退化为无 CFG 情况（方程（9））。  

为了使方程（17）中的网络 u_\theta^{\text{cfg}} 暴露于类无条件输入，我们以 10% 的概率丢弃类条件，遵循 [18]。出于类似动机，我们还可以使方程（19）中的 u_\theta^{\text{cfg}}(z_t, t, t) 暴露于类无条件和类条件版本：细节在附录 B.1 中。

带 CFG 的单 NFE 采样  

4.3 设计决策

损失度量  

5、实验

实验设置  

我们在 ImageNet [7] 上以 256×256 分辨率进行主要实验。我们在 50K 张生成图像上评估 Fréchet Inception Distance（FID）[17]。我们检查函数评估次数（NFE），并默认研究 1-NFE 生成。遵循 [34, 13, 52]，我们在预训练变分自编码器（VAE）分词器的潜在空间上实现我们的模型。对于 256×256 图像，分词器产生 32×32×4 的潜在空间，这是模型的输入。我们的模型全部从头训练。实现细节在附录 A 中。  

在我们的消融研究中，我们使用 ViT-B/4 架构（即“基础”大小，补丁大小为4）[9]，如 [34] 开发，训练 80 个 epoch（400K 次迭代）。作为参考，[34] 中的 DiT-B/4 具有 68.4 FID，[33] 中的 SiT-B/4（我们的重现）具有 58.9 FID，两者均使用 250-NFE 采样。

5.1 消融研究

表1：ImageNet 256×256 数据集上 1 - NFE 生成的消融实验 说明：评估指标为 FID - 50K。默认配置用灰色标注：使用 B/4 骨干网络，从头开始训练 80 个 epoch。

我们在表1中研究模型属性，分析如下：  

从流匹配到均值流  

JVP 计算  

条件于 (r, t)  

时间采样器  

损失度量  

据报告 [43]，损失度量的选择显著影响少步/单步生成的性能。我们在表1e中研究这一方面。我们的损失度量通过自适应损失加权 [15] 实现，幂为 p（4.3节）。表1e 显示，p=1 实现最佳结果，而 p=0.5（类似于 [43] 中的伪 Huber 损失）也具有竞争力。标准平方 L2 损失（此处 p=0）相较于其他设置表现不佳，但仍能产生有意义的结果，与 [43] 的观察一致。

指导尺度  

表1f 报告了带 CFG 的结果。与多步生成中的观察一致 [34]，CFG 在我们的 1-NFE 设置中也显著提高了生成质量。我们强调，我们的 CFG 公式（4.2节）天然支持 1-NFE 采样。

可扩展性  

图4：均值流模型在 ImageNet 256×256 上的可扩展性。  

报告了 1-NFE 生成 FID。所有模型从头训练。应用 CFG，同时保持 1-NFE 采样行为。我们的方法在模型大小方面表现出良好的可扩展性。

图4：MeanFlow在不同模型规模和训练时长下的1-NFE FID结果。该图展示了使用不同模型规模和训练时长训练的MeanFlow模型在ImageNet 256×256上的1-NFE FID性能。每个点代表一个单独训练的模型。 

5.2 与先前工作的比较

ImageNet 256×256 比较  

在图1中，我们与先前的一步扩散/流模型进行比较，总结在表2（左）中。总体而言，均值流在同类方法中大幅优于先前方法：它实现了 3.43 FID，相较于 IMM 的一步结果 7.77 [52] 有超过 50% 的相对改进；如果仅比较 1-NFE（不仅仅是一步）生成，均值流相较于之前最先进的 10.60（Shortcut [13]）有近 70% 的相对改进。我们的方法大幅缩小了一步与多步扩散/流模型之间的差距。  

在 2-NFE 生成中，我们的方法实现了 2.20 FID（表2，左下）。此结果与领先的多步扩散/流模型基准相当，即 DiT [34]（FID 2.27）和 SiT [33]（FID 2.15），两者均具有 250×2 NFE（表2，右），在相同的 XL/2 骨干下。我们的结果表明，少步扩散/流模型可以与多步前驱竞争。正交改进，如 REPA [51]，是适用的，留待未来工作。  

值得注意的是，我们的方法是自包含的，完全从头训练。它在不使用任何预训练、蒸馏或 [43, 15, 31] 中采用的课程学习的情况下实现了强大结果。

表2：ImageNet-256×256上的类条件生成

所有条目在适用时均报告 CFG。左：从头训练的 1-NFE 和 2-NFE 扩散/流模型。右：作为参考的其他生成模型家族。在两表中，“×2”表示 CFG 每采样步引入 2 NFE。我们的均值流模型均训练 240 个 epoch，除了“MeanFlow-XL+”训练更多 epoch 并使用为长期训练选择的配置，详见附录。†：iCT [43] 结果由 [52] 报告。

CIFAR-10 比较  

我们在表3中报告了 CIFAR-10 [25]（32×32）上的无条件生成结果。报告了 50K FID，1-NFE 采样。所有条目使用相同的 U-net [38]，由 [44] 开发（55M），直接应用于像素空间。所有其他竞争者使用 EDM 风格预处理器 [22]，而我们没有预处理器。实现细节在附录中。在此数据集上，我们的方法与先前方法具有竞争力。

表3：CIFAR-10 无条件生成

6、结论  

我们提出了均值流，一个原则性且高效的一步生成框架。广义上讲，本工作考虑的场景与物理中的多尺度模拟问题相关，可能涉及空间或时间上的多种尺度、长度和分辨率。进行数值模拟本质上受到计算机分辨尺度范围能力的限制。我们的公式涉及在较粗糙的粒度级别描述底层量，这是物理中许多重要应用的共同主题。我们希望我们的工作将连接生成模型、模拟和相关领域的动态系统研究。

本文转载自​​​AIRoobt​​ ，作者：Zhengyang Geng等