步数蒸馏详解（二）：DMD2 如何移除回归损失并引入多步生成

上一篇介绍了 DMD 的核心思路：用真实分布与学生分布的 score 差来训练生成器，再用回归损失稳定训练。DMD2 没有推翻这条路线，而是集中解决它在数据、稳定性和生成步数上的限制。

DMD2 的三项关键变化是：让 fake score 网络以更快的节奏追踪学生分布，从而移除回归损失；加入基于真实数据的 GAN 监督；把一步生成扩展为多步生成，并在训练中模拟真实推理轨迹。

DMD 的瓶颈在哪里

DMD 的分布匹配梯度依赖两个 score：冻结的教师网络估计 \(s_{\mathrm{real}}\)，在线训练的 fake score 网络估计 \(s_{\mathrm{fake}}\)。问题在于，生成器每更新一次，学生分布都会改变。如果 fake score 网络追踪得不够快，它给出的方向就会滞后。

DMD 用回归损失缓解这一问题，但代价是需要教师提前生成大量配对样本。数据准备和存储成本都很高，而且逐点回归会限制学生偏离教师轨迹的自由度。

DMD2 的出发点是：如果 fake score 网络能够更充分地拟合学生的最新分布，是否就能去掉这根“训练拐杖”？

双时间尺度更新：让 fake score 先跟上

DMD2 采用 two time-scale update rule。在一次生成器更新之前，先对 fake score 网络执行多次更新；论文配置中常见的比例是 5 次 fake score 更新对应 1 次生成器更新。

用 \(\phi\) 表示 fake score 网络参数、\(\theta\) 表示生成器参数，可以把节奏写成：

\[ \underbrace{\phi\leftarrow\phi-\eta_\phi\nabla_\phi \mathcal{L}_{\mathrm{denoise}}}_{\text{重复多次}} \quad\Longrightarrow\quad \theta\leftarrow\theta-\eta_\theta\nabla_\theta \mathcal{L}_{\mathrm{DM}}. \]

这样做的目的不是简单增加训练量，而是减少分布估计的时延：生成器更新时，\(s_{\mathrm{fake}}\) 应尽可能描述当前学生分布，而不是若干步之前的学生。

DMD2 通过增加 fake score 更新稳定训练

在这个训练节奏下，DMD2 可以移除 DMD 中依赖预生成教师样本的回归损失，训练数据链路也随之简化。

加入 GAN 损失：让学生直接看到真实数据

分布匹配使用教师扩散模型提供的 score，本质上仍然受到教师能力限制。DMD2 又加入一个判别器，让真实图像和学生生成图像直接参与对抗训练：

\[ \mathcal{L}_{D} =-\mathbb{E}_{x\sim p_{\mathrm{data}}}[\log D(x)] -\mathbb{E}_{z}[\log(1-D(G_\theta(z)))], \]
\[ \mathcal{L}_{\mathrm{GAN}} =-\mathbb{E}_{z}[\log D(G_\theta(z))]. \]

生成器的目标由分布匹配与对抗监督共同组成：

\[ \mathcal{L}_{G} =\mathcal{L}_{\mathrm{DM}}+\lambda_{\mathrm{GAN}}\mathcal{L}_{\mathrm{GAN}}. \]

GAN 损失并不是所有步数蒸馏方法在理论上的必选项，但在 DMD2 以及多种近期方法中，它是一种有效的真实数据监督。因为学生不再只模仿教师，特定数据集和指标下出现超过教师评测结果的情况并不矛盾；这仍应结合评测配置理解，不能笼统解释为学生在所有方面都优于教师。

为什么从一步扩展到多步

一步模型追求极致速度，但所有生成工作都压在一次网络调用中，质量上限和训练难度都更受约束。DMD2 支持少步生成，用少量额外计算换取更好的样本质量。

设推理时间序列为 \(t_K>t_{K-1}>\cdots>t_0\)。每一步先由生成器预测干净样本，再按下一个时间步重新加噪：

\[ \hat{x}_0^{(k)}=G_\theta(x_{t_k},t_k), \]
\[ x_{t_{k-1}} =\alpha_{t_{k-1}}\hat{x}_0^{(k)} +\sigma_{t_{k-1}}\epsilon_k, \qquad \epsilon_k\sim\mathcal{N}(0,I). \]

以四步生成为例，输入会依次经过 \(t_{999}\rightarrow t_{749}\rightarrow t_{499}\rightarrow t_{249}\rightarrow t_0\)。每一步的输入必须来自前一步结果，而对应的时间条件也必须与当前噪声状态一致。原始笔记中四步示例的后两步存在复制形成的索引错位，这里已按连续推理轨迹修正。

训练与推理必须使用相同的输入分布

多步模型还有一个容易忽略的问题：训练时若总是给生成器输入由真实图像直接加噪得到的 \(x_t\)，推理时它看到的却是自己上一轮输出重新加噪的结果。前者来自真实数据轨迹，后者来自学生自己的生成轨迹，两者存在 train-inference mismatch。

DMD2 在训练中模拟少步推理。先从纯噪声出发，让学生执行若干前置步骤，再把得到的中间状态用于当前训练步骤。这样，模型在训练阶段就能看到推理时真正会遇到的误差和输入分布。

DMD2 的整体算法

完整训练结构

DMD2 可以理解成三个相互配合的更新环：

  1. Fake score 更新：从学生生成样本构造带噪数据，连续更新多次 fake score 网络，使其追踪当前学生分布。
  2. 生成器更新：用教师 score 与 fake score 的差异计算分布匹配梯度，同时加入 GAN 损失。
  3. 判别器更新：用真实数据和学生生成数据训练判别器，为生成器提供直接的数据监督。

多步版本还会在生成器更新前模拟部分推理轨迹。训练完成后，教师、fake score 网络和判别器都不参与部署，线上只保留学生生成器。

DMD 与 DMD2 的区别

维度 DMD DMD2
核心目标 真假 score 的分布匹配 保留分布匹配并改善训练机制
稳定方式 教师预生成样本上的回归损失 fake score 多次更新
真实数据监督 间接来自教师模型 增加 GAN 判别器直接使用真实数据
生成步数 重点展示一步生成 同时支持一步与少步生成
训练/推理分布 主要围绕一步训练 多步训练中显式模拟推理轨迹

DMD2 的价值不只是“多加一个 GAN”。它把分布估计速度、训练数据来源和少步推理轨迹放在同一个框架里处理，使 DMD 从具有启发性的单步方法，变成更完整的快速生成方案。

论文报告的实验结果

论文作者报告，DMD2 在 ImageNet 64×64 上取得 FID 1.28，在 COCO 2014 文本到图像任务上报告 FID 19.01。下面是论文中的对比结果与样例；数字仅代表论文所述数据集、分辨率和评测协议。

DMD2 在 ImageNet 64×64 上的论文结果

DMD2 的消融与对比实验

DMD2 文本到图像生成样例

从工程视角看，少步生成并不是越少越好。真正部署时仍需在模型规模、单步延迟、总调用次数、分辨率和质量目标之间权衡。DMD 与 DMD2 解决的是算法层面的步数压缩，因此本系列归入 AIGC 理论；如果进一步讨论具体硬件上的吞吐、延迟、显存与部署参数，则更适合归入 推理优化。

参考资料

  • Tianwei Yin 等，One-step Diffusion with Distribution Matching Distillation，CVPR 2024。
  • Tianwei Yin 等，Improved Distribution Matching Distillation for Fast Image Synthesis，NeurIPS 2024。