本篇分享 NeurIPS 2024 论文 One-Step Diffusion Distillation through Score Implicit Matching ，单步生成：让扩散模型实现高速无损的内容生成。

• 团队成员：Weijian Luo, Zemin Huang, Zhengyang Geng, J. Zico Kolter, Guo-jun Qi
• 相关论文：https://arxiv.org/abs/2410.16794
• 项目网站：https://github.com/maple-research-lab/SIM

引言

扩散模型（Diffusion Models, DMs） 是一类通过逐步添加噪声来破坏数据分布并逆向生成数据的模型。这类模型在图像生成、文本到图像生成等领域表现出色，但其生成过程往往需要几十步的采样步骤，造成高昂的计算成本和较长的响应时间。这在实际应用中，尤其是在实时生成或移动设备上，是一个重要的限制。

随着对生成模型效率的需求不断增加，研究人员开始探索如何将预训练的扩散模型转化为高效的单步或少步生成模型。我们的研究聚焦于扩散模型蒸馏（Diffusion Distillation），通过优化策略实现 单步生成图像 ，同时尽量保留与原始模型相近的生成能力。

单步扩散算法

我们提出的 隐式得分匹配 算法可以将多步推理的扩散模型蒸馏为单步推理的生成模型，并且在蒸馏后没有任何图像生成质量上的损失。这一算法的关键思想是推导出一个可优化的训练目标，同时其梯度与不可直接优化的基于得分函数的散度相同。一般的基于得分函数的散度定义如下，它可以用于将单步生成器的分布和预训练教师模型的分布进行匹配以优化单步生成器：

但实际由于这个损失函数很难直接优化，因此我们通过推导出 得分散度梯度定理 获得了一种等效但易于处理的训练目标，用于优化单步生成器，这一等效的训练目标可以表示为：

我们最终的训练算法包含了如下两个阶段的交替训练：

1. 在线扩散模型更新： 在这一阶段，我们固定单步生成模型的参数，利用扩散模型的损失来优化在线模型。
2. 学生模型参数更新： 在这一阶段，我们固定在线扩散模型的参数，通过上文提出的损失函数来优化单步生成器。

实验结果

为了验证 隐式得分匹配 这一蒸馏算法的有效性，我们在多个设置上进行了大量实验，包括小规模的CIFAR10 数据集和开放数据集上训练的文生图模型 PixArt-α 实验结果显示，我们的单步生成模型在无条件和有条件的生成中均显著优于其他蒸馏模型，达到2.06/1.96的 FID 值，显著优于其他蒸馏方法。

在文本到图像生成的实验中，我们选择了 DiT 架构的 PixArt-α 模型作为教师模型，应用我们的算法蒸馏出了 SIM-DiT-600M 。与其他经过少步蒸馏的扩散模型相比，我们只需一步就能生成高质量的图像，并在多个基准测试中取得了优异的性能。这些结果展示了我们的蒸馏算法在快速生成和高质量样本方面的潜力。

下图我们对比了 SIM-DIT-600M 和其他单步与少步生成的蒸馏工作，最左侧是我们生成的图像，可以看出我们的蒸馏模型与其他的模型相比具有更强的生成质量以及美学效果。

下图最右侧是 LCM-PixArt-α生成的结果，左边两侧的图由教师模型 PixArt-α和我们蒸馏后的模型生成并打乱排列顺序。我们的蒸馏后的模型有几乎和原模型相当的生成质量，并且将 推理步数从几十步降低到只需要单步推理 。