Elucidating the Design Space of Diffusion-Based Generative Models

此篇内容源于论文：Elucidating the Design Space of Diffusion-Based Generative Models
EDM 构建了一个清晰的扩散模型设计空间，其中各种具体的设计选项解耦开来，能够更精确地分析不同设计选项对于采样、训练各过程的实际影响。

Expressing diffusion models in a common framework

• DDPM/DDIM

一步加噪公式：
$$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,$$

写成概率分布形式：
$$p(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)I)$$

• Score Matching

一步加噪公式：
$$x_t=x_0+{\sigma }_tϵ$$

写成概率分布形式：
$$p(x_t|x_0)=\mathcal{N}(x_t;x_0,{\sigma }_t^{2}I)$$

• Flow Matching

一步加噪公式：
$$x_t=(1-t)x_0+tϵ$$

写成概率分布形式：
$$p(x_t|x_0)=\mathcal{N}(x_t;(1-t)x_0,t^{2}I)$$

---

• 通用形式 —— ODE / SDE 公式

$$p_{0t}(x(t)|x(0))=\mathcal{N}\left(x(t);s(t)x(0),s(t{)}^2\sigma (t{)}^{2}I\right),$$

---

Song 等人 将他们的前向随机微分方程（SDE）定义为：

$$dx=f(x,t)dt+g(t)d{\omega }_t,(9)$$

其中，${\omega }_t$ 是标准的 Wiener 过程（布朗运动），$f(\cdot ,t):{\mathbb{R}}^d\to {\mathbb{R}}^d$ 和 $g(\cdot ):\mathbb{R}\to \mathbb{R}$ 分别是漂移系数和扩散系数，$d$ 是数据集的维度。这些系数在方差保持（Variance Preserving, VP）和方差爆炸（Variance Exploding, VE）公式中选择不同，且 $f(\cdot )$ 始终具有形式 $f(x,t)=f(t)x$，其中 $f(\cdot ):\mathbb{R}\to \mathbb{R}$。因此，SDE 可以等价地写为：

$$dx=f(t)xdt+g(t)d{\omega }_t.(10)$$

该 SDE 的扰动核具有以下一般形式：

$$p_{0t}(x(t)|x(0))=\mathcal{N}\left(x(t);s(t)x(0),s(t{)}^2\sigma (t{)}^{2}I\right),(11)$$

其中，$\mathcal{N}(x;\mu ,\Sigma )$ 表示均值为 $\mu$、协方差为 $\Sigma$ 的高斯分布在 $x$ 处的概率密度函数，

$$s(t)=\exp \left({\int }_0^{t}f(\xi )d\xi \right),\sigma (t)=\sqrt{{\int }_0^t\frac{g(\xi {)}^2}{s(\xi {)}^2}d\xi }.(12)$$

边缘分布 $p_t(x)$ 通过对 $x(0)$ 积分扰动核得到：

$$p_t(x)={\int }_{{\mathbb{R}}^d}{p}_{0t}(x|x_0)p_{\text{data}}(x_0)d{x}_0.(13)$$

Song 等人定义了概率流常微分方程（ODE），使其遵循相同的 $p_t(x)$：

$$dx=\left[f(t)x-\frac{1}{2}g(t{)}^2{\nabla }_x\log p_t(x)\right]dt.(14)$$

---

从式子（14）中可以看到，Song 的 Probability Flow ODE 由 $f(\cdot )$ 和 $g(\cdot )$ 决定。然而，边缘分布 { p t ( x ) } t = 0 T \{p_t(x)\}_{t=0}^T {pt​(x)}t=0T​ 却直接由 $s(\cdot )$ 和 $\sigma (\cdot )$ 决定。因此，作者希望将 Probability Flow ODE 改写为由 $s(\cdot )$ 和 $\sigma (\cdot )$ 直接控制的表达式。

可以反写出：

$$f(t)=\frac{\dot{s}(t)}{s(t)},(6)$$

$$g(t)=s(t)\sqrt{2\dot{\sigma }(t)\sigma (t)}.(7)$$

将 $f(t)$ 和 $g(t)$ 的表达式代入(14)可以得到新的 Probability Flow ODE 为：

$$\begin{array}{rl}dx& =\left[\left(\dot{s}(t)/s(t)\right)x-\frac{1}{2}{\left(s(t)\sqrt{2\dot{\sigma }(t)\sigma (t)}\right)}^2{\nabla }_x\log p\left(x/s(t);\sigma (t)\right)\right]dt\\ & =\left[\left(\dot{s}(t)/s(t)\right)x-\frac{1}{2}\left(2s(t{)}^2\dot{\sigma }(t)\sigma (t)\right){\nabla }_x\log p\left(x/s(t);\sigma (t)\right)\right]dt\\ & =\left[\frac{\dot{s}(t)}{s(t)}x-s(t{)}^2\dot{\sigma }(t)\sigma (t){\nabla }_x\log p\left(\frac{x}{s(t)};\sigma (t)\right)\right]dt.\end{array}$$

$$dx=-\dot{\sigma }(t)\sigma (t){\nabla }_x\log p\left(x;\sigma (t)\right)dt.(4)$$

我们的公式强调了以下事实：概率流 ODE 的每个实现都是同一规范 ODE 的重参数化；改变 $\sigma (t)$ 对应于对 $t$ 的重新参数化，而改变 $s(t)$ 对应于对 $x$ 的重新参数化。

我们可以将不同的扩散模型解构为不同的 Schedule $\sigma (t)$
以及 Scaling $s(t)$，如下图所示：

---

Improvements to deterministic sampling

本文假设，与采样过程相关的选择在很大程度上独立于其他组件，例如网络架构和训练细节。

换句话说，$D_{\theta }$ 的训练过程不应决定 $\sigma (t)$、$s(t)$ 和 { t i } \{t_i\} {ti​}，反之亦然；从采样器的角度来看，$D_{\theta }$ 只是一个黑箱。

我们通过在三个预训练模型上评估不同的采样器来验证这一点，每个模型代表不同的理论框架和模型家族。我们首先使用这些模型的原始采样器实现测量基线结果，然后使用表 1 中的公式将这些采样器引入我们的统一框架，接着应用我们的改进。这使得我们能够评估不同的实际选择，并提出适用于所有模型的采样过程的通用改进。

在这里，作者比较了三种 pre-trained model，分别为"DDPM++cont. (VP)"，“NCSN++ cont. (VE)” 以及 “DDIM” 模型。下图中的横轴为 neural function evaluation（NFE），即生成一副图像所需调用 $D_{\theta }$的次数；纵轴为 FID，即模型的表现：

Discretization and higher-order integrators

数值求解 ODE 的方式需要我们跟踪真实的求解轨迹。在每一步迭代的过程中，求解器都会带来 truncation error，并且这种错误会不断累积。想要提升求解精度，就不得不增大求解步数（使每次的求解步长更小）。

现在最经常使用的是 Euler’ method 来求解 ODE，这是一种一阶 ODE 求解器，会带来 $O(h^2)$ 的 truncation error（h 为步长）。本文作者发现 Heun’s 二阶 ODE 求解器可以在 truncation error 与 NFE 之间取得较好的平衡，算法如下：

Trajectory curvature and noise shcedule

ODE 方程的解的轨迹取决于 $\sigma (t)$与 $s(t)$的选择；如果这两个函数的选择得当，则可以降低 truncation error。作者认为，最优的函数选择为$\sigma (t)=t$ 以及 $s(t)=1$，这也是 DDIM 中的选择。

在上图中，作者在一维数据分布上对比了三种不同的 ODE 选择。作者发现，无论是 VP ODE 还是 VE ODE，其迭代方向都会偏离数据均值，因而引入了较大的 truncation error；而使用本文的方式则不会偏离均值。

Stochastic sampling

Background.

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

其中，${\omega }_t$ 是标准的 Wiener 过程。$d{x}_{+}$ 和 $d{x}_{-}$ 现在是分别用于时间向前和向后移动的独立 SDE，它们通过 Anderson的时间反转公式相关联。

Langevin 项可以进一步看作是一个确定性的基于分数的去噪项和一个随机的噪声注入项的组合，其净噪声水平贡献相互抵消。

因此，$\beta (t)$ 有效地表达了现有噪声被新噪声替换的相对速率。Song 等人 的 SDE 可以通过选择 $\beta (t)=\dot{\sigma }(t)/\sigma (t)$ 恢复，此时前向 SDE 中的分数项消失。

这一视角揭示了为什么随机性在实践中是有帮助的：隐式的 Langevin 扩散将样本驱动到给定时间所需的边缘分布，主动纠正早期采样步骤中的任何错误。另一方面，用离散的 SDE 求解器步骤近似 Langevin 项本身会引入误差。先前的结果 [3, 24, 47, 49] 表明，非零的 $\beta (t)$ 是有帮助的，但据我们所知，Song 等人中对 $\beta (t)$ 的隐式选择并不具有特殊性质。因此，最佳的随机性量应通过实验确定。

stochastic sampler.

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

Preconditioning and training

作者认为，直接训练神经网络来建模 $D$ 远非理想选择
由于输入 $x=y+n$ 是干净信号 $y$ 和噪声 $n\sim \mathcal{N}(0,{\sigma }^{2}I)$ 的组合，其幅度会因噪声水平 $\sigma$ 的不同而发生巨大变化。因此，通常的做法不是直接将 $D_{\theta }$ 表示为神经网络，而是训练一个不同的网络 $F_{\theta }$，并从 $F_{\theta }$ 导出 $D_{\theta }$。

之前的方法 [37, 47, 49] 通过依赖于 $\sigma$ 的归一化因子来解决输入缩放问题，并尝试通过训练 $F_{\theta }$ 来预测缩放为单位方差的噪声 $n$，然后通过 $D_{\theta }(x;\sigma )=x-\sigma F_{\theta }(\cdot )$ 重建信号。这种方法的缺点是，当 $\sigma$ 较大时，网络需要仔细微调其输出，以精确抵消现有噪声 $n$ 并以正确的尺度给出输出；需要注意的是，网络的任何错误都会被放大 $\sigma$ 倍。在这种情况下，直接预测期望输出 $D(x;\sigma )$ 似乎要容易得多。

本文提出了一种带有依赖于 $\sigma$ 的跳跃连接的神经网络预条件方法，使其能够估计 $y$ 或 $n$，或者介于两者之间的某种状态。因此，我们将 $D_{\theta }$ 写成以下形式：

$$D_{\theta }(x;\sigma )=c_{\text{skip}}(\sigma )x+c_{\text{out}}(\sigma )F_{\theta }\left(c_{\text{in}}(\sigma )x;c_{\text{noise}}(\sigma )\right),(7)$$

其中：

• $F_{\theta }$ 是要训练的神经网络；
• $c_{\text{skip}}(\sigma )$ 调节跳跃连接；
• $c_{\text{in}}(\sigma )$ 和 $c_{\text{out}}(\sigma )$ 分别缩放输入和输出的幅度；
• $c_{\text{noise}}(\sigma )$ 将噪声水平 $\sigma$ 映射为 $F_{\theta }$ 的条件输入。

通过对噪声水平 $\sigma$ 的加权期望，得到整体训练损失：

$${\mathbb{E}}_{\sigma ,y,n}\left[\lambda (\sigma )\Vert D(y+n;\sigma )-y{\Vert }_2^2\right],$$

其中：

• $\sigma \sim p_{\text{train}}$，$y\sim p_{\text{data}}$，$n\sim \mathcal{N}(0,{\sigma }^{2}I)$；
• $p_{\text{train}}(\sigma )$ 是采样噪声水平 $\sigma$ 的概率；
• $\lambda (\sigma )$ 是对应的权重。

我们可以将这一损失等价地表示为关于原始网络输出 $F_{\theta }$ 的形式（公式 7）：

$${\mathbb{E}}_{\sigma ,y,n}\left[\lambda (\sigma )c_{\text{out}}(\sigma {)}^2{\Vert F_{\theta }\left(c_{\text{in}}(\sigma )\cdot (y+n);c_{\text{noise}}(\sigma )\right)-\frac{1}{c_{\text{out}}(\sigma )}\left(y-c_{\text{skip}}(\sigma )\cdot (y+n)\right)\Vert }_2^2\right].(8)$$

这种形式揭示了 $F_{\theta }$ 的有效训练目标，使我们能够从基本原理出发确定合适的预条件函数选择。