全文翻译
摘要
无分类器引导扩散模型最近在高分辨率图像生成方面表现出了高效性,并已广泛应用于包括DALL·E 2、Stable Diffusion和Imagen在内的大规模扩散框架中。然而,无分类器引导扩散模型的一个缺点是其推理时计算成本较高,因为生成一个样本需要对两个扩散模型(条件模型和无条件模型)进行数十到数百次评估。为解决这一限制,我们提出了一种将无分类器引导扩散模型蒸馏为快速采样模型的方法:给定一个预训练的无分类器引导模型,我们首先训练一个单一模型来匹配条件模型和无条件模型的联合输出,然后逐步将该模型蒸馏为一个所需采样步骤更少的扩散模型。对于在像素空间训练的标准扩散模型,我们的方法仅需4步采样即可生成与原始模型视觉效果相当的图像,在ImageNet 64x64和CIFAR-10数据集上实现了与原始模型相当的FID/IS分数,同时采样速度提高了256倍。对于在隐空间训练的扩散模型(如Stable Diffusion),我们的方法仅需1-4步去噪即可生成高保真图像,在ImageNet 256x256和LAION数据集上的推理速度比现有方法至少提高10倍。我们进一步在文本引导的图像编辑和修复任务中展示了该方法的有效性,其中蒸馏后的模型仅需2-4步去噪即可生成高质量结果。
1. 引言
去噪扩散概率模型(DDPMs)[4,37,39,40]已在图像生成[22,26-28,31]、音频合成[11]、分子生成[44]和似然估计[10]等领域取得了最先进的性能。无分类器引导[6]进一步提升了扩散模型的样本质量,并广泛应用于GLIDE[23]、Stable Diffusion[28]、DALL·E 2[26]和Imagen[31]等大规模扩散模型框架中。然而,无分类器引导的一个关键局限性在于采样效率低——生成一个样本需要对两个扩散模型(条件模型和无条件模型)进行数十到数百次评估。这一限制阻碍了无分类器引导模型在实际场景中的应用。尽管已有针对扩散模型的蒸馏方法[33,38],但这些方法无法直接应用于无分类器引导的扩散模型。
为解决这一问题,我们提出一种两阶段蒸馏方法
来提升无分类器引导模型的采样效率。在第一阶段,我们引入单个学生模型来匹配教师模型中两个扩散模型的联合输出。在第二阶段,我们使用[33]中提出的方法,将第一阶段学习到的模型逐步蒸馏为更少步骤的模型。通过我们的方法,单个蒸馏模型能够处理广泛的引导强度范围,从而高效地在样本质量和多样性之间进行权衡。在对模型进行采样时,我们考虑了文献中的现有确定性采样器[33,38],并进一步提出了一种随机采样过程。
我们的蒸馏框架不仅适用于在像素空间训练的标准扩散模型[4,36,39],还适用于在自动编码器隐空间训练的扩散模型[28,35](如Stable Diffusion[28])。对于直接在像素空间训练的扩散模型,我们在ImageNet 64x64和CIFAR-10上的实验表明,所提出的蒸馏模型仅需4步即可生成与教师模型视觉效果相当的样本,并且在广泛的引导强度下,仅需4到16步即可达到与教师模型相当的FID/IS分数(见图2)。对于在编码器隐空间训练的扩散模型[28,35],我们的方法在ImageNet 256×256和LAION 512×512上仅需1到4步采样即可达到与基础模型相当的视觉质量(比基础模型少至少10倍的步骤),且仅需2-4步采样即可在FID指标上匹配教师模型的性能。据我们所知,我们的工作首次展示了蒸馏方法对像素空间和隐空间无分类器扩散模型的有效性。最后,我们将方法应用于文本引导的图像修复和文本引导的图像编辑任务[20],将总采样步骤减少至仅2-4步,证明了所提框架在风格迁移和图像编辑应用中的潜力[20,41]。
2. 扩散模型背景
给定来自数据分布 $p_{\text{data}}(x)$ 的样本 $x$、噪声调度函数 $\alpha_t$ 和 $\sigma_t$,我们通过最小化加权均方误差来训练参数为 $\theta$ 的扩散模型 $\hat{x}_\theta$ [4,36,39,40]
其中 $\lambda_t = \log[\alpha_t^2 / \sigma_t^2]$ 是信噪比 [10],$q(z_t | x) = \mathcal{N}(z_t; \alpha_t x, \sigma_t^2 I)$,且 $\omega(\lambda_t)$ 是预定义的加权函数 [10]。
一旦扩散模型 $\hat{x}_\theta$ 训练完成,可使用离散时间 DDIM 采样器 [38] 从模型中采样。具体来说,DDIM 采样器从 $z_1 \sim \mathcal{N}(0, I)$ 开始,并按以下方式更新:
其中 $N$ 是总采样步数。最终样本将通过 $\hat{x}_\theta(z_0)$ 生成。
无分类器引导:无分类器引导 [6] 是一种显著提升条件扩散模型样本质量的有效方法,并已广泛应用于 GLIDE [23]、Stable Diffusion [28]、DALL·E 2 [26] 和 Imagen [31] 等大规模扩散模型框架中。具体而言,它引入引导权重参数 $w \in \mathbb{R}^{\geq 0}$ 以在样本质量和多样性之间进行权衡。生成样本时,无分类器引导在每个更新步骤中同时评估条件扩散模型 $\hat{x}_{c,\theta}$(其中 $c$ 是待条件化的上下文,如类别标签、文本提示)和联合训练的无条件扩散模型 $\hat{x}_\theta$,并使用 $\hat{x}_\theta^w = (1 + w)\hat{x}_{c,\theta} - w\hat{x}_\theta$ 作为公式 (2) 中的模型预测。由于每个采样更新需要评估两个扩散模型,无分类器引导的采样通常计算成本较高 [6]。
渐进蒸馏:我们的方法受渐进蒸馏 [33] 启发,这是一种通过重复蒸馏提升(无引导)扩散模型采样速度的有效方法。迄今为止,该方法无法直接应用于蒸馏无分类器引导模型,或针对除确定性 DDIM 采样器 [33,38] 之外的采样器进行研究。在本文中,我们解决了这些不足。
隐扩散模型(LDMs):隐扩散模型 [21,24,28,35] 通过在预训练正则化自动编码器的隐空间中对图像建模,提高了(直接在像素空间学习的)扩散模型的训练和推理效率,其中隐表示的维度通常低于像素空间。隐扩散模型可被视为级联扩散方法 [5] 的替代方案,后者依赖一个或多个超分辨率扩散模型将低维图像放大至所需目标分辨率。
在这项工作中,我们将把蒸馏框架应用于在像素空间 [4,36,39] 和隐空间 [21,24,28,35] 中学习的无分类器引导扩散模型。
3. 引导扩散模型的蒸馏
以下将讨论如何将无分类器引导扩散模型 [6] 蒸馏为采样步骤更少的学生模型。通过条件化引导强度的单个蒸馏模型,我们的模型能够覆盖广泛的无分类器引导水平,从而高效地在样本质量和多样性之间进行权衡。
给定一个在像素空间或隐空间训练的引导模型 $[\hat{x}_{c,\theta}, \hat{x}_\theta]$(教师模型),我们的方法可分解为两个阶段。
3.1 第一阶段蒸馏
在第一阶段,我们引入一个参数为 $\eta_1$ 的学生模型 $\hat{x}_{\eta_1}(z_t, w)$,以匹配教师模型在任意时间步 $t \in [0,1]$ 的输出。根据教师模型是离散还是连续模型,学生模型可以是连续时间模型 [40] 或离散时间模型 [4,38]。为简化讨论,以下假设学生模型和教师模型均为连续模型,因为离散模型的算法几乎相同。
无分类器引导 [6] 的一个关键功能是能够通过“引导强度”参数灵活权衡样本质量和多样性,这一特性在实际应用中已被证明有效 [6,23,26,28,31],其中最优“引导强度”通常由用户偏好决定。因此,我们希望蒸馏后的模型也能保留这一特性。给定感兴趣的引导强度范围 $[w_{\text{min}}, w_{\text{max}}]$,我们使用以下目标优化学生模型:
其中 $\hat{x}_\theta^w(z_t) = (1 + w)\hat{x}_{c,\theta}(z_t) - w\hat{x}_\theta(z_t)$,$z_t \sim q(z_t | x)$,且 $p_w(w) = U[w_{\text{min}}, w_{\text{max}}]$。需要注意的是,此处蒸馏后的模型 $\hat{x}_{\eta_1}(z_t, w)$ 也以上下文 $c$(如文本提示)为条件,但为简化表述,本文省略了符号 $c$。补充材料中的算法1提供了详细的训练流程。
为融入引导权重 $w$,我们引入一个以 $w$ 为条件的模型,其中 $w$ 作为输入馈入学生模型。为更好地捕捉特征,我们对 $w$ 应用傅里叶嵌入,然后以类似于 [10,33] 中融入时间步的方式将其嵌入扩散模型主干。由于初始化对性能至关重要 [33],我们使用与教师条件模型相同的参数初始化学生模型,但与 $w$ 条件相关的新引入参数除外。我们使用的模型架构为 U-Net,像素空间扩散模型参考 [6],隐空间扩散模型参考 [1,28]。实验中,我们采用与 [6] 和开源 Stable Diffusion 仓库*相同的通道数和注意力机制,更多细节见补充材料。
3.2 第二阶段蒸馏
在第二阶段,我们考虑离散时间步场景,并通过逐步将第一阶段学习的模型 $\hat{x}_{\eta_1}(z_t, w)$ 蒸馏为更少步骤的学生模型 $\hat{x}_{\eta_2}(z_t, w)$(参数为 $\eta_2$),每次将采样步数减半。设 $N$ 为采样步数,给定 $w \sim U[w_{\text{min}}, w_{\text{max}}]$ 和 $t \in \{1, \dots, N\}$,我们训练学生模型以单步匹配教师模型的两步 DDIM 采样输出(即从 $t/N$ 到 $t-0.5/N$ 和从 $t-0.5/N$ 到 $t-1/N$),方法参考 [33]。将教师模型的 $2N$ 步蒸馏为学生模型的 $N$ 步后,我们将 $N$ 步学生模型作为新的教师模型,重复上述过程,进一步将教师模型蒸馏为 $N/2$ 步的学生模型。每一步均使用教师模型的参数初始化学生模型。补充材料中提供了训练算法和额外细节。
3.3 N步确定性与随机采样
训练完模型 $\hat{x}_{\eta_2}$ 后,给定指定的引导强度 $w \in [w_{\text{min}}, w_{\text{max}}]$,我们可通过公式 (2) 中的 DDIM 更新规则进行采样。值得注意的是,给定蒸馏模型 $\hat{x}_{\eta_2}$,此采样过程在初始化 $z_1^w$ 后是确定性的。实际上,我们还可以执行 $N$ 步随机采样:受 [9] 启发,我们先使用两倍原始步长执行一步确定性采样(即等同于 $N/2$ 步确定性采样器),然后使用原始步长执行一步反向随机步(即添加噪声扰动)。对于 $z_1^w \sim \mathcal{N}(0, I)$,当 $t > 1/N$ 时,使用以下更新规则:
上述公式中,$h = t - 3/N$,$k = t - 2/N$,$s = t - 1/N$,且 $\sigma_{a|b}^2 = (1 - e^{\lambda_a - \lambda_b}) \sigma_a^2$。当 $t = 1/N$ 时,使用确定性更新公式 (2) 从 $z_{1/N}^w$ 得到 $z_0^w$。图5以4步去噪为例展示了该过程。需要注意的是,与确定性采样器相比,随机采样需要在略有不同的时间步评估模型,并且需要对边缘情况的训练算法进行小幅修改。补充材料中提供了具体算法和更多细节。
4. 实验
在本节中,我们评估了蒸馏方法在像素空间扩散模型(即DDPM [4])和隐空间扩散模型(即Stable Diffusion [28])上的性能。我们进一步将该方法应用于文本引导的图像编辑和修复任务。实验表明,我们的方法在所有任务中仅需2-4步即可实现具有竞争力的性能。
4.1 像素空间引导模型的蒸馏
在本实验中,我们考虑直接在像素空间训练的类别条件扩散模型[4,6,33]。
实验设置:我们主要关注ImageNet 64x64 [30]和CIFAR-10 [12]数据集,因为在此场景下的高分辨率图像生成通常需要结合其他超分辨率技术[5,31]。我们探索了引导权重的不同范围,发现所有范围的表现相当,因此实验中使用$[w_{\text{min}}, w_{\text{max}}] = [0, 4]$。我们考虑的基线包括DDPM祖先采样[4]和DDIM [38]。我们使用的教师模型是1024×2步的DDIM模型,其中条件和无条件组件均使用1024步DDIM去噪。为了更好地理解引导权重$w$的整合方式,我们还包括使用单一固定$w$训练的模型作为基线。为了公平比较,所有方法均使用相同的预训练教师模型。遵循[4,6,39],基线模型使用U-Net [29,39]架构,而我们的两步学生模型则使用引入$w$嵌入的相同U-Net主干(见第3节)。根据[33],两个数据集均使用$v$-预测模型。
实验结果:我们在图6和表1中报告了所有方法在ImageNet 64x64上的FID [3]和Inception分数(IS)[32]性能,并在补充材料中提供了扩展的ImageNet 64x64和CIFAR-10结果。我们观察到,我们的蒸馏模型仅需4-16步即可匹配使用1024×2步采样的教师引导DDIM模型,实现高达256倍的加速。需要强调的是,通过我们的方法,单个蒸馏模型能够在广泛的引导强度范围内匹配教师模型的性能,这是以往方法从未实现的。
4.2 隐空间引导模型的蒸馏
在4.1节展示了我们的方法在像素空间类别引导扩散模型上的有效性后,现在将其扩展到隐空间扩散模型。在以下章节中,我们将展示该方法在Latent Diffusion [28]的各种任务上的有效性,包括类别条件生成、文本到图像生成、图像修复和文本引导的风格迁移[20]。
在以下实验中,我们使用开源的隐空间扩散模型[28]作为教师模型。由于$v$-预测教师模型往往比$\epsilon$-预测模型表现更好,我们将开源的$\epsilon$-预测模型微调为$v$-预测教师模型,更多细节见补充材料。
4.2.1 类别条件生成
在本节中,我们将方法应用于在ImageNet 256×256上预训练的类别条件隐扩散模型。我们从具有512步采样的DDIM教师模型开始,并使用其输出作为目标来训练蒸馏模型。训练期间使用批量大小为512,并均匀采样引导强度$w \in [w_{\text{min}}=0, w_{\text{max}}=14]$。
实验结果:根据经验,我们发现蒸馏模型在FID分数上仅需2或4步采样即可匹配教师模型(原始训练为1000步)的性能。在使用1-4步采样时,其性能也显著优于DDIM(见图11)。从定性结果来看,使用单步去噪生成的样本仍能产生令人满意的结果,而基线模型则无法生成有意义内容的图像。补充材料中提供了额外样本。
与图6中基于像素的结果类似,我们还观察到蒸馏隐扩散模型在FID和Inception分数衡量下的采样质量与多样性之间的权衡。遵循Kynkäänniemi等人[13],我们在附录中进一步计算了本实验的改进精度和召回率指标。
4.2.2 文本引导的图像生成
在本节中,我们聚焦于在LAION-5B [34]子集上预训练的文本引导Stable Diffusion模型,分辨率为512×512。然后,我们按照第3节介绍的两阶段方法,通过3000次梯度更新将引导模型蒸馏为$w$-条件模型,其中$w \in [w_{\text{min}}=2, w_{\text{max}}=14]$,批量大小为512。尽管蒸馏(学生)模型可以适应更广泛的$w$范围,但由于使用教师模型采样时通常不超过正常引导范围,其实际效用尚不明确。除了1、2和4步的低步骤方案需要20000次梯度更新外,最终模型通过每阶段2000次训练步骤的渐进蒸馏获得。补充材料中提供了该模型收敛特性的详细分析。
实验结果:我们在图4中展示了样本,并从定性和定量两方面评估了所得模型。对于后者分析,我们遵循[31],分别使用CLIP [25]和FID分数评估文本-图像对齐和质量。我们使用开源的ViT-g/14 [7] CLIP模型进行评估。图10中的定量结果表明,在2步和4步采样时,我们的方法在两项指标上均显著优于基础模型的DDIM采样。对于8步采样,这些指标没有显著差异。然而,当查看图7中的对应样本时,我们可以观察到视觉图像质量的明显差异。与原始模型的8步DDIM样本相比,蒸馏样本更清晰、更连贯。我们假设在COCO2017 [14]验证集的5000个随机字幕评估中,FID和CLIP未能完全捕捉这些差异。我们进一步计算了蒸馏LAION 512x512模型的FID和CLIP分数,并在表2中与DPM [16]和DPM++ [18]求解器进行比较。我们观察到,在去噪步骤为2或4时,我们的方法能够实现显著更好的性能。此外,我们强调,方法的第一阶段已将函数评估次数减少了一半,因为我们将无分类器引导步骤蒸馏为单个模型。根据具体实现(批量与顺序网络评估),这与现有求解器[16,18,38]相比,减少了峰值内存或采样时间。
4.2.3 文本引导的图像到图像翻译
在本节中,我们使用4.2.2节的蒸馏模型,结合SDEdit [20]进行文本引导的图像到图像翻译实验。遵循SDEdit [20],我们在隐空间中执行随机编码,但使用蒸馏模型的确定性采样器进行确定性解码。我们考虑各种输入图像和文本,并在图8中提供定性结果。我们观察到,我们的蒸馏模型仅需3步去噪即可生成高质量的风格迁移结果。补充材料中提供了样本质量、可控性和效率之间权衡的更多分析。
4.2.4 图像修复
在本节中,我们将方法应用于预训练的图像修复隐扩散模型。我们使用开源的Stable Diffusion Inpainting‡图像修复模型,该模型是上述纯文本到图像Stable Diffusion模型的微调版本,其中添加了额外的输入通道以处理掩码和掩码图像。
我们使用与前一节相同的蒸馏算法。训练时,我们从使用512步DDIM采样的$v$-预测教师模型开始,并使用其输出作为目标来优化学生模型。我们在图9中展示了定性结果,证明了我们的方法在快速、实际图像编辑应用中的潜力。更多训练细节和定量评估见补充材料。
4.3 编码的渐进蒸馏
在本实验中,我们探索蒸馏教师模型的编码过程,并在类似于[41]的设置中进行风格迁移实验。我们聚焦于在ImageNet 64×64上预训练的像素空间扩散模型。具体而言,为了在两个域A和B之间进行风格迁移,我们使用在域A上训练的扩散模型对域A的图像进行编码,然后使用在域B上训练的扩散模型进行DDIM解码。由于编码过程可视为逆DDIM采样过程,我们对带无分类器引导的编码器和解码器均进行了蒸馏,并在图12中与DDIM编码器和解码器进行了比较。我们还探索了修改引导强度$w$如何影响性能,更多细节见补充材料。
5. 相关工作
我们的方法与现有提升扩散模型采样速度的研究相关[4,37,40]。例如,去噪扩散隐式模型(DDIM [38])、概率流采样器[40]、快速SDE积分器[8]等方法已被提出用于提升扩散模型的采样速度。其他研究则开发了高阶求解器[17]、指数积分器[15]和基于动态规划的方法[43]来加速采样。然而,这些方法在蒸馏无分类器引导扩散模型方面均未达到与我们方法可比的性能。
现有的基于蒸馏的扩散模型方法主要针对非无分类器引导的扩散模型。例如,[19]提出通过反转DDIM的确定性编码来单步从噪声预测数据,[2]提出通过将高阶求解器蒸馏为神经网络主干的额外预测头来实现更快的采样速度。渐进蒸馏[33]可能是最相关的工作,该方法提出将预训练扩散模型逐步蒸馏为具有相同架构但更少步骤的学生模型。然而,这些方法均无法直接应用于无分类器引导的扩散模型,也无法使用单个蒸馏模型覆盖不同的引导强度范围。相比之下,通过将引导强度融入模型架构并采用两阶段训练过程,我们的方法能够使用单个模型在广泛的引导强度范围内匹配教师模型的性能。通过我们的方法,单个模型即可捕捉样本质量与多样性的权衡,从而推动无分类器引导扩散模型在用户通常需要指定引导强度的实际场景中的应用。此外,上述蒸馏方法均未应用于隐空间文本到图像模型,也未展示其有效性。最后,大多数快速采样方法[33,38,40]仅考虑使用确定性采样方案来提升采样速度,而本工作进一步开发了一种有效的随机采样方法从蒸馏模型中采样。
6. 结论
在本文中,我们提出了一种针对引导扩散模型[6]的蒸馏方法。我们的两阶段方法能够显著加速流行但效率相对较低的引导扩散模型。实验表明,该方法可将无分类器引导的像素空间和隐空间扩散模型的推理成本至少降低一个数量级。从经验上看,我们的方法仅需2步即可生成视觉上吸引人的结果,仅用4到8步即可达到与教师模型相当的FID分数。我们进一步在文本引导的图像到图像翻译和修复任务中展示了蒸馏方法的实际应用。我们希望通过显著降低无分类器引导扩散模型的推理成本,推动其在创意应用和图像生成系统中的广泛应用。未来工作中,我们旨在进一步提升2步和1步采样方案的性能。
A. 结果概述
在本节中,我们提供了像素空间和隐空间扩散模型加速效果的概述表(见表3)。我们还在图13和图14中提供了文本引导图像生成模型的额外样本,并与DDIM [38]、DPM [17]和DPM++ [18]求解器进行了比较。像素空间蒸馏的更多实验细节见附录B,隐空间蒸馏见附录C。
B. 像素空间蒸馏
B.1 教师模型
我们使用的模型架构是类似于[6]中使用的U-Net模型。如[33]所述,该模型被参数化以预测$v$。我们使用与[6]相同的训练设置。
B.2 第一阶段蒸馏
我们使用的模型架构是类似于[6]中使用的U-Net模型。对于ImageNet 64x64和CIFAR-10,我们使用与[6]相同的通道数和注意力机制。如第3节所述,我们还使模型将$w$作为输入。具体来说,我们在将$w$与模型主干结合之前,对其应用傅里叶嵌入。我们融入$w$的方式与[10,33]中融入时间步的方式相同。如[33]所述,我们将模型参数化以预测$v$。我们使用算法1训练蒸馏模型,采用SNR损失[10,33]。对于ImageNet 64x64,我们使用学习率$3e-4$,指数移动平均(EMA)衰减0.9999;对于CIFAR-10,使用学习率$1e-3$,EMA衰减0.9999。除了与$w$-嵌入相关的参数外,我们使用教师模型的参数初始化学生模型。
B.3 确定性采样器的第二阶段蒸馏
我们使用与第一阶段(见附录B.2)相同的模型架构。我们使用算法2训练蒸馏模型。我们首先将第一阶段的学生模型作为教师模型。从1024步DDIM采样开始,逐步将第一阶段的学生模型蒸馏为单步模型。在将采样步数减半并将学生模型作为新的教师模型之前,我们对学生模型进行50,000次参数更新,对于采样步数等于1或2的情况,则进行100,000次参数更新。在每个采样步骤,我们使用教师模型的参数初始化学生模型。我们使用SNR截断损失[10,33]训练模型。在每个步骤中,我们在参数更新过程中从$1e-4$到0线性退火学习率。训练时不使用EMA衰减,我们的训练设置与[33]非常接近。
B.4 随机采样的第二阶段蒸馏
我们使用算法3训练蒸馏模型。对于ImageNet 64x64和CIFAR-10,我们使用与附录B.3中确定性采样的第二阶段蒸馏相同的模型架构和训练设置:这里的主要区别在于,我们的蒸馏目标对应于采用比确定性采样器大两倍的采样步长。我们在图15中提供了不同引导强度$w$的样本可视化。
B.5 基线样本
我们在图16和图17中提供了DDIM基线的额外样本。
B.6 额外蒸馏结果
我们在图22b、图22a和表4中提供了我们的方法和基线在ImageNet 64x64和CIFAR-10上的FID和IS结果。我们还在图18和图19中可视化了两个数据集的FID和IS权衡曲线,其中我们为ImageNet 64x64选择引导强度$w = \{0, 0.3, 1, 2, 4\}$,为CIFAR-10选择$w = \{0, 0.1, 0.2, 0.3, 0.5, 0.7, 1, 2, 4\}$。
B.7 风格迁移
我们在ImageNet 64x64上进行此实验。如[41]所述,通过使用在域A上训练的扩散模型对图像进行编码(执行反向DDIM),然后使用在域B上训练的扩散模型通过DDIM进行解码,可以在域A和域B之间进行风格迁移。我们使用算法4训练模型,采用与附录B.3中讨论的相同的$w$-条件模型架构和训练设置。
对无分类器引导模型进行渐进蒸馏[33]的一种自然方法是使用与教师结构相同的蒸馏学生模型——即具有联合训练的蒸馏条件和无条件扩散组件。记预训练教师模型为$[\hat{x}_{c,\theta}, \hat{x}_\theta]$,学生模型为$[\hat{x}_{c,\eta}, \hat{x}_\eta]$,我们在算法5中提供了训练算法。要从训练好的模型中采样,可以使用DDIM确定性采样器[38]或提出的随机采样器。我们遵循附录B.3中的训练设置,使用$w$-条件模型,并训练模型以适应引导强度$[0, 4]$。我们观察到,使用算法5蒸馏的模型在采样步数较少时无法生成合理的样本。我们在图23中提供了使用DDIM采样器在CIFAR-10上生成的样本,并在表5中提供了FID/IS分数。
B.8 朴素蒸馏方法
对无分类器引导模型进行渐进蒸馏 [33] 的一种自然方法是使用与教师结构相同的蒸馏学生模型 —— 即具有联合训练的蒸馏条件和无条件扩散组件。记预训练教师模型为$[\hat{x}_{c,\theta}, \hat{x}_\theta]$,学生模型为$[\hat{x}_{c,\eta}, \hat{x}_\eta]$,我们在算法 5 中提供了训练算法。要从训练好的模型中采样,可以使用 DDIM 确定性采样器 [38] 或提出的随机采样器。我们遵循附录 B.3 中的训练设置,使用w- 条件模型,并训练模型以适应引导强度$[0, 4]$。我们观察到,使用算法 5 蒸馏的模型在采样步数较少时无法生成合理的样本。我们在图 23 中提供了使用 DDIM 采样器在 CIFAR-10 上生成的样本,并在表 5 中提供了 FID/IS 分数。
C. 隐空间蒸馏
C.1 类别条件生成
C.1.1 训练细节
在本实验中,我们考虑ImageNet 256×256上的类别条件生成任务。我们首先将原始的ϵ-预测模型微调为v-预测模型,然后从具有512步采样的DDIM教师模型开始,使用其输出作为目标来训练我们的蒸馏模型。在第一阶段,我们使用常数损失[10,33]对模型进行2000次梯度更新训练。在第二阶段,除了采样步数为1、2或4时需要20000次梯度更新外,我们对模型进行2000次梯度更新训练。第二阶段模型使用SNR截断损失[10,33]进行训练。在两个阶段中,我们都进行了额外的500步学习率热身,其中学习率从0线性增加到目标学习率。训练期间使用批量大小为2048,并均匀采样引导强度$w \in [w_{\text{min}}=0, w_{\text{max}}=14]$。
补充结果:我们在图25中提供了通过精度和召回率评估的定量结果。这些结果证实了我们的方法在少步数场景下的显著性能提升,尤其是在1-4步采样时。我们的2步和4步采样蒸馏隐扩散模型在精度上接近32步DDIM的性能,在低步数时的召回率显著优于DDIM。更多定性结果见图25,其中我们展示了1步和2步模型的随机样本,并与DDIM采样进行了对比。
C.2 文本引导的图像生成
C.2.1 训练细节
在本实验中,我们考虑分辨率为256×256和512×512的LAION-5B数据集。
LAION-5B 256×256:与附录C.1类似,我们首先将原始的ϵ-预测模型微调为v-预测模型。我们从具有512步采样的DDIM教师模型开始,使用其输出作为目标来训练蒸馏模型。在第一阶段,我们使用常数损失[10,33]对模型进行2000-5000次梯度更新训练。在第二阶段,除了采样步数为1、2或4时需要10000-50000次梯度更新外,我们对模型进行2000-5000次梯度更新训练。第二阶段模型使用SNR截断损失[10,33]进行训练。在两个阶段中,我们都进行了额外的100-1000步学习率热身。训练期间使用批量大小为1024,并均匀采样引导强度$w \in [w_{\text{min}}=2, w_{\text{max}}=14]$。
图26提供了上述不同训练设置的收敛性分析。我们观察到,我们的方法在几千次训练迭代后接近基础模型的DDIM采样性能,并在1步和2步场景中迅速超越。然而,为了达到最佳性能,需要更长时间的训练。
LAION-5B 512×512:类似地,我们首先将原始的ϵ-预测模型微调为v-预测模型。我们从具有512步采样的DDIM教师模型开始,使用其输出作为目标来训练蒸馏模型。在第一阶段,我们使用常数损失[10,33]对模型进行2000-5000次梯度更新训练。在第二阶段,除了采样步数为1、2或4时需要10000-50000次梯度更新外,我们对模型进行2000-5000次梯度更新训练。第二阶段模型使用SNR截断损失[10,33]进行训练。在两个阶段中,我们都进行了额外的1000步学习率热身。训练期间使用批量大小为512,并均匀采样引导强度$w \in [w_{\text{min}}=2, w_{\text{max}}=14]$。
补充结果:除了DDIM,我们还在此处将我们的方法与DPM++-Solver[16,18]进行了比较,这是一种先进的采样器,无需额外训练,且在隐扩散模型的≥10步采样中取得了良好结果。与我们的蒸馏模型不同,该方法与DDIM类似,必须使用无分类器引导才能获得良好结果,这使得U-Net评估次数比我们的w-条件化方法增加了一倍。
我们在图28中提供了这些采样方法的定性比较,从中可以清楚地看到我们的蒸馏方法在少步数时的优势:我们的方法比无需训练的采样器生成的结果更清晰、更连贯。这种表现反映在图27和表6、表7的定量FID和CLIP分析中。尽管此处的加速效果不如像素空间显著,但我们的方法在2步或4步采样时仍能取得非常好的结果。由于w-条件化(因为此处我们不再需要为无分类器引导评估无条件模型和条件模型,只需要一个蒸馏的w-条件模型),我们的方法与现有方法相比,进一步将峰值内存或去噪步骤减少了一半。我们希望我们的工作将推动通用高分辨率文本到图像系统在实时应用中的进展。
C.3 文本引导的图像到图像翻译
C.3.1 训练细节
我们使用为文本引导图像生成训练的模型,训练细节见附录C.2。
C.3.2 额外分析
我们在图30和图31中提供了样本质量、可控性和效率之间权衡的更多分析。与[20]类似,我们还观察到,随着初始扰动噪声水平的增加,真实性、可控性和忠实性之间存在权衡:添加的噪声越多,图像与文本提示的对齐越好,但对输入图像的忠实度越低(见图30和图31)。
C.4 图像修复
C.4.1 训练细节
与我们之前的实验类似,我们使用LAMA[42]建议的大掩码生成方案将ϵ-预测模型微调为v-预测模型,并在512×512分辨率的LAION-5B上进行训练。我们从具有512步采样的DDIM教师模型开始,使用其输出作为目标来训练蒸馏模型。在第一阶段,我们使用常数损失[10,33]对模型进行2000次梯度更新训练。在第二阶段,除了采样步数为1或2时需要5000次梯度更新外,我们对模型进行10000次梯度更新训练。第二阶段模型使用SNR截断损失[10,33]进行训练。在两个阶段中,我们都进行了额外的1000步学习率热身。训练期间使用批量大小为512,并均匀采样引导强度$w \in [w_{\text{min}}=2, w_{\text{max}}=14]$。
补充评估结果:与低采样步数的DDIM采样的定量比较见Tab.8,额外样本见图32。
D. 像素空间蒸馏的额外样本
在本节中,我们为像素空间蒸馏模型提供了额外的样本。我们使用确定性采样器(见算法2)和随机采样器(见算法3)生成样本。