0. TL; DR

这项工作提出了改进校正流训练的技术，使其即使在低函数评估次数（Number of Function Evaluations, NFE）设置下也能与知识蒸馏方法竞争。在实际设置下，用于训练校正流的Reflow算法单次迭代就足以学习到近乎笔直的轨迹；因此，当前使用多次Reflow迭代的做法是不必要的。

因此，作者提出了改进单轮校正流训练的技术，包括一个U形的时间步长分布和LPIPS-Huber预度量。借助这些技术，作者在CIFAR-10的1 NFE设置下，将先前2-rectified flow的FID提高了多达75%。在ImageNet 64×64上，作者改进的校正流在单步和两步设置中都优于最先进的蒸馏方法，如一致性蒸馏和渐进式蒸馏，并在FID上与改进的一致性训练（iCT）相媲美。

1. 背景介绍

扩散模型通过模拟一个随机去噪过程来生成数据，即逐渐将噪声转化为数据。为了高效地从扩散模型中采样，去噪过程通常被转换为其对应的常微分方程（Ordinary Differential Equations, ODEs），称为概率流ODE（probability flow ODEs, PF-ODEs）。

尽管使用PF-ODEs的扩散模型取得了成功，但绘制高质量样本需要用小步长对PF-ODE进行数值积分，这在计算上是昂贵的。目前，解决这个问题有两大类方法：

知识蒸馏（Knowledge Distillation）：例如一致性蒸馏（consistency distillation）和渐进式蒸馏（progressive distillation）。这类方法训练一个“学生”模型来直接预测PF-ODE的解。它们目前在低函数评估次数（low NFE，例如1-4步）的场景下处于领先地位。
无模拟流模型（Simulation-Free Flow Models）：例如校正流（rectified flows）和流匹配（flow matching）。校正流通过学习更平滑的ODE轨迹来减少截断误差，从而用更少的NFE生成高质量样本。它在中到高NFE场景下已显示出优于扩散模型的性能，但在低NFE场景下仍落后于蒸馏方法。

与知识蒸馏方法相比，校正流有几个优势：它可以推广到任意两个分布之间的映射，支持从数据到噪声的可逆映射（这在图像编辑等应用中很有用），并且可以评估模型的似然。目前最先进的训练校正流的技术是Reflow算法，它通过迭代训练来拉直ODE轨迹。然而，Reflow通常需要至少两轮迭代才能获得合理的单步生成性能，这个过程计算成本高昂且会累积误差。

本文的核心贡献在于，作者发现单次Reflow迭代在理论上就足以学习到近乎笔直的轨迹。基于这一发现，作者提出了一系列新的训练技术，显著提升了单轮训练的校正流（即2-rectified flow）的性能，使其在低NFE场景下能够与最先进的蒸馏方法相媲美。

2. 方法介绍

本章将首先回顾校正流和Reflow算法的基础，然后阐述作者的核心见解——即单次Reflow迭代的充分性，最后介绍基于此见解提出的几项关键训练改进技术。

2.1 校正流 (Rectified Flow)

校正流是一种通过求解ODE在两个分布$p_x$和$p_z$之间进行平滑过渡的生成模型。对于从$p_x$采样的$x$和从$p_z$采样的$z$，它们之间的线性插值被定义为$x_t = (1-t)x + tz$，其中$t \in [0, 1]$。

对于从$p_x$采样的$z_0$，以下ODE的解$z_t$与$x_t$具有相同的边际分布：

\[\frac{dz_t}{dt} = v_t(z_t) := \frac{1}{t}(z_t - \mathbb{E}[x|x_t = z_t])\]

在训练时，作者用一个神经网络$x_\theta$来估计条件期望$\mathbb{E}[x|x_t = z_t]$，并最小化一个加权的均方误差损失：

\[\min_\theta \mathbb{E}_{x,z \sim p_{xz}} \mathbb{E}_{t \sim p_t} \left[ \omega(t) \|x - x_\theta(x_t, t)\|_2^2 \right]\]

其中，$p_{xz}$是$x$和$z$的联合分布（耦合），$p_t$是时间$t$的采样分布，$\omega(t)$是一个权重函数。在最优情况下，$x_\theta$成为最小均方误差（MMSE）估计器。

2.2 Reflow算法

初始的独立耦合$p_{xz}(x, z) = p_x(x)p_z(z)$会导致弯曲的ODE轨迹，需要大量的NFE才能生成高质量样本。Reflow是一个递归的训练算法，旨在通过找到更好的耦合来拉直ODE轨迹。

Reflow的流程如下：

第一次迭代（训练1-rectified flow）：使用独立耦合$p^0_{xz}$训练一个模型$\theta_1$。
后续迭代 $k=1, \dots, K-1$：
- 使用上一阶段的模型$\theta_k$和其生成的ODE轨迹，来生成一个新的、更优的耦合$p^k_{xz}(x, z)$。具体来说，对于从$p_z$采样的$z$，其配对的$x$被定义为从$t=1$到$t=0$反向求解ODE所得到的结果，即$x = T^k(z)$。
- 使用这个新的合成耦合$p^k_{xz}$来训练下一个阶段的模型$\theta_{k+1}$（即(k+1)-rectified flow）。

当迭代次数$K \to \infty$时，Reflow的轨迹会收敛到直线。然而，在实践中，这需要多次昂贵的迭代。

2.3 一次Reflow就足够了

作者指出：在实际设置下，由训练良好的1-rectified flow生成的配对$(x’, z’)$，其线性插值轨迹几乎不会相互交叉。

如果两个不同的轨迹相交，例如$(x’, z’)$和$(x’’, z’’)$在某个时刻$t$相交，那么根据几何关系，必须满足$z’’ = z’ + \frac{1-t}{t}(x’ - x’’)$。然而，对于一个训练良好的1-rectified flow来说，这个构造出来的$z’‘$通常是一个“非典型”的噪声（例如，它可能具有非零的自相关性，或者其范数过大/过小，不在典型高斯噪声所在的“环带”上）。由于模型主要是在典型的高斯噪声上训练的，它很难将这样一个非典型的$z’‘$映射到数据流形$M_x$上。

这个观察表明，在训练2-rectified flow时，轨迹交叉的事件非常罕见。这意味着，最优的2-rectified flow的轨迹本身就应该近乎是笔直的。因此，作者得出结论：进行更多轮次的Reflow是不必要的，我们应该专注于如何更好地训练2-rectified flow。

2.4 改进Reflow的训练技术

基于上述见解，作者提出了几项关键的训练改进技术。

2.4.1 时间步长分布 (Timestep distribution)

作者分析了2-rectified flow在不同时间步$t$上的训练损失。与1-rectified flow的损失主要集中在中间部分不同，2-rectified flow的损失在$t$接近0和1的两端都很大。这是因为在$t=1$时，模型需要从噪声直接预测数据；在$t=0$时，模型需要从数据直接预测噪声，这都是非平凡的任务。

因此，作者提出使用一个U形的时间步长分布来代替均匀分布，从而让模型在训练时更多地关注两端这些更具挑战性的区域。实验表明，仅此一项改进就能在多个数据集上带来超过20%的FID提升。

2.4.2 损失函数

由于最优的2-rectified flow轨迹是近乎笔直的，其训练损失的下界也近乎为零。这给了我们使用除了均方误差（squared ℓ2 distance）之外的其他预度量（premetric）的自由。

伪-Huber损失 (Pseudo-Huber loss)：对离群点不那么敏感，可以减少梯度方差，使训练更容易。
LPIPS-Huber损失：将Pseudo-Huber损失与LPIPS（一种衡量图像感知相似度的度量）相结合。损失函数的形式为：

\[m_{lp-hub} = (1-t)m_{hub} + \text{LPIPS}(x, \hat{x}_\theta)\]

其中$\hat{x}\theta = x_t - t \cdot v\theta(x_t, t)$是模型预测的$x_0$。这个损失函数在$t$接近1时（即接近生成最终图像时），会更多地依赖LPIPS，从而强制模型关注生成图像的感知质量。

实验表明，引入LPIPS-Huber损失可以带来巨大的性能提升，在CIFAR-10上将FID从5.17降低到3.42。

2.4.3 使用预训练的扩散模型进行初始化

从头开始训练1-rectified flow的计算成本很高。作者证明了可以用一个预训练好的扩散模型（如EDM）来近似1-rectified flow的条件期望$\mathbb{E}[x|x_t]$。作者推导了如何将校正流的时间和尺度与VP和VE扩散模型的时间和尺度进行精确转换。

通过使用预训练的EDM模型来初始化Reflow的第一步（即生成合成数据对），可以极大地节省计算资源，并显著提升基线性能。

2.4.4 结合真实数据

虽然训练2-rectified flow可以完全在合成数据上进行（即data-free），但如果真实数据可用，也可以将其整合进来。作者通过在训练好的2-rectified flow++ 模型上，使用（真实数据，合成噪声）对进行短时间的微调，进一步提升了性能，在CIFAR-10上将FID从3.38降低到3.07。

3. 实验分析

作者将上述所有改进技术结合起来，称之为2-rectified flow++，并在多个图像生成任务上进行了广泛的评估。

3.1 无条件和类条件图像生成

表3（CIFAR-10无条件生成）：2-rectified flow++ 在单步生成（NFE=1）时取得了3.07的FID，优于所有已发表的蒸馏方法，如CD (3.55), PD (9.12), DFNO (3.78)等。在两步生成（NFE=2）时，FID进一步降低到2.40，性能接近最先进的一致性模型iCT-deep (2.24)。与原始的2-rectified flow (12.21) 相比，性能提升了惊人的75%。
表4（ImageNet 64x64类条件生成）：在单步生成时，2-rectified flow++ 的FID为4.31，同样优于所有蒸馏方法，并缩小了与iCT (4.02) 的差距。在两步生成时，FID为3.64，性能同样具有很强的竞争力。