Zepeng Lin

Zepeng Lin

Sophomore at South China University of Technology, diving into the research of Embodied AI, especially focusing on how to narrow down the sim-to-real gap and scale up simulation data for robotic learning.

Efficient Diffusion（五）：对抗性蒸馏 —— ACM、Hyper-SD 与 DMD2

4 minute read

Published:

书接上文，SiD 和 sCM 分别在 DMD 和 CM 两条路线上做了结构性的简化。但一个更深层的问题仍然悬而未决：预训练扩散模型自身的信息是否足以支撑极少数步（甚至单步）的高质量生成？ 2024 年的一系列工作给出了同一个答案——不够。加上 GAN 的对抗性判别信号，才能把质量真正推上去。本文讲三条将对抗性思路融入蒸馏的代表性路线。

目录

对抗一致性训练 (ACT)

CT 的误差累积与 Wasserstein 视角

ACT（Adversarial Consistency Training，Kong et al., CVPR 2024）的出发点是对 Consistency Training 损失函数的一个理论重解读。作者证明：优化 CT 的 MSE 自洽性损失，本质上等价于在每一步最小化生成分布与目标分布之间的 Wasserstein 距离。

这听起来像是一个好性质——Wasserstein 距离是衡量分布差异的合理度量。但问题出在”累积”上。CT 在 $t_n$ 处的训练目标是 $f_{\theta^-}(x_{t_{n-1}}, t_{n-1})$，而这个 EMA 目标的准确性本身又依赖于更小 $t$ 处的训练质量。因此，任意时刻 $t_n$ 的 CT 损失的 Wasserstein 上界 不仅包含当前步的误差，还包含所有更小 $t$ 处误差的总和：

\[\mathcal{W}(p_{\text{fake}}^{(t_n)}, p_{\text{real}}^{(t_n)}) \leq \sum_{i=1}^{n} \mathcal{L}_{\text{CT}}^{(t_i)}.\]

这意味着，要控制 $t_{\max}$ 端的分布偏差，必须在所有 $n$ 个时间步上都保持极低的训练误差。而达成这一点的唯一方法是使用超大的 batch size——CT 原作在 ImageNet 64×64 上需要 64 块 A100——因为大 batch 能提供更稳定的梯度估计，减少每一步的方差。

但这个开销显然不现实。ACT 提供了一个根本性的替代方案。

ACT 的思想并不复杂：既然 CT 的 Wasserstein 距离会在时间轴上累积误差，那为什么不直接在每个时间步上独立地最小化一个不具备累积性的分布距离？最经典的选择就是 GAN 中的 Jensen-Shannon（JS）散度——判别器在每个 $t$ 上的训练是独立的，不依赖其他 $t$ 处的模型质量。

ACT 的算法流程

ACT 在 CT 的框架上增加了一个条件判别器 $D(x_0, x_t)$，将纯粹的自洽性训练改造为对抗 + 自洽的联合训练。

网络结构：

  • 生成器 $G_\theta(x_t, t)$：即一致性模型 $f_\theta$ 本身，结构与 CT 中的 UNet 一致。
  • 判别器 $D(x_0, x_t)$：复用扩散模型 UNet 的下采样部分（down-sampling blocks 到 mid-block），后接一个线性层输出标量 logit。判别器同时接收生成/真实样本 $x_0$ 和对应噪声水平的带噪输入 $x_t$ 作为条件——它学会判断”在噪声水平 $t$ 下，给定的干净样本 $x_0$ 和带噪样本 $x_t$ 是否来自同一扩散轨迹”。

训练流程（每轮迭代）：

  1. 采样干净样本 $x_0 \sim p_{\text{data}}$，噪声 $\varepsilon \sim \mathcal{N}(0, I)$，以及时间步 $t \sim \mathcal{U}[t_{\min}, T]$。

  2. 构造带噪样本 $x_t = \alpha_t x_0 + \sigma_t \varepsilon$。

  3. 用生成器一步输出干净样本的估计 $\hat{x}0 = G\theta(x_t, t)$。

  4. 计算条件判别器损失（hinge loss 形式）：

\[\mathcal{L}_D = \mathbb{E}\Bigl[ \max(0, 1 - D(x_0, x_t)) + \max(0, 1 + D(\hat{x}_0, x_t)) \Bigr].\]
  1. 计算生成器的对抗损失（非饱和形式）:
\[\mathcal{L}_{\text{adv}} = -\mathbb{E}\Bigl[ D(\hat{x}_0, x_t) \Bigr].\]
  1. 总生成器损失为对抗损失与原始 CT 自洽性损失的加权和：
\[\mathcal{L}_G = \mathcal{L}_{\text{adv}} + \lambda \cdot \mathcal{L}_{\text{CT}}.\]

$\mathcal{L}{\text{CT}}$ 的作用在此变成了一种正则项——它帮助维护轨迹自洽性，避免纯对抗训练可能导致的模式坍塌。而 $\mathcal{L}{\text{adv}}$ 承担了主要的分布对齐任务，因为它的梯度在每个 $t$ 处独立，不受其他时间步误差的影响。

这里的关键在于判别器 $D$ 的条件化设计。如果判别器只看 $\hat{x}_0$，那它只能在”图片是否真实”这一个维度上给信号——相当于一个普通 GAN 的判别器，丢失了扩散模型最重要的逐级去噪结构。而条件化 $x_t$ 后，判别器学会回答的是一个更精细的问题：”给定这团特定噪声水平的输入，模型输出的干净样本应该长什么样？”这使得对抗信号能按噪声水平提供差异化的指导——$t$ 大时，关注全局结构；$t$ 小时，关注细节纹理。

判别器训练与生成器训练交替进行（判别器更新 1 次，生成器更新 1 次），这一比一的更新频率与 CT 的简单性保持一致。

ACT 的意义与局限

ACT 最直接的收益是训练成本的大幅下降。CT 在 ImageNet 64×64 上需要 64 块 A100，而 ACT 用 4 块 A100（batch size 320）就能达到更好的 FID。在 CIFAR-10 上甚至只需 1 块 RTX 3090。作为参考，这大概是 CT 原版所需 batch size 的 1/6。

但 ACT 并非没有局限。首先，引入判别器意味着新增一组超参数（判别器学习率、对抗损失权重 $\lambda$、梯度惩罚系数等），这些在不同数据集上需要独立调参。其次，判别器与生成器的交替训练本身也可能不稳定——虽然 CT 的 EMA 目标网络提供了一定的稳定性基础，但对抗训练的 mode collapse 风险并未完全消除。ACT 能保留零样本 inpainting 能力（一致性模型的关键特性），但在其他 zero-shot 编辑任务上的表现尚待验证。

Hyper-SD：轨迹分段一致性蒸馏

两种路线之争：轨迹保真 vs 轨迹重建

在理解 Hyper-SD（Ren et al., NeurIPS 2024）之前，需要先看清 2024 年初蒸馏方法的两个阵营：

轨迹保真派（Progressive Distillation、Consistency Distillation、CTM）：目标是让学生模型的输出轨迹尽可能贴合教师 PF-ODE 的原始轨迹。优点是与教师行为一致，zero-shot 能力保留较好；缺点是拟合误差难以避免——少步模型的容量限制意味着它不可能完美复制教师多步轨迹的每一段弯曲，保真度越高，对模型表达力的要求越高。

轨迹重建派（DMD、SDXL-Lightning、ADD）：不关心学生轨迹是否贴近教师轨迹，只关心最终输出的分布是否与教师一致。优点是可以”抄近路”——学生不需要走教师的路，可以直接学习更高效的生成路径；缺点是可能偏离原始模型的行为特性（domain shift），且在没有轨迹约束时训练可能不稳定。

两个阵营看似对立，但 Hyper-SD 的思路是：为什么不先分段保真，再全局重建？

TSCD 的算法流程

TSCD（Trajectory Segmented Consistency Distillation）是 Hyper-SD 的核心算法，它将时间轴 $[0, T]$ 划分为 $K$ 个连续段落（segment），在每个段内独立做一致性蒸馏，然后逐步合并段落。

阶段划分： 将全时间域 $[0, T]$ 均匀分为 $K$ 个段落。$K$ 的取值从大到小渐进缩减——典型设置为 $K = 8 \to 4 \to 2 \to 1$。当 $K=1$ 时，TSCD 退化为标准的轨迹一致性蒸馏（CTM）。

阶段 $K$ 的训练流程：

  1. 采样干净样本 $x_0 \sim p_{\text{data}}$，噪声 $\varepsilon \sim \mathcal{N}(0, I)$，以及当前段落内的时间起始点 $t_n \sim \mathcal{U}[(k-1) \cdot T/K, k \cdot T/K]$。

  2. 带噪：$x_{t_n} = \alpha_{t_n} x_0 + \sigma_{t_n} \varepsilon$。在同一段落内，按照 CTM 的方式构造相邻时间步的带噪样本 $x_{t_{n-1}}$。

  3. 用 ODE 求解器 $\Psi$ 将学生的预测分别从 $x_{t_n}$ 和 $\hat{x}{t{n-1}}$ 推进到该段落的终点 $t_{\text{end}}$（即段内最小的时间步）：

\[\mathcal{L}_{\text{TSCD}} = d\Bigl( \Psi(x_{t_n}, f_\theta(x_{t_n}, t_n, c), t_{\text{end}}),\; \Psi(\hat{x}_{t_{n-1}}, f_{\theta^-}(\hat{x}_{t_{n-1}}, t_{n-1}, c), t_{\text{end}}) \Bigr),\]

其中 $f_{\theta^-}$ 是 EMA 副本（stop-gradient），$c$ 是文本条件。

  1. 距离度量 $d$ 采用动态混合策略：当 $K$ 较大（如 $K=8,4$）时以 MSE 为主——段内时间跨度小，$f_\theta$ 和 $f_{\theta^-}$ 的预测差异不大，MSE 的平滑梯度更适合精细对齐。当 $K$ 较小（如 $K=2,1$）时逐步引入对抗损失（借鉴 SDXL-Lightning 的 adversarial loss）——时间跨度增大后，MSE 的”回归到均值”效应会导致模糊，对抗损失提供分布层面的判别信号来维持清晰度。

  2. 当前 $K$ 训练收敛后，减少段落数至 $K/2$，将相邻段落合并。合并后每个段落的时间跨度翻倍，一致性约束的难度增加，但此时模型已经从上一阶段获得了较好的段落内自洽性作为起点。

辅助组件：Human Feedback 与统一 LoRA

除了 TSCD 核心算法，Hyper-SD 还引入了两个辅助组件：

Human Feedback Learning：在蒸馏训练完成后，用 LAION 美学预测器（Aesthetic Score Predictor）和 ImageReward（人类偏好模型）作为额外的奖励信号进行 RL 微调。这一步直接优化人类偏好，弥补蒸馏过程中可能的性能损失，在极低步数（1-2 步）下效果尤为显著。

Score Distillation + 统一 LoRA：将 DMD 的分布匹配损失作为额外的监督信号，进一步优化单步生成质量。最关键的是，Hyper-SD 训练出的是一组统一的 LoRA 权重——同一个 LoRA 支持从 1 步到 8 步的所有推理配置。这意味着部署时不需要为不同步数维护不同模型，只需一个基础模型加一组 LoRA 即可灵活切换速度与质量的权衡。

从效果上看，Hyper-SDXL 在 1 步生成下的 CLIP Score 比 SDXL-Lightning 高 0.68，Aesthetic Score 高 0.51，在整个 1-8 步范围内都达到了 SOTA。更重要的是，它证明了轨迹保真和轨迹重建两条路线并非互斥——通过分段渐进策略，可以同时收获保真的稳定性和重建的高质量。

DMD2：分布匹配蒸馏的进化

DMD 的遗留问题

回顾第三篇中 DMD 的训练流程，有三个痛点：

  1. 回归损失的高昂代价。 DMD 需要额外的回归损失 $\mathcal{L}_{\text{reg}}$——在大量噪声-图像对上让生成器学习教师的确定性映射。这需要预先采样海量的教师轨迹数据，计算量大且存储开销高。

  2. 假 score 模型的训练不稳定。 虽然 SiD 后来用 score identity 消去了假 score 模型，但 DMD 原版仍然需要交替更新 $s_\psi$。单时间尺度（生成器和假 score 模型以相同频率更新）下，如果生成器更新过快，假 score 模型来不及适应新的假分布，梯度方向就会出现系统性偏差。

  3. 训练-推理的不匹配。 训练时生成器以纯噪声 $z_T \sim \mathcal{N}(0, I)$ 为输入做单步生成，但如果想在推理时做多步（比如 4 步），中间步骤的输入并不完全是高斯噪声——它们来自前一步的部分去噪结果，与训练分布存在差异。

DMD2（Yin et al., NeurIPS 2024 Oral）针对这三个问题做了系统性的改进。

DMD2 的算法流程

DMD2 的核心改进可以归纳为四项，每一项对应解决一个问题。

改进一：去掉回归损失。

DMD2 发现回归损失可以被完全移除，理由是分布匹配损失 $\mathcal{L}_{\text{DMD}}$ 本身已经提供了足够强的学习信号。关键在于：原版 DMD 中的回归损失主要是为了在训练初期稳定生成器输出，防止其偏离教师轨迹太远。DMD2 通过下面改进二（多时间尺度）来替代这种稳定性功能——假 score 模型更新更频繁后，分布匹配的梯度本身就足够可靠，不再需要回归损失做”拐杖”。这大大简化了训练流程，不再需要预先采样教师轨迹。

改进二：双时间尺度（Two Time-Scale）更新规则。

这是 DMD2 最关键的改进。生成器 $G_\theta$ 和假 score 模型 $s_\psi$ 使用不同的更新频率：假 score 模型更新 $N$ 次后，生成器才更新 1 次（$N$ 的典型取值为 5-10）。

为什么这样做？假 score 模型的目标是准确估计 $p_{\text{fake}}$ 在当前生成器参数下的 score。如果生成器每步都在变，$p_{\text{fake}}$ 就在持续漂移，假 score 模型永远追不上。让假 score 模型更新更频繁，相当于让它在每个”准静态”的 $p_{\text{fake}}$ 下有足够时间逼近最优——类似于强化学习中 critic 更新得比 actor 快，以确保 critic 提供的是当前 actor 策略的可靠价值估计。

从收敛性角度看，这实际上构造了一个双时间尺度动力系统：快时间尺度（$s_\psi$ 的 denoising score matching）跟踪慢时间尺度（$G_\theta$ 的分布匹配）的瞬时状态。当两个时间尺度分离充分大时，系统等价于生成器在真实分布匹配梯度方向上做确定性的梯度下降，而不是在随机噪声中摸索。

改进三：GAN 损失。

DMD2 在分布匹配损失的基础上增加了一个 GAN 判别器。这与 ACT 的动机一致：预训练扩散模型”见过”的只是教师生成的图像分布，而不是真实图像分布。依赖教师 score 的分布匹配本质上是在让学生逼近教师——学生的极限就是教师的质量。

而 GAN 判别器直接使用真实数据训练，它的对抗性信号告诉学生：”即使教师认为这样已经够了，但真实图像还有更丰富的细节。”这使得 DMD2 的生成质量可以超越教师模型——ImageNet 64×64 上学生 FID = 1.28，而教师（EDM2 的 512 步）FID 约 1.33。这是极少数蒸馏工作中学生真正超越教师的案例。

GAN 损失的另一个贡献是缓解了模式坍塌——分布匹配损失中 KL 散度的”mode-seeking”倾向在有判别器监督时被 JS 散度的”mode-covering”倾向平衡。

改进四：多步生成训练。

为解决训练-推理不匹配，DMD2 在训练时模拟推理时的多步采样过程。具体做法是：在训练循环中，除了从纯噪声 $z_T$ 做单步生成外，还间歇性地用生成器自身执行 $m$ 步”伪推理”——将前一步的部分去噪结果作为当前步的输入，让模型学会处理中间噪声水平的输入。这个机制确保模型在训练中就见过推理时多步采样中可能遇到的各种中间状态，消除了训练-推理之间的分布偏移。

此外，DMD2 通过一个预训练的 VAE 将方法扩展到高分辨率（如 1024×1024 的 SDXL 蒸馏），实现了百万像素级别的单步/少步生成。

DMD2 的结果与定位

DMD2 在几个关键基准上的结果：

  • ImageNet 64×64：FID 1.28（单步），超越教师，推理速度约为教师的 500 倍。
  • COCO 2014 zero-shot：FID 8.35。
  • SDXL 蒸馏：百万像素级别的高质量少步生成。

DMD2 的定位是独特的：ACT 用对抗信号加固了一致性训练，Hyper-SD 用分段策略调和了轨迹保真和重建，而 DMD2 则表明——即使在纯分布匹配的框架下，引入 GAN 损失和多时间尺度训练也足以让单步生成超越教师。三条路线共同的结论是：2024 年以后，对抗性信号不再是少步蒸馏的可选项，而是必需品。

总结一下第五篇：ACT 从理论层面揭示了 CT 的误差累积问题，用 JS 散度在每个时间步上独立地做分布对齐，大幅降低了训练成本；Hyper-SD 通过轨迹分段蒸馏（TSCD）渐进合并时间段落，在轨迹保真和轨迹重建之间找到了平衡点，并贡献了实用的统一 LoRA 方案；DMD2 去掉回归损失、引入双时间尺度更新和 GAN 判别器，让分布匹配蒸馏首次实现了学生超越教师的质量。

三项工作从不同角度切入，但共同标志着一个转折：少步扩散模型的训练范式从”小心翼翼地模仿教师”转变为”借用对抗信号实现独立生成”。下一篇，我们将进入 2025 年的最新进展，看看这一趋势如何进一步演化。