去噪扩散概率模型

Title:Denoising Diffusion Probabilistic Models

paper是UC Berkeley发表在NIPS 2020的工作

paper地址

Abstract

我们使用扩散概率模型（一类受非平衡热力学启发的潜在变量模型）展示了高质量的图像合成结果。根据扩散概率模型与采用Langevin动力学的去噪得分匹配之间的新颖联系设计的加权变分下界进行训练，我们获得了最佳结果。我们的模型自然而然地使用一种渐进的有损解压方案，这可以解释为自回归解码的推广。在无条件CIFAR10数据集上，我们取得了9.46的Inception得分和3.17的最先进的FID得分。在 $256\times 256$ 的LSUN数据集上，我们获得了与ProgressiveGAN相似的样本质量。代码。

1 Introduction

各种深度生成模型最近在多种数据模态中展示了高质量的样本。生成对抗网络（GANs）、自回归模型、流模型和变分自编码器（VAEs）已经合成了引人注目的图像和音频样本【14, 27, 3, 58, 38, 25, 10, 32, 44, 57, 26, 33, 45】。同时，并且基于能量的建模和得分匹配方面的显著进展也产生了与GANs可媲美的图像【11, 55】。

图 1：在 CelebA-HQ 256 × 256（左）和无条件 CIFAR10（右）上生成的样本

图 2：本工作中考虑的有向图模型。

本文介绍了扩散概率模型的进展【53】。扩散概率模型（简称为“扩散模型”）是一种参数化的马尔可夫链，通过变分推断进行训练，在有限时间内生成与数据匹配的样本。这个链的转移过程被学习来逆转扩散过程，即一个逐渐向数据中添加噪声的马尔可夫链，直到信号被完全破坏。当扩散过程由少量高斯噪声组成时，只需将采样链的转移过程也设置为条件高斯分布，从而允许采用特别简单的神经网络参数化方式。

扩散模型定义简单且训练高效，但据我们所知，尚未有研究表明它们能够生成高质量样本。我们展示了扩散模型实际上能够生成高质量样本，有时甚至优于其他类型生成模型的已发表结果（第4节）。此外，我们展示了一种特定的扩散模型参数化方法，这种方法揭示了在训练过程中扩散模型与多个噪声级别上的去噪得分匹配，以及在采样过程中与退火Langevin动力学的等价性（第3.2节）【55, 61】。我们使用这种参数化方法获得了最好的样本质量结果（第4.2节），因此我们认为这种等价性是我们主要的贡献之一。

尽管我们的模型样本质量较高，但与其他基于似然的模型相比，我们的模型在对数似然上并不具有竞争力（不过，我们的模型在对数似然上优于已报告的退火重要性采样对能量基模型和得分匹配产生的大量估计【11, 55】）。我们发现，大多数模型的无损编码长度主要用于描述难以察觉的图像细节（第4.3节）。我们以有损压缩的语言对这一现象进行了更精细的分析，并展示了扩散模型的采样过程是一种渐进解码形式，类似于按照位排序进行的自回归解码，这大大扩展了自回归模型通常可能实现的范围。

2 Background

扩散模型【53】是一种以 $p_{\theta }\left({\mathbf{x}}_0\right):=\int p_{\theta }\left({\mathbf{x}}_{0:T}\right)d{\mathbf{x}}_{1:T}$ 形式的潜在变量模型，其中 ${\mathbf{x}}_1,\dots ,{\mathbf{x}}_T$ 是与数据 ${\mathbf{x}}_0\sim q\left({\mathbf{x}}_0\right)$ 同维度的潜变量。联合分布 $p_{\theta }\left({\mathbf{x}}_{0:T}\right)$ 被称为反向过程，其定义为从 $p\left({\mathbf{x}}_T\right)=\mathcal{N}\left({\mathbf{x}}_T;0,\mathbf{I}\right)$ 开始的学习高斯过渡的马尔科夫链：

$$p_{\theta }\left({\mathbf{x}}_{0:T}\right):=p\left({\mathbf{x}}_T\right)\prod _{t=1}^T{p}_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right),p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right):=\mathcal{N}\left({\mathbf{x}}_{t-1};{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t,t\right),{\mathbf{\Sigma }}_{\theta }\left({\mathbf{x}}_t,t\right)\right)(1)$$

将扩散模型与其他类型的潜在变量模型区别开来的是，近似后验 $q\left({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0\right)$（被称为前向过程或扩散过程）固定为一个马尔科夫链，该链根据方差调度 ${\beta }_1,\dots ,{\beta }_T$ 逐渐向数据添加高斯噪声：

$$q\left({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0\right):=\prod _{t=1}^{T}q\left({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1}\right),q\left({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1}\right):=\mathcal{N}\left({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I}\right)(2)$$

训练过程主要是通过优化变分下界（也称作证据下界，或ELBO）的负对数似然性来完成的：

$$\mathbb{E}\left[-\log p_{\theta }\left({\mathbf{x}}_0\right)\right]\le {\mathbb{E}}_q\left[-\log \frac{p_{\theta }\left({\mathbf{x}}_{0:T}\right)}{q\left({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0\right)}\right]={\mathbb{E}}_q\left[-\log p\left({\mathbf{x}}_T\right)-\sum _{t\ge 1}\log \frac{p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)}{q\left({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1}\right)}\right]=:L(3)$$

可以通过重参数化【33】学习前向过程的方差 ${\beta }_t$，或者将其视为恒定的超参数，并通过选择高斯条件下的 $p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)$，确保反向过程的表达力，因为当 ${\beta }_t$ 很小的时候，这两个过程具有相同的函数形式【53】。前向过程的一个显著特性是，它允许在任意时间步长 $t$ 以封闭形式采样 ${\mathbf{x}}_t$：使用符号 ${\alpha }_t:=1-{\beta }_t$ 和 ${\bar{\alpha }}_t:={\prod }_{s=1}^t{\alpha }_s$，我们有：

$$q\left({\mathbf{x}}_t\mid {\mathbf{x}}_0\right)=\mathcal{N}\left({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,\left(1-{\bar{\alpha }}_t\right)\mathbf{I}\right)(4)$$

因此，我们可以通过使用随机梯度下降优化$L$的随机项进行高效训练。进一步的改进来自于通过将$L$（3）重写为以下格式来降低方差：
$${\mathbb{E}}_q[\underset{L_T}{\underbrace{D_{\mathrm{KL}}\left(q\left({\mathbf{x}}_T\mid {\mathbf{x}}_0\right)\Vert p\left({\mathbf{x}}_T\right)\right)}}+\sum _{t>1}\underset{L_{t-1}}{\underbrace{D_{\mathrm{KL}}\left(q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)\Vert p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)\right)}}\underset{L_0}{\underbrace{-\log p_{\theta }\left({\mathbf{x}}_0\mid {\mathbf{x}}_1\right)}}](5)$$
（详情请参见附录A。这些项的标签用于第3节。）公式（5）使用KL散度直接将$p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)$与条件于${\mathbf{x}}_0$时可以处理的前向过程后验进行比较：
$$\begin{array}{rl}q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)& =\mathcal{N}\left({\mathbf{x}}_{t-1};{\tilde{\mathbf{\mu }}}_t\left({\mathbf{x}}_t,{\mathbf{x}}_0\right),{\tilde{\beta }}_t\mathbf{I}\right)(6)\\ \text{ where }{\tilde{\mathbf{\mu }}}_t\left({\mathbf{x}}_t,{\mathbf{x}}_0\right)& :\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{\mathbf{x}}_0+\frac{\sqrt{{\alpha }_t}\left(1-{\bar{\alpha }}_{t-1}\right)}{1-{\bar{\alpha }}_t}{\mathbf{x}}_t\text{ and }{\tilde{\beta }}_t:=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t(7)\end{array}$$
因此，公式（5）中的所有KL散度都是高斯间的比较，所以它们可以以Rao-Blackwellized的方式通过封闭形式的表达式计算，而不是使用高方差的蒙特卡洛估计。

前向加噪过程的推导，即式（2）和式（4）
前向加噪（${\mathbf{x}}_{t-1}\to {\mathbf{x}}_t$）：
$${\mathbf{x}}_t=a_t{\mathbf{x}}_{t-1}+b_t{\epsilon }_t,{\epsilon }_t\sim \mathcal{N}(0,\mathbf{I})(a.1)$$
其中$a_t$和$b_t$表示图像与噪声的混合比例，${\epsilon }_t$是噪声，满足$a_t^2+b_t^2=1$的约束。
展开上式，得：
$$\begin{array}{rl}{\mathbf{x}}_t& =a_t{\mathbf{x}}_{t-1}+b_t{\epsilon }_t\\ & =a_t\left(a_{t-1}{\mathbf{x}}_{t-2}+b_{t-1}{\epsilon }_{t-1}\right)+b_t{\epsilon }_t\\ & =a_t{a}_{t-1}{\mathbf{x}}_{t-2}+a_t{b}_{t-1}{\epsilon }_{t-1}+b_t{\epsilon }_t\\ & =\dots \\ & =\left(a_t\dots a_1\right){\mathbf{x}}_0+\underset{\text{多个相互独立的正态噪声之和 }}{\underbrace{\left(a_t\dots a_2\right)b_1{\epsilon }_1+\left(a_t\dots a_3\right)b_2{\epsilon }_2+\cdots +a_t{b}_{t-1}{\epsilon }_{t-1}+b_t{\epsilon }_t}}\end{array}(a.2)$$
根据独立高斯的可加性：
假设 $X$和 $Y$是独立的随机变量，并且服从正态分布（因此也是联合正态分布），那么它们的和也是正态分布的。即：
如果
$$\begin{array}{rl}& X\sim N\left({\mu }_X,{\sigma }_X^2\right)\\ & Y\sim N\left({\mu }_Y,{\sigma }_Y^2\right)\\ & Z=X+Y\end{array}$$
然后
$$Z\sim N\left({\mu }_X+{\mu }_Y,{\sigma }_X^2+{\sigma }_Y^2\right).$$
这意味着两个独立的正态分布随机变量的和是正态的，其均值是两个均值的和，其方差是两个方差的和。
因此上式可以写为：
$$\begin{array}{c}{\mathbf{x}}_t=\left(a_t\dots a_1\right){\mathbf{x}}_0+\sqrt{{\left(a_t\dots a_2\right)}^2{b}_1^2+{\left(a_t\dots a_3\right)}^2{b}_2^2+\cdots +a_t^2{b}_{t-1}^2+b_t^2}{\overline{\mathbf{\epsilon }}}_t,{\overline{\mathbf{\epsilon }}}_t\sim \mathcal{N}(0,\mathbf{I})\end{array}(a.3)$$
注意到式(a.1)满足$a_t^2+b_t^2=1$的约束，下面给出推导。
计算式(a.2)系数的平方和：
$$\begin{array}{rl}& {\left(a_t\dots a_1\right)}^2+{\left(a_t\dots a_2\right)}^2{b}_1^2+{\left(a_t\dots a_3\right)}^2{b}_2^2+\cdots +a_t^2{b}_{t-1}^2+b_t^2\\ =& {\left(a_t\dots a_2\right)}^2{a}_1^2+{\left(a_t\dots a_2\right)}^2{b}_1^2+{\left(a_t\dots a_3\right)}^2{b}_2^2+\cdots +a_t^2{b}_{t-1}^2+b_t^2\\ =& {\left(a_t\dots a_2\right)}^2\left(a_1^2+b_1^2\right)+{\left(a_t\dots a_3\right)}^2{b}_2^2+\cdots +a_t^2{b}_{t-1}^2+b_t^2\\ =& {\left(a_t\dots a_3\right)}^2\left(a_2^2\left(a_1^2+b_1^2\right)+b_2^2\right)+\cdots +a_t^2{b}_{t-1}^2+b_t^2\\ =& a_t^2\left(a_{t-1}^2\left(\dots \left(a_2^2\left(a_1^2+b_1^2\right)+b_2^2\right)+\dots \right)+b_{t-1}^2\right)+b_t^2\end{array}(a.3)$$
如果满足$a_t^2+b_t^2=1$的约束，那么式(a.3)会大大简化。即式(a.2)的系数平方和也为1。记${\bar{a}}_t={\left(a_t\dots a_1\right)}^2$，式(a.3)可以写成：
$${\mathbf{x}}_t=\sqrt{{\bar{a}}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{a}}_t}{\bar{\epsilon }}_t,{\bar{\epsilon }}_t\sim \mathcal{N}(0,\mathbf{I})(a.4)$$
令$a_t=\sqrt{{\alpha }_t}$,代入式(a.4)可得：
$${\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}{\bar{\epsilon }}_t,{\bar{\epsilon }}_t\sim \mathcal{N}(0,\mathbf{I})$$
对应式(4)
由${\alpha }_t:=1-{\beta }_t$，重写式(a.1)可得：
$${\mathbf{x}}_t=\sqrt{(1-{\beta }_t)}{\mathbf{x}}_{t-1}+\sqrt{{\beta }_t}{\epsilon }_t,{\epsilon }_t\sim \mathcal{N}(0,\mathbf{I})$$
对应式(2)
上述高斯相加的过程，也可以从重参数化来理解：若希望从高斯分布 $N\left(\mu ,{\sigma }^2\right)$ 中采样，可以先从标准分布 $N(0,1)$ 采样出 $z$ ，再得到 $\sigma \ast z+\mu$ ，这就是我们想要采样的结果。这样做的好处是将随机性转移到了 $z$ 这个常量上，而 $\sigma$ 和 $\mu$ 则当作仿射变换网络的一部分。
此外，$a_t^2+b_t^2=1$根本原因在于这个约束保证了信号的总能量不变，同时确保了生成过程中的稳定性和归一化特性。

• 在信号传输和处理过程中，保持信号的原始能量有助于保留信号的主要特征，从而提高处理结果的保真度。
• 数值稳定性：归一化处理有助于避免计算过程中出现过大的数值，防止溢出或下溢。
• $a_t^2+b_t^2=1$描述的是一个单位圆上的点。在单位圆上，所有点到原点的距离都是1，表示信号和噪声在每一步的组合仍然在一个固定的能量水平上。
• 信号和噪声的加权可以看作是两个向量的合成，这两个向量在每一步的组合结果仍然保持在单位圆上，确保总的“长度”或能量不变。
• 数学推导的简化
  从方差合成归一化的视角理解：
  在前向加噪过程中，每一步加噪的形式为：
  $${\mathbf{x}}_t=\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1}+\sqrt{{\beta }_t}{\epsilon }_t,$$
  其中${\epsilon }_t\sim \mathcal{N}(0,\mathbf{I})$，并且${\beta }_t$控制噪声的比例。
  我们来看每一步的方差归一化过程：
  初始时刻$t=0$：
  $${\mathbf{x}}_0\sim \mathcal{N}(0,{\sigma }_0^2\mathbf{I}).$$
  第$t$步：
  $${\mathbf{x}}_t=\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1}+\sqrt{{\beta }_t}{\epsilon }_t.$$
  在这里，${\epsilon }_t$的方差为$\mathbf{I}$。
  为了计算${\mathbf{x}}_t$的方差，我们需要考虑两个独立正态分布变量的加权和：
  $$\begin{array}{rl}\text{Var}({\mathbf{x}}_t)& =\text{Var}(\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1}+\sqrt{{\beta }_t}{\epsilon }_t)\\ & =(1-{\beta }_t)\text{Var}({\mathbf{x}}_{t-1})+{\beta }_t\text{Var}({\epsilon }_t)\\ & =(1-{\beta }_t)\text{Var}({\mathbf{x}}_{t-1})+{\beta }_t\cdot \mathbf{I}.\end{array}$$
  累积方差的归一化
  假设在每一步中，${\mathbf{x}}_{t-1}$的方差为$\mathbf{I}$（这表示信号保持归一化方差）。则我们希望通过选择合适的${\beta }_t$保持方差的归一化，即每一步的${\mathbf{x}}_t$也具有单位方差。
  在每一步中：
  $(1-{\beta }_t)\mathbf{I}+{\beta }_t\mathbf{I}=\mathbf{I}.$
  这意味着无论我们选择的${\beta }_t$是多少，只要满足${\beta }_t\in [0,1]$，前向加噪过程中的信号和噪声的加权组合将始终保持单位方差$\mathbf{I}$。这就是方差合成归一化的意义。
  累积效应
  考虑从初始状态${\mathbf{x}}_0$到第$t$步的累积效应，我们有：
  $${\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}{\bar{\epsilon }}_t,$$
  其中${\bar{\alpha }}_t={\prod }_{i=1}^t(1-{\beta }_i)$是所有${\alpha }_i$的累积乘积，${\bar{\epsilon }}_t$是独立正态噪声的加权和。
  再次计算方差：
  $$\begin{array}{rl}\text{Var}({\mathbf{x}}_t)& =\text{Var}(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}{\bar{\epsilon }}_t)\\ & ={\bar{\alpha }}_t\text{Var}({\mathbf{x}}_0)+(1-{\bar{\alpha }}_t)\text{Var}({\bar{\epsilon }}_t)\\ & ={\bar{\alpha }}_t\mathbf{I}+(1-{\bar{\alpha }}_t)\mathbf{I}\\ & =\mathbf{I}.\end{array}$$