步数蒸馏详解（一）：DMD 如何用分布匹配实现一步生成

扩散模型的生成质量很高，但它通常需要反复调用去噪网络。步数蒸馏要解决的问题很直接：能否把教师模型几十次甚至上百次的迭代，压缩到学生模型的一步或少数几步？

DMD（Distribution Matching Distillation，分布匹配蒸馏）的关键变化，是不再要求学生逐点复刻教师的某一条采样轨迹，而是直接让学生的输出分布逼近真实数据分布。理解这一点，也就抓住了 DMD 与传统轨迹蒸馏的根本区别。

本文依次说明：DMD 为什么采用分布匹配、梯度如何转化为 score 差、为什么还要训练一个 fake score 网络，以及回归损失在训练中扮演什么角色。

从轨迹匹配到分布匹配

传统蒸馏常把教师在某个噪声状态下得到的结果当作监督目标，要求学生输出与它尽可能接近。这类逐点目标容易实现，但它隐含了一个很强的约束：相同输入最好对应相同输出。

生成模型真正关心的却不是某一个样本是否逐像素一致，而是整体分布是否正确。对同一文本提示，只要生成结果合理、清晰且具有多样性，就不必沿着教师的固定轨迹到达同一个终点。

DMD 因此直接优化学生生成分布 \(p_{\mathrm{fake}}\) 与真实数据分布 \(p_{\mathrm{real}}\) 之间的差异：

\[ \mathcal{L}_{\mathrm{DM}} =D_{\mathrm{KL}}\left(p_{\mathrm{fake}}\,\|\,p_{\mathrm{real}}\right) =\mathbb{E}_{x\sim p_{\mathrm{fake}}} \left[\log p_{\mathrm{fake}}(x)-\log p_{\mathrm{real}}(x)\right]. \]

学生生成器记为 \(x=G_\theta(z)\)。如果直接对生成器参数求导，概率密度本身很难计算；DMD 的做法是把梯度改写成两个分布的 score 之差。这里 score 对样本 \(x\) 求梯度，而不是对模型参数求梯度：

\[ s_{\mathrm{real}}(x)=\nabla_x\log p_{\mathrm{real}}(x),\qquad s_{\mathrm{fake}}(x)=\nabla_x\log p_{\mathrm{fake}}(x). \]

通过重参数化，可以得到对生成器参数的实用梯度形式：

\[ \nabla_\theta \mathcal{L}_{\mathrm{DM}} =\mathbb{E}_{z} \left[ \left(s_{\mathrm{fake}}(x)-s_{\mathrm{real}}(x)\right) \frac{\partial G_\theta(z)}{\partial\theta} \right]. \]

直观上，真实数据的 score 把样本推向真实分布的高密度区域，学生分布自己的 score 则提供反向校正。两者缺一不可：只保留真实 score，无法正确表达当前学生分布与目标分布之间的差距。

为什么要在带噪空间里比较

真实图像空间中的概率密度高度复杂，直接估计 score 并不稳定。DMD 沿用扩散模型的思路，随机采样时间步 \(t\)，把学生输出加噪：

\[ x_t=\alpha_t x+\sigma_t\epsilon,qquad \epsilon\sim\mathcal{N}(0,I). \]

教师扩散模型可以估计真实分布在 \(t\) 时刻的 score。与此同时，DMD 额外训练一个扩散网络 \(\mu_{\mathrm{fake}}^\phi\)，让它拟合学生当前生成样本的带噪分布。这样便得到：

\[ s_{\mathrm{real}}(x_t,t) \approx -\frac{\epsilon_{\mathrm{real}}(x_t,t)}{\sigma_t}, \qquad s_{\mathrm{fake}}(x_t,t) \approx -\frac{\epsilon_{\mathrm{fake}}^\phi(x_t,t)}{\sigma_t}. \]

训练时，fake score 网络持续追踪不断变化的学生分布；生成器则根据真假 score 的差异更新。二者形成一个交替过程，而不是只训练一个学生网络。

为了平衡不同噪声强度下的梯度尺度，实际目标还会加入与时间步有关的权重 \(w_t\)：

\[ \nabla_\theta \mathcal{L}_{\mathrm{DM}} =\mathbb{E}_{z,t,\epsilon} \left[ w_t\left(s_{\mathrm{fake}}(x_t,t)-s_{\mathrm{real}}(x_t,t)\right) \frac{\partial G_\theta(z)}{\partial\theta} \right]. \]

这也是 DMD 能借用预训练扩散模型知识的核心：教师网络提供真实分布方向，在线训练的 fake score 网络描述学生此刻站在哪里。

DMD 的分布匹配训练结构

回归损失为什么仍然存在

只使用分布匹配目标，训练早期容易出现 mode collapse 或 mode dropping：生成结果集中到少数模式，看起来可能清晰，却丢失了多样性。

仅使用分布匹配时可能出现模式坍缩

DMD 为此保留一项回归损失。它先用教师模型为一组固定噪声生成参考结果 \(y\)，再约束学生从同一输入得到相近输出：

\[ \mathcal{L}_{\mathrm{reg}} =\mathbb{E}_{z}\left[\left\|G_\theta(z)-y\right\|_2^2\right]. \]

最终目标可以概括为：

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

这里两部分并不重复。分布匹配负责让整体输出接近真实数据分布；回归损失则像训练护栏，在 fake score 尚未跟上学生变化时提供稳定的逐点监督。

完整训练流程

一次典型的 DMD 训练迭代包含三条数据流：

  1. 采样噪声或条件输入，通过学生生成器得到 \(x=G_\theta(z)\)。
  2. 随机采样 \(t\) 并构造 \(x_t\)，分别送入冻结的教师 score 网络和在线更新的 fake score 网络。
  3. 使用 score 差更新生成器，并用学生样本训练 fake score 网络；同时加入教师预生成样本上的回归损失。

学生一般从教师模型参数或其适配版本初始化，而不是从零训练。训练结束后只保留学生生成器，因此在线推理不需要 fake score 网络，也不需要完整教师模型。

论文结果应该怎样理解

DMD 论文展示了 ImageNet 和大规模文本到图像任务上的一步生成结果。下面两张图均为论文作者报告的实验和样例，不应脱离具体数据集、分辨率与评测配置横向比较。

DMD 在 ImageNet 上的论文实验结果

DMD 文本到图像生成样例

DMD 最重要的贡献不是某个单独指标，而是证明了“直接匹配生成分布”可以成为扩散步数蒸馏的有效路线。不过，它也留下两个明显问题：回归损失需要预先生成并保存大量教师样本；学生和 fake score 网络的更新速度不匹配时，训练仍可能不稳定。

这两个问题正是 DMD2 的起点。下一篇将继续讨论：DMD2 如何通过双时间尺度更新移除回归损失，引入真实数据的 GAN 监督，并把一步生成扩展为更实用的少步生成。

参考资料

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