扩散模型设计选项的全面拆解与分析：EDM

扩散模型设计选项的全面拆解与分析：EDM

EDM 是扩散模型形式化研究中的一篇重要工作。在当时，DDPM、SMLD、SDE 等扩散模型工作陆续被提出，它们的建模角度不尽相同，却又有着千丝万缕的联系。EDM 指出现有的扩散模型设计过于杂糅，并构建了一个清晰的扩散模型设计空间，其中各种具体的设计选项解耦开来，从而能够更精确地分析不同设计选项对于采样、训练各过程的实际影响。

统一的框架

我们首先基于 Song 等人提出的 SDE 形式，对已有的扩散模型理论进行整理，构建一个统一的理论框架，其中各种扩散模型设计中的各种选项相互解耦开来，便于我们进一步通过推导或实验确定通用的最佳选项。

之前工作中的 ODE 形式

记数据分布为 $p_{\text{data}}(\mathbf{x})$，其方差为 ${\sigma }_{\text{data}}^2$。通过向数据添加方差为 ${\sigma }^2$ 的不同强度的高斯噪声，可以得到一系列加噪分布 $p(\mathbf{x};\sigma )$。当噪声的 sigma 足够大，${\sigma }_{\text{max}}>>{\sigma }_{\text{data}}$ ，此时 $p(\mathbf{x};{\sigma }_{\text{max}})$ 中就几乎看不出数据，成为完全的的噪声分布了。扩散模型的想法就是，随机采样一个起始纯噪声图 $\mathbf{x}\sim \mathcal{N}(0,{\sigma }_{\text{max}}^2\mathbf{I})$，并逐步对其进行去噪，每一步得到噪声稍小的 ${\mathbf{x}}_i\sim p({\mathbf{x}}_i;{\sigma }_i)$，最终就能得到服从 $p_{\text{data}}$ 的干净数据了。

在 Song 等人的工作中，扩散模型 SDE 定义为如下形式：
$$\begin{gathered} d\mathbf{x}=f(\mathbf{x},t)dt+g(t)d\mathbf{w},\text{forward} \\ d\mathbf{x}=[f(\mathbf{x},t)-g(t){\nabla }_{\mathbf{x}}\log p(\mathbf{x})]dt+g(t)d\tilde{\mathbf{w}},\text{backward}\text{\hspace{1em}(1)} \end{gathered}$$
一般来说，$f(\cdot )$ 取关于 $\mathbf{x}$ 的线性形式，即 $f(\mathbf{x},t)=f(t)\mathbf{x}$，上式的前向过程可写为 $d\mathbf{x}=f(t)\mathbf{x}dt+g(t)d\mathbf{w}$。

这个 SDE 对应了一个具有完全相同边缘分布 $p_t(\mathbf{x})$ 的概率流 ODE 形式：
$$d\mathbf{x}=\left[f(t)\mathbf{x}-\frac{1}{2}{\nabla }_{\mathbf{x}}\log p(\mathbf{x})\right]dt\text{\hspace{1em}(2)}$$
Song 等人还给出了从时刻 $0$ 到 $t$ 加噪过程中扰动核的一般形式：
$$\begin{gathered} p_{0t}({\mathbf{x}}_t|{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_t;s(t){\mathbf{x}}_0,s^2(t){\sigma }^2(t)\mathbf{I}) \\ i.e.{\mathbf{x}}_t=s(t){\mathbf{x}}_0+s(t)\sigma (t)ϵ \\ 其中s(t)=\exp \left({\int }_0^{t}f(\xi )d\xi \right),\sigma (t)=\sqrt{{\int }_0^t\frac{g^2(\xi )}{s^2(\xi )}}d\xi \text{\hspace{1em}(3)} \end{gathered}$$

可以看到，SDE/ODE 的形式是由 $f,g$ 定义的。但是实际上，边缘分布 $p_t(\mathbf{x})$ 才是我们在扩散模型形式设计中最关心的东西，它关系到模型训练、采样加速以及如何理解 ODE 的实际行为。而 $f,g$ 与边缘分布的关系并不直接。因此，EDM 中作者认为直接从 $s,\sigma$ 的角度出发来定义扩散模型的规格是更合理、更有实际意义的。接下来，我们将推到从 $s,\sigma$ 的角度来从新推导 ODE 形式。

ODE 形式重新推导

将边缘分布写成积分形式并进行推导（原文式 15-19），有如下结果：
$$\begin{array}{rl}{p}_t(\mathbf{x})& ={\int }_{{\mathbb{R}}^d}{p}_{0t}(\mathbf{x}|{\mathbf{x}}_0)p_{\text{data}}({\mathbf{x}}_0)d{\mathbf{x}}_0\\ & =...\\ & =s(t{)}^{-d}[p_{\text{data}}\ast \mathcal{N}(0,{\sigma }^2(t)\mathbf{I})](\mathbf{x}/s(t))\end{array}$$
其中 $\ast$ 表示两个概率密度函数的卷积操作，$d$ 是数据集的维度。方括号内的卷积操作相当于是对数据分布进行加噪扰动，我们将这样得到的加噪分布记为 $p(\mathbf{x};\sigma )=p_{\text{data}}\ast \mathcal{N}(0,{\sigma }^2(t)\mathbf{I})$。从而可以用加噪分布 $p(\mathbf{x};\sigma )$ 来表示边缘分布 $p_t(\mathbf{x})$：
$$p_t(\mathbf{x})=s(t{)}^{-d}p(\mathbf{x}/s(t);\sigma (t))$$
将这个关系带入到概率流 ODE （式 2）中，并整理得到：
$$d\mathbf{x}=...=[f(t)\mathbf{x}-\frac{1}{2}{g}^2(t){\nabla }_{\mathbf{x}}\log (p(\mathbf{x}/s(t),\sigma (t))]dt\text{\hspace{1em}(4)}$$
基于式 3，我们可以反写出用 $s,\sigma$ 表示 $f,g$ 的形式：
$$f(t)=\dot{s}(t)/s(t)g(t)=s(t)\sqrt{2\dot{\sigma }(t)\sigma (t)}$$

其中上面的点表示对时间 $t$ 的微分。现在，可以将式 (4) 中的 $f,g$ 也替换成 $s,\sigma$：
$$d\mathbf{x}=...=\left[\frac{\dot{s}(t)}{s(t)}\mathbf{x}-s^2(t)\dot{\sigma }(t)\sigma (t){\nabla }_{\mathbf{x}}\log p(\frac{\mathbf{x}}{s(t)};\sigma (t))\right]dt\text{\hspace{1em}(5)}$$
这就得到了 EDM 中基于 $s,\sigma$ 的 ODE 方程。我们再来理解一下 $s,\sigma$ 的含义，$s(t)$ 实际上是对于原始 $\hat{\mathbf{x}}$ 的一个关于时间 $t$ 的缩放因子 $\mathbf{x}=s(t)\hat{\mathbf{x}}$，$\sigma (t)$ 则是在 $t$ 时刻加噪分布的标准差，也就是我们常说的噪声计划表 noise schedule。注意这里显式地撤销了 $s(t)$ 的缩放，从而保持 $p(\mathbf{x};\sigma )$ 是独立于 $s(t)$ 的。 可以看到，在这个式子中，我们唯一未知的就是得分 ${\nabla }_{\mathbf{x}}\log p(\mathbf{x};\sigma )$，这正是我们的神经网络需要去学习估计的东西。

训练目标：去噪得分匹配

在训练时，扩散模型一般采用去噪得分匹配（denoising score matching）的目标函数。这种方法绕开了对无法计算的归一化常数 $Z$，直接计算加噪分布的得分。具体来说，如果我们记 $D(\mathbf{x};\sigma )$ 是一个去噪函数，它能够最小化下面的 L2 误差：
$${\mathbb{E}}_{\mathbf{y}\sim p_{\text{data}}}{\mathbb{E}}_{\mathbf{n}\sim \mathcal{N}(0,{\sigma }^2\mathbf{I})}||D(\mathbf{y}+\mathbf{n};\sigma )-\mathbf{y}|{|}_2^2$$
经推导（附录 B.3 ），有：
$${\nabla }_{\mathbf{x}}\log p(\mathbf{x};\sigma )=(D(\mathbf{x};\sigma )-\mathbf{x})/{\sigma }^2$$
其中 $\mathbf{y}$ 是干净数据，$\mathbf{n}$ 是噪声。从上式的角度来看，得分函数的实际上在做的事情就是将噪声从干净数据中分离出来，即所谓去噪。扩散模型中，我们将上式作为目标函数，训练一个神经网络 $D_{\theta }(\mathbf{x};\sigma )$ 来估计这个去噪函数 $D$。

ODE 离散化与求解

现在，我们去噪得分匹配的目标函数训练了一个神经网络 $D_{\theta }$，估计出了加噪分布的得分，将这个得分代回到式 5，我们就得到了一个可求解的 ODE。接下来的采样过程，就是要在给定初值 ${\mathbf{x}}_T$ 的条件下，求解出 ${\mathbf{x}}_0$，这样就得到了符合数据分布 $p_{\text{data}}$ 的干净数据。

注意，我们现在的采样过程完全就是在解一个已知 ODE 的初值问题，只不过每一步 ODE 中的某些项（得分）需要我们跑一下神经网络估计出来（这就是一次所谓的 NFE Network Function Evaluation）。既然是求解 ODE，那么采样过程就我们需要确定两件事情：一是采用何种求解器（ODE Solver），二是采用什么样的求解步长离散策略（Schedule）。

preconditioning和损失权重

在训练神经网络时，有一个非常有用的经验是：我们最好保持网络的输入输出在一个稳定的数值范围内（比如具有单位方差）。这样可以避免梯度的数值有大范围的波动，从而稳定训练。而在训练扩散模型来拟合去噪函数 $D$ 时，由于我们需要向干净数据中加入方差不等的高斯噪声，因此是无法满足这一要求的。我们一般考虑改为训练神经网络 $F_{\theta }$（而非直接训练 $D_{\theta }$），再通过适当的形式化调整计算出去噪函数 $D$。之前的工作一般是引入一个关于 $\sigma$ 的归一化项，将噪声 $\mathbf{n}$ 归一化到单位方差，让模型来预测这个单位方差的噪声，再计算出去噪函数 $D_{\theta }=\mathbf{x}-\sigma F_{\theta }(\cdot )$。作者认为，这样做的问题是当 $\sigma$ 非常大时，网络预测的一点微小错误都会被这个很大的 $\sigma$ 进一步放大，网络的学习比较困难，此时反而是直接预测 $D$ 对模型来说比较简单。

为了改善这个问题，作者参考之前自适应混合信号噪声的方法，提出了一种新的扩散模型形式化。通过引入一个关于 $\sigma$ 的跳跃连接来对模型进行 preconditioning，这样模型可以预测 $\mathbf{y}$，也可以预测 $\mathbf{n}$，也可以预测二者之间。具体来说，作者将 $D_{\theta }$ 构造成如下形式：
$$D_{\theta }(\mathbf{x};\sigma )=c_{\text{skip}}(\sigma )\mathbf{x}+c_{\text{out}}(\sigma )F_{\theta }\left(c_{\text{in}}(\sigma )\mathbf{x};c_{\text{noise}}(\sigma )\right)\text{\hspace{1em}(6)}$$
其中 $F_{\theta }$ 是实际训练的神经网络，$c_{\text{skip}}(\sigma )$ 用于调制跳跃连接，$c_{\text{in}}(\sigma ),c_{\text{out}}(\sigma )$ 分别用于对输入和输出进行缩放，$c_{\text{noise}}(\sigma )$ 将噪声等级 $\sigma$ 转换为网络的输入条件。

我们之前提到 DDPM 通过修改原扩散模型的形式化，从预测原图改为预测噪声，大大提升了出图的质量。这里 EDM 实际上是将扩散模型的形式化拓展成了一种更一般的形式。其中 $c_{\text{skip}},c_{\text{out}},c_{\text{in}},c_{\text{noise}}$ 的选取，扩散模型训练过程设计选项的关键。

此外，在训练过程中，损失的加权方式 $\lambda (\sigma )$ 和采样 $\sigma$ 的分布，也很重要。

总结

以上我们已经将整个扩散模型的设计空间（包括训练过程、采样过程）的所有规格可选项全部梳理出来了。具体包括如下内容：

采样

• $s(t)$：缩放因子
• $\sigma (t)$：噪声的标准差
• { t i } \{t_i\} {ti​}：离散化的步长
• ODE Solver：所选用的 ODE 求解器

训练

• preconditioning：$c_{\text{skip}}(\sigma )$：用于调制跳跃连接；$c_{\text{out}}(\sigma )$：输出缩放因子；$c_{\text{in}}(\sigma )$：输入缩放因子；$c_{\text{noise}}(\sigma )$：噪声条件化

• 损失加权：$\lambda (\sigma )$

• 采样 $\sigma$ 的分布

在 EDM 论文中，作者选取了 DDPM（VP SDE）、SMLD（VE SDE）、improved DDPM training + DDIM sampling 三种方法作对比分析通过理论分析和实验，设计出了自己的 EDM，所有这些方法在 EDM 框架下的设计选项重新推到结果如下表所示。之所以选择这几项工作作为对比，一是因为它们都是当时 SOTA 的扩散生成模型，另一点原因，是它们分别来自不同的扩散模型理论，通过将他们全都在上述提出的扩散模型设计空间中进行统一的分析、对比，也验证了本文所提框架对不同扩散模型理论的统一性和有效性。

本文所提的扩散模型分析框架，将之前工作中耦合在一起的扩散模型规格设计解耦独立开来，各组件之间不再有隐式依赖关系，可以进行独立的优化选择，而不会影响彼此。

确定性采样

我们先来看如何基于新的框架，如何优化扩散模型的采样过程。这里我们先看 ODE 确定性采样。采样过程有两个优化目标。一是提升出图的质量，由 FID 指标来评估，越低越好。二是提高采样速度，采样过程时间消耗的大头是在模型推理上，因此可以用模型推理的次数 NFE 来反映采样速度。这两者之间存在权衡关系，我们需要在尽可能快的采样速度下得到尽量好的结果。

作者认为，采样过程中的各项参数选择，即缩放因子 $s(t)$，噪声方差 $\sigma (t)$ 和 ODE 离散化的策略 { t i } \{t_i\} {ti​}，与其他组件（模型结构和训练细节）应当是相互独立的。模型 $D_{\theta }$ 对于采样过程来说，可以看作是一个黑盒的得分估计器，其结构和训练方式不会影响采样参数的选择，采样参数也不应该决定模型的训练过程。作者对来自不同理论的扩散模型方法进行了统一的实验，测试各种采样选项对不同模型的影响，并提出通用的最优选项。

离散化和求解器

首先我们来看 ODE 的离散化方式和求解器的选择。

ODE 的数值求解实际上是对真实解轨迹的离散化近似。每一步都会存在误差，并且会逐步积累。常用的欧拉法一阶 ODE 求解器，其单步误差与步长的关系是 $\mathcal{O}(h^2)$。更高阶的龙格库塔等方法精度更高，但高阶求解器会需要更多的推理次数 NFE。作者通过广泛的实验来权衡 NFE 和误差，最终推荐选用二阶 Heun 方法。

二阶 Heun 求解器的具体算法步骤如图所示。其实是在欧拉方法的基础上加了一个校正步（6,7,8 三行，去掉这三行就是欧拉法），这样需要一次额外的模型推理，但是能将单步误差降低到 $\mathcal{O}(h^3)$，从而在采样速度和误差之间达到最好的权衡。

离散化方式，也就是采样时选取的时间步 { t i } \{t_i\} {ti​}，决定了每步的步长 $h$，从而决定了不同噪声程度下的误差分布。作者的结论是，步长应该随着 $\sigma$ 的降低单调递减，并且没有必要根据不同样本采用不同的步长。作者根据 $\sigma$ 构造了形如一个 ${\sigma }_{i<N}=(Ai+B{)}^{\rho }$ 参数化的步长策略，考虑到我们需要保证 ${\sigma }_0={\sigma }_{\text{max}},{\sigma }_{n-1}={\sigma }_{\text{min}}$，具体公式写为：
$${\sigma }_{i<N}={\left({\sigma }_{\text{max}}^{\frac{1}{\rho }}+\frac{i}{N-1}({\rho }_{\text{min}}^{\frac{1}{\rho }}-{\sigma }_{\text{max}}^{\frac{1}{\rho }})\right)}^{\rho },{\sigma }_N=0\text{\hspace{1em}(8)}$$
参数 $\rho$ 越大，在 $\sigma$ 较小处的步长越小，但同时 $\sigma$ 较大处的步长也会越大。这样相当于让采样过程在 $\sigma$ 较小处的求解更精细。当设置 $\rho =3$ 时，在每一步的单步误差基本是相等的，但是当 $\rho \in [5,10]$ 时出图质量更高。这表明在 $\sigma$ 较小时的误差影响相对比较大。作者最终推荐的值是 $\rho =7$。

轨迹曲率和噪声计划表

然后我们再来看 $\sigma (t)$ 和 $s(t)$，这两项的选择决定了 ODE 解轨迹的形状。上面提到的 ODE 离散数值求解误差与 $d\mathbf{x}/dt$ 的曲率是成比例的，因此选择一个更好的$s,\sigma$ 能够降低误差。作者直接给出最优的选择是取：
$$\sigma (t)=t,s(t)=1$$
这与 DDIM 采样的选择是一样的。这样一来，ODE 式就可以简化为：
$$d\mathbf{x}/dt=(\mathbf{x}-D(\mathbf{x},t))/t$$
注意现在 $\sigma$ 和 $t$ 是等价的了。这样选择一个直接的好处是在任何一个时刻 $t$，解轨迹的切线都指向起点 ${\mathbf{x}}_0$。当模型训练地足够好时，在任何时刻 $t$ 都能一步回到干净图片 ${\mathbf{x}}_0$。顺便提一句，沿着这个分析走下去，我们马上就能得到大名鼎鼎的一致性模型（Consistency Models）了。

下图中对比了三种不同的 ODE 在一维数据上的解轨迹可视化结果。图 a 中的 VP-ODE 的解轨迹，在 $\sigma $ 较大时一直是水平的，直到 $\sigma$ 很小之后才开始指向数据分布；图 b 中的 VE-ODE 的解轨迹，一直都是弯曲的；图 c 中是本文 edm 的（也是 ddim 的）解轨迹，随着 $\sigma$ 增大，解轨迹接近指向数据均值的直线。

下图对比了基于不同理论方法的扩散模型（VP、VE 和 DDIM）在采用了 EDM 中提出的采样选项后，采样速度和出图质量都有所提升。

随机性采样

确定性采样有很多好处，比如可以对图片作 inversion 得到 latent，然后进行图像编辑。但是，实际中确定性采样的出图质量通常不如随机性采样。

Song 等人的 SDE 可以推广为如下形式：
$$d{\mathbf{x}}_{\pm }=\underset{\text{PF ODE}}{\underbrace{-\dot{\sigma }(t)\sigma (t){\nabla }_{\mathbf{x}}\log p(\mathbf{x};\sigma (t))dt}}\pm \underset{\text{朗之万扩散 SDE}}{\underbrace{\underset{确定性的噪声衰减}{\underbrace{\beta (t){\sigma }^2(t){\nabla }_{\mathbf{x}}\log (\mathbf{x};\sigma (t))dt}}+\underset{随机噪声注入}{\underbrace{\sqrt{2\beta (t)}\sigma (t)d{\omega }_t}}}}\text{\hspace{1em}(9)}$$

其中 $\omega (t)$ 是标准维纳过程，$d{\mathbf{x}}_{+},d{\mathbf{x}}_{-}$ 分别表示 SDE 前向和反向移动。该 SDE 形式可以看做在 PF-ODE 的基础上添加了朗之万扩散 SDE 项，而该项又可被拆解分一个确定性的去噪项和一个随机的噪声注入项，它们对净噪声水平的贡献相互抵消。$\beta (t)$ 用于控制现有的噪声被替代为新噪声的比率，当取 $\beta (t)=\dot{\sigma }(t)\sigma (t)$ 时，前向过程中的 score 项就没了，这就回到了 Song 等人提出的 SDE 形式。

从这个角度，我们可以看到为什么引入随机性对出图质量是有帮助的：隐式的朗之万扩散将样本推向特定时刻的目标边缘分布，积极纠正采样早期出现的误差。另一方面，用离散的 SDE 求解器步骤近似朗之万项本身会引入误差。已有工作证明确定性的噪声衰减这一项（$\beta (t)\ne 0$）是有帮助的，$\beta (t)$ 实际如何确定，也应该通过实验来分析。

本文提出了的随机采样器（如下算法所示），在二阶确定性 ODE 积分器的基础上结合了一个显式的再添加和移除噪声，作者将这个过程形象地称为“搅拌（churn）” 。具体来说，在第 $i$ 步，样本 ${\mathbf{x}}_i$ 的噪声水平 $t_i$（由于 $\sigma =t$，所以其实也就是 $\sigma (t_i)$ 了），我们执行两个子步骤。首先，根据一个因子 ${\gamma }_i\ge 0$ 向样本添加噪声，先达到更高的噪声水平 ${\hat{t}}_i=t_i+{\gamma }_i{t}_i$。第二步，基于加过噪声的样本 ${\mathbf{x}}_{i+1}$，我们进行一步 ODE 求解，从 ${\hat{t}}_i$ 回到 $t_{i+1}$。此时我们就得到了噪声强度为 $t_{i+1}$ 下的样本 ${\mathbf{x}}_{i+1}$​。然后再进行下一步的迭代。

引入随机性有助于纠正早期采样步骤中产生的错误，但它也有自身的缺陷。作者观察到过度的噪声添加和移除过程会导致生图结果逐渐丢失细节。在噪声水平特别高和特别低时，还会出现颜色过饱和的漂移。作者猜测这是因为实际的去噪网络 $D_{\theta }$ 在估计噪声时诱导了一个轻微的非守恒向量场，破坏了朗之万扩散的前提条件，从而导致了这些问题。

如果这种退化是由网络本身 $D_{\theta }(\mathbf{x};\sigma )$ 的缺陷引起的，那么只能在采样过程中根据经验来补救。具体来说，作者限制仅在特定的噪声水平范围 $[S_{\text{tmin}},S_{\text{tmax}}]$ 内引入随机性，来解决向过饱和颜色的漂移；定义 ${\gamma }_i=S_{\text{churn}}/N$，其中 $S_{\text{churn}}$ 用来控制引入控制随机性的总量。并且限制 ${\gamma }_i$ 的值，以确保引入的新噪声不会超过图像中已有的噪声。此外，作者还发现通过设置 $S_{\text{noise}}$ 略大于 1，可以部分抵消除去细节的损失，以增加新添加噪声的标准差。这表明猜测的 $D_{\theta }(\mathbf{x};\sigma )$ 导致的非守恒性的主要组成部分是倾向于去除稍多的噪声，这很可能是由于回归均值的倾向，曾有研究指出这是任何 L2 训练的去噪器都可能发生的情况。

模型训练

上面，我们通过引入 preconditioning $c_{\text{skip}},c_{\text{out}},c_{\text{in}},c_{\text{noise}}$，将扩散模型的训练过程重新形式化为式 (6)。现在，我们要根据神经网络训练过程中的一些经验，确定下这几个设计选项。

首先，我们基于式 6 将损失重写为关于 $F_{\theta }$ 的形式：
$${\mathbb{E}}_{\sigma ,\mathbf{y},\mathbf{n}}\left[\underset{权重}{\underbrace{\lambda (\sigma )c_{\text{out}}^2(\sigma )}}||\underset{网络输出}{\underbrace{F_{\theta }(c_{\text{in}}(\sigma )\cdot (\mathbf{y}+\mathbf{n});c_{\text{noise}}(\sigma ))}}-\underset{训练目标}{\underbrace{\frac{1}{c_{\text{out}}}(\mathbf{y}-c_{\text{skip}}(\sigma )\cdot (\mathbf{y}+\mathbf{n})|{|}_2^2}}\right]\text{\hspace{1em}(10)}$$
写成这样的形式之后，我们就能根据输入输出具有单位方差、尽量少地放大网络 $F_{\theta }$ 的偏差、权重在不同噪声层级之间尽可能均匀分布这几条网络训练的经验，来分别推导出 $c_{\text{in}},c_{\text{out}}$ 、$c_{\text{skip}}$ 和 $\lambda (\sigma )$。至于 $c_{\text{noise}}$，作者是通过实验确定下来的。具体来说，我们希望

1 网络输入具有单位方差，即：
$$\begin{array}{rl}{\text{Var}}_{\mathbf{y},\mathbf{n}}[c_{\text{in}}(\sigma )(\mathbf{y}+\mathbf{n})]& =1\\ c_{\text{in}}^2(\sigma )\text{Var}(\mathbf{y}+\mathbf{n})& =1\\ c_{\text{in}}^2(\sigma )({\sigma }_{\text{data}}^2+{\sigma }^2)& =1\\ c_{\text{in}}(\sigma )& =\frac{1}{{\sigma }_{\text{data}}^2+{\sigma }^2}\end{array}$$
2 训练目标具有单位方差，即：
$$\begin{array}{rl}{\text{Var}}_{\mathbf{y},\mathbf{n}}\left[\frac{1}{c_{\text{out}}}(\mathbf{y}-c_{\text{skip}}(\sigma )\cdot (\mathbf{y}+\mathbf{n}))\right]& =1\\ \frac{1}{c_{\text{out}}^2(\sigma )}{\text{Var}}_{\mathbf{y},\mathbf{n}}[\mathbf{y}-c_{\text{skip}}(\sigma )\cdot (\mathbf{y}+\mathbf{n})]& =1\\ c_{\text{out}}^2(\sigma )& ={\text{Var}}_{\mathbf{y},\mathbf{n}}[(1-c_{\text{skip}}(\sigma ))\mathbf{y}+c_{\text{skip}}(\sigma )\mathbf{n}]\\ c_{\text{out}}^2(\sigma )& =(1-c_{\text{skip}}(\sigma ){)}^2{\sigma }_{\text{data}}^2+c_{\text{skip}}^2(\sigma ){\sigma }^2\end{array}$$
3 我们想要尽可能减小网络的误差被放大的程度，为此，我们需要使得 $c_{\text{out}}(\sigma )$ 最小的 $c_{\text{skip}}(\sigma )$，即：$c_{\text{skip}}=\arg {\min }_{c_{\text{skip}}(\sigma )}{c}_{\text{out}}(\sigma )$，由于 $c_{\text{out}}\ge 0$，因此可以等价地写为 $c_{\text{skip}}=\arg {\min }_{c_{\text{skip}}(\sigma )}{c}_{\text{out}}^2(\sigma )$。这是一个凸优化问题，当关于 $c_{\text{skip}}$ 的导数为 0 时取得唯一解，再将上面的关系带入，有：
$$\begin{array}{rl}d[c_{\text{out}}^2(\sigma )]/d{c}_{\text{skip}}(\sigma )& =0\\ d[(1-c_{\text{skip}}(\sigma ){)}^2{\sigma }_{\text{data}}^2+c_{\text{skip}}^2(\sigma ){\sigma }^2]/d{c}_{\text{skip}}(\sigma )& =0\\ {\sigma }_{\text{data}}^{2}d[(1-c_{\text{skip}}(\sigma ){)}^2]/d{c}_{\text{skip}}(\sigma )+{\sigma }^{2}d[c_{\text{skip}}^2(\sigma )]/d{c}_{\text{skip}}(\sigma )& =0\\ {\sigma }_{\text{data}}^2[2c_{\text{skip}}(\sigma )-2]+{\sigma }^2[2c_{\text{skip}}(\sigma )]& =0\\ ({\sigma }^2+{\sigma }_{\text{data}}^2)c_{\text{skip}})(\sigma )-{\sigma }_{\text{data}}^2& =0\\ c_{\text{skip}}(\sigma )& =\frac{{\sigma }_{\text{data}}^2}{{\sigma }^2+{\sigma }_{\text{data}}^2}\end{array}$$

再将这个结果代回，可以求出 $c_{\text{out}}(\sigma )$：
$$c_{\text{out}}(\sigma )=\sigma \cdot {\sigma }_{\text{data}}/\sqrt{{\sigma }^2+{\sigma }_{\text{data}}^2}$$
4 我们希望损失权重在不同的噪声层级之间是均匀分布的，即：
$$\begin{gathered} w(\sigma )=\lambda (\sigma )c_{\text{out}}^2(\sigma )=1 \\ ... \\ \lambda (\sigma )=({\sigma }^2+{\sigma }_{\text{data}}^2)/(\sigma \cdot {\sigma }_{\text{data}}{)}^2 \end{gathered}$$
5 我们还想通过调整训练时采样 $\sigma$ 的策略 $p_{\text{train}}(\sigma )$，尽可能高效地进行训练。作者进行了相关的实验，发现只有在中等强度的噪声处，损失才可能降得很小。作者认为，在噪声强度较低时，识别出图像中的微弱噪声很困难，而且对最终的出图结果也没有特别大的帮助；而在噪声强度较高时，训练目标又与数据集平均的正确结果非常不同。因此，在中等噪声强度处的去噪能力是最有意义的，需要尽可能多的将 $\sigma$ 采样在这个区间内。为此，作者采用了一种简单地 log-normal 分布来用作 $p_{\text{train}}(\sigma )$。

至此，我们已经将训练阶段作者推荐的配置全部或推导，或分析，或实验出来了：

preconditioning：

$$\begin{array}{r}\{\begin{array}{ll}{c}_{\text{skip}}(\sigma )& ={\sigma }_{\text{data}}^2/({\sigma }^2+{\sigma }_{\text{data}}^2)\\ c_{\text{out}}(\sigma )& =\sigma \cdot {\sigma }_{\text{data}}/\sqrt{{\sigma }_{\text{data}}^2+{\sigma }^2}\\ c_{\text{in}}(\sigma )& =1/\sqrt{{\sigma }_{\text{data}}^2+{\sigma }^2}\\ c_{\text{noise}}(\sigma )& =\frac{1}{4}\ln (\sigma )\end{array}\end{array}$$

加权：
$$\lambda (\sigma )=({\sigma }^2+{\sigma }_{\text{data}}^2)/(\sigma \cdot {\sigma }_{\text{data}}{)}^2$$
采样：
$$\ln (\sigma )\sim \mathcal{N}(P_{\text{mean}},P_{\text{std}}^2),P_{\text{mean}}=-1.2,P_{\text{std}}=1.2$$
另外，在训练时，为了避免在小数据集上出现过拟合的现象，作者还引入了数据扩增管线，并将数据扩增方法作为额外的条件输入给模型以避免泄露。

总结

EDM 是扩散模型发展中非常重要、非常具有实际意义的一篇工作。它构建了一个清晰、统一的扩散模型框架，将之前的几种经典方法都囊括其中，在这个框架下，不同的扩散模型设计选项相互解耦，并且对训练、采样各过程的设计选项进行了清晰的拆分。随后，针对各个设计选项，或通过详尽的理论推导，或通过扎实的实验分析，得出了最优的扩散模型设计形式和典型值。从各大扩散模型代码库中带有 “Karras” 字样的采样和训练设计选项可以得知，这些结果直到今天都还非常有价值。

不过毕竟实验只是做在 32、64 这种小数据集上的，模型规模也不大，笔者比较好奇的是这些结论在高分辨率图像、大数据大模型上的应用价值如何呢？