扩散模型=流匹配？谷歌DeepMind博客深度详解这种惊人的等价性

扩散模型和流匹配实际上是同一个概念的两种不同表达方式吗？

从表面上看，这两种方法似乎各有侧重：扩散模型专注于通过迭代的方式逐步去除噪声，将数据还原成清晰的样本。

而流匹配则侧重于构建可逆变换系统，目标是学习如何将简单的基础分布精确地映射到真实数据分布。

因为流匹配的公式很简单，并且生成样本的路径很直接，最近越来越受研究者们的欢迎，于是很多人都在问：

「到底是扩散模型好呢？还是流匹配好？」

现在，这个困扰已得到解答。Google DeepMind 的研究团队发现，原来扩散模型和流匹配就像一枚硬币的两面，本质上是等价的 (尤其是在流匹配采用高斯分布作为基础分布时)，只是不同的模型设定会导致不同的网络输出和采样方案。

这无疑是个好消息，意味着这两种框架下的方法可以灵活搭配，发挥组合技了。比如在训练完一个流匹配模型后，不必再局限于传统的确定性采样方法，完全可以引入随机采样策略。

链接：https://diffusionflow.github.io

在这篇博客的开头，作者们写道：「我们的目标是帮助大家能够自如地交替使用这两种方法，同时在调整算法时拥有真正的自由度 —— 方法的名称并不重要，重要的是理解其本质。」

扩散模型与流匹配

采样

看到这里，需要思考一下。人们常说流匹配会产生直线路径，但在上图中，其采样轨迹看起来是弯曲的。

在下面的交互式图表中，我们可以通过滑块更改右侧数据分布的方差。

不过，在像图像这样的真实数据集上找到这样的直线路径要复杂得多。但结论仍然是相同的：最优的积分方法取决于数据分布。

我们可以从确定性采样中得到的两个重要结论：

1. 采样器的等价性：DDIM 与流匹配采样器等价，并且对噪声调度的线性缩放不变。
2. 对直线性的误解：流匹配调度仅在模型预测单个点时才是直线。

训练

对于扩散模型，学习模型是通过最小化加权均方误差（MSE）损失来完成的：

流匹配也符合上述训练目标：

网络应该输出什么

如何选择加权函数

加权函数是损失函数中最重要的部分，它平衡了图像、视频和音频等数据中高频和低频分量的重要性。这一点至关重要，因为这些信号中的某些高频分量是人类无法感知的。如果通过加权情况来查看损失函数，可以得出以下结果：

即公式 (7) 中的条件流匹配目标与扩散模型中常用的设置相同。下面绘制了文献中常用的几个加权函数。

流匹配加权（也称为 v-MSE + 余弦调度加权）会随着 λ 的增加而呈指数下降。该团队在实验中发现了另一个有趣的联系：Stable Diffusion 3 加权 [9]（这是流匹配的一种重新加权版本）与扩散模型中流行的 EDM 加权 [10] 非常相似。

如何选择训练噪声调度？

最后讨论训练噪声调度，因为在以下意义上，它对训练的重要程度最低：

1. 训练损失不会随训练噪声调度变化。具体来说，损失函数可以重写为  它只与端点（λ_max, λ_min）有关，但与中间的调度 λ_t 无关。在实践中，应该选择合适的 λ_max, λ_min，使得两端分别足够接近干净数据和高斯噪声。λ_t 可能仍然会影响训练损失的蒙特卡洛估计量的方差。一些文献中提出了一些启发式方法来在训练过程中自动调整噪声调度。这篇博文有一个很好的总结：https://sander.ai/2024/06/14/noise-schedules.html#adaptive

2. 类似于采样噪声调度，训练噪声调度不会随线性扩展（linear scaling）而变化，因为人们可以轻松地将线性扩展应用于 z_t，并在网络输入处进行 unscaling 以获得等价性。噪声调度的关键定义属性是对数信噪比 λ_t。

3. 人们可以根据不同的启发式方法为训练和采样选择完全不同的噪声调度：对于训练，最好有一个噪声调度来最小化蒙特卡洛估计量的方差；而对于采样，噪声调度与 ODE / SDE 采样轨迹的离散化误差和模型曲率更相关。

总结

下面给出了训练扩散模型 / 流匹配的一些要点：

1. 加权中的等价性：加权函数对于训练很重要，它平衡了感知数据不同频率分量的重要性。流匹配加权与常用的扩散训练加权方法相同。
2. 训练噪声调度的不重要性：噪声调度对训练目标的重要性要小得多，但会影响训练效率。
3. 网络输出的差异：流匹配提出的网络输出是新的，它很好地平衡了 

更深入地理解采样器

这一节将更详细地介绍各种不同的采样器。

回流算子

流匹配中的回流（Reflow）运算是使用直线将噪声与数据点连接起来。通过基于噪声运行一个确定性的采样器，可以得到这些 (数据，噪声) 对。然后，可以训练模型，使之可以根据给定噪声直接预测数据，而无需采样。在扩散技术的相关文献中，这同样的方法是最早的蒸馏技术之一。

确定性采样器与随机采样器

此前已经讨论了扩散模型或流匹配的确定性采样器。另一种方法是使用随机采样器，例如 DDPM 采样器。

执行一个从 λ_t 到 λ_t+Δλ 的 DDPM 采样步骤完全等价于执行一个到 λ_t+2Δλ 的 DDIM 采样步骤，然后通过执行前向扩散重新噪声化到 λ_t+Δλ。也就是说，通过前向扩散重新噪声化恰好逆转了 DDIM 所取得的一半进展。为了理解这一点，让我们看一个 2D 示例。从相同的高斯分布混合开始，我们可以执行一个小的 DDIM 采样步骤，左图带有更新反转的符号，右图则是一个小的前向扩散步骤：

对于单个样本而言，这些更新的行为完全不同：反转的 DDIM 更新始终将每个样本推离分布模式，而扩散更新完全是随机的。但是，在汇总所有样本时，更新后得到的分布是相同的。因此，如果执行 DDIM 采样步骤（不反转符号），然后执行前向扩散步骤，则整体分布与更新之前的分布保持不变。

通过重新加噪来撤消的 DDIM 步骤的比例是一个超参数，并且可以自由选择（即不必一定是 DDIM 步骤的一半）。这个超参数在《Elucidating the design space of diffusion-based generative models》中被称为 level of churn，可译为「搅动水平」。有趣的是，将搅动添加到采样器的效果是：减少采样过程早期做出的模型预测对最终样本的影响，并增加对后续预测的权重。如下图所示：

SDE 和 ODE 视角

前面，我们已经观察到扩散模型和流匹配算法之间的等价性。下面将使用 ODE 和 SDE 来形式化地描述正向过程和采样的等价性，以实现理论上的完整性。

扩散模型

扩散模型的前向过程涉及到随时间推移逐渐破坏一个数据，而该过程可使用以下随机微分方程（SDE）来描述：

其中 dz 是无穷小的高斯（即布朗运动）。f_t 和 g_t 决定了噪声调度。其生成过程由前向过程的逆过程给出，其公式为：

其中 ∇log p_t 是前向过程的分数。

请注意，这里引入了一个附加参数 η_t，它控制的是推理时的随机性。这与之前介绍的搅动（churn）参数有关。当离散化后向过程时，如果 η_t=0，则是恢复 DDIM；如果 η_t=1，则是恢复 DDPM。

流匹配

两个框架的等价性

结语

读到这里，想必你已经理解了扩散模型和高斯流匹配的等价性。不过，文中重点介绍的是流匹配为该领域带来的两个新模型规范：

• 网络输出：流匹配提出了一种网络输出的向量场参数化方案，并且其不同于扩散文献中使用的方案。当使用高阶采样器时，网络输出可能会有所不同。它也可能影响训练动态。
• 采样噪声调度：流匹配利用了简单的采样噪声调度 α_t = 1-t 和 σ_t = t，并且更新规则与 DDIM 相同。

该团队最后表示：「如果能通过实证方式研究这两个模型规范在不同的真实应用中的重要性，那一定会很有趣。我们将此留给了未来的工作。」

参考文献

1. Flow matching for generative modeling
2. Lipman, Y., Chen, R.T., Ben-Hamu, H., Nickel, M. and Le, M., 2022. arXiv preprint arXiv:2210.02747.
3. Flow straight and fast: Learning to generate and transfer data with rectified flow
4. Liu, X., Gong, C. and Liu, Q., 2022. arXiv preprint arXiv:2209.03003.
5. Building normalizing flows with stochastic interpolants
6. Albergo, M.S. and Vanden-Eijnden, E., 2022. arXiv preprint arXiv:2209.15571.
7. Stochastic interpolants: A unifying framework for flows and diffusions
8. Albergo, M.S., Boffi, N.M. and Vanden-Eijnden, E., 2023. arXiv preprint arXiv:2303.08797.
9. Denoising diffusion implicit models
10. Song, J., Meng, C. and Ermon, S., 2020. arXiv preprint arXiv:2010.02502.
11. Score-based generative modeling through stochastic differential equations
12. Song, Y., Sohl-Dickstein, J., Kingma, D.P., Kumar, A., Ermon, S. and Poole, B., 2020. arXiv preprint arXiv:2011.13456.
13. Understanding diffusion objectives as the elbo with simple data augmentation
14. Kingma, D. and Gao, R., 2024. Advances in Neural Information Processing Systems, Vol 36.
15. Diffusion is spectral autoregression [HTML]
16. Dieleman, S., 2024.
17. Scaling rectified flow transformers for high-resolution image synthesis
18. Esser, P., Kulal, S., Blattmann, A., Entezari, R., Muller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F. and others,, 2024. Forty-first International Conference on Machine Learning.
19. Elucidating the design space of diffusion-based generative models
20. Karras, T., Aittala, M., Aila, T. and Laine, S., 2022. Advances in neural information processing systems, Vol 35, pp. 26565—26577.
21. Knowledge distillation in iterative generative models for improved sampling speed [PDF]
22. Luhman, E. and Luhman, T., 2021. arXiv preprint arXiv:2101.02388.
23. Denoising diffusion probabilistic models
24. Ho, J., Jain, A. and Abbeel, P., 2020. Advances in neural information processing systems, Vol 33, pp. 6840—6851.
25. Progressive Distillation for Fast Sampling of Diffusion Models
26. Salimans, T. and Ho, J., 2022. International Conference on Learning Representations.
27. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models
28. Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C. and Zhu, J., 2022. arXiv preprint arXiv:2211.01095.

文章来自微信公众号“机器之心”