[学习笔记]Diffusion Probabilistic Model Made Slim

Diffusion Probabilistic Model Made Slim

原文链接：原文链接

“本作品仅供学习交流使用，严禁用于商业用途。如果保存，请在24小时内删除。作者不承担任何责任。”

尽管最近取得了视觉上令人愉悦的结果，但扩散概率模型（DPM）的计算成本一直是其长期存在的缺陷，这反过来又大大限制了其在资源有限的平台上的应用。然而，之前对高效DPM的方法大多集中在加速测试，而忽视了它们的巨大复杂性和规模。在这篇论文中，我们致力于减轻DPM的负担，同时努力保持其有利的性能。**我们首先从头开始训练一个小型的潜在扩散模型（LDM），但观察到合成图像的显著失真。通过全面评估，我们发现DPM本质上对高频生成有偏见，并学习在不同的时间步长恢复不同的频率成分。**这些特性使得紧凑的网络无法准确估计高频的频率动态。为此，我们引入了一个定制的设计，用于轻量级的图像合成，我们称之为谱扩散（SD）。SD在其架构中引入了小波门控，以在每个反向步骤中实现频率动态特征提取，并通过谱感知蒸馏来促进高频恢复，通过根据谱幅度反向加权目标来实现。实验结果表明，SD在一系列条件和无条件图像生成任务上实现了与潜在扩散模型相比8-18倍的计算复杂度降低，同时保持了具有竞争力的图像保真度。

然而，现有的高效DPM的努力集中在模型加速上，但在很大程度上忽略了模型的轻量级。一些方法力求更快的采样，而其他则依赖于减少输入大小。诚然，所有这些方法都缩短了训练或推理时间，但由于规模过大，它们无法在许多现实世界的应用场景中使用。

在本文中，我们致力于建立紧凑的DPM。首先，我们通过减少通道大小来训练流行的潜在扩散模型（LDM）[48]的精简版本。我们在图1中示出了由原始DPM和精简DPM生成的图像。虽然精简LDM勾画了面部的整体结构，但不幸的是，诸如皮肤和头发纹理的高频分量恢复得很差。这种现象实际上可以通过右列所示的离散傅里叶变换（DFT）系数来揭示，这表明当模型变得苗条时，用于DPM的常规设计导致高频不足。

我们在图1中显示了原始和精简DPM生成的图像。虽然lite LDM绘制了人脸的整体结构，但不幸的是，高频分量，如皮肤和头发纹理，恢复得很差。这一现象实际上可以从离散傅立叶变换(DFT)系数中揭示出来，这表明传统的DPM设计在模型变细时会导致高频不足。

然后，我们通过频率的视角对DPM进行了深入分析，得到了两个关键观察结果。

1. 频率演化。在一些假设下，我们通过数学证明了DPM在去噪过程的不同阶段学习了不同的功能。具体来说，我们展示了最优去噪器实际上可以归结为一个具有增长带宽的wiener滤波器的级联。在恢复低频成分后，高频特征在后期去噪阶段逐渐添加。因此，这种进化特性使得小型DPM无法学习具有有限参数的动态带宽。
2. 频率偏见。DPM对数据分布的主导频率成分有偏见。当噪声幅度小时，这种偏见最为明显，导致反向过程结束时噪声预测不准确。因此，小型DPM难以恢复高频带和图像细节。

受这些观察结果的启发，我们提出了一种新的光谱扩散（SD）模型，专门用于轻量级图像合成。我们的核心思想是将频率动力学和先验引入小型DPM的架构设计和训练目标中，以明确地保留高频细节。拟议的解决方案由两部分组成，每个部分都说明了上述一项意见。对于频率演化，我们提出了一种小波门控操作，使网络能够动态适应不同时间步长的频谱响应。在上采样和下采样阶段，首先通过小波变换对输入特征进行分解，并通过可学习的门控函数对系数进行重新加权。它显著降低了在反向过程中表示频率演变的参数要求。

我们建立了一个名为SD的轻量级潜在扩散模型，它在很大程度上保留了LDM的性能。值得注意的是，SD继承了DPM的优点，包括优越的样本多样性、训练稳定性和可处理的参数化。我们的模型比原模型小8∼18倍，运行2∼快5倍，同时获得了具有竞争力的图像保真度。

本研究的贡献有三方面：

1. 本研究调查了扩散模型瘦身的任务，这是以前很少探索的。
2. 我们发现关键的挑战在于它对高频成分的不现实恢复。通过从频率的角度探测DPM，我们表明在不同的去噪步骤中存在频谱演化，而稀有的频率不能被小型模型准确估计。
3. 我们提出了SD，一个有效地通过增强高频生成性能来恢复图像纹理的轻量级DPM。SD在低成本下在图像生成任务上取得了令人满意的表现。

Related Work

• Diffusion Probabilistic Models

• Efficient Diffusion

  有效扩散

• Knowledge Transfer and Distillation

  知识转移与提炼

• Frequency Analysis for Generative Model

  生成模型的频率分析

在本文中，我们研究了DPM的频率行为。利用其频率特性，我们的SD以低成本实现了逼真的图像生成。

Background

• Denoising Diffusion Probabilistic Models
• Frequency Domain Representation of Images

频域分析根据一组基函数分解图像。我们关注两个离散变换：傅立叶和小波变换。

Frequency Perspective for Diffusion

在一般的信号处理中，去噪通常在频率空间中进行。表1比较了FFHQ数据集上不同DPM的低频和高频误差2。Lite LDM由于缺乏高频生成而表现不佳。

Spectrum Evolution over Time

DPM优化时间条件网络以适应多个尺度的噪声，这产生了随时间步长的去噪轨迹。我们从频率的角度仔细研究了这一轨迹。当假设网络是线性滤波器时，我们根据其在每个时间步长的频谱响应给出最优滤波器。这种滤波器通常被称为维纳滤波器。

命题1。假设$x_0$是一个广义平稳信号，而$ϵ$是方差${\sigma }_2=1$的白噪声。对于$x_t=\sqrt{\bar{\alpha }}{x}_0+\sqrt{1-\bar{\alpha }}ϵ$，在时间$t$最小化$J_t={\Vert h_t\ast x_t-ϵ\Vert }^2$的最优线性去噪滤波器ht具有闭式解

其中$|{\mathcal{X}}_0(f){|}^2$是$X_0$的功率谱，${\mathcal{H}}_t^{\ast }（f）$是H*t的频率响应。

尽管线性假设对模型架构构成了强烈的限制，但我们相信它为如何执行反向过程提供了有价值的见解。DPM从结构到细节。在这项研究中，我们对自然图像的功率谱做出了一个被广泛接受的假设，即遵循幂律$\mathbb{E}[|X_0（f）{|}^2]=A_s（\theta ）/f^{\alpha S(\theta )}$。$A_s(\theta )$被称为振幅比例因子，$\alpha S(\theta )$频率指数。如果我们设置$A_s(\theta )=1$和$\alpha S(\theta )=2$，则信号重构滤波器$1-\sqrt{1-\bar{\alpha }}h$的频率响应如图所示。

在相反的过程中，t从T→0开始，而α从0→1开始增加。因此，随着时间的推移，DPM表现出一种光谱变化的行为。在开始时，我们有一个窄带滤波器(α=0.1andα=0.01)，它只恢复控制粗糙结构的低频分量。T下降，α逐渐增加，图像中恢复了更多细节和高频成分，如人的头发、皱纹和毛孔

我们在图2中使用预训练的LDM[48]绘制了不同步骤的去噪预测的图，这表明DPM首先生成低频并转变为高频。[6，18，39，48]中显示了DPM从粗略到细节的相同经验观察结果，而我们是第一个给出其数值解的人。

Frequency Bias in Diffusion Model

基于扩散的模型中的另一个挑战是低密度区域中不准确的噪声估计。它产生于损失函数中对$p(x_0)$的期望。通过$p(x_0)$对去噪目标进行加权可能会在训练的扩散模型中引入偏差，导致其优先考虑高密度区域，而忽略罕见模式。

图像生成任务中长尾模式的一个例子是频率偏移，其中高频分量很少。因此，基于有偏分布训练基于扩散的小模型可能会使生成这种高频模式变得具有挑战性，因为该模型往往过于强调低频图像。这个问题会严重影响生成图像的质量。

Spectral Diffusion Model

如上所述，我们的目标是通过将频率动态和先验纳入架构设计和训练目标来减少DPM的大小。我们从LDM作为基线开始，然后设计了一个小波门控模块，该模块能够对轻量级模型进行时间动态推理。应用频谱感知蒸馏来提高高频生成性能。这两种修改都使我们能够以最小的模型大小和计算工作量实现照片逼真的图像生成。

Dynamic Wavelet Gating 动态小波门控

如第4.1节所述，反向过程需要具有动态频率响应的级联滤波器。虽然Vanilla UNet[50]在重建图像细节方面是有效的，但它不能将动态频谱合并到一组参数中。因此，小型DPM无法补偿不断变化的带宽。

为了解决这个问题，我们建议在网络中插入小波门控（WG）模块，使其自动适应变化的频率响应。WG将特征图分解为小波带，并在不同的反向步骤中选择性地关注适当的频率，这是为扩散模型唯一定制的。

对小波系数的门控。我们用DWT和IDWT[12，69]替换UNet中的所有下采样和上采样，并对小波系数进行软门控操作，以促进阶跃自适应图像去噪。我们称它们为WGDown和WGUp，如图5所示。

类似于通道注意力，来自输入特征**$X$**的信息被聚合以产生软门控
g { L L ， L H ， H L ， H H } ＝ S i g m o i d ( F F N ( A v g p o o l ( X ) ) ) g_{\{LL，LH，HL，HH\}}＝Sigmoid(FFN(Avgpool(X))) g{LL，LH，HL，HH}​＝Sigmoid(FFN(Avgpool(X)))
其中$g_i$是每个小波带的门控分数；$FFN$是一个2层前馈网络，$Avgpool$代表平均池。然后用$g_i$对系数进行门控以产生输出$X\prime$。

在WG-Down中，我们在DWT运算后应用WG来融合具有加权求和$X\prime =\sum {}_{i\in ｛LL，LH，HL，HH｝}gi\odot Xi$的子带系数，其中$\odot$是全元素乘法。在WG-Up中，输入特征被分割成4个块作为小波系数。然后，对$X\prime =IDWT（gLL\odot XLL，gLH\odot XLH，gHL\odot XHL，gHH\odot XHH）$之前的每个子带进行WG重新加权。在本文中，我们默认使用Haar小波。

Spectrum-Aware Knowledge Distillation

扩散模型在高频组件建模方面遇到了挑战（见第4.2节），尤其是在高效需求方面。在与图像生成中的频谱不足作斗争时，我们提取了一个大型预先训练的教师模型对一个紧凑的WG-Unet学生的预测。除了空间输出匹配之外，我们还应用SpectrumAware蒸馏来指导学生模型合成自然图像细节。我们的方法包括根据光谱大小重新加权蒸馏损失。我们增加了低幅度分量（如高频带）的误差惩罚，同时减少了低频元素的权重。

给定一个教师扩散模型$s_T(\cdot ；{\theta }_T)$，我们希望通过模拟输出和特征来提取学生的$s_T(\cdot ；{\theta }_T)$。在时间步长$t$，扰动图像$x_t$被馈送到两个网络中以产生输出和特征。用L2损失量化它们的空间距离

其中$X_T^{(i)}$和$X_S^{(i)}$代表一对教师/学生的输出特征或相同规模的输出。单个1×1 CONV层用于对齐尺寸。

除了空间蒸馏，我们还从不平衡学习中获得灵感，设计了一种蒸馏损失，以促进少数频率的恢复。所提出的方法包括采用一对模型预测和一个干净的图像$x_0$，将$x_0$插值到与特征图相同的大小，然后计算它们的二维DFT

然后应用$X_0$来调制$X_T^{(i)}$和$X_S^{(i)}$之间的差

在比例因子α＜0（在我们的实验中，α＝−1）的情况下，${\mathcal{L}}_{freq}$促使学生学习少数频率，但对大多数分量进行加权。与方程3中的DDPM目标一起，我们的训练目标变为$\mathcal{L}={\mathcal{L}}_{DDPM}+{\lambda }_s{\mathcal{L}}_{spatial}+{\lambda }_f{\mathcal{L}}_{freq}$，加权因子${\lambda }_s$=0.1和${\lambda }_f$=0.1。

注意，我们的方法的目的是在每个去噪步骤学习准确的分数预测，这与现有的关于采样步长减少的蒸馏[40，52]是正交的。

Experiments

数据集和评估。我们在4个无条件和2个条件Benckmark上评估了我们的模型。特别是，我们在LSUNChurches/Bedrooms[75]、FFHQ[27]和CelebAHQ[26]上训练无条件SD模型。此外，模型的性能也在类条件ImageNet[7]和MS-COCO[31]文本到图像生成的背景下进行了评估。对于文本图像任务，我们首先在LAION-400M[53]上进行训练，并直接在MS-COCO上进行测试。

培训和评估详细信息。我们在LDM[48]框架3上构建我们的模型。为了进行公平比较4，我们实现了LDM的精简版本，通道维度为64作为我们的基线模型。我们称之为Lite LDM。

我们提出的SD是在4个无条件基准上训练的，持续时间为150k次迭代，使用512或256的小批量大小。我们使用AdamW[36]优化器，初始学习率为1.024×10−3，线性学习率衰减。对于类和文本条件生成，我们将初始学习率设置为5.12×10−4，同时保持其他参数不变。基于FID评分[15]，使用分辨率为256的50k个生成样本来评估合成图像质量。默认情况下，我们使用200步DDIM[56]采样。我们根据参数计数和乘加累积（MAC）来比较模型大小和计算成本，并报告速度比较的吞吐量。所有实验均在8x NVIDIA特斯拉V100 GPU上进行。其他详细信息可在补充材料中找到。