引言

随着扩散模型的出现，生成式 AI 经历了巨大的变革。这些模型为 Stable Diffusion 和 DALL-E 等工具提供了动力，能够通过逐步从信号中去除噪声来生成令人惊叹的图像。然而，它们面临着一个众所周知的瓶颈: 速度。生成单张图像通常需要数十甚至数百个连续步骤。

为了解决这个问题，研究人员引入了一致性模型 (Consistency Models, CMs) 。 一致性模型的前景非常诱人: 它的目标是通过学习将噪声轨迹上的任意一点直接映射回其干净的起始点，从而在单步 (或极少步数) 内生成高质量数据。

然而，与从现有的慢速扩散模型中蒸馏 (“一致性蒸馏”) 相比，从头开始训练这些模型 (“一致性训练”) 已被证明十分困难。两者之间存在理论和实践上的差距。

在这篇文章中，我们将深入探讨论文 《利用生成器增强流改进一致性模型》 (Improving Consistency Models with Generator-Augmented Flows) 。 我们将探讨为什么标准的一致性训练在数学上存在缺陷，作者如何识别出即使在连续时间极限下仍然存在的“差异”，以及他们提出的新颖解决方案: 生成器增强流 (Generator-Augmented Flows, GC) 。 该方法利用模型自身的预测来拉直生成路径，从而大幅提高训练速度和质量。

如上图 图 1 所示，标准方法 (独立耦合，IC) 导致了混沌的轨迹 (左图) ，而本文提出的生成器增强流 (右图) 创建了更直、更连贯且与真实速度场一致的路径。

让我们来详细拆解它是如何工作的。

背景: 一致性的机制

要理解这项创新，我们首先需要了解基准。扩散模型可以使用概率流常微分方程 (Probability Flow Ordinary Differential Equation, PF-ODE) 来描述。

PF-ODE

扩散过程逐渐向数据 \(\mathbf{x}_{\star}\) 添加噪声，直到其变成纯噪声 \(\mathbf{z}\)。这个过程描述了一条轨迹。我们可以逆转这条轨迹来生成图像。沿该轨迹的移动由速度场 \(\mathbf{v}_t(\mathbf{x})\) 定义。控制这种移动的方程为:

这里，\(\mathbf{v}_t(\mathbf{x})\) 是速度场。在标准扩散中，它是从分数函数 (对数密度的梯度) 推导出来的。

一致性模型

逐步求解 ODE 的速度很慢。一致性模型试图学习一个函数 \(f_{\theta}(\mathbf{x}_t, \sigma_t)\)，该函数将时间 \(t\) 时的任意点 \(\mathbf{x}_t\) 映射回原始数据 \(\mathbf{x}_0\)。

如果模型是完美的，那么将 \(f_{\theta}\) 应用于轨迹上的任何点都会得到相同的结果: 干净的图像。这就是自洽性 (self-consistency property) 。

为了确保模型遵守边界条件 (即在时间 0 时，输出即为输入本身) ，通常使用跳跃连接 (skip connections) 来参数化架构:

两种训练方法的故事

学习这个函数 \(f_{\theta}\) 主要有两种方法:

1. 一致性蒸馏 (Consistency Distillation, CD) : 你有一个预训练的教师模型 (标准的扩散模型) ，它为你提供“真实”的速度场 \(\mathbf{v}_t\)。你利用这个教师来模拟一步，并训练你的学生模型 \(f_{\theta}\) 去匹配它。
2. 一致性训练 (Consistency Training, CT) : 你没有教师模型。你想从头开始训练模型。

本文重点关注一致性训练面临的挑战。

在蒸馏中，我们最小化模型在当前步骤的预测与下一步预测之间的差异，其中下一步是使用教师模型的速度计算出来的。

在这个方程中，\(\mathbf{x}_{t_i}^{\Phi}\) 代表使用教师模型的精确速度场走出的“基本真值 (ground truth) ”步:

然而，在一致性训练 (CT) 中，我们无法获得真实速度 \(\mathbf{v}_{t_{i+1}}\)。我们用单样本蒙特卡洛估计来代替它。我们取一个数据点 \(\mathbf{x}_{\star}\) 和一个噪声向量 \(\mathbf{z}\)，混合它们得到 \(\mathbf{x}_{t_i}\) 和 \(\mathbf{x}_{t_{i+1}}\)，并将它们之间的差值视为方向。

问题就在这里。我们正在用单个带噪声的样本来代替真实的速度场 (一个期望值) 。

差异: 为什么训练落后于蒸馏

直观地说，人们可能会认为用噪声估计代替真实速度只会让训练过程更嘈杂，但随着时间的推移应该会平均化。先前的研究表明，当时间步数 \(N \to \infty\) (连续时间极限) 时，一致性训练和蒸馏将变得等价。

本文作者证明了这是错误的。

定理 1: 差距确实存在

研究人员提供了严格的证明 (定理 1) ，表明即使有无限的时间步数，训练和蒸馏的损失函数也会收敛到不同的值。

存在一个持久的差距，记为 \(\mathcal{R}(\theta)\)。这个项严格为正，代表了由于使用单样本而非真实向量场而引入的根本偏差。

为什么存在这个差距？归根结底在于梯度。蒸馏中使用的“真实”梯度 (\(\partial_{\mathrm{CD}}\)) 依赖于期望速度 \(\mathbf{v}_t\)。而训练中使用的梯度 (\(\partial_{\mathrm{CT}}\)) 依赖于瞬时样本速度 \(\dot{\mathbf{x}}_t\)。

当我们观察距离度量为平方欧几里得距离 (\(\alpha=2\)) 的具体情况时，残差项 \(\mathcal{R}(\theta)\) 呈现出非常清晰的形式:

这个方程是“确凿的证据”。误差取决于样本速度 \(\dot{\mathbf{x}}_t\) 与真实速度 \(\mathbf{v}_t(\mathbf{x}_t)\) 之间差值的平方。

通俗地说: 因为标准训练随机耦合数据和噪声 (独立耦合) ，单一样本采取的路径 (\(\dot{\mathbf{x}}_t\)) 通常与概率分布的平均流向 (\(\mathbf{v}_t\)) 大相径庭。这种方差在损失景观中造成了永久性的差异。

梯度指向了不同的最优解:

核心方法: 生成器增强流

作者意识到，要修复一致性训练，我们必须减小方差项: \(\|\dot{\mathbf{x}}_t - \mathbf{v}_t(\mathbf{x}_t)\|^2\)。

在没有预训练教师的情况下，我们无法神奇地知道真实速度 \(\mathbf{v}_t\)。但是，我们可以构建“更聪明”的轨迹，使样本路径 \(\dot{\mathbf{x}}_t\) 自然地更好地与平均流向对齐。

超越独立耦合 (IC)

标准训练使用独立耦合 (Independent Coupling, IC) : 选取一个随机图像 \(\mathbf{x}_{\star}\)，选取随机噪声 \(\mathbf{z}\)，然后进行插值。这很简单但效率低下。噪声向量可能“想要”去往一个与它被配对的图像完全不同的图像，从而产生高曲率路径和高方差。

最优传输 (Optimal Transport, OT) 是另一种选择。它将特定的噪声向量与特定的图像配对，以最小化传输距离。然而，OT 计算成本高昂 (\(O(N^3)\)) ，并且必须在每个小批量 (mini-batch) 中进行近似，这又重新引入了误差。

引入生成器增强耦合 (GC)

作者提出了一个绝妙的启发式方法: 使用一致性模型本身来决定耦合。

如果一致性模型 \(f_{\theta}\) 经过了部分训练，它对噪声点 \(\mathbf{x}_t\) 应该映射到哪里已经有了一些概念。我们可以利用这个预测来创建一个“合成”数据目标。

过程:

1. 使用独立耦合采样一个标准的随机点 \(\mathbf{x}_{t_i}\)。
2. 让模型预测干净的数据: \(\hat{\mathbf{x}}_{t_i} = f(\mathbf{x}_{t_i}, \sigma_{t_i})\)。
3. 关键步骤: 创建一条新的轨迹，将这个预测数据 \(\hat{\mathbf{x}}_{t_i}\) 与第 1 步中使用的同一个噪声 \(\mathbf{z}\) 连接起来。

现在，我们使用这个新的配对定义生成器增强耦合 (GC) 点 \((\tilde{\mathbf{x}}_{t_i}, \tilde{\mathbf{x}}_{t_{i+1}})\):

然后，我们在这个新的、更聪明的轨迹上计算一致性损失。

为什么这样更好？

通过将噪声 \(\mathbf{z}\) 与模型从该噪声预测出的数据点 \(\hat{\mathbf{x}}\) 进行耦合，我们实际上是在构建一条模型认为“自然”的轨迹。这将单样本路径方向 (\(\dot{\mathbf{x}}_t\)) 与向量场方向 (\(\mathbf{v}_t\)) 对齐得更加紧密，有效地缩小了差距 \(\mathcal{R}(\theta)\)。

理论验证

作者证明了这种新流的两个关键属性。

1. 减小差异 他们分析了差异项 \(\mathcal{R}(\theta)\) 的一个代理指标，记为 \(\tilde{\mathcal{R}}_t\)。

理论分析 (定理 2) 和实证测量表明，GC 的这一项显著低于 IC。

我们可以在图 2 中通过实证观察到这一点。蓝线 (GC) 始终低于红线 (IC) ，甚至低于橙线 (Batch-OT) 。

2. 降低传输成本 传输成本衡量路径的“笔直”程度。笔直的路径更容易学习和模拟。

作者定义了传输成本 \(c(t)\)，并证明对于生成器增强流，该成本随着时间的推移而降低 (具体来说，导数为负) 。

在实证中，这导致了更直、更高效的路径。请看图 3 : 与 IC (红圈) 相比，GC (蓝星) 大幅降低了传输成本，甚至在许多区域优于批量最优传输 (Batch-OT，橙色/紫色) 。

算法: 联合学习

这里存在一个先有鸡还是先有蛋的问题 (Catch-22) 。为了生成好的 GC 轨迹，我们需要一个好的一致性模型 \(f_{\theta}\) 来预测 \(\hat{\mathbf{x}}\)。但为了得到一个好的模型，我们需要训练它。

解决方案是联合学习 (Joint Learning) 。 我们从头开始训练单个模型 \(f_{\theta}\)。在每个训练步骤中，我们混合数据:

• 以 \((1-\mu)\) 的概率，我们使用标准的独立耦合 (IC) 。
• 以 \(\mu\) 的概率，我们使用基于模型当前预测的生成器增强耦合 (GC) 。

这使得模型能够自我引导 (bootstrap) 。IC 样本确保模型看到真实的数据分布，而 GC 样本则细化梯度并拉直流动。

作者发现混合率 \(\mu = 0.5\) (各占一半) 通常效果最好。

实验与结果

研究人员在标准的图像生成基准测试上测试了这种方法: CIFAR-10、ImageNet、CelebA 和 LSUN Church。他们将自己的方法 (iCT-GC) 与标准改进的一致性训练 (iCT-IC) 以及带最优传输的一致性训练 (iCT-OT) 进行了比较。

收敛速度

最重要的发现之一是收敛速度。由于梯度估计器的方差较低，模型学习得更快。

图 5 显示了训练迭代期间的 FID (Fréchet Inception Distance，越低越好) 。紫色曲线 (\(\mu=1.0\) 的 GC) 在早期下降得非常快，尽管混合策略 (如 \(\mu=0.5\)，浅蓝色) 具有最佳的长期稳定性。所有 GC 变体的收敛速度都快于基线 (红色) 。

定量性能

最终结果总结在表 1 中。所提出的方法 (iCT-GC) 在几乎所有数据集上都取得了最低的 FID 分数，击败了标准基线和最优传输变体。

例如，在 CIFAR-10 上，FID 从 7.42 (基线) 提高到了 5.95 (GC) 。在 CelebA 上，从 15.82 提高到了 11.74 。

视觉质量

定量的改进转化为了视觉质量的提升。 图 8 比较了不同方法生成的样本。由 GC 训练的模型 (图 c) 生成的人脸看起来比基线 (图 a) 更清晰、更一致。

简易一致性微调 (ECT)

作者还在“简易一致性微调” (Easy Consistency Tuning) 设置中测试了他们的方法 (这是一种微调预训练扩散模型的最新技术) 。如表 2 所示，GC 继续优于标准方法，证明了其多功能性。

结论

论文《利用生成器增强流改进一致性模型》强调了一致性模型训练方式中一个微妙但关键的缺陷。通过依赖于速度场的单样本估计以及随机的数据-噪声配对，标准训练会产生永久性的误差偏差和高方差。

提出的解决方案——生成器增强流——因其自力更生而显得优雅。通过利用模型自身对数据逐渐形成的理解来指导耦合过程，作者实现了:

1. 更低的差异: 训练目标与理想的蒸馏目标更加一致。
2. 更低的传输成本: 从噪声到数据的轨迹变得更直，更容易学习。
3. 更好的性能: 收敛速度更快，图像生成质量更高。

这项工作提醒我们，在深度学习中，如何配对输入和目标 (耦合) 与模型架构本身一样重要。通过从随机耦合转向“智能的”、模型引导的耦合，我们可以缩小理论理想与实践现实之间的差距。