一、介绍
扩散模型虽然能够产生高质量的样本，但这方法仍然需要迭代推理过程，即需要神经网络的数十次到数百次前向传递，使得生成样本的速度缓慢且昂贵。而以往的加速采样方法需要复制的训练机制，例如多个训练阶段、多个网络等等。于是，本文作者考虑到端到端设置，单个模型在单次训练运行中获得一步去噪。
二、方法——少步生成的快捷模型
扩散和流匹配模型，通过学习将噪声转换为数据的常微分方程 (ODE) 来解决生成建模问题。在这项工作中，为了简单起见，作者采用最佳传输流匹配目标。我们将 \(\textcolor{blue}{x_t}\) 定义为相同维度的数据点 \(\textcolor{blue}{x_1 \sim D}\) 和噪声点 \(\textcolor{blue}{x_0 ∼ N (0, I)}\) 之间的线性插值。速度 \(\textcolor{blue}{v_t}\) 是从噪声到数据点的方向：
给定 \(\textcolor{blue}{x_0}\) 和 \(\textcolor{blue}{x_1}\)，速度 \(\textcolor{blue}{v_t}\) 就完全确定了。但仅给定 \(\textcolor{blue}{x_t}\)，就有多个可能的对 \(\textcolor{blue}{(x_0, x_1)}\)，因此速度可以采用不同的值，从而使 \(\textcolor{blue}{v_t}\) 成为随机变量。流模型学习神经网络来估计期望值 \(\textcolor{blue}{\bar{v}_t = E[v_t | x_t]}\) 是 \(\textcolor{blue}{x_t}\) 处所有可能速度的平均值。可以通过对噪声 \(x_0\) 和数据 \(x_1\) 对的随机采样对的经验速度进行回归来优化流模型：
快捷模型，这是一类端到端生成模型，可以在任何推理预算下（包括在单个采样步骤中）生成高质量的生成模型。该模型的关键思想是不仅根据噪声水平还根据所需的步长来调节神经网络，使其能够在去噪过程中准确地向前跳跃。快捷模型可以被视为在训练期间执行自蒸馏，因此不需要单独的蒸馏步骤，并且可以在单次运行中进行训练。不需要时间表或仔细的热身。快捷模型训练起来非常高效，只需要比基础扩散模型多 16% 的计算量。
图 2：朴素的扩散和流量匹配模型在几个步骤的生成中失败了。左：训练路径是通过随机配对数据和噪声创建的。请注意，路径重叠；仅给定 \(x_t\) 时，\(v_t\) 到数据点的方向存在固有的不确定性。右图：虽然流匹配模型学习确定性 ODE，但其路径不是笔直的，必须严格遵循。预测方向 \(v_t\) 指向合理数据点的平均值。推理步骤越少，各代对数据集平均值的偏差就越大，导致它们偏离轨道。在第一个采样步骤中，模型指向数据集均值，因此无法在单个步骤中生成多模态数据（参见红色圆圈）。
如图 2 所示，流匹配学习一个 ODE，将噪声沿弯曲路径映射到数据。天真地采取大的采样步骤会导致大的离散化误差，并且在单步情况下，会导致灾难性的失败。对 \(\textcolor{blue}{d}\) 进行调节允许快捷模型考虑未来的曲率，并跳转到正确的下一个点而不是偏离轨道。作者将从 \(\textcolor{blue}{x_t}\) 到正确的下一个点 \(\textcolor{blue}{x^′_{t+d}}\) 的归一化方向称为快捷方式 \(\textcolor{blue}{s(x_t, t, d)}\)：
本文的目标是训练一个快捷模型 \(\textcolor{blue}{s_\theta(x_t, t, d)}\) 以学习 \(\textcolor{blue}{x_t}\)、\(\textcolor{blue}{t}\) 和 \(\textcolor{blue}{d}\) 所有组合的快捷方式。因此，快捷模型可以被视为流匹配模型对更大步长的推广：流匹配模型仅学习瞬时速度，而快捷模型还学习做出更大的跳跃。在 \(\textcolor{blue}{d → 0}\) 时，捷径相当于流。
计算训练 \(\textcolor{blue}{s_\theta(x_t, t, d)}\) 目标的一种简单方法是使用足够小的步长完全模拟 ODE 前向。然而，这种方法的计算成本很高，尤其是对于端到端训练。相反，我们利用快捷方式模型固有的自洽属性，即一个快捷方式步骤等于两个连续的一半大小的快捷方式步骤：
图 3：快捷模型训练概述。在 \(d ≈ 0\) 时，快捷目标相当于流匹配目标，并且可以通过回归到经验 \(E[v_t|x_t]\) 样本来进行训练。较大 \(d\) 快捷方式的目标是通过连接两个 \(d/2\) 快捷方式的序列来构造的。两个目标可以联合训练；捷径模型不需要两阶段过程或离散化时间表。
这使我们能够使用 \(\textcolor{blue}{d > 0}\) 的自洽目标并使用流匹配损失（方程 2）作为 \(\textcolor{blue}{d = 0}\) 的基本情况来训练快捷模型。原则上，可以在 \(\textcolor{blue}{d ∼ p(d)}\)任何分步上训练模型。在实践中，我们将目标分成两部分，一部分是使用 \(\textcolor{blue}{d = 0}\) 进行训练，另一部分是使用随机采样的 \(\textcolor{blue}{d > 0}\) 目标进行训练。因此，我们得到组合的捷径模型损失函数：
直观地说，上述目标学习了从噪声到数据的映射，该映射在任何步长序列下（包括直接在单个步骤中）查询时都是一致的。目标的流量匹配部分以小步长建立快捷模型，以匹配经验速度样本。这确保了快捷模型在使用多个步骤进行查询时能够开发基本生成功能，就像等效的流匹配模型一样。在自洽部分，通过连接两个较小的快捷方式的序列来构造较大步长的适当目标。这将生成能力从多步传播到几步再到一步。可以使用单个模型并在单个端到端训练运行中联合训练组合目标。
三、实验
3.1 快捷模型与之前的一步生成方法相比如何？
在本节中，我们通过使用相同的模型架构和代码库从头开始训练每个目标，仔细地将我们的方法与许多先前的方法进行比较。我们利用 DiT-B 扩散变压器架构。我们考虑将 CelebAHQ-256 用于无条件生成，将 Imagenet-256 用于类条件生成。对于所有运行，我们使用 AdamW 优化器，其恒定学习率为 \(1 * 10^{−4}\)，权重衰减为 0.1。所有运行都使用 sd-vae-ft-mse 自动编码器的潜在空间。
与之前的工作进行比较。比较了两类基于扩散的现有方法：两阶段方法，即预先训练扩散模型，然后单独提取它；以及端到端方法，即在单次训练中从头开始训练一步模型。
- Diffusion (DiT-B)：代表一个标准扩散模型
- Flow Matching：用最佳传输目标取代了扩散目标。 这两种方法为迭代多步去噪模型的性能提供了基线。其余方法均以流量匹配目标为基础，如果适用，则使用标准流量匹配模型作为教师。
- Reflow：提供了与标准两阶段蒸馏方法的比较，该方法全面评估教师模型以生成合成 \((x_0, x_1)\) 对。
- Progressive Distillation（渐进蒸馏）提供了与两阶段二元时间蒸馏方法的比较，这与我们提出的方法一致。从预训练的教师模型开始，提炼出一系列学生模型，每个模型的步长都比之前的大 2 倍。
- Consistency Distillation（一致性蒸馏）：与基于一致性模型的两阶段蒸馏策略进行比较。 \((x_t, x_{t+d})\) 对是通过教师扩散模型生成的，并且训练单独的学生一致性模型以在这些对之间强制执行一致的 \(x_1\) 预测。
- Consistency Training（一致性训练）：提供了与之前的端到端方法的比较，该方法从头开始训练一步模型，最接近我们的设置。一致性模型是根据数据集中的经验 \((x_t, x_{t+d})\) 样本进行训练的。
- Live Reflow：是我们提出的另一种端到端方法，其中模型同时针对流匹配和回流蒸馏目标进行训练。该模型分别以每种目标类型为条件。蒸馏目标是使用完全去噪的每个训练步骤自行生成的，因此这种方法的计算成本相当高。我们将其包含在内是为了进行比较。
表 1 强调了快捷模型在少步采样和一步采样下保持准确生成的能力。除了两阶段渐进蒸馏之外，快捷模型优于所有现有方法，且无需多个训练阶段。请注意，渐进蒸馏模型失去了多步采样的能力，而快捷模型保留了这种能力。毫不奇怪，扩散模型和流量匹配模型在 4 步和 1 步生成中表现不佳。有趣的是，在我们的实验中，捷径模型显示出比可比较的流量匹配模型稍好的 FID。一个盲目的假设是，自一致性损失充当模型隐式正则化的一种形式，尽管我们将这项研究留给未来的工作。
3.2 不同推理预算下的快捷模型的行为是什么？
现在，我们分析以不同步数生成图像时快捷模型的行为。虽然简单的流匹配模型以切线速度近似去噪 ODE，但快捷模型经过训练以近似积分去噪 ODE，从而在低步数下提供更准确的行为。图 4 显示，虽然更多的步数总是有帮助的，但在低步数下的退化在原始流匹配模型中比在捷径模型中更为明显。在我们的实验中，我们发现多步性能没有被牺牲，并且即使在高步数下，快捷模型也能保持基线流模型的性能。
图 4：流匹配和快捷模型在去噪步骤数量减少时的行为。虽然原始的流匹配模型会导致退化和模式崩溃，但快捷模型能够在很少的步数生成中保持类似的样本分布。在大量推理预算下，此功能不会影响生成质量。定性示例见图 1。
图 1：针对不同推理预算的几代流程匹配模型和快捷模型。快捷模型可在各种推理预算范围内生成高质量图像，包括使用单个前向传递，与扩散和流匹配模型相比，可将采样时间大幅缩短高达 128 倍。相比之下，当在少步设置中查询时，扩散和流量匹配模型会迅速恶化。每列中使用相同的起始噪声，并且在 CelebA-HQ 和 Imagenet-256（类条件）上训练两个模型。
流模型中的少步伪影类似于模糊和模式崩溃（图 1，左）。快捷模型通常不会遇到此类问题，并且全局匹配由相同初始噪声生成的等效多步图像。快捷模型中的伪影类似于高精度细节中的错误。因此，一步快捷模型可以用作代理：如果下游用户希望细化一步生成的图像，他们可以简单地从相同的初始噪声中重新生成图像，但具有更多的生成步骤。
3.3 快捷模型性能如何随着模型规模的扩大而提高？
深度学习的一个强大趋势是，大型过度参数化模型在扩大规模时能力会增加。这种趋势在基于引导的方法中不太明显，例如学习用于强化学习的 Qfunctions。常见问题类似于某种形式的排名崩溃：通过对自身生成的输出进行重新训练，模型的表达能力可能会受到限制。由于快捷模型本质上是基于引导程序的，因此我们检查快捷模型是否可以从模型规模中获益。如图 5 所示，即使在一步设置中，随着模型大小的增加，我们基于 Transformer 的快捷模型架构也可以实现越来越准确的生成。这些结果表明，尽管快捷模型依赖于引导，但它们仍然避免了剧烈崩溃，并继续随着模型参数计数而扩展。使用 DiT-XL 大小训练的快捷模型能够在 Imagenet-256 上实现 10.6 的单步 FID 和 3.8 的 128 步 FID，如附录中的表 2 所示。
图 5：随着模型参数数量的增加，一步生成的质量不断提高。虽然生成模型往往会随着模型规模的扩大而不断改进，但 Q-learing等基于引导的方法已被证明会失去这一特性。我们证明，快捷模型虽然是一种基于引导的方法，但保留了随模型大小缩放精度的能力。
3.4 快捷模型可以为我们提供可插值的潜在空间吗？
快捷模型表示从噪声到数据的确定性映射。为了检查该映射的结构，我们可以在初始噪声样本之间进行插值并查看相应的生成。图 6 显示了以这种方式生成的插值示例。对一对噪声点 \((x^0_0, x^1_0)\) 进行采样，然后以保留方差的方式进行插值：
尽管没有对噪声到数据映射进行明确的正则化，但所得到的插值代显示出定性的平滑过渡。中间图像在语义上似乎是合理的。请注意，在可视化生成中，所有图像均已生成。虽然我们没有探索现有图像之间的插值，但有可能向现有图像添加噪声并插值这些中间点。
图 6：两个采样噪声点之间的插值。所有显示的图像都是模型生成的。每个水平集表示通过两个高斯噪声样本之间的方差保留插值的一步去噪生成的图像。
四、总结
这项工作引入了快捷模型，这是一种新型的基于扩散的生成模型，支持少步和一步生成。快捷模型背后的关键思想是学习一个步长条件模型，该模型以小步长的数据为基础，并通过大步长的自引导进行训练。快捷模型提供了一个简单的方法，无需调度，并且可以在单次运行中进行端到端训练。与之前的蒸馏方法不同，快捷模型不需要单独的预训练和蒸馏阶段。与一致性训练相比，快捷模型不需要训练计划，仅使用标准 L2 回归，并且需要更少的引导步骤。虽然快捷模型为训练一步生成模型提供了一个简单的框架，但一个关键的固有限制是噪声和数据之间的映射完全依赖于对数据集的期望。在其他生成模型风格中（例如 GAN、VAE），可以调整这种映射，这有可能简化学习问题。其次，在我们的快捷模型实现中，多步生成质量和一步生成质量之间仍然存在差距。快捷模型开辟了一系列未来的研究方向。在表 1 中，快捷模型显示的多步生成效果略好于基本流模型 — 有没有办法让单步生成效果提高到多步生成效果（而不是相反）？也许可以扩展快捷模型的表述，迭代地调整噪声到数据的映射（例如，使用类似回流的过程），或者减少多步生成和一步生成之间的差距。这项工作为开发理想的生成建模目标奠定了基础——一个满足快速采样、模式覆盖和高质量生成三重效果的简单方法。