【论文阅读】Elucidating the Design Space of Diffusion-Based Generative Models

Elucidating the Design Space of Diffusion-Based Generative Models

引用： Karras T, Aittala M, Aila T, et al. Elucidating the design space of diffusion-based generative models[J]. Advances in Neural Information Processing Systems, 2022, 35: 26565-26577.

论文链接： （NIPS 2022）https://proceedings.neurips.cc/paper_files/paper/2022/hash/a98846e9d9cc01cfb87eb694d946ce6b-Abstract-Conference.html

代码链接： https://github.com/NVlabs/edm

简介

作者认为，基于扩散的生成模型的理论和实践目前过于复杂，并通过提出一个清晰分离具体设计选择的设计空间来解决这个问题。这使我们能够识别出采样过程和训练过程中的几项改进，以及对分数网络的预处理。我们的改进得到了新的最先进FID（1.79）在CIFAR-10的类条件设置中，以及在无条件设置中的1.97，采样速度比以前的设计快得多（每张图像35次网络评估），将一个预先训练的ImageNet-64模型的FID从2.07提高到接近最先进的1.55，并且在经过改进后重新训练达到了新的最先进水平1.36。

1. 从实践的角度来研究这些模型背后的理论，更多地关注在训练和采样阶段出现的“有形”对象和算法，而不是可能从中得出它们的统计过程。目标是更好地了解这些组件是如何连接在一起的，以及在整个系统的设计中有哪些自由度。专注于一大类模型，其中神经网络用于模拟被高斯噪声破坏的训练数据的噪声水平相关边际分布的分数。因此，该工作是在去噪分数匹配的背景下进行的。

2. 针对扩散模型合成图像的采样过程，确定了采样的最佳时间离散化，在采样过程中应用了高阶 Runge-Kutta 方法，评估了不同的采样器计划，并分析了随机性在采样过程中的有用性。改进后，合成过程中所需的采样步骤数量显著减少，改进后的采样器可以作为几种广泛使用的扩散模型的直接替代品。

3. 侧重于评分建模神经网络的训练。虽然继续依赖常用的网络架构（DDPM [1]，NCSN [2]），但在扩散模型设置中对网络的输入、输出和损失函数的预处理进行了第一个原则性分析，并得出了改进训练动态的最佳实践。在训练期间改善噪声水平的分布，并且通常与GAN一起使用的无泄漏增强[3]也有利于扩散模型。

Method

用标准差 ${\sigma }_{data}$ 表示 $p_{data}(x)$ 的数据分布，并考虑通过添加如标准差为$\sigma$的高斯噪声数据得到的软化分布 $p（x;\sigma ）$。。对于${\sigma }_{max}\gg {\sigma }_{data}$，$p(x;{\sigma }_{max})$与纯高斯噪声几乎没有区别。扩散模型的思想是随机采样噪声图像 $x_0\sim N(0，{\sigma }_{max}^{2}I)$，并依次将其降噪为噪声水平 ${\sigma }_0={\sigma }_{max}>{\sigma }_1>\cdot \cdot \cdot >{\sigma }_N=0$ 的图像，因此在每个噪声水平下 $x_i\sim p(x_i;{\sigma }_i)$。因此，该过程的终点 $x_N$ 根据数据进行分配。Song等[4]提出了一个随机微分方程（SDE），该方程在样本$x$随时间变化时保持所需的分布$p$。这允许使用随机求解器实现上述过程，该求解器在每次迭代时都会消除和添加噪声。它们还给出了相应的“概率流”常微分方程 （ODE），其中随机性的唯一来源是初始噪声图像 $x_0$。

概率流ODE在时间上分别向前或向后移动时不断增加或减少图像的噪声水平。要指定常微分方程，必须首先选择一个时间表 $\sigma (t)$，该时间表定义了时间 t 所需的噪声水平。然而，时间表的选择具有重大的实际意义，不应基于理论上的方便而作出。

表 1 给出了在论文框架中重现三种早期方法的确定性变体的公式。之所以选择这些方法，是因为它们被广泛使用并实现了最先进的性能，但也因为它们来自不同的理论基础。一些公式看起来与原始论文有很大不同，因为间接和递归已被删除。这种重构的主要目的是揭示所有在以前的工作中经常纠缠在一起的独立组件。在框架中，组件之间没有隐含的依赖关系——原则上，对各个公式的任何选择（在合理范围内）都会导致一个有效的模型。换言之，更改一个组件并不需要在其他地方进行更改，例如，保持模型收敛到极限中数据的属性。在实践中，一些选择和组合当然会比其他选择和组合效果更好。

Improvements to deterministic sampling

论文假设，与采样过程相关的选择在很大程度上独立于其他组件，例如网络架构和训练细节。换句话说，$D_{\theta }$ 的训练过程不应规定 $\sigma (t)$、$s(t)$ 和 $t_i$，反之亦然;从采样器的角度来看，$D_{\theta }$只是一个黑匣子。通过在三个预训练模型上评估不同的采样器来测试这一点，每个模型代表不同的理论框架和模型系列。首先使用这些模型的原始采样器实现来测量这些模型的基线结果，然后使用表 1 中的公式将这些采样器引入统一框架，然后进行改进。这能够评估不同的实际选择，并提出适用于所有模型的抽样过程的一般改进。

论文评估了在32×32的无条件CIFAR-10上训练的“DDPM++ cont. (VP)”和“NCSN++ cont. (VE)”模型，对应于方差保持（VP）和方差爆炸（VE）公式，最初受到DDPM和SMLD的启发。还评估了在 64×64 的类条件 ImageNet 上训练的“ADM(dropout)”模型，对应于改进的 DDPM(iDDPM) 公式。该模型使用一组离散的 M = 1000 噪声水平进行训练。

根据 50,000 张生成图像和所有可用真实图像之间计算的 Fréchet 初始距离 （FID） 评估结果质量。图 2 显示了 FID 作为神经功能评估 （NFE） 的函数，即评估 $D_{\theta }$ 的次数以产生单个图像。鉴于采样过程完全由 $D_{\theta }$ 的成本主导，NFE 的改进直接转化为采样速度。原始的确定性采样器以蓝色显示，在统一框架（橙色）中重新实现这些方法会产生类似但始终更好的结果。这些差异的原因是原始实现中的某些疏忽，以及在DDIM的情况下对离散噪声水平的更仔细处理;注意，重新实现完全由算法 1 和表 1 指定，即使原始代码库的结构彼此非常不同。

Discretization and higher-order integrators：以数值方式求解常微分方程必然是遵循真实解轨迹的近似值。在每个步骤中，求解器都会引入截断误差，该误差会在 N 个步骤的过程中累积。局部误差通常相对于步长呈超线性缩放，因此增加 N 可以提高解的精度。常用的欧拉方法是一阶常微分方程求解器，其相对于步长 h 具有 $O(h^2)$ 局部误差。 高阶 Runge-Kutta 方法的扩展性更有利，但每步需要对 $D_{\theta }$ 进行多次评估。最近还提出了用于采样扩散模型的线性多步方法。通过广泛的测试，发现Heun的二阶方法（又名改进的欧拉，梯形规则）——之前由Jolicoeur-Martineau等人在扩散模型的背景下探索过——在截断误差和NFE之间提供了一个很好的权衡。如算法 1 所示，它为 $t_{i+1}$ 引入了一个额外的校正步骤，以考虑 $t_i$ 和 $t_{i+1}$ 之间 $dx/dt$ 的变化。这种校正导致 $O(h^3)$ 局部误差，代价是每一步额外评估一次 $D_{\theta }$ 。

时间步长 $t_i$ 决定了步长和截断误差在不同噪声水平之间的分布方式，实验结论是步长应该随着σ的减小而单调减小，并且不需要根据每个样本而变化。采用参数化方案，其中时间步长根据噪声水平 ${\sigma }_i$ 的序列定义，例如 $t_i={\sigma }^{-1}({\sigma }_i)$。设置 ${\sigma }_{i<N}=(Ai+B{)}^{\rho }$ 并选择常数 A 和 B，使 ${\sigma }_0={\sigma }_{max}$ 和 ${\sigma }_{N-1}={\sigma }_{min}$，即

其中， $\rho$ 控制 ${\sigma }_{min}$ 附近的步长缩短了多少，而 ${\sigma }_{max}$ 附近的步长却缩短了多少。分析表明，设置 $\rho =3$ 几乎可以平衡每个步骤的截断误差，但 5 到 10 范围内的 $\rho$ 在采样图像中表现更好。这表明${\sigma }_{min}$附近的误差有很大的影响。将本文的其余部分设置为 $\rho =7$。

Trajectory curvature and noise schedule：常微分方程解轨迹的形状由函数 $\sigma (t)$ 和 $s(t)$ 定义。选择这些函数提供了一种减少上述截断误差的方法，因为它们的大小可以预期与 $dx/dt$ 的曲率成正比。这些函数的最佳选择是 $\sigma (t)=t$ 和 $s(t)=1$，这也是 DDIM 中的选择。一个直接的结果是，在任何 $x$ 和 $t$ 处，到 $t=0$ 的单个欧拉步长会产生去噪图像 $D_{\theta }(x;t)$。因此，解轨迹的正切始终指向降噪器输出。预计这只会随着噪声水平而缓慢变化，这在很大程度上对应于线性求解轨迹。图3c的一维常微分方程草图支持这种直觉;求解轨迹在大噪声和小噪声水平下都接近线性，并且仅在两者之间的一个小区域内具有较大的曲率。在图1b中的真实数据中也可以看到同样的效果，其中不同去噪器目标之间的变化发生在相对较窄的$\sigma$范围内。对于建议的时间表，这对应于高ODE曲率被限制在相同的范围内。设置 $\sigma (t)=t$ 和 $s(t)=1$ 的效果如图 2 中的红色曲线所示。由于 DDIM 已采用这些相同的选择，因此红色曲线与 ImageNet-64 的绿色曲线相同。然而，VP 和 VE 从原来的时间表中获益匪浅。

Stochastic sampling

确定性采样具有许多优点，例如，能够通过反转常微分方程将真实图像转换为相应的潜在表示。然而，与随机采样相比，它往往会导致更差的输出质量，随机采样在每个步骤中都会向图像注入新噪声。鉴于常微分方程和标准差别在理论上恢复相同的分布，那么随机性的作用究竟是什么？

Song等的SDE可以推广为方程1的概率流常微分方程和时变Langevin扩散SDE的总和：

其中 ${\omega }_t$ 是标准的 Wiener 过程。$d{x}_{+}$ 和 $d{x}_{-}$ 现在是独立的 SDE，用于在时间上向前和向后移动，与 Anderson 的时间反转公式相关。Langevin 项可以进一步看作是确定性基于分数的去噪项和随机噪声注入项的组合，其净噪声水平贡献抵消。因此，$\beta （t）$有效地表示了现有噪声被新噪声取代的相对速率。

这一观点揭示了为什么随机性在实践中是有帮助的：隐含的Langevin扩散在给定时间将样本推向所需的边缘分布，积极纠正早期采样步骤中的任何错误。另一方面，近似Langevin项。这一观点揭示了随机性在实践中有帮助的原因：隐含的Langevin扩散在给定时间将样本推向所需的边缘分布，积极纠正早期采样步骤中的任何错误。另一方面，用离散SDE求解器步骤近似Langevin项本身会引入误差。先前的结果表明非零β（t）是有帮助的，但据我们所知，Song等人中对β（t）的隐式选择没有特殊性质。因此，随机性的最佳量应该根据经验来确定。

论文提出了一个随机采样器，它将二阶确定性 ODE 积分器与显式 Langevin 式的添加和去除噪声的“搅动”相结合。算法 2 中给出了一个伪代码，这不是一个通用的SDE求解器，而是为特定问题量身定制的采样程序。它的正确性源于两个子步骤的交替，每个子步骤都保持正确的分布（直到 ODE 步骤中的截断误差）。增加随机性可以有效地纠正早期采样步骤产生的误差，但它也有其自身的缺点。过度的 Langevin 样噪声添加和去除会导致所有数据集和降噪网络生成的图像中的细节逐渐丢失。在非常低和高的噪点水平下，颜色也会偏向过度饱和。我们怀疑实际降噪器在方程 3 中诱导出一个略微非保守的向量场，违反了 Langevin 扩散的前提并导致这些有害影响。值得注意的是，用分析降噪器进行的实验（如图1b）没有显示出这种降解。

如果退化是由 $D_{\theta }(x;\sigma )$ 中的缺陷引起的，则只能在采样期间使用启发式方法进行补救。通过仅在特定范围内的噪声水平 $t_i\in [S_{tmin}，S_{tmax}]$ 范围内启用随机性来解决向过度饱和的颜色漂移。对于这些噪声水平，定义 ${\gamma }_i=S_{churn}/N$，其中 $S_{churn}$ 控制随机性的总量。我们进一步钳制${\gamma }_i$，使其永远不会引入比图像中已经存在的噪声更多的新噪声。最后发现，通过将 $S_{noise}$ 设置为略高于 1 以夸大新添加噪声的标准偏差，可以部分抵消细节损失。这表明，假设的$D_{\theta }(x;\sigma )$非保守性的一个主要组成部分是去除略微过多的噪声的倾向——这很可能是由于向任何L2训练的降噪器都可以预期的均值回归。

Preconditioning and training

训练神经网络直接对 D 进行建模远非理想——例如，由于输入 $x=y+n$ 是干净信号 $y$ 和噪声 $n\sim N(0，{\sigma }^{2}I)$的组合，其幅度因噪声水平 $\sigma$ 而变化很大。出于这个原因，通常的做法是不直接将 $ 表示为神经网络，而是训练一个不同的网络 $F_{\theta }$，从中派生 $D_{\theta }$。

以前的方法通过σ相关归一化因子来解决输入缩放问题，并尝试通过训练$F_{\theta }$来预测$n$缩放到单位方差，然后通过$D_{\theta }(x;\sigma )=x-\sigma F_{\theta }(\cdot )$对输出进行预处理。这样做的缺点是，对大$\sigma$，网络需要仔细微调其输出，以精确地抵消现有的噪声 $n$，并以正确的比例提供输出;但网络所犯的任何错误都会放大 $\sigma$ 倍。在这种情况下，直接预测预期输出 $D_{\theta }(x;\sigma )$ 似乎要容易得多。本着与之前自适应混合信号和噪声的参数化相同的精神，建议使用与σ相关的跳跃连接来预处理神经网络，使其能够估计y或n，或介于两者之间。因此，用以下形式写出 $D_{\theta }$：

其中 $F_{\theta }$ 是要训练的神经网络，$c_{skip}(\sigma )$ 调制跳跃连接，$c_{in}(\sigma )$ 和 $c_{out}(\sigma )$ 缩放输入和输出幅度，$c_{noise}(\sigma )$ 将噪声水平$\sigma$映射到 $F_{\theta }$ 的调节输入中。可以等效地表示损失相对于原始网络输出 $F_{\theta }$：

该形式揭示了$F_{\theta }$的有效训练目标，使我们能够从第一性原理确定预处理函数的合适选择。通过要求网络输入和训练目标具有单位方差 $(c_{in}，c_{out})$ 并尽可能少地放大 $F_{\theta }$ 中的误差 ($c_{skip}$) 来推导出表 1 所示的选择。$c_{noise}$的公式是根据经验选择的。

表 2 显示了一系列训练设置的 FID，使用确定性采样器进行评估。从Song等人的基线训练设置开始，VP和VE之间差异很大;为每个（配置 A）提供单独的结果。为了获得更有意义的比较点，重新调整了基本超参数（配置 B）并通过删除最低分辨率层并将最高分辨率层的容量加倍来提高模型的表现力（配置 C）。然后，我们将 {$c_{in}，c_{out}，c_{noise}，c_{skip}$} 的原始选择替换为论文提出的预处理（配置 D），这将使结果基本保持不变——除了 VE 在 64×64 分辨率下显着改善。预处理的主要好处不是改善FID本身，而是使训练更加稳健，能够将注意力转向重新设计损失函数，而不会产生不利影响。

为了平衡有效损失权重，设置 $\lambda (\sigma )=1/c_{out}(\sigma {)}^2$，这也均衡了整个σ范围内的初始训练损失，如图 5a（绿色曲线）所示。最后，需要选择 $p_{train}(\sigma )$，即在训练过程中如何选择噪声水平。检查训练后的每σ损耗（蓝色和橙色曲线）表明，只有在中等噪声水平下才有可能显着降低;在非常低的水平上，辨别出消失的小噪声分量既困难又无关紧要，而在高水平上，训练目标总是与接近数据集平均值的正确答案不同。因此，使用 $p_{train}(\sigma )$ 的简单对数正态分布将训练工作定位到相关范围，如表 1 中详述，如图 5a（红色曲线）所示。

为了防止潜在的过拟合，这种过拟合经常困扰具有较小数据集的扩散模型，从GAN文献中借用了一个增强管道。该管道由各种几何变换组成，在添加噪点之前将其应用于训练图像。为了防止增强泄漏到生成的图像中，提供了增强参数作为 $F_{\theta }$ 的调节输入;在推理过程中，将它们设置为零，以保证仅生成非增强图像。表 2 显示，数据增强提供了一致的改进（配置 F），为有条件和无条件 CIFAR-10 产生了 1.79 和 1.97 的最先进的新 FID，超过了之前的 1.85 和 2.10 的记录。

随机抽样的相关性似乎随着模型本身的改进而减弱，如图5b，c所示。当在CIFAR-10中使用论文的训练设置时（图5b），通过确定性抽样获得了最佳结果，任何数量的随机抽样都是有害的。

参考文献

[1] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. In Proc. NeurIPS, 2020.

[2] Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distribution. In Proc. NeurIPS, 2019.

[3] T. Karras, M. Aittala, J. Hellsten, S. Laine, J. Lehtinen, and T. Aila. Training generative adversarial networks with limited data. In Proc. NeurIPS, 2020.

[4] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. In Proc. ICLR, 2021.