扩散模型(DDPM)介绍

文章主要介绍了DDPM扩散模型，包含了详细的数学推导，可能公式有点多，但是只要能够耐心看完，相信会有一些收获的。好了下面进入正题。

扩散模型

扩散模型是一种新的生成模型，可以生成高分辨率、高质量的图像。目前像谷歌，OpenAI，Nvidia等大厂都在持续投入研究，并推出自家的扩散模型，包括谷歌的Imagen、OpenAI的DALLE-2等，还有像近期大热的Midjourney和StableDiffusion，背后的原理均是基于扩散模型。但是扩散模型背后的原理到底是什么呢？今天笔者将聚焦于Denoising Diffusion Probabilistic Models (DDPM)这篇文章，带大家一探扩散模型的原理。

扩散过程

扩散过程主要分为两个阶段，前向过程和反向过程。当给定一张图像$x_0$时，前向过程就是经过$T$步，给输入图像$x_0$逐步加上高斯噪声，前向过程被用来生成训练过程中的目标数据；反向过程则与前向过程相反，是逐步去掉添加到$x_0$上的噪声，最后生成一张无噪声的图像。在实际操作中，我们通过一个神经网络（unet）来预测噪声。首先在训练阶段，通过前向过程我们知道了添加到图像上的噪声$s$（因为$s$是我们生成的），此噪声$s$即作为unet的目标输出，输入则为在$x_0$上添加噪声后的图像$x_s$，通过（$x_s$，$s$）数据对，我们可以训练一个具备预测噪声能力的unet，该unet具备能力：当输入含噪图像时，可以预测出该含噪图像上添加的噪声$s$，这样就可以计算出不含噪的图像。当然这里只介绍了一步，实际上反向过程是反复执行上面预测噪声，去掉噪声的多步去噪的过程。

前向过程

扩散模型的前向过程其实就是一个马尔科夫链。马尔科夫链到底是什么，这里我们不必深究，只需要知道它有一个特性就是链上新产生的一个数只依赖于这个新产生数的前一个数。前向过程的理解可以结合下图：$$\begin{array}{c}q(x_t|x_{t-1})=\mathcal{N}(x_t;{\mu }_t=\sqrt{1-{\beta }_t}{x}_{t-1},{\Sigma }_t={\beta }_t\mathbf{I})\end{array}$$

$x_0$是一张干净的没有噪声的图像，采样自真实的数据分布$q(x)(x_0\sim q(x))$，在马尔科夫链的每一步我们将方差为${\beta }_t$的高斯噪声添加到$x_{t-1}$上，得到了一个新的含噪图像$x_t$，其分布满足$q(x_t|x_{t-1})$。最终经过$T$步后，得到了一个纯噪声图像$x_T$。公式(1)中，$\mathbf{I}$是一个单位矩阵。我们注意到，$q(x_t|x_{t-1})$是一个由均值${\mu }_t$和方差${\Sigma }_t$定义的正太分布。
由$q(x_t|x_{t-1})$可以推导出
$$\begin{array}{c}q(x_{1:T}|x_0)={\prod }_{t=1}^{T}q(x_t|x_{t-1})\end{array}$$
上式中的$t$表示时间戳，$q(x_{1:T}|x_0)$表示从1到$T$，共应用$T$次$q$函数。那么；例如当$t=500<T$时，我们需要应用$q$函数500次，才能得到最终想要的$x_T$，这个需要重复计算的次数太多了，且按照上面的方式，梯度是没有办法进行反向传播的，因此这里需要用到一个重参数的技巧。

重参数技巧就是将相邻的两个变量建模成一个线性关系：
$$\begin{array}{c}{x}_t=\sqrt{1-{\beta }_t}{x}_{t-1}+\sqrt{{\beta }_t}{ϵ}_{t-1},其中{ϵ}_{t-1}\sim \mathcal{N}(0,\mathbf{I})\end{array}$$
这里我们设${\alpha }_t=1-{\beta }_t$，那么(3)式可写为：
$$\begin{array}{c}\begin{array}{rl}{x}_t& =\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}\\ & =\sqrt{{\alpha }_t}(\sqrt{{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{ϵ}_{t-2})+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{{\alpha }_t(1-{\alpha }_{t-1})}{ϵ}_{t-2}+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}\end{array}\end{array}$$
由于多个高斯分布相加，结果还是高斯分布，所以（4）式的后两项可以进行合并，合并过程如下：
$$\begin{array}{rl}\mathcal{N}(0,{\sigma }_1^2\mathbf{I})+\mathcal{N}(0,{\sigma }_2^2\mathbf{I})& \sim \mathcal{N}(0,({\sigma }_1^2+{\sigma }_2^2)\mathbf{I})\\ \sqrt{1-{\alpha }_t}{ϵ}_{t-1}& \sim \mathcal{N}(0,(1-{\alpha }_t)\mathbf{I})\\ \sqrt{{\alpha }_t(1-{\alpha }_{t-1})}{ϵ}_{t-2}& \sim \mathcal{N}(0,{\alpha }_t(1-{\alpha }_{t-1})\mathbf{I})\\ \sqrt{{\alpha }_t(1-{\alpha }_{t-1})}{ϵ}_{t-2}+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}& \sim \mathcal{N}(0,[{\alpha }_t(1-{\alpha }_{t-1})+(1-{\alpha }_t)]\mathbf{I})=\mathcal{N}(0,(1-{\alpha }_t{\alpha }_{t-1})\mathbf{I})\end{array}$$
那么经过合并后（4）式变为：
$$\begin{array}{c}{x}_t=\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{ϵ}}_{t-2},其中{\overline{ϵ}}_{t-2}\sim \mathcal{N}(0,\mathbf{I})\end{array}$$

以此类推，再将$x_{t-2}=\sqrt{{\alpha }_{t-2}}{x}_{t-3}+\sqrt{1-{\alpha }_{t-2}}{ϵ}_{t-3}$继续带入，这样不断代入，最终我们可以得到
$$\begin{array}{c}{x}_t=\sqrt{{\alpha }_t{\alpha }_{t-1}...{\alpha }_1}{x}_0+\sqrt{1-{\alpha }_t{\alpha }_{t-1}...{\alpha }_1}ϵ\end{array}$$
另${\overline{\alpha }}_t={\alpha }_t{\alpha }_{t-1}...{\alpha }_1$,那么式（6）就会变成
$$\begin{array}{c}\begin{array}{rl}{x}_t& =\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ\\ & =\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)\mathbf{I})\end{array}\end{array}$$
其中$ϵ$依然为一标准高斯分布。可以知道，当$t$趋于无穷时，由于${\alpha }_t=1-{\beta }_t<1$，因此${\overline{\alpha }}_t$趋近于0，此时$x_t$趋向于一个标准的高斯分布，也就是纯噪声了。这也就是如此构造式（3）的原因。

逆向过程

逆向过程就是逐步去噪的过程。在前向过程中我们要计算的是$q(x_t|x_{t-1})$，那么在逆向过程中就是要计算$q(x_{t-1}|x_t)$。因为我们无法直接得到$q(x_{t-1}|x_t)$，但是加上$x_0$的后验分布$q(x_{t-1}|x_t,x_0)$是可以计算的。根据贝叶斯公式
$$\begin{array}{c}q(x_{t-1}|x_t,x_0)=q(x_t|x_{t-1},x_0)\frac{q(x_{t-1}|x_0)}{q(x_t|x_0)}\end{array}$$
根据马尔科夫链特性，
$$\begin{array}{c}q(x_t|x_{t-1},x_0)=q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_t\mathbf{I})（这里x_0是多余的）\end{array}$$
由公式（7）可以直接计算出$q(x_{t-1}|x_0)$和$q(x_t|x_0)$，即
$$\begin{array}{c}\begin{array}{rl}q(x_{t-1}|x_0)& =\mathcal{N}(x_{t-1};\sqrt{{\overline{\alpha }}_{t-1}}{x}_0,(1-{\overline{\alpha }}_{t-1})\mathbf{I})\\ q(x_t|x_0)& =\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)\mathbf{I})\end{array}\end{array}$$
将公式(9)，(10)代入公式(8)中，经过一系列推导（这里暂时省略），就可以得到$q(x_{t-1}|x_t,x_0)$的均值和方差：
$$\begin{array}{c}\begin{array}{rl}{\beta }_t& =\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}{\beta }_t\\ {\mu }_t(x_t,x_0)& =\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}{x}_0+\frac{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{x}_t\end{array}\end{array}$$
根据前向过程的式（7），可以知道：
$$\begin{array}{c}{x}_0=\frac{1}{\sqrt{{\overline{\alpha }}_t}}(x_t-\sqrt{1-{\overline{\alpha }}_t}ϵ)\end{array}$$
将式（12）和${\beta }_t=1-{\alpha }_t$代入式（11）中，那么
$$\begin{array}{c}\begin{array}{rl}{\mu }_t(x_t)& =\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}\frac{1}{\sqrt{{\overline{\alpha }}_t}}(x_t-\sqrt{1-{\overline{\alpha }}_t}ϵ)+\frac{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{x}_t\\ & =\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}\frac{1}{\sqrt{{\overline{\alpha }}_t}}{x}_t+\frac{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{x}_t-\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}\frac{\sqrt{1-{\overline{\alpha }}_t}}{\sqrt{{\overline{\alpha }}_t}}ϵ\\ & =\frac{{\beta }_t\frac{1}{{\sqrt{\alpha }}_t}+\sqrt{{\alpha }_t}-\sqrt{{\alpha }_t}{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}{x}_t-\frac{{\beta }_t}{\sqrt{{\alpha }_t}\sqrt{1-{\overline{\alpha }}_t}}ϵ\\ & =\frac{{\beta }_t+{\alpha }_t-{\alpha }_t{\overline{\alpha }}_{t-1}}{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_t)}{x}_t-\frac{{\beta }_t}{\sqrt{{\alpha }_t}\sqrt{1-{\overline{\alpha }}_t}}ϵ\\ & =\frac{1-{\alpha }_t+{\alpha }_t-{\alpha }_t{\overline{\alpha }}_{t-1}}{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_t)}{x}_t-\frac{{\beta }_t}{\sqrt{{\alpha }_t}\sqrt{1-{\overline{\alpha }}_t}}ϵ\\ & =\frac{1-{\overline{\alpha }}_t}{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_t)}{x}_t-\frac{{\beta }_t}{\sqrt{{\alpha }_t}\sqrt{1-{\overline{\alpha }}_t}}ϵ\\ & =\frac{1}{\sqrt{{\alpha }_t}}{x}_t-\frac{{\beta }_t}{\sqrt{{\alpha }_t}\sqrt{1-{\overline{\alpha }}_t}}ϵ\\ & =\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}ϵ)\end{array}\end{array}$$
这个时候，我们就可以知道，在每一个时间戳都会有一个均值${\mu }_t(x_t)$依赖于$x_t$和$ϵ$，这里$x_t$是网络的输入，所以只要能够知道$ϵ$，那么我们就能得到${\mu }_t(x_t,x_0)$，在DDPM论文中，作者为了简化优化过程，保持方差固定，仅仅让网络学习均值，在有了均值${\mu }_t(x_t,x_0)$和方差${\beta }_t$后，我们就可以得到$q(x_{t-1}|x_t,x_0)$，那么一步逆向过程就完成了。因此损失函数为
$$\begin{array}{c}\begin{array}{rl}{\mathbf{L}}_t& ={\mathbb{E}}_{x_0,t,ϵ}\left[\frac{1}{2||{\Sigma }_{\theta }(x_t,t)|{|}_2^2}||{\mu }_t-{\mu }_{\theta }(x_t,t)|{|}_2^2\right]\\ & ={\mathbb{E}}_{x_0,t,ϵ}\left[\frac{{\beta }_t^2}{2{\alpha }_t(1-{\overline{\alpha }}_t)||{\Sigma }_{\theta }|{|}_2^2}||{ϵ}_t-{ϵ}_{\theta }(\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ,t)|{|}^2\right]\end{array}\end{array}$$
式（14）等效于网络直接预测噪声$ϵ$，我们用神经网络${ϵ}_{\theta }(x_t,t)$来近似噪声$ϵ$，因此简化后的损失函数为：
$$\begin{array}{c}{\mathbf{L}}_t^{simple}={\mathbb{E}}_{x_0,t,ϵ}\left[||ϵ-{ϵ}_{\theta }(\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ,t)|{|}^2\right]\end{array}$$
下图是DDPM中训练和采样的整体过程

在训练过程中，DDPM的过程为：
1、取一张无噪声图像$x_0$，在$0-T$范围内随机选择一个采样时间戳$t$
2、采样一个标准高斯分布$ϵ$作为噪声，将噪声$ϵ$加到时间戳$t$上，得到$x_t$
3、2中采样的噪声为网络的目标输出，$x_t$为网络输入，计算网络预测噪声和目标噪声之间的loss和梯度

网络结构

前面我们说过，用来预测噪声的是一个神经网络（unet），unet的特点是输出与输入具有相同的形状，unet是一个encoder-decoder结构，encoder和decoder中间用skip connection相连。结构如下图：