【论文笔记】Adversarial Diffusion Distillation

Abstract

本文提出了一种新的训练方法，在保持较高图像质量的前提下，仅用1~4步就能有效地对大规模传统图像扩散模型进行采样，使用分数蒸馏(score distillation)，来利用大规模现成的图像扩散模型作为教师信号，并结合对抗损失，以确保即使在1~2个采样步骤的低步骤机制下也能获得高图像保真度。

分析表明，ADD在单个步骤中明显优于现有的少步方法(GANs)，并且仅在4个步骤中达到了最先进的扩散模型(SDXL)的性能。ADD是第一种利用基础模型解锁单步实时图像合成的方法。

1 Introduction

扩散模型的迭代推理过程需要大量的采样步骤，阻碍了其实时应用。
GANs的特点是单步生成，速度固定，但是即便推广到大型数据集，GANs的样本质量方面仍然比不过Diffusion Model。

本文方法很简单，提出Adversarial Diffusion Distillation (ADD)，将预训练的扩散模型的推理步数降低到1~4步，同时保持较高的采样保真度，并有可能进一步提高模型的整体性能。
为实现上述目标，提出两种训练目标：

• 对抗损失：强制模型在每个前向通道上直接生成真实图像样本，避免其他蒸馏方法中看到的模糊和伪影。
• 蒸馏损失：使用另一个扩散模型作为教师，有效地利用预训练模型的知识，保留大型扩散模型观察到的强组合性。
  推理过程中，本方法不使用classifier-free guidance，减少内存需求，通过迭代细化保留模型生成能力。

3 Method

本文目标：在尽可能少的采样步骤中升成高保真度的样本。

对抗损失支持快速生成，因为在一步前向步骤中训练一个在图像流形中输出样本的模型。
将GANs扩展到大型数据集不仅依赖于判别器，还使用预训练的分类器或CLIP网络来改善文本对齐，过度使用判别器会引入伪影，使图像质量受损。

本文通过分数蒸馏目标，利用预训练扩散模型的梯度，来提高文本对齐和采样质量。
本文还使用预训练的扩散模型权重初始化模型，因为预训练的生成网络可以提高带有对抗损失的训练过程。

3.1 Training Procedure

图2：ADD。ADD-student训练成一个denoiser，接受扩散输入图像$x_s$，输出样本${\hat{x}}_{\theta }(x_s,s)$，然后优化两个目标：

• 对抗损失(adversarial loss)，模型旨在误导判别器，判别器经过训练，用于从真实图像集$x_0$中分辨生成样本${\hat{x}}_{\theta }$。
• 蒸馏损失(distillation loss)，模型经过训练，用于匹配冻结的扩散模型教师的去噪目标${\hat{x}}_{\psi }$。

训练过程如图2所示，包含3个网络。

• ADD-student：借助预训练的UNet-DM初始化
• 判别器：带有可训练权重$\varphi$
• DM教师模型：带有冻结参数$\psi$的。

训练过程中，ADD-student从ADD-student从噪声数据$x_s$中生成样本${\hat{x}}_{\theta }(x_s,s)$。噪声数据是从真实图像$x_0$通过前向扩散过程$x_s={\alpha }_s{x}_0+{\sigma }_sϵ$产生的。

本实验中，使用了和学生DM模型相同的${\alpha }_s$和${\sigma }_s$系数，从集合 T student = { τ 1 , ⋯   , τ n } T_\text{student}=\{\tau_1,\cdots,\tau_n\} Tstudent​={τ1​,⋯,τn​}中，用$N$个学生选定的timestep均匀采样$s$个样本（实践中$N=4$，${\tau }_n=1000$，强制模型推理时从纯噪声开始）。

对于对抗目标，生成的样本${\hat{x}}_{\theta }$和真实图像$x_0$送入判别器中区分他们。为了从DM教师处蒸馏知识，将学生样本${\hat{x}}_{\theta }$使用教师模型的前向过程扩散得到${\hat{x}}_{\theta ,t}$，再使用DM教师的${\hat{x}}_{\psi }({\hat{x}}_{\theta ,t},t)$去噪预测过程作为蒸馏损失${\mathcal{L}}_{\text{distill}}$的引导目标：
$$\mathcal{L}={\mathcal{L}}_{\text{adv}}^{\text{G}}({\hat{x}}_{\theta }(x_s,s),\varphi )+\lambda {\mathcal{L}}_{\text{distill}}({\hat{x}}_{\theta }(x_s,s),\psi )\text{\hspace{1em}(1)}$$
上述过程是基于像素域描述的，但在潜变量域上操作的过程是一样的。

3.2 Adversarial Loss

本文使用冻结的预训练特征网络$F$和一些轻量判别器头${\mathcal{D}}_{\varphi ,k}$。
对于特征网络$F$，之前的工作提出ViTs做的很好。可训练的判别器头被放到$F_k$上（特征网络上的不同层）。
实践中，使用额外的特征网络提取图像嵌入$c_{\text{img}}$。
这里使用hinge loss作为对抗目标函数。ADD学生的对抗目标${\mathcal{L}}_{\text{adv}}({\hat{x}}_{\theta }(x_s,s),\varphi )$写作：
$${\mathcal{L}}_{\text{adv}}^{\text{G}}({\hat{x}}_{\theta }(x_s,s),\varphi )=-{\mathbb{E}}_{s,ϵ,x_0}[\sum _k{\mathcal{D}}_{\varphi ,k}(F_k({\hat{x}}_{\theta }(x_s,s)))]\text{\hspace{1em}(2)}$$
判别器头通过最小化下式训练：
$$\begin{array}{rl}{\mathcal{L}}_{\text{adv}}^{\text{D}}({\hat{x}}_{\theta }(x_s,s),\varphi )& ={\mathbb{E}}_{x_0}[\sum _k\max (0,1-{\mathcal{D}}_{\varphi ,k}(F_k(x_0)))+\gamma \text{R1}(\varphi )]\\ & +{\mathbb{E}}_{{\hat{x}}_{\theta }}[\sum _k\max (0,1+{\mathcal{D}}_{\varphi ,k}(F_k({\hat{x}}_{\theta })))]\end{array}\text{\hspace{1em}(3)}$$

3.3 Score Distillation Loss

蒸馏损失可写为
$${\mathcal{L}}_{\text{distill}}({\hat{x}}_{\theta }(x_s,s),\psi )={\mathbb{E}}_{t,{ϵ}^{\prime }}[c(t)d({\hat{x}}_{\theta },{\hat{x}}_{\psi }(\text{sg}({\hat{x}}_{\theta ,t});t))]\text{\hspace{1em}(4)}$$
其中$\text{sg}$表示中断梯度操作。
$d$衡量ADD学生生成的样本$x_{\theta }$和DM教师的输出${\hat{x}}_{\psi }({\hat{x}}_{\theta ,t},t)=({\hat{x}}_{\theta ,t}-{\sigma }_t{\widehat{ϵ}}_{\psi }({\hat{x}}_{\theta ,t},t))/{\alpha }_t$之间的距离。
教师不是直接用在生成图像${\hat{x}}_{\theta }$上，而是扩散后的输出${\hat{x}}_{\theta ,t}={\alpha }_t{\hat{x}}_{\theta }+{\sigma }_t{ϵ}^{\prime }$，因为对于教师而言和生成图像的数据分布不同。

定义距离函数$d(x,y)=||x-y|{|}_2^2$。
关于加权函数$c(t)$，考虑两个选择：

• 指数加权(exponential weighting)，即$c(t)={\alpha }_t$
• 分数蒸馏采样加权(score distillation sampling, SDS)。

实现过程中，$d(x,y)=||x-y|{|}_2^2$，且对$c(t)$有特定的选择时，蒸馏损失和SDS目标${\mathcal{L}}_{\text{SDS}}$等价。SDS损失为：
$${\mathcal{L}}_{\text{SDS}}(\varphi ,x=g(\theta ))={\nabla }_{\theta }{\mathbb{E}}_t[{\sigma }_t/{\sigma }_{t}w(t)\text{KL}(q(z_t|g(\theta );y,t)||p_{\theta }(z_t;y,t))]$$
本文公式的优点是它能够使重建目标直接可视化，并且它自然地促进了连续几个去噪步骤的执行。