扩散模型去伪求真，Straightness Is Not Your Need？！

23年以来，基于rectified flow[1]和flow-matching[2]的扩散模型在生成领域大展异彩，许多基于flow-matching的模型和文章如雨后春笋般涌现，仿佛“直线”已经成为一种政治正确。

该blog的动机，是源于网络平台，诸如知乎，小红书，乃至许多论文中都出现了很多对于rectified flow (flow-matching) 的错误理解和解读。本文希望能够提供一个相对合理的视角来重新思考关于rectified flow相关的话题， 抛砖引玉。

该blog基于最近的论文：

Rectified Diffusion: Straightness is Not Your Need

单位: MMLab-CUHK, Peking University, Princeton University

https://arxiv.org/pdf/2410.07303

https://github.com/G-U-N/Rectified-Diffusion/tree/master

https://huggingface.co/wangfuyun/Rectified-Diffusion

直观对比，FMs和一般的DMs有本质区别吗？

早在21年，Variational Diffusion Models[3] 就已经使用了一种general的扩散模型表达式 其中 关于时间 t 的严格正函数。许多的后续研究，例如非常出名的DPM-Solver[4]，都是基于这种general的扩散模型的表达式。

这种表达式的好处，在于他提供了一个统一的form来看待很多diffusion的形式。下面我们举几个例子：

• 的表达式，也被称为variance preserving (VP)[6]的表达式。其只需要满足 。在这种情况下，我们很容易看到 ，这也是该表达式被称为 variance preserving的原因。所以，我们这样理解DDPM就会发现他的form也并不复杂，之所以 DDPM论文的原始推导显得如此复杂，就在于他额外引入了一个 序列，然后用 序列的累乘来计算得到 的序列。
• variance exploding (VE) [6]的表达式，也常常被认为是EDM[7]的form，只要满足 即可。这个时候， 的值域就不会仅仅限制在0-1之间，而是会有 以满足图像的原始信号在noise较大时，被近似完全抹除。
• 我们以同样的视角来看flow-matching的表达式，只要设置 就可以了。要说这种form的优点，大概就是他简单，不过从笔者的视角看，他也并不比variance exploding (VE) 的form简单, 因为VE中甚至只需要考虑唯一的变量 。

以这种视角我们看到，flow-matching，它只是general 扩散模型表达式的一种特例，他并不比其他的形式包括VP, VE, Sub-VP更加特殊。要说为什么有很多文章诸如sd3，flow-matching等文章中的效果要比其他form好，我个人认为更多是超参数导致的问题，例如如何分配时间的采样，时间的weigthing，还有prediction type的差异等等。论文中通常为了表现自己方法的优越性，都会对自己的方法进行比较精细的超参数搜索。而对于一种diffusion form的最优超参数设置，未必适用于其他的diffusion forms。

FMs的轨迹真的直吗？

关于这点的误解是最多的，其实所有的diffusion模型训练结束后，它们的采样轨迹几乎必然是curved。我们来看flow- matching的形式

因此

从上面的式子可能会误认为, 永远都是，因此是一条直线。但是由于我们diffusion的训练， 和 都是随机配对的，所以实际上 的最终的收敛目标会是

并且注意到, 对于不同时间刻, 这种期望并不相同, 因此不同时间刻 对应的 总是不同的。因此, 虽然flow-matching的每次denoise一小步的范围可以认为是以直线前进, 但是由于这些小的直线的方向并不相同, 总体的采样曲线就是曲线 (或者说很多小折线近似出来的曲线)。

我们用下面一张图更加形象的说明这个例子, 红色的点, 代表离散采样过程中每一步的 , 虽然每步模型denoise都是走直线, 但是由于 并不相同, 因此整体的轨迹为曲线。如果你训练一个flowmatching模型, 打印他的轨迹, 大概率会发现几乎所有的采样轨迹都是曲线（折线）, 而非直线。

Rectified Flow是怎么让轨迹变直的？

通读rectified flow[1] [8] [9]的相关文章，我们可以看到rectified flow相较于一般的DDPM，主要有三点核心的观点：

1. Flow-matching。Rectified Flow提出上述的的扩散形式。中间的噪声状态 定义为 , 其中 为干净数据, 是标准正态噪声, t in[0,1]$ 是时间步长。与原始DDPM的semi-linear形式相比, 这种设计更为简单。
2. -Prediction。Rectified Flow提出采用 -预测, 即模型学习预测 。这使得去噪形式变得简单。例如, 可以基于 预测 , 公式为 , 其中 表示模型参数，^表示预测值。此外, 它避免了在 时使用 -预测时出现的数值问题。例如, , 这是invalid的。
3. Rectification。Rectification（也称为Reflow）是Rectified Flow中提出的重要技术。它是一种渐进式的重训练方法, 能够显著提高低步长下的生成质量, 并保持标准扩散模型的多步预测的灵活性。具体来说, 它将标准扩散训练中采用的任意塊合 (真实数据) 和 (噪声)转变为一个新的确定性塊合 (生成数据) 和 (预先收集的噪声)。简单来说,它将 替换为 , 其中 是真实数据, 是由预训练扩散模型 生成的数据, 是随机采样的噪声, 是生成 所使用的噪声。

Rectified flow的采样轨迹，只有在执行了多次（一次）的rectification的操作之后，才会慢慢的变为直线，这也正是rectified flow中实现单步生成的重要操作。

Rectified flow包括后续的诸多工作，都强调Rectification这个操作，仅适用于flow-matching形式的diffusion模型。也就是说他们认为前两点~(1和2)是采用Rectification并实现采样加速的的基础，并强调修正过程（Rectification）将ODE路径“拉直”。如果是其他的diffusion forms，如VP, VE, sub-VP[6]等，则无法采用rectification的操作来实现加速。

这就导致，在InstaFlow[10]中 (Rectified Flow的后续工作，其作者尝试将rectified flow拓展到stable diffusion上的文生图任务)，rectified flow的作者使用stable diffusion的原始权重初始化，但是将SD转变为flow-matching的form，并采用-prediction来进行重新训练。

所以，rectified flow中的这种观点真的正确吗？在加速sd的过程中，真的有必要首先将其转化为flow-matching的形式和-prediction吗? Flow-matching的diffusion form真的显著比别的形式好吗？

Recfified Diffusion: 本质是使用配对的噪声样本对重训练。

Rectification实现采样加速的成功之处在于使用成对的噪声-样本进行训练。 为了清晰地展示差异，我们可视化了标准的flow matching训练和rectified flow的rectification训练的过程，分别对应算法1和算法2。差异部分以红色标出。一个关键的观察点是，在标准流匹配训练中，表示从训练集中随机采样的真实数据，而噪声 是从高斯分布中随机采样的。这导致了噪声和样本之间的随机配对。相对的，在rectification的训练中，噪声是预先从高斯分布中采样的，图像则是通过预先采样的噪声由前一轮重整（预训练模型）生成的，形成了确定性的噪声与样本的配对。

Flow-matching Training是标准diffusion training的子集。此外，算法2 可视化了更一般的扩散模型的训练过程，与算法1的差异以蓝色和橙色标出。值得注意的是，流匹配是我们讨论的扩散形式的一个特殊情况。从算法中可以看出，它们之间的唯一区别在于扩散形式和预测类型。因此，流匹配训练只是特定扩散形式和预测类型下的标准扩散训练的特殊情况。

通过比较算法2 和算法3 与算法1，可以自然而然的推导出算法4。本质上，通过引入预训练模型来收集噪声-样本对，并在标准的扩散训练中用这些预先收集的配对替换随机采样的噪声和真实样本，我们就得到了Rectified Diffusion的训练算法。

以此为基础，我们提出rectified diffusion。如图下图所示，我们的总体设计非常简单。我们保留了预训练扩散模型的所有内容，包括噪声调度器、预测类型、网络架构，甚至训练和推理代码。唯一的不同在于用于训练的噪声 和数据 是由预训练扩散模型预先收集和生成的，而不是独立地从高斯分布和真实数据集中采样。也正因为，我们不改变任何原始diffusion的设置，我们在stable diffusion上的实验，具有更小的gap，我们只使用了Rectified flow 8%的trained images，就在一步生成的性能上获得了远超rectified flow的性能。

训练目标是轨迹一阶化

对于上述讨论的通用扩散形式 ，存在一个精确的ODE解形式[4]，

其中, 是 的反函数。左侧项 是一个预定义的确定性缩放。右侧项是关于 预测的指数加权积分。如果满足一阶ODE，上述任意 和 的积分始终等价于

在文章中, 我们证明了任意 和 的上述两个方程等价当且仅当ODE轨迹上的 预测是常数。

一阶ODE与预定义扩散形式具有相同的形式

简而言之, 我们假设ODE轨迹是一阶ODE, 且存在一个解点 。因此, ODE轨迹上解点 对应的 预测是常数, 我们将其记作 。将 和 代入方程

这与预定义的前向过程形式完全相同。因此, 我们可以得出一阶ODE正是数据和噪声按预定义的前向扩散形式进行加权插值。唯一的区别在于, 上述方程中的 和 是在同一ODE轨迹上的确定性对, 而在标准扩散训练中, 和 是随机采样的。这意味着, 如果我们在训练中实现了数据 和噪声 的完美塊合, 并且不同路径之间不存在交叉 (否则 预测可能是不同路径的 预测期望）, 则在理想情况下, 训练出的扩散模型将获得一阶ODE。

一阶化轨迹可能是弯曲的

对于一阶ODE, 虽然基于流匹配的方法的轨迹是直线, 但其他形式的扩散模型的轨迹可以本质上是弯曲的。但如果我们定义 , 我们将从方程

得到 。我们可以很容易地观察到 的轨迹是一条从初始点 朝向 方向的直线（即一阶轨迹可以转换为直线）。我们在下图中展示了我们的发现：我们选择 和 。图 (a) 和图 (b) 展示了流匹配和EDM的一阶轨迹, 它们都是直线, 但EDM的一阶轨迹和flowmatching完全不同。图 (c) 和图 (d) 展示了DDPM和Sub-VP的一阶轨迹, 它们的一阶轨迹本质上是弯曲的。图 (e) 展示了 的轨迹, 表明所有的一阶轨迹都可以通过简单的时间步依赖缩放转换为直线。

实验验证

我们进行了广泛的实验验证和方法对比，我们的方法取得了一致超越rectified flow相关方法的性能，并且与最先进的蒸馏加速算法[11] [12]也取得了comparable的结果。

最后欢迎大家关注我们的论文：

参考

1. ^abFlow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow https://arxiv.org/abs/2209.03003

2. ^Flow Matching for Generative Modeling https://arxiv.org/abs/2210.02747

3. ^Variational Diffusion Models https://arxiv.org/abs/2107.00630

4. ^abDPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps https://arxiv.org/abs/2206.00927

5. ^Denoising Diffusion Probabilistic Models https://arxiv.org/abs/2006.11239

6. ^abcdScore-Based Generative Modeling through Stochastic Differential Equations https://arxiv.org/abs/2011.13456

7. ^Elucidating the Design Space of Diffusion-Based Generative Models https://arxiv.org/abs/2206.00364

8. ^PeRFlow: Piecewise Rectified Flow as Universal Plug-and-Play Accelerator https://arxiv.org/abs/2405.07510

9. ^Improving the Training of Rectified Flows https://arxiv.org/abs/2405.20320

10. ^InstaFlow: One Step is Enough for High-Quality Diffusion-Based Text-to-Image Generation https://arxiv.org/abs/2309.06380

11. ^Phased Consistency Model https://arxiv.org/abs/2405.18407

12. ^Score identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation https://arxiv.org/abs/2404.04057

公众号后台回复“数据集”获取100+深度学习各方向资源整理