MAPLE 2024-11-19 17:23 江苏
关注公众号，发现CV技术之美
本篇分享 NeurIPS 2024 论文One-Step Diffusion Distillation through Score Implicit Matching，单步生成：让扩散模型实现高速无损的内容生成。
团队成员：Weijian Luo, Zemin Huang, Zhengyang Geng, J. Zico Kolter, Guo-jun Qi
引言
扩散模型（Diffusion Models, DMs）是一类通过逐步添加噪声来破坏数据分布并逆向生成数据的模型。这类模型在图像生成、文本到图像生成等领域表现出色，但其生成过程往往需要几十步的采样步骤，造成高昂的计算成本和较长的响应时间。这在实际应用中，尤其是在实时生成或移动设备上，是一个重要的限制。
随着对生成模型效率的需求不断增加，研究人员开始探索如何将预训练的扩散模型转化为高效的单步或少步生成模型。我们的研究聚焦于扩散模型蒸馏（Diffusion Distillation），通过优化策略实现单步生成图像，同时尽量保留与原始模型相近的生成能力。
单步扩散算法
我们提出的隐式得分匹配算法可以将多步推理的扩散模型蒸馏为单步推理的生成模型，并且在蒸馏后没有任何图像生成质量上的损失。这一算法的关键思想是推导出一个可优化的训练目标，同时其梯度与不可直接优化的基于得分函数的散度相同。一般的基于得分函数的散度定义如下，它可以用于将单步生成器的分布和预训练教师模型的分布进行匹配以优化单步生成器：
但实际由于这个损失函数很难直接优化，因此我们通过推导出得分散度梯度定理获得了一种等效但易于处理的训练目标，用于优化单步生成器，这一等效的训练目标可以表示为：
我们最终的训练算法包含了如下两个阶段的交替训练：
在线扩散模型更新：在这一阶段，我们固定单步生成模型的参数，利用扩散模型的损失来优化在线模型。
学生模型参数更新：在这一阶段，我们固定在线扩散模型的参数，通过上文提出的损失函数来优化单步生成器。
实验结果
为了验证隐式得分匹配这一蒸馏算法的有效性，我们在多个设置上进行了大量实验，包括小规模的CIFAR10 数据集和开放数据集上训练的文生图模型 PixArt-α 实验结果显示，我们的单步生成模型在无条件和有条件的生成中均显著优于其他蒸馏模型，达到2.06/1.96的 FID 值，显著优于其他蒸馏方法。
在文本到图像生成的实验中，我们选择了 DiT 架构的 PixArt-α 模型作为教师模型，应用我们的算法蒸馏出了 SIM-DiT-600M。与其他经过少步蒸馏的扩散模型相比，我们只需一步就能生成高质量的图像，并在多个基准测试中取得了优异的性能。这些结果展示了我们的蒸馏算法在快速生成和高质量样本方面的潜力。
下图我们对比了 SIM-DIT-600M 和其他单步与少步生成的蒸馏工作，最左侧是我们生成的图像，可以看出我们的蒸馏模型与其他的模型相比具有更强的生成质量以及美学效果。
下图最右侧是 LCM-PixArt-α生成的结果，左边两侧的图由教师模型 PixArt-α和我们蒸馏后的模型生成并打乱排列顺序。我们的蒸馏后的模型有几乎和原模型相当的生成质量，并且将推理步数从几十步降低到只需要单步推理。
结论
本文介绍了 Score Implicit Matching(SIM) 方法在扩散模型蒸馏中的应用，展示了如何通过新颖的基于得分函数的散度训练来实现单步生成。实验结果验证了 SIM 在多个生成任务 中的优越性能，标志着在快速生成和高质量样本方面的一个重要进展。未来，随着对 SIM 方法的进一步研究和优化，我们相信其将在更多领域展现出更大的应用潜力，推动生成模型的发展。
撰稿：黄泽敏，罗维俭编辑：姜天欣
最新 AI 进展报道
请联系：amos@52cv.net
END
欢迎加入「图像生成」交流群?备注：生成
