Elucidated Diffusion Models

Introduction

扩散模型的火热吸引了众多研究者，相关工作百花齐放，但这样繁荣的景象也就带来了一个问题——每篇论文都有一套自己的体系，缺乏一个统一的框架；并且扩散模型的各个组件相互耦合，不便于单独研究。终于，NVIDIA 的研究人员在论文^[1]中详细阐述了扩散模型的设计空间。其实扩散模型无非由以下几部分构成：

• Formulation：扩散和去噪过程可以用 SDE/ODE 描述，方程的系数有多种选取方案；
• Sampling：可以自定义采样器种类、时间步离散化方式和采样轨迹；
• Training：可以将网络参数化为不同形式，网络的输入输出可以进行各种预处理，可以调整各时间步权重，可以运用各种训练技巧。

将这些部分解耦并分别研究后，研究者将各部分的最优实践组合起来，提出了 EDM，达到了扩散模型的 SOTA.

Formulation

在扩散模型的SDE与ODE描述一文中我们看到，Song 等人^[2]通过将 DDPM 和 SMLD 连续化，指出扩散过程可以描述为如下的 SDE 形式： $$ 由于无论是 VP SDE (对应 DDPM) 还是 VE SDE (对应 SMLD)，$$ 其实都是 $$ 的形式，所以不妨将该 SDE 改写作： $$ 另一方面，在以 Variational Diffusion Models^[3]为代表的一系列基于贝叶斯框架推导扩散模型的论文中，人们常用如下马尔可夫过程表示扩散过程： $$ 为方便后续推导，将其改写作： $$ 事实上，$$ 式可以视作 $$ 式的扰动核 (perturbation kernel)，因此二者本质其实是一样的。鉴于 SDE 形式的灵活性，现阶段的研究倾向于用 SDE 描述扩散模型，但是 $$ 式中的参数 $$ 本身没有什么实际意义，反而 $$ 式的参数 $$ 有直观的含义，即加噪强度。因此，我们希望结合二者的优点，将扩散过程写作以 $$ 为参数的 SDE 形式。经过推导，可以证明 $$ 与 $$ 之间存在如下的关系： $$ 于是将 $$ 式代入 $$ 式就得到了我们想要的扩散过程的 SDE 描述： $$

推导过程：由 $$ 式可知： $$ 即： $$ 整理一下： $$ 取 $$： $$ 令 $$，得： $$ 整理得： $$ 对比 $$ 式可知： $$ 反过来用 $$ 表示 $$ 即为： $$

根据论文^[2]，$$ 式对应如下的与之具有相同边缘分布的 reverse SDE 和 probability flow ODE： $$

其中 $$ 表示数据在 $$ 时刻加噪后的分布。考虑到 ODE 比 SDE 更容易分析，我们从 ODE 下手，将 $$ 式代入 $$ ​式得： $$

进一步地，注意到 $$ 这种表示方式并不能直观地告诉我们 $$ 时刻加了多少噪声。换句话说，上式与 noise schedule 耦合起来了。为了解耦，定义 $$ 表示数据加上方差为 $$ 的高斯噪声后的分布，那么可以推导： $$

代入 $$​​ 式得： $$ 特别地，若 $$，则有： $$ 上述 $$ 式即是 EDM 论文采用的 ODE 形式化。

为了获得 score function $$，我们利用 denoising score matching 训练模型拟合之。在不同的工作中，模型有不同的参数化形式，如拟合 $$、$$、$$ 或 score function 等，它们可以容易地互相转换。这里统一用 $$ 表示“去噪模型”（即拟合 $$ 的模型）而暂时不讨论其具体形式，那么有： $$

于是： $$ 代入 $$​ 式得到： $$

生成过程（采样过程）就是使用数值方法迭代求解该 ODE 的过程。

Deterministic Sampling

ODE Solver

采样器，即 reverse SDE 或 probability flow ODE 的数值求解器，是扩散模型的一个重点研究方向。采样器用迭代前进的方式近似 SDE/ODE 真实的轨迹，因此每一步都会带来截断误差。越高阶的采样器允许人们用越少的迭代步数达到相同的效果。然而，高阶采样器的每一步迭代往往也需要更多次的模型推断（即神经网络的前向传播），而扩散模型的时耗基本都是前向传播造成的，所以在扩散模型的语境下，并不一定是越高阶的采样器就越高效。事实上正因如此，人们倾向于用 NFE (Neural Function Evaluation)，即网络前向传播次数而非求解器的采样步数作为采样时间来比较不同采样器的效率。

最简单常用的 ODE 求解器无疑是 Euler 法，即在每一个时间步上用矩形近似积分，这是一个一阶求解器，即若设步长为 $$，则其局部截断误差是 $$ 的。另一些常用的求解器包括高阶 Runge-Kutta 方法和线性多步法。EDM 的作者通过大量实验发现，二阶的 Heun 方法（即改进 Euler 法，用梯形近似积分）的效率最高，其算法如下所示：

如果去掉 6-8 行则 Heun 方法退化为 Euler 法，即 $$ 式的简单离散化。

Discretization

除了 ODE 求解器的算法以外，还有一个因素会影响采样质量——时间步离散化的方式，即如何选取序列 $$. 由于 $$ 单调增加存在反函数，这等价于选取序列 $$，然后令 $$. 作者决定采用 $$ 的幂函数形式并选择 $$ 以使得 $$，可解得： $$ 超参数 $$ 越大，则接近 $$ 处的步长越小（采得越密），接近 $$ 处的步长越大（采得越密）；而步长越小，求解器的局部截断误差就越小。实验发现当 $$ 时各 $$ 处的局部截断误差比较均衡，但对图像而言 $$ 时生成质量更好，这说明接近 $$ 处的误差对图像质量影响更大，因此作者最终采用 $$.

Trajectory

ODE 解的轨迹形状由 $$ 和 $$，即 noise schedule 决定，下图展示了 VP ODE、VE ODE 和 DDIM 的轨迹：

它们对应的 noise schedule 分别为：

• VP ODE: $$
• VE ODE: $$
• DDIM: $$

作者认为，DDIM 的设置最好，因为对任意的 $$，执行一个 Euler 步就可以得到去噪结果 $$. 从图中也可以看见 DDIM 的轨迹曲率最小，任意噪声下轨迹的切线都基本指向去噪结果。并且，实验也发现在 VP ODE 和 VE ODE 的预训练模型的基础上使用 DDIM 的设置，图像质量相比原设置有着明显的提升。在 DDIM 的 noise scheduler 设置下，对应的 ODE 方程 $$ 式简化为： $$

Stochastic Sampling

确定性采样有非常多的好处，例如可以反向采样 ODE 从而将图像转换回它的隐变量表示，但是其采样质量往往不如随机性采样。

可以证明，ODE $$ 式（取 $$ 的特殊情况）对应一族与之具有相同边缘分布的 SDE： $$ 其中，$$ 和 $$ 分别是正向和逆向的 SDE. 事实上，论文^[2]提出的 reverse SDE $$ 式其实只是上述 SDE 取 $$ 的一个特殊情形。上式表明，SDE 是在 ODE 的基础上添加了 Langevin diffusion 过程作为修正，从而相比 ODE 有更好的采样效果。

基于上述理论，当取 $$ 时，作者在针对 ODE 的 Heun 采样器的基础上设计了如下针对 SDE 的采样方法：

其中 4、5、6 行在执行 Euler 步前添加一定量的噪声来模拟 Langevin dynamics 过程。这不是一个通用的 SDE 求解器，其正确性保证来源于 Langevin 步在理论上是保持边缘分布不变的，因此可以在 ODE 求解过程中引入一些 Langevin 步。不过，由于模型具有预测误差，Langevin 步本身也不是精确的，因此作者对引入的噪声做了限制，即上图中 $$ 的作用。这与论文^[2]的 Predictor-Corrector 采样器有着类似的思想。

Training

Network

在实际训练时，我们不一定用神经网络直接对 $$ 建模，例如 DDPM 原文预测的就是噪声，而 Progressive Distillation^[4]提出预测 $$. 为将不同的设计统一起来，作者将 $$ 建模为： $$ 其中 $$ 才是真正要训练的神经网络。于是根据 denoising score matching，损失函数为： $$ 为了让网络的输入和输出在不同大小的噪声保持单位方差，同时尽可能小地放大网络的误差，作者最终选择： $$

Loss weighting

为了平衡损失函数的各项，取 $$，于是： $$ 训练时还需按照 $$ 采样 $$，作者选择对数正态分布： $$ 其中 $$.

Augmentation

当在小数据集上训练时，扩散模型有可能会过拟合。为此，受到 GANs 相关文献的启发，作者在训练时采用了一系列数据增强。为了避免模型生成增强的结果，作者在训练时将增强信息作者为条件给到网络 $$，并在推断时将其置零。

Unified Framework

在作者的统一框架下，原来的各个经典扩散模型对应的设置如下表所示：

下面我们推导一下表中 Sampling 一块的设置。

VP Sampling

VP SDE 是 DDPM 的连续化版本，对应 $$ 式中取： $$ 在 DDPM 中，$$ 是从 $$ 到 $$ 的线性插值： $$ 所以： $$ 设 $$ 是 $$ 的积分： $$ 其中 $$. 则有： $$ 类似地，有： $$ 又 DDPM++（DDPM 的连续时间版本）的时间步是 $$ 之间的均匀采样，所以： $$ 这就得到了表中 VP Sampling 的部分。

VE formulation

VE SDE 是 SMLD 的连续化版本，对应 $$ 式中取： $$ 于是： $$ 在 SMLD 中，$$ 被设置为几何级数： $$ 但这里 EDM 作者设置 $$，然后将时间步写作： $$ 从而达到一样的效果。

iDDPM + DDIM

在 DDIM与加速采样一文中我们看到，DDIM 可以写作如下的 ODE 形式： $$ 用本文中的 $$ 替换 $$，再替换为 score function 表示： $$ 因此得到 ODE： $$ 取 $$，则： $$

说实话这里没看懂作者的脑回路，感觉反而弄得更复杂了。

References

1. Karras, Tero, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. Advances in Neural Information Processing Systems 35 (2022): 26565-26577. ↩︎
2. Song, Yang, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations. In International Conference on Learning Representations. 2020. ↩︎
3. Kingma, Diederik, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems 34 (2021): 21696-21707. ↩︎
4. Salimans, Tim, and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512 (2022). ↩︎
5. Diffusion Models 基础知识总结回顾 - 董子斌的文章 - 知乎 https://zhuanlan.zhihu.com/p/682151286 ↩︎