Improved Distribution Matching Distillation for Fast Image Synthesis
论文的核心贡献
- 消除回归损失，实现更廉价的大规模训练： 原始 DMD 依赖预计算的“噪声-图像”配对数据来计算回归损失，以稳定训练。本文直接去除这一目标，避免用教师确定性采样器生成数百万配对样本，使蒸馏更容易扩展到 SDXL 这类大规模文本到图像模型。相对旧方法，这保留了分布匹配“不绑定单条采样轨迹”的优势，也减少了训练前的数据构造成本。
- 用双时间尺度更新规则稳定纯分布匹配： 去掉回归损失后，生成器分布持续变化，fake score function 如果学习滞后，会给生成器提供偏差梯度。作者让 fake score function 比生成器更新更频繁，例如多次 fake score 更新对应一次生成器更新，使其更好地追踪当前生成分布。这个机制替代了原先回归损失的稳定器角色。
- 引入 GAN 损失和真实数据监督： 仅依赖教师模型 real score function 时，学生可能继承教师近似误差。DMD2 在分布匹配框架中加入 GAN 判别器，用真实图像和生成图像进行对抗训练，让学生除向教师分布靠近外，还能直接接收真实数据分布的反馈。论文报告称，这有助于学生生成器在部分设置下超过教师质量。
- 用 backward simulation 缓解多步生成的训练-推理偏差： 多步蒸馏时，如果训练输入来自真实图像加噪，而推理输入来自学生上一步生成结果，就会出现域偏差。Backward simulation 在训练中让学生模型先模拟推理过程，生成带有自身误差特征的中间输入，再训练下一步去噪，从而更贴近实际推理分布。
复杂 Pipeline 深度解析
图中元素对照解读
- 图中左侧 input 图像与黑色横向箭头： 对应论文中的生成器输入，单步时是噪声输入，多步时可代表中间步输入；它在 pipeline 中的作用是启动学生模型的少步生成过程。
- 图中左中部粉红色竖块并标注 few-step generator / Gθ： 对应论文中的 one-step 或 few-step student generator；它在 pipeline 中的作用是把输入快速映射成学生当前分布下的 fake image。
- 图中中部上方标注 fake image 的生成图： 对应论文中的学生生成样本；它在 pipeline 中的作用是同时送入分布匹配、fake score 训练和 GAN 判别这些训练信号。
- 图中 fake image 右侧黑色虚线箭头并标注 diffusion： 对应论文中对 fake image 加噪得到扩散中间态的步骤；它在 pipeline 中的作用是为 real score function 和 fake score function 提供同一类噪声状态输入。
- 图中右上灰色方块、锁图标与标签 real score function： 对应论文中的冻结教师/真实分布 score 模型；它在 pipeline 中的作用是提供目标分布方向的 score 估计，不随学生训练更新。
- 图中右下蓝色方块与标签 fake score function： 对应论文中的可训练 fake score function；它在 pipeline 中的作用是拟合当前学生生成分布的 score，并通过右侧蓝色 diffusion loss 更新。
- 图中最右侧红色减号圆圈、computed gradient 图像与顶部红色回路箭头： 对应论文中的隐式 distribution matching gradient ∇θDKL；它在 pipeline 中的作用是把 real score 与 fake score 的差异转换成反馈到 Gθ 的生成器更新方向。
- 图中下方真实松鼠图、绿色虚线 diffusion 路径、绿色 discriminator 方块与 GAN loss 箭头： 对应论文中的真实图像监督和 GAN 判别器分支；它在 pipeline 中的作用是用 real image 与 fake image 的判别信号补充分布匹配梯度，给生成器提供来自真实数据的对抗反馈。
1. few-step generator 产生 fake image
- 输入： 噪声输入；在多步模型中，也可以是 intermediate step input。
- 操作： 学生生成器用极少步数完成从噪声到图像的映射。
- 输出： fake image。
- 作用： 提供当前学生分布的样本，后续分布匹配梯度、fake score 训练和 GAN 判别都围绕它展开。
2. real score function 和 fake score function 估计两类分数
- 输入： fake image 经过 diffusion 后得到的噪声中间态。
- 操作： 冻结的教师模型提供 real score function；可训练的 fake score function 学习当前生成器输出分布的 score，并通过 diffusion loss 更新。
- 输出： real score 和 fake score。
- 作用： 两个 score 的差构成分布匹配目标的隐式梯度。real score 指向教师/目标分布，fake score 刻画当前学生分布。
3. computed gradient 更新生成器
- 输入： real score 与 fake score。
- 操作： 图中红色路径计算分布匹配梯度。
- 输出： 指向生成器参数更新方向的 gradient。
- 作用： 让生成器的输出分布整体向目标分布移动，而不是逐样本复刻教师采样轨迹。
4. discriminator 和 GAN loss 引入真实数据反馈
- 输入： fake image 与 real image。
- 操作： discriminator 区分真实图像和学生生成图像，并产生 GAN loss。
- 输出： 对抗损失及其对生成器的反馈。
- 作用： 补充教师 score 可能带来的近似误差，让训练信号直接接触真实数据分布。
这篇论文的关键价值在于，它把扩散模型蒸馏从“强依赖教师采样轨迹的配对回归”推进到“无配对数据的分布匹配与对抗学习融合”。DMD2 的核心目标不是逐步复刻教师模型的去噪路径，而是让极少步学生生成器的输出分布对齐目标分布，同时降低大规模文本到图像蒸馏中的数据构造成本。论文报告称，这套改动在 1 步或 4 步生成设置下取得了优于原教师模型的若干指标结果；这些属于论文报告结果，未在此处做跨协议可比性复核。
核心方法图是 Figure 3。该图展示 DMD2 如何把昂贵的扩散教师模型蒸馏为 one-step 或 few-step generator，并把训练拆成两类交替优化：一类更新生成器，另一类更新 fake score function 和 GAN discriminator。
关键术语
- Distribution Matching Distillation (DMD)： 本文的技术基础。它关注学生最终输出分布是否匹配教师/目标分布，而不是要求学生沿着教师的采样轨迹做逐点回归，因此理论上给学生更大的映射自由度。
- Regression loss： 原始 DMD 中用于稳定训练的附加项。本文认为它既昂贵，又把学生绑定到教师采样路径上，因此成为 DMD2 首先移除的组件。
- Two Time-scale Update Rule： DMD2 的稳定训练机制。它让 fake score function 更快更新，以便追踪非平稳的学生生成分布，降低 fake score 滞后造成的梯度偏差。
- GAN loss： DMD2 用来引入真实数据监督的补充目标。它不是替代分布匹配，而是与 DMD 梯度并行更新生成器。
- Backward simulation： 多步学生模型的训练策略。它用学生自身生成的中间状态模拟推理时输入，避免训练时“加噪真图”和推理时“学生生成中间态”之间的分布错位。
为什么 Figure 3 是核心图
Figure 3 同时串联了 DMD2 的三个关键设计：去掉回归损失后的分布匹配主路径、用于稳定 fake score 的 diffusion loss 路径，以及引入真实图像监督的 GAN loss 路径。它不是结果图或示例图，而是完整展示了学生生成器、教师 real score、可训练 fake score、判别器和各类梯度如何相互作用的训练框架图。