扩散模型详解——DDPM

前言

扩散模型是一种新型生成式模型，专攻于图像生成任务。其特点是把图像视为张量空间中离散分布的点，每个点之间由大量高斯随机噪声淹没，图像与图像之间存在相当的距离，因此，能够设计一种噪声去除的方法，将不同张量点之间的噪声去除，即可得到离散的张量点值，即需要生成的图像。

Denoising Diffusion Probabilistic Model

去噪扩散概率模型（以下简称扩散模型），英文缩写为DDPM，是20年以降兴起的一种基于马尔可夫链生成式概率空间模型，它的特点是将噪声去除的过程建模为离散马尔可夫链，使用条件概率逐项传播迭代，且计算过程完全位于图像空间中，计算量极大，理解上比较直观。

真是蚌埠住了，VAE和DDPM的贝叶斯理论符号假设正好相反，可千万别搞混了。

前提与假设

扩散模型的数学底层工具是马尔可夫链，首先假设图像空间为$\mathbf{x}\in\Omega(\mathbf{x})$，把一幅完整美观的图像视作为一个低熵体，记为$\mathbf{x}_0$，将其放于封闭空间中自行扩散Diffusion，文章假设：扩散过程是不断增加高斯白噪声的过程，直到达到热平衡，即图像完全随机化destoryed，扩散过程增加的噪声服从参数为$\boldsymbol{\theta}$的高斯分布。称扩散过程为正过程。

从热平衡图像恢复至完全低熵图像的过程为逆过程（或反过程）。称达到热平衡的图像为起始图像（即生成式模型最初使用的图像），记为 $\mathbf{x}_T$，其服从标准白噪声分布，即 $\mathbf{x}_T\sim\mathcal{N}(\mathbf{0,I})$，在文中记为 $p=\mathcal{N}(\mathbf{x}_T;\mathbf{0,I})$，将该图像恢复至 $\mathbf{x}_0$ 的过程即为逆过程。

扩散过程假设为马尔可夫链，即每步扩散所增加的白噪声都是服从高斯分布（不是定常的）。记扩散过程中所涉及到的真实参数和真实概率分布为 $q_{\boldsymbol{\phi}}$，而该分布$q$和参数$\boldsymbol{\phi}$实际上是无法求解的，因此设计一个合理的分布去逼近扩散分布使得热平衡图像的恢复成为可能，这个逼近的分布就被称为 $p_{\boldsymbol{\theta}}$ 实际参数和实际分布（和VAE论文是反着的，注意啦！）。

优化目标

根据马尔可夫链，正反过程都是马尔可夫的，因此其正过程可以表示为：

$$q(\mathbf{x}_{1:T}|\mathbf{x}_0)=\prod_{t=1}^Tq(\mathbf{x}_t|\mathbf{x}_{t-1})$$

逆过程可以表示为：

$$p(\mathbf{x}_{0:T-1}|\mathbf{x}_T)=\prod_{t=1}^Tp(\mathbf{x}_{t-1}|\mathbf{x}_{t})$$

根据假设，扩散过程中每一步都是一次高斯白噪声，因此可写出正过程的单步条件概率分布：

$$q(\mathbf{x}_t|\mathbf{x}_{t-1})=\mathcal{N}[\mathbf{x}_t;\boldsymbol{\mu}_{\boldsymbol{\phi}}(\mathbf{x}_{t-1}),\boldsymbol{\Sigma}_{\boldsymbol{\phi}}\mathbf{I}]$$

同理，逆过程的单步条件概率分布：

$$p(\mathbf{x}_{t-1}|\mathbf{x}_t)=\mathcal{N}[\mathbf{x}_{t-1};\boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{x}_t),\boldsymbol{\Sigma}_{\boldsymbol{\theta}}(\mathbf{I})]$$

贝叶斯算法的优化目标通常是缩小实际分布与真实分布的KL散度，即：

$$\begin{aligned} &\argmin_{\boldsymbol{\theta}}D_\mathrm{KL}\left[\frac{\log q(\mathbf{x}_0)}{\log p_{\boldsymbol{\theta}}(\mathbf{x}_0)}\right]\\ &=\argmin_{\boldsymbol{\theta}}\mathbb{E}_{q(\mathbf{x}_0)}[\log q(\mathbf{x}_0)-\log p_{\boldsymbol{\theta}}(\mathbf{x}_0)] \end{aligned}$$

其中，$q(\mathbf{x}_0)$是一个确定的值，即逆过程的初始图像，即为热平衡图像，因此优化目标可简化为：

$$\argmin_{\boldsymbol{\theta}}\mathbb{E}_q[-\log p_{\boldsymbol{\theta}}(\mathbf{x}_0)]$$

文章认为，正过程中施加的高斯白噪声的均值和方差与上一个时刻的图像有关且增量极小，图像各个维度之间的噪声都是互不相关的。

化简

以上就是要求的全部内容了，本质上就是求使得目标函数最小的参数值$\boldsymbol{\theta}$。但是直接求解解不出来。所以作者对上述目标作了一个放缩，改为求放缩后函数的变分上界。

$$\begin{aligned} &\argmin_{\boldsymbol{\theta}}\mathbb{E}_q[-\log p_{\boldsymbol{\theta}}(\mathbf{x}_0)]\\ &\le\argmin_{\boldsymbol{\theta}}\mathbb{E}_q\left[-\log \frac{p_{\boldsymbol{\theta}}(\mathbf{x}_0)}{q(\mathbf{x}_{1:T}|\mathbf{x}_0)}\right]\\ \end{aligned}$$

这一放缩过程可谓是巧夺天工，为了推出它，我不知道了花了多少的时间，主要一个核心在于：

期望的表示法！

$$\begin{aligned} \log p_{\boldsymbol{\theta}}(\mathbf{x}_0)&=\int p_{\boldsymbol{\theta}}(\mathbf{x}_{0:T})\mathrm{d}\mathbf{x}_{1:T}\\ &=\sum_{\mathbf{x}_{1:T}}q(\mathbf{x}_{1:T}|\mathbf{x}_0)\frac{p_{\boldsymbol{\theta}}(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}|\mathbf{x}_0)}\\ &=\mathbb{E}_{q(\mathbf{x}_{1:T}|\mathbf{x}_0)}\left[\frac{p_{\boldsymbol{\theta}}(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}|\mathbf{x}_0)}\right] \end{aligned}$$

然后就可以推导得到后续的目标函数表达式：

$$\mathbb{E}_q\left\{\underbrace{D_\mathrm{KL}\left[q(\mathbf{x}_T|\mathbf{x}_0)||p(\mathbf{x}_T)\right]}_{L_T} +\sum_{t=2}^T\underbrace{D_\mathrm{KL}\left[q(\mathbf{x}_{t-1}|\mathbf{x}_0,\mathbf{x}_t)||p_{\boldsymbol{\theta}}(\mathbf{x}_{t-1}|\mathbf{x}_t)\right]}_{L_{t-1}}\underbrace{-\log p_{\boldsymbol{\theta}}(\mathbf{x}_0|\mathbf{x}_1)}_{L_0}\right\}$$

目标函数最优化

目标函数未和项包含$q(\mathbf{x}_t|\mathbf{x}_{0})$类型、$q(\mathbf{x_{t-1}}|\mathbf{x}_t,\mathbf{x}_0)$以及$p_{\boldsymbol{\theta}}(\mathbf{x}_{t-1}|\mathbf{x}_t)$ 这三种类型，其中含$p$的项在训练中根据施加的噪声情况是已知量，在此不必讨论。前两项可分别称为逆扩散过程项和逆扩散过程联合分布项，两项的求解如下所示：

逆过程项

扩散过程施加的热噪声在训练过程中可人为控制，但是在推理过程中处于完全未知的状态。

文章假设：单步扩散的高斯白噪声与上一步的白噪声参数相差很小，且高度相关（以符合热力学定律）。

因此将扩散过程的条件概率分布写为：

$$q(\mathbf{x}_t|\mathbf{x}_{t-1})=\mathcal{N}(\mathbf{x}_t;\sqrt{1-\beta_t}\mathbf{x}_{t-1},\beta_t\mathbf{I})$$

其中$\beta_0,\beta_1,...,\beta_T\approx 0$都是很小的数，保证扩散过程是近似均匀的（小量假设）。为了求解目标函数中的前两项$L_T,L_{t-1}$，需要求解条件概率$q(\mathbf{x}_t|\mathbf{x}_0)$。因此，使用换元法：$\alpha_t=1-\beta_t$

注意，带有$\beta$的乘积项可以看作为均值为0的高斯分布。

对待求的条件概率，直接递推如下：

$$\begin{aligned} \mathbf{x}_t&=\sqrt{1-\beta_t}\mathbf{x}_{t-1}+\sqrt{\beta_t}\boldsymbol{\epsilon}_{t-1}\\ &=\sqrt{1-\beta_t}(\sqrt{1-\beta_{t-1}}\mathbf{x}_{t-2}+\sqrt{\beta_{t-1}}\boldsymbol{\epsilon}_{t-2})+\sqrt{\beta_t}\boldsymbol{\epsilon}_{t-1}\\ &=\sqrt{\alpha_t}(\sqrt{\alpha_{t-1}}\mathbf{x}_{t-2}+\sqrt{1-\alpha_{t-1}}\boldsymbol{\epsilon}_{t-2})+\sqrt{\beta_t}\boldsymbol{\epsilon}_{t-1}\\ &=\sqrt{\alpha_t}\sqrt{\alpha_{t-1}}\mathbf{x}_{t-2}+\sqrt{\alpha_t-\alpha_t\alpha_{t-1}}\boldsymbol{\epsilon}_{t-2}+\sqrt{\beta_t}\boldsymbol{\epsilon}_{t-1}\\ \end{aligned}$$

其中的隐变量$\boldsymbol{\epsilon}$也是服从高斯分布的随机矢量，因此该条件概率可以看作是高斯分布随机矢量的线性组合。使用换元法reparameterization，可把含有隐变量的两项合并为：

$$\tilde{\boldsymbol{\epsilon}}_{t-2}=\sqrt{\alpha_t-\alpha_t\alpha_{t-1}}\boldsymbol{\epsilon}_{t-2}+\sqrt{\beta_t}\boldsymbol{\epsilon}_{t-1}\sim\mathcal{N}(\tilde{\boldsymbol{\epsilon}}_{t-2};\mathbf{0},(1-\alpha_t\alpha_{t-1})\mathbf{I})$$

这里使用到了之前的小量假设，对每一项推导式均使用一次换元法，则可得最终的单项表达式：

$$\begin{aligned} \mathbf{x}_t&=\sqrt{\alpha_t\alpha_{t-1}\cdots\alpha_1}\mathbf{x}_{0}+\sqrt{1-\alpha_t\alpha_{t-1}\cdots\alpha_1}\tilde{\boldsymbol{\epsilon}}_t\\ &=\sqrt{\bar{\alpha}_t}\mathbf{x}_0+\sqrt{1-\bar{\alpha}_t}\tilde{\boldsymbol{\epsilon}}_t\end{aligned}\tag{1}$$

逆过程联合分布项

该项不再引入另外的假设条件，直接使用条件概率定义式展开得：

$$\begin{aligned} q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0)&=\frac{q(\mathbf{x}_{t-1},\mathbf{x}_t|\mathbf{x}_0)}{q(\mathbf{x}_t|\mathbf{x}_0)}\\ &=\frac{q(\mathbf{x}_t|\mathbf{x}_{t-1},\mathbf{x}_0)q(\mathbf{x}_{t-1}|\mathbf{x}_0)}{q(\mathbf{x}_t|\mathbf{x}_0)}\\ &=\frac{q(\mathbf{x}_t|\mathbf{x}_{t-1})q(\mathbf{x}_{t-1}|\mathbf{x}_0)}{q(\mathbf{x}_t|\mathbf{x}_0)} \end{aligned}$$

所有的项都是已知项，因此将全部项的高斯分布概率密度函数代入，忽略归一化因子（不影响结果，因为是常数，在优化时随对数部分转移为定值）。可得：

$$\begin{aligned} q(x_{t-1}|x_t,x_0)&=\frac{q(x_t|x_{t-1})q(x_{t-1}|x_0)}{q(x_t|x_0)}\\ &=\frac{\displaystyle\exp\left[-\frac{\left(x_t-\sqrt{1-\beta_t}x_{t-1}\right)^2}{2\beta_t}\right]\times\exp\left[-\frac{\left(x_{t-1}-\sqrt{\bar{\alpha}_{t-1}}x_0\right)^2}{2(1-\bar{\alpha}_{t-1})}\right]}{\displaystyle\exp\left[-\frac{\left(x_t-\sqrt{\bar{\alpha}_{t}}x_0\right)^2}{2(1-\bar{\alpha}_{t})}\right]} \end{aligned}$$

这里，需要将上述式子整理成以$\mathbf{x}_{t-1}$
为中心变量的高斯函数表达式，其中系数部分就是方差，偏移部分就是均值，略去不含 $\mathbf{x}_{t-1}$ 的项，只看指数内部分，可简化为：

$$-\left[\frac{\alpha_t}{2\beta_t}+\frac{1}{2(1-\bar{\alpha}_{t-1})}\right]x^2_{t-1}+\left(\frac{\sqrt{1-\beta_t}x_t}{\beta_t}+\frac{\sqrt{\bar{\alpha}_{t-1}}x_0}{1-\bar{\alpha}_{t-1}}\right)x_{t-1}$$

比较高斯分布的系数，（注意小量假设的应用）可得：

$$\left\{\begin{aligned} \frac{1}{2\sigma^2}=\frac{\alpha_t}{2\beta_t}+\frac{1}{2(1-\bar{\alpha}_{t-1})}\\\\ \frac{\mu}{2}=\frac{\displaystyle\frac{\sqrt{1-\beta_t}x_t}{\beta_t}+\frac{\sqrt{\bar{\alpha}_{t-1}}x_0}{1-\bar{\alpha}_{t-1}}}{\displaystyle\frac{\alpha_t}{2\beta_t}+\frac{1}{2(1-\bar{\alpha}_{t-1})}} \end{aligned}\right.\implies \left\{\begin{aligned} \bm{\tilde{\mu}_t}(x_t,x_0)&=\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}\mathbf{x}_0+\frac{\alpha(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}\mathbf{x}_t\\\\ \sigma^2_t&=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\beta_t \end{aligned}\right.\tag{*}$$

注意估计的参数$\bm{\tilde{\mu}}_t$需要使用当前时刻的图像和热平衡图像，而当前时刻的图像$\mathbf{x}_t$必须通过上一时刻图像$\mathbf{x}_{t-1}$经去噪UNet $\bm{\epsilon_\theta}$作用后才能得到，因此公式(*)的计算将涉及到去噪UNet，具体过程参考章节训练过程。

模型设计

扩散模型的整体结构可以如下图所示：


说明：图中使用到Latent Diffusion的方法，即在VAE描述提取的隐空间下实现扩散，而原始的DDPM只在图像空间下扩散，计算效率不高，现已经被LDM取代。

如果刨去VAE的部分不看，模型整体上只由两个部分构成，即去噪UNet和Embedding模块（语义嵌入模块）去噪UNet接受带噪图像的输入，在论文中以$\bm{\epsilon_\theta}$的形式出现，产生去噪输出。语义嵌入模块在原始的论文中并未设计，该部分实现上使用的是条件扩散模型和空间注意力完成。不管模型如何设计，其都可以使用模型参数$\bm{\epsilon_\theta}(\mathbf{x}_t,t)$来表征，参考公式(1)，网络实现的输出就是对$\tilde{\epsilon}_t$的估计值，即有：

$$\tilde{\bm{\epsilon}}_t\approx\bm{\epsilon_\theta}(\mathbf{x}_t,t)$$

根据公式(1)反解：

$$\begin{aligned} \mathbf{x}_0&=\frac{1}{\sqrt{\bar{\alpha}_t}}\left(\mathbf{x}_t-\sqrt{1-\bar{\alpha}_t}\tilde{\bm{\epsilon}}_t\right)\\ \implies\mathbf{\hat{x}}_0&=\frac{1}{\sqrt{\bar{\alpha}_t}}\left[\mathbf{x}_t-\sqrt{1-\bar{\alpha}_t}\bm{\epsilon_\theta}(\mathbf{x}_t,t)\right] \end{aligned}$$

和公式(*)联立，可得：

$$\tilde{\bm{\mu}}_t(\mathbf{x}_t,\mathbf{\hat{x}}_0)=\frac{1}{\sqrt{\alpha_t}}\left[\mathbf{x}_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\bm{\epsilon_\theta}(\mathbf{x}_t,t)\right]$$

训练过程

DDPM的训练是完成的马尔可夫链，即首先正向扩散加入足够多的噪声，分别记录下正向扩散各个时间段得到的带噪图像。然后进行逆向去噪，每步去噪结果与记录的带噪图像作差计算损失函数完成对UNet的优化。去噪步长通常较大，$T=1000$。
利用换元法，在已知当前图像$\mathbf{x}_t$的情况下，可以推断逆过程前一个时刻的噪声图像：

$$\mathbf{\hat{x}}_t=\frac{1}{\sqrt{\alpha_t}}\left[\mathbf{x}_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\bm{\epsilon_\theta}(\mathbf{x}_t,t)\right]+\sigma_t\mathbf{z},\quad\mathbf{z}\sim\mathcal{N}(\mathbf{0,I})$$

由此可以完成从$t=T$向$t=0$的逆过程训练。

扩散编码器和解码器

编码器和解码器是基于VAE模型提出的特征提取与反演网络，本质上是网络而不是算法。

前提假设

文章认为，扩散过程的高斯分布参数都是事先选好的定值，因此不需要学习，损失函数中的$L_T$项可以被忽略。

另一个重要的损失函数部分就是$L_{t-1}$，在通式中其表现为KL散度，描述的是分布$q_{\boldsymbol{\phi}}$和$p_{\boldsymbol{\theta}}$ 的差距，而在文献中却变成了关于两个高斯分布的均值的RMSE。

KL散度描述的是两个分布之间的差距，当两个分布是同型分布时，完全可以使用参数的差距代替KL散度的计算。

而文章中假定扩散过程的$\beta_t\mathbf{I}$和逆过程的$\tilde{\beta}_t=\frac{1-\tilde{\alpha}_{t-1}}{1-\tilde{\alpha}_t}\beta_t$是几乎相等的。因此，仅考虑均值的差距就可以度量两个分布之间的差距了。有公式：

$$L_{t-1}=\mathbb{E}_q\left[\frac{1}{2\sigma_t^2}||\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t,\mathbf{x}_0)-\boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{x}_t,t)||^2\right]+C$$

其中，扩散过程的$\boldsymbol{\theta}$在训练中选定，因此其均值$\boldsymbol{\mu}_{\boldsymbol{\theta}}$也是已知量。另一项$\tilde{\boldsymbol{\mu}}_t$是本次训练过程中解出来的逆过程参数。如果将上面计算得到的$\mathbf{x}_t=\sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t}\tilde{\epsilon}_0$表达，则可以具体化该损失函数：

$$L_{t-1}-C=\mathbb{E}_{\mathbf{x}_0,\boldsymbol{\epsilon}}\left[\frac{1}{2\sigma_t^2}\left|\left|\frac{1}{\sqrt{\alpha_t}}\left(\mathbf{x}_t(\mathbf{x}_0,\boldsymbol{\epsilon})-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\boldsymbol{\epsilon}\right)-\boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{x}_t(\mathbf{x}_0,\boldsymbol{\epsilon}),t)\right|\right|\right]$$

换元法

换元法的思想来自于VAE论文中对真实分布的先验性假设，即$\tilde{\mathbf{x}}=\mathbf{x}+\mathbf{z}$；而本文是高斯类型的变量，因此换元方式适用于$\tilde{x}=x+y\sigma$，其中$y\sim\mathcal{N}(0,1)$。在逆过程推断过程时，在已知逆过程后一个时刻的噪声图像$\mathbf{x}_t$，要根据去噪UNet $\bm{\epsilon_\theta}$完成对前一个时刻的噪声图像$\mathbf{x}_{t-1}$的推断时，就采用了这样的换元方法：

$$\boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{x}_t(\mathbf{x}_0,\boldsymbol{\epsilon}),t)=\frac{1}{\sqrt{\alpha_t}}\left(\mathbf{x}_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t1}}\boldsymbol{\epsilon}_{\boldsymbol{\theta}}\right)$$

从而完成了对真实分布的量化，进一步简写损失函数可得：

$$\mathbf{E}_{\mathbf{x}_0,\boldsymbol{\epsilon}}\left[\frac{\beta_t^2}{2\sigma^2_t\alpha_t(1-\bar{\alpha}_t)}||\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\sqrt{\bar{\alpha}_t}\mathbf{x}_0-\sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon},t)||\right]$$

进一步简化甚至可以把前面的系数也简化掉，那个似乎也是常数。

扩散自编码器

重新看图，扩散正过程与逆过程是不是与VAE的编码器与解码器有异曲同工之处，二者都是将图像空间的张量经过一系列运算得到某个隐藏的变量，然后又可以通过一系列反运算将隐藏的变量恢复于原始的图像张量，即整个扩散过程就是从图像空间$\mathcal{X}$向图像空间$\mathcal{X}$的满射，而中间产物$\mathbf{x}_T$就可以作为图像张量的一种描述，即编码结果，因此扩散模型也可以用作为一种编码-解码器，文中称为去噪自编码器Denoising Auto-encoder。

去噪编码器

编码过程就是加噪声的过程，得到的热平衡图像就是图像的描述隐变量。

去噪解码器

解码过程就是去噪的过程，注意到解码过程中引入了去噪UNet，因此需要对该网络作优化，优化过程即使用上述的损失函数$L_{t-1}$，对UNet的参数进行优化，简单地可以写为：

$$L_\mathrm{simple}(\bm{\theta})=\mathbb{E}_{t,\mathbf{x}_0,\bm{\epsilon}}\left[\|\bm{\epsilon}-\bm{\epsilon_\theta}(\sqrt{\bar{\alpha}_t}\bm{x}_0+\sqrt{1-\bar{\alpha}_t}\bm{\epsilon},t)\|^2\right]$$

建议译名

DDPM：去噪扩散概率模型

参考文献

[1] Ho J. , Jain A. , Abbeel P. .Denoising Diffusion Probabilistic Models[M/OL].arXiv,2020