ADD：具有对抗学习和知识提炼功能的扩散模型

生成图像 21/05/2024

三个要点
✔️ 扩散模型的推理速度非常慢，因此难以实时使用
✔️ 引入 ADD 后，在保持质量的同时显著提高了估算速度
✔️ ADD 是首个使用对抗学习和知识提炼的单步扩散模型。

Adversarial Diffusion Distillation
written by Axel Sauer, Dominik Lorenz, Andreas Blattmann, Robin Rombach
(Submitted on 28 Nov 2023)
Comments: Published on arxiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV)

code：

本文所使用的图片要么来自论文、介绍性幻灯片，要么是参考这些图片制作的。

介绍

扩散模型作为生成模型受到了广泛关注，最近在生成高质量图像和视频方面取得了显著进展。高图像质量和高多样性是扩散模型的两大优势。然而，在生成图像时，需要成百上千个采样步骤，而且估计速度非常慢。

另一方面，生成对抗网络（GANs）的特点是单步表述和快速采样。然而，尽管人们试图将其扩展到大型数据集，但 GANs 在样本质量上往往不如扩散模型。另一个弱点是生成图像的多样性较低。

本评论文章的目标是将扩散模型的出色样本质量与 GAN 的快速采样相结合。为此，本文引入了两种训练目标的组合。

敌对损失
与分馏取样 (SDS) 相对应的蒸馏损失

逆向损失通过一个鉴别器对真实图像和生成图像进行比较，从而避免了其他蒸馏方法中经常出现的模糊和其他伪影现象。蒸馏损耗使用另一个预先训练好的（固定的）扩散模型作为教师，有效地利用了预先训练好的扩散模型的广泛知识。

所提出的方法优于扩散模型 SOTA SDXL，只需一到四个采样步骤即可生成高保真实时图像。

建议方法

培训程序

训练过程如图 1 所示，其中主模型 ADD-student由三个预先训练过的权重为 θ 的扩散模型（UNet-DM）、一个可训练权重为 $j$ 的判别器和一个权重为 $ψ$ 的 DM-Teacher （扩散模型）组成。使用的模型

对于敌意损失，生成的样本 $(\hat{x}_\theta) $ 和实际图像 $( x_0) $ 被传递给一个判别器来区分它们。下一节将详细介绍判别器和敌意损失的设计。为了从 DM-Teacher 中提炼知识，ADD-学生样本 $(\hat{x}_{\theta))$ 被扩散到教师（DM-Teacher）的前瞻过程 $(\hat{x}_{\theta,t})$ 和提炼损失 $( L_{text{distill})$，并使用教师的去噪预测 $(\hat{x}_\psi(\hat{x}_{\theta,t},t)) $ 作为重建目标。详情将在下一节给出。

整体损失函数如下

敌对损失和标识符

在判别器方面，我们使用了 Stylegan-t 的结构和设置（Sauer et al, 2023）。我们使用一个固定的预训练特征网络 F 和一组可训练的轻量级判别头 $ D_{(j, k)} $。对于特征网络 F，我们将在下一节中研究视觉转换器（ViT）和模型大小的不同选择，因为 Sauer 等人发现视觉转换器（ViT）效果很好。可训练的判别头应用于特征网络不同层中的特征 Fk。

判别器 $ L_{adv}^D $ 和主要模型 $ L_{adv}^G $ 的损失如下。

其中 $ R1 $ 表示 $ R1 $ 梯度惩罚。我们不计算像素值的梯度惩罚，而是在每个判别头 ǫ( D_{(ϕ, k)} \) 的输出端计算梯度惩罚。如果输出分辨率大于 $128 \times 128 $ 像素，R1 惩罚就特别有用。

分馏损失

馏分损失的计算公式如下

$ sg $ 表示停止梯度操作。分数蒸馏损失使用的是距离指标（d），它计算的是ADD-学生生成的样本（x_\theta \）与 DM-教师输出之间的差异。为了找到一个合适的 $ d$ ，我们在实验中测试了许多函数，其中平均平方误差（MSE）是最有效的。

试验

生成模型与 SOTA 的定量比较

在本实验中，通过用户偏好调查，而不是常用的自动计算评价指标，更客观地检验了建议方法的有效性。用户选择了两个评价指标中较好的一个：提示符合性（输入提示是否正确反映在输出图像中）和图像质量。图 2 和图 3 对结果进行了总结。只需几个采样步骤（1-4），提议的方法就能超越生成模型的代表模型，并实现 SOTA 结果，尤其是在第 4 步。

定性结果和比较

图 4 显示了 SDXL 和建议方法之间的定性比较。可以看出，拟议方法只需四个步骤就能生成与 SDXL 相同或更好的图像质量。此外，还可以确认输入的质子也能正确反映在生成的结果中。特别是，如图 4左下方图像所示，可以确认所提出的方法比 SDXL 产生的噪声和伪影更少。包括定量实验结果在内，可以看出所提出的方法在质量和及时一致性方面都优于扩散模型的 SOTASDXL，而且采样步骤更少。

摘要

在这篇文章中，我们介绍了对抗性扩散蒸馏（ADD），用于将预先训练好的扩散模型蒸馏为快速、低步骤的图像生成模型。所提出的方法结合了对抗性蒸馏和分数蒸馏损失，利用来自判别器的真实数据和来自扩散教师的结构理解，对稳定扩散和 SDXL 等训练有素的模型进行蒸馏。所提出的方法在进行一到两步的超快速采样时表现尤为出色，实验结果表明，它在很多情况下都优于之前的研究。另一方面，进一步增加步数会产生更好的结果，优于常用的多步扩散模型，如 SDXL、IF 和 OpenMUSE。不过，在图像质量和与 pronto 的一致性方面，单步采样生成模型仍有改进的余地。如果能做出更多改进，所提出的方法可能会成为第一个可实时使用的扩散模型。