浅析扩散模型与图像生成【理论篇】（一）——DDPM

写在最前面：近几年扩散模型已经逐步取代GAN、VAE等模型成为图像生成领域最炙手可热的研究方向，本文从基础理论的角度出发，简要的介绍一下扩散模型在图像生成领域的几个代表性工作，包括DDPM、DDIM和NCSN等。为保证文章的可读性，本文将尽可能避免具体的公式推导和证明，使用更加通俗易懂的方式来介绍相关内容。

• DDPM
• NCSN
• SDE
• DDIM
• IDDPM

   Ho et al. [1] 和 Song et al. [2] 分别从扩散模型（DDPM）和分数模型（NCSN）的角度出发阐述了如何实现图像生成，后面我们会知道二者本质上是等价的，只相差一个系数，而且可以用随机微分方程的形式进行统一（SDE）[3]。下面，我们将从DDPM开始，介绍一下扩散模型是如何实现图像生成的。

1. DDPM

   在DDPM中，图像生成被分为两个阶段：扩散和采样。所谓扩散阶段$q(x_T|x_0)$，就是对一幅图像$x_0$通过逐步增加噪声$ϵ$，使其最终变成一个随机噪声$x_T$；而采样阶段$p_{\theta }(x_0)$，则是从一个随机噪声$x_T$中通过逐步去噪，最终得到一张生成的图像$x_0$。无论是扩散还是采样过程，都是一个马尔可夫链，每一步的预测都依赖于上一步的输出结果。在扩散阶段我们将会训练一个神经网络（UNet），让其根据当前的时间步数$t$和输入的噪声图像$x_t$，预测一个噪声${ϵ}_{\theta }(x_t,t)$，用于采样阶段中对图像进行去噪。下图很形象的表述了扩散和采样过程（图源：B站博主VictorYuki的视频“从零开始了解 Diffusion Model”），图中使用$\tilde{z}$表示预测噪声${ϵ}_{\theta }$。

扩散

采样

   有了一个大致的概念后，我们来具体介绍一下扩散和采样的实现过程。首先给定一幅图像$x_0$，根据上面的介绍我们得知，我们可以对其添加噪声得到$x_1$，再对$x_1$添加噪声得到$x_2$，重复这个过程直至得到纯噪声$x_T$，该过程可以表示为$$q(x_{1:T}|x_0):=\prod _{t=1}^{T}q(x_t|x_{t-1})\text{\hspace{1em}(1-1)}$$
其中
$$q(x_t|x_{t-1}):=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)\text{\hspace{1em}(1-2)}$$
作者将每一步的扩散过程描述成从以$\sqrt{1-{\beta }_t}{x}_{t-1}$为均值，以${\beta }_t$为方差的高斯分布中采样的过程，利用重参数技巧可转化为
$$x_t=\sqrt{1-{\beta }_t}{x}_{t-1}+\sqrt{{\beta }_t}ϵ，ϵ\backsim \mathcal{N}(0,I)\text{\hspace{1em}(1-3)}$$
可以看到每一步的$x_t$都是由上一步的输出结果$x_{t-1}$与噪声$ϵ$进行加权求和得到的，权重系数${\beta }_t$随着$t$的增加而逐步增长，直至$x_T$变成一个完全的随机噪声。这个过程通常要迭代成百上千次，在DDPM中$T$取值为1000，这就导致扩散的过程十分漫长。为了解决这个问题，作者提出一种快速扩散的方法，可以由$x_0$直接生成任意时刻$t$的噪声图$x_t$，如下式$$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)I)\text{\hspace{1em}(1-4)}$$
即
$$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ，ϵ\backsim \mathcal{N}(0,I)\text{\hspace{1em}(1-5)}$$
其中${\alpha }_t=1-{\beta }_t,{\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$。
   介绍完扩散过程，我们再来看一下采样过程。采样过程（也叫重建过程）是一个逆扩散过程，即从随机采样得到的噪声$X_T$逐步去噪直至生成一个图像$x_0$，表示为
$$p_{\theta }(x_{0:T}):=p(x_T)\prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t)\text{\hspace{1em}(1-6)}$$
其中$p(x_T)=\mathcal{N}(0,I),p_{\theta }(x_{t-1}|x_t):=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))$。分布$p_{\theta }(x_{t-1}|x_t)$实际上是未知的，而逆转后的扩散过程的分布$q(x_{t-1}|x_t,x_0)$是可以利用贝叶斯公式计算出来的，因此作者使用神经网络模型去拟合一个分布，使其逼近$q(x_{t-1}|x_t,x_0)$。
$$q(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};{\tilde{\mu }}_t(x_t,x_0),{\tilde{\beta }}_{t}I)\text{\hspace{1em}(1-8)}$$
其中
$${\tilde{\mu }}_t(x_t,x_0):=\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{x}_0+\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t,{\tilde{\beta }}_t:=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t\text{\hspace{1em}(1-9)}$$
训练的目标函数为
$${\mathcal{L}}_{VLB}={\mathbb{E}}_q\left[\underset{L_T}{\underbrace{D_{KL}(q(x_T|x_0)||p(x_T))}}+\sum _{t>1}\underset{L_{t-1}}{\underbrace{D_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{x-1}|x_t))}}\underset{L_0}{\underbrace{-\log p_{\theta }(x_0|x_1)}}\right]\text{\hspace{1em}(1-10)}$$
该目标函数也就是所谓的变分下界（variational lower bound,VLB），其中第一项$L_T$不包含可训练的参数，因此在训练时通常忽略，而对于第三项$L_0$，作者构建了一个离散化的分段积分累乘方法来计算，此处不再展开介绍，感兴趣的可阅读原文[1]。我们重点关注第二项$L_{t-1}$。
   对于分布$p_{\theta }(x_{x-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))$，当$0<t\le T$时，作者令方差项为一个无需训练只与时间$t$有关的常数，即${\Sigma }_{\theta }(x_t,t)={\sigma }_t^{2}I$，其中${\sigma }_t^2={\beta }_t$或${\sigma }_t^2={\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t$。实验结果表明两种取值方式并没有显著差别，${\Sigma }_{\theta }(x_t,t)$也可以作为一个可训练参数，但本文的实验表示这会使得训练不稳定。利用多元高斯分布的KL散度求解公式，$L_{t-1}$可以写为
$$L_{t-1}={\mathbb{E}}_q\left[\frac{1}{2{\sigma }_t^2}{\Vert {\tilde{\mu }}_t(x_t,x_0)-{\mu }_{\theta }(x_t,t)\Vert }^2\right]+C\text{\hspace{1em}(1-11)}$$
则训练目标转化为训练一个网络根据输入的噪声图$x_t$和时间$t$输出对应的均值${\mu }_{\theta }(x_t,t)$使其接近${\tilde{\mu }}_t(x_t,x_0)$。利用公式(1-5)，我们可以得到
$$x_0=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}ϵ)\text{\hspace{1em}(1-12)}$$
将其带入${\tilde{\mu }}_t(x_t,x_0)$的表达式中可得
$${\tilde{\mu }}_t(x_t,x_0)=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}ϵ)\text{\hspace{1em}(1-13)}$$其中$x_t,{\alpha }_t,{\beta }_t$都是已知量，只有噪声$ϵ$是一个随机变量，因此预测${\mu }_{\theta }(x_t,t)$进一步转化为根据当前的输入$x_t$和时间$t$预测一个噪声${ϵ}_{\theta }(x_t,t)$，即
$${\mu }_{\theta }(x_t,t)={\tilde{\mu }}_t\left(x_t,\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t))\right)=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t,t)\right)\text{\hspace{1em}(1-14)}$$
则训练目标可以进一步的写为
$$L_{t-1}={\mathbb{E}}_{x_0,ϵ}\left[\frac{{\beta }_t^2}{2{\sigma }_t^2{\alpha }_t(1-{\bar{\alpha }}_t)}{\Vert ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t)\Vert }^2\right]\text{\hspace{1em}(1-15)}$$
结合$L_{t-1}$和$L_0$，最终的目标函数${\mathcal{L}}_{VLB}$简化为
$$L_{simple}(\theta ):={\mathbb{E}}_{t,x_0,ϵ}\left[{\Vert ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t)\Vert }^2\right]\text{\hspace{1em}(1-16)}$$
利用扩散过程中输入的$x_0$和每次随机采样的噪声$ϵ$就能训练得到一个噪声预测模型${ϵ}_{\theta }(x_t,t)$。最后在采样阶段，根据$x_{t-1}\backsim p_{\theta }(x_{t-1}|x_t)$可得
$$x_{t-1}={\mu }_{\theta }(x_t,t)+{\sigma }_{t}z=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t,t)\right)+{\sigma }_{t}z,z\backsim \mathcal{N}(0,I)\text{\hspace{1em}(1-17)}$$
当模型训练完成时，${\mu }_{\theta }(x_t,t)$是一个确定过程，只与输入的$x_t$有关，而在每次重建时都会随机采样一个$z$来增加一定的随机性，使得生成样本具有更大的多样性。我们再看一下作者给出的训练和采样阶段的算法流程

  训练阶段，首先采样得到一个图像样本$x_0$，然后从 { 1 , . . . , T } \{1,...,T\} {1,...,T}中均匀采样一个时间$t$，并按照标准正态分布采样一个噪声$ϵ$，利用目标函数(1-16)训练一个UNet模型用于预测${ϵ}_{\theta }(x_t,t)$。在采样阶段，在标准正态分布中采样得到初始的噪声图$x_T$，并且在每一步都采样一个随机变量$z$，然后按照公式(1-17)，由$T$到1逐步预测$x_{t-1}$，直至得到$x_0$。值得注意的是当$t=1$时，也就是最后一步预测时，随机变量$z=0$。
  至此，我们对DDPM的扩散和采样过程，以及训练目标进行了一个相对完整的介绍。我们只给出了结论性的公式，如果要了解具体推导和证明步骤，可参考原文及附录，这篇博客和这个视频也给出了更加详细的分析。受篇幅限制，本文仅对DDPM进行了介绍，其他的方法将在后续的文章中分别介绍。

参考文献
[1] Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in neural information processing systems, 33, 6840-6851.
[2] Song Y, Ermon S. Generative modeling by estimating gradients of the data distribution[J]. Advances in neural information processing systems, 2019, 32.
[3] Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., & Poole, B. (2020). Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456.