扩散模型（Diffusion Model）——由浅入深的理解

Diffusion model 是一种图片生成的范式，大量的数学公式让许多同学望而却步，但实际研究下来，它的公式推导其实大部分都在射程范围之内。

本文在概览中对Diffusion model抛去细节做一个整体的梳理，而细节的推导会在下文的扩散过程、逆扩散过程、损失函数中展示。如果只想对Diffusion model有一个定性的了解而不关系推导的话，只看概览就可以了。

概览

扩散模型有两个过程，分别为扩散过程和逆扩散过程。

如上图所示，扩散过程为从右到左（$X_0\to X_T$）的过程，表示对图片逐渐加噪，且$X_{t+1}$是在$X_t$上加躁得到的，其只受$X_t$的影响，因此扩散过程是一个马尔科夫过程。$X_0$表示从真实数据集中采样得到的一张图片，对$X_0$添加$T$次噪声，图片逐渐变得模糊，当$T$足够大时，$X_T$为标准正态分布。在训练过程中，每次添加的噪声是已知的，即$q(X_t|X_{t-1})$是已知的，根据马尔科夫过程的性质，我们可以递归得到$q(X_t|X_0)$，即$q(X_t|X_0)$是已知的。扩散过程最主要的就是$q(X_t|X_0)$和$q(X_t|X_{t-1})$的推导，推导细节见下文的扩散过程。

如上图所示，逆扩散过程为从左到右（$X_T\to X_0$）的过程，表示从噪声中逐渐复原出图片。如果我们能够在给定$X_t$条件下知道$X_{t-1}$的分布，即如果我们可以知道$q(X_{t-1}|X_t)$，那我们就能够从任意一张噪声图片中经过一次次的采样得到一张图片而达成图片生成的目的。显然我们很难知道$q(X_{t-1}|X_t)$，因此我们才会用$p_{\Theta }(X_{t-1}|X_t)$来近似$q(X_{t-1}|X_t)$，$p_{\Theta }(X_{t-1}|X_t)$就是我们要训练的网络，在原文中就是个U-Net。而很妙的是，虽然我们不知道$q(X_{t-1}|X_t)$，但是$q(X_{t-1}|X_t{X}_0)$却是可以用$q(X_t|X_0)$和$q(X_t|X_{t-1})$表示的，即$q(X_{t-1}|X_t{X}_0)$是可知的，因此我们可以用$q(X_{t-1}|X_t{X}_0)$来指导$p_{\Theta }(X_{t-1}|X_t)$进行训练。逆扩散过程最主要的就是$q(X_{t-1}|X_t{X}_0)$的推导，推导细节见下文的逆扩散过程。

我们已经明确了要训练$p_{\Theta }(X_{t-1}|X_t)$，那要怎么确定目标函数呢？有两个很直接的想法，一个是负对数的最大似然概率，即$-log{p}_{\Theta }(X_0)$，另一个是真实分布与预测分布的交叉熵，即$-E_{q(X_0)}log{p}_{\Theta }(X_0)$，而显然这两种都不好搞，因此他参考了VAE，不去优化这两个东西，而是优化他们的变分上界(variational lower bound)，定义$L_{VLB}$如下:
$$\begin{gathered} L_{VLB}=E_{q(x_{0:T})}[log\frac{q(X_{1:T}|X_0)}{p_{\Theta }(X_{0:T})}] \\ 可证：L_{VLB}\ge -log{p}_{\Theta }(X_0)\&L_{VLB}\ge -E_{q(X_0)}log{p}_{\Theta }(X_0)\text{\hspace{1em}(1)} \end{gathered}$$
即$L_{VLB}$减小就代表$-log{p}_{\Theta }(X_0)$和$-E_{q(X_0)}log{p}_{\Theta }(X_0)$的上界减小。
且经过推导，$L_{VLB}$可写成如下形式：
$$L_{VLB}=L_T+L_{T-1}+...+L_0\text{\hspace{1em}(2)}$$$$\begin{gathered} where:L_T=D_{KL}(q(X_T|X_0)||p_{\Theta }(X_T)) \\ {L}_t=D_{KL}(q(X_t|X_{t+1}{X}_0)||p_{\Theta }(X_t|X_{t+1})),1\le t\le T \\ {L}_0=-log{p}_{\Theta }(X_0|X_1) \end{gathered}$$
由上式不难发现，$L_t$就是逆扩散过程中$q(X_t|X_{t+1}{X}_0)$和$p_{\Theta }(X_t|X_{t+1})$的KL散度，这也就是我上面说的，用$q(X_t|X_{t+1}{X}_0)$来指导$p_{\Theta }(X_t|X_{t+1})$进行训练。这部分主要就是(1)式和(2)式的推导，细节部分见下文的损失函数。

扩散过程

如上图所示，扩散过程为从右到左（$X_0\to X_T$）的过程，表示对图片逐渐加噪，且$X_{t+1}$是在$X_t$上加躁得到的，其只受$X_t$的影响，因此扩散过程是一个马尔科夫过程。且每一步扩散的步长受变量$\{{\beta }_t\in (0,1){\}}_{t=1}^T$的影响。$q(X_t|X_{t-1})$可写为如下形式，即给定$X_{t-1}$的条件下，$X_t$服从均值为$\sqrt{1-{\beta }_t}{X}_{t-1}$，方差为${\beta }_{t}I$的正态分布:
$$q(X_t|X_{t-1})=N(X_t;\sqrt{1-{\beta }_t}{X}_{t-1},{\beta }_{t}I)\text{\hspace{1em}(3)}$$
用重参数化技巧表示$X_t$，令${\alpha }_t=1-{\beta }_t$，令$Z_t\sim N(0,I),t\ge 0$，即:
$$X_t=\sqrt{{\alpha }_t}{X}_{t-1}+\sqrt{1-{\alpha }_t}{Z}_{t-1}\text{\hspace{1em}(4)}$$
写多几行:
$$\begin{gathered} X_{t-1}=\sqrt{{\alpha }_{t-1}}{X}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{Z}_{t-2} \\ {X}_{t-2}=\sqrt{{\alpha }_{t-2}}{X}_{t-3}+\sqrt{1-{\alpha }_{t-2}}{Z}_{t-3} \\ ... \\ {X}_1=\sqrt{{\alpha }_1}{X}_0+\sqrt{1-{\alpha }_1}{Z}_0 \end{gathered}$$
易归纳得，令${\bar{\alpha }}_t={\prod }_{i=1}^t{\alpha }_i$：
$$X_t=\sqrt{{\bar{\alpha }}_t}{X}_0+\frac{\sqrt{{\bar{\alpha }}_t}}{\sqrt{{\alpha }_1}}\sqrt{1-{\alpha }_1}{Z}_0+\frac{\sqrt{{\bar{\alpha }}_t}}{\sqrt{{\bar{\alpha }}_2}}\sqrt{1-{\alpha }_2}{Z}_1+\frac{\sqrt{{\bar{\alpha }}_t}}{\sqrt{{\bar{\alpha }}_3}}\sqrt{1-{\alpha }_3}{Z}_2+...+\sqrt{1-{\alpha }_t}{Z}_{t-1}$$
设随机变量${\bar{Z}}_{t-1}$为：
$${\bar{Z}}_{t-1}=\frac{\sqrt{{\bar{\alpha }}_t}}{\sqrt{{\alpha }_1}}\sqrt{1-{\alpha }_1}{Z}_0+\frac{\sqrt{{\bar{\alpha }}_t}}{\sqrt{{\bar{\alpha }}_2}}\sqrt{1-{\alpha }_2}{Z}_1+\frac{\sqrt{{\bar{\alpha }}_t}}{\sqrt{{\bar{\alpha }}_3}}\sqrt{1-{\alpha }_3}{Z}_2+...+\sqrt{1-{\alpha }_t}{Z}_{t-1}$$
则${\bar{Z}}_{t-1}$的期望和方差如下：
$$\begin{gathered} E({\bar{Z}}_{t-1})=0 \\ D({\bar{Z}}_{t-1})=\frac{{\bar{\alpha }}_t}{{\alpha }_1}(1-{\alpha }_1)+\frac{{\bar{\alpha }}_t}{{\bar{\alpha }}_2}(1-{\alpha }_2)+\frac{{\bar{\alpha }}_t}{{\bar{\alpha }}_3}(1-{\alpha }_3)+...+\frac{{\bar{\alpha }}_t}{{\bar{\alpha }}_t}(1-{\alpha }_t)=1-{\bar{\alpha }}_t \end{gathered}$$
所以，
$$\begin{gathered} X_t=\sqrt{{\bar{\alpha }}_t}{X}_0+{\bar{Z}}_{t-1}=\sqrt{{\bar{\alpha }}_t}{X}_0+\sqrt{1-{\bar{\alpha }}_t}Z,Z\sim N(0,I) \\ q(X_t|X_0)=N(X_t;\sqrt{{\bar{\alpha }}_t}{X}_0,(1-{\bar{\alpha }}_t)I) \end{gathered}$$
至此，我们推出了$q(X_t|X_{t-1})$和$q(X_t|X_0)$。

逆扩散过程

如果我们能够在给定$X_t$条件下知道$X_{t-1}$的分布，即如果我们可以知道$q(X_{t-1}|X_t)$，那我们就能够从任意一张噪声图片中经过一次次的采样得到一张图片而达成图片生成的目的。显然我们很难知道$q(X_{t-1}|X_t)$，因此我们才会用$p_{\Theta }(X_{t-1}|X_t)$来近似$q(X_{t-1}|X_t)$，$p_{\Theta }(X_{t-1}|X_t)$就是我们要训练的网络。而很妙的是，虽然我们不知道$q(X_{t-1}|X_t)$，但是$q(X_{t-1}|X_t{X}_0)$却是可以用$q(X_t|X_0)$和$q(X_t|X_{t-1})$表示的，即$q(X_{t-1}|X_t{X}_0)$是可知的。
下面对$q(X_{t-1}|X_t{X}_0)$进行推导:
$$\begin{gathered} q(X_{t-1}|X_t{X}_0)=\frac{q(X_0{X}_{t-1}{X}_t)}{q(X_0{X}_t)}=\frac{q(X_0{X}_{t-1}{X}_t)}{q(X_0{X}_{t-1})}\frac{q(X_0{X}_{t-1})}{q(X_0{X}_t)}=q(X_t|X_{t-1}{X}_0)\ast \frac{q(X_{t-1}|X_0)}{q(X_t|X_0)} \\ \because 扩散过程是马尔科夫过程 \\ \therefore q(X_t|X_{t-1}{X}_0)=q(X_t|X_{t-1}) \\ \therefore q(X_{t-1}|X_t{X}_0)=q(X_t|X_{t-1})\ast \frac{q(X_{t-1}|X_0)}{q(X_t|X_0)} \end{gathered}$$
至此，已经把$q(X_{t-1}|X_t{X}_0)$用$q(X_t|X_0)$和$q(X_t|X_{t-1})$进行表示，下面对$q(X_{t-1}|X_t{X}_0)$的表达式进行推导:
$$\begin{gathered} q(X_t|X_{t-1})=N(X_t;\sqrt{1-{\beta }_t}{X}_{t-1},{\beta }_{t}I)=\frac{1}{\sqrt{2\pi (1-{\alpha }_t)}}exp(-\frac{1}{2}\frac{(X_t-\sqrt{{\alpha }_t}{X}_{t-1}{)}^2}{1-{\alpha }_t}) \\ q(X_t|X_0)=N(X_t;\sqrt{{\bar{\alpha }}_t}{X}_0,(1-\overline{{\alpha }_t})I)=\frac{1}{\sqrt{2\pi (1-{\bar{\alpha }}_t)}}exp(-\frac{1}{2}\frac{(X_t-\sqrt{{\bar{\alpha }}_t}{X}_0{)}^2}{1-{\bar{\alpha }}_t}) \\ q(X_{t-1}|X_0)=N(X_{t-1};\sqrt{{\bar{\alpha }}_{t-1}}{X}_0,(1-{\bar{\alpha }}_{t-1})I)=\frac{1}{\sqrt{2\pi (1-{\bar{\alpha }}_{t-1})}}exp(-\frac{1}{2}\frac{(X_{t-1}-\sqrt{{\bar{\alpha }}_{t-1}}{X}_0{)}^2}{1-{\bar{\alpha }}_{t-1}}) \end{gathered}$$
$$\begin{gathered} q(X_{t-1}|X_t{X}_0)=\frac{1}{\sqrt{2\pi \frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}}{\beta }_t}exp(-\frac{1}{2\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t}(X_{t-1}^2-2(\frac{(1-{\bar{\alpha }}_{t-1})\sqrt{{\alpha }_t}{X}_t}{1-{\bar{\alpha }}_t}+\frac{{\beta }_t\sqrt{{\bar{\alpha }}_{t-1}}{X}_0}{1-{\bar{\alpha }}_t})X_{t-1}+C(X_0,X_t)) \\ q(X_{t-1}|X_t{X}_0)=N(X_{t-1};\frac{(1-{\bar{\alpha }}_{t-1})\sqrt{{\alpha }_t}{X}_t}{1-{\bar{\alpha }}_t}+\frac{{\beta }_t\sqrt{{\bar{\alpha }}_{t-1}}{X}_0}{1-{\bar{\alpha }}_t},\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t) \\ \because X_t=\sqrt{{\bar{\alpha }}_t}{X}_0+\sqrt{1-{\bar{\alpha }}_t}Z,Z\sim N(0,I) \\ \therefore q(X_{t-1}|X_t{X}_0)=N(X_{t-1};\frac{1}{{\sqrt{\alpha }}_t}{X}_t-\frac{{\beta }_t}{\sqrt{{\alpha }_t(1-{\bar{\alpha }}_t)}}Z,\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t),Z\sim N(0,I) \end{gathered}$$
至此，得到了$q(X_{t-1}|X_t{X}_0)$的分布表达式。
在下文的损失函数中，会介绍我们要怎么用$q(X_{t-1}|X_t{X}_0)$来监督$p_{\Theta }(X_{t-1}|X_t)$进行训练。

损失函数

我们已经明确了要训练$p_{\Theta }(X_{t-1}|X_t)$，那要怎么确定目标函数呢？有两个很直接的想法，一个是负对数的最大似然概率，即$-log{p}_{\Theta }(X_0)$，另一个是真实分布与预测分布的交叉熵，即$-E_{q(X_0)}log{p}_{\Theta }(X_0)$，然而，类似于VAE，由于我们很难对噪声空间进行积分，因此直接优化$-log{p}_{\Theta }(X_0)$或$-E_{q(X_0)}log{p}_{\Theta }(X_0)$是很困难的，因此我们不会直接优化它们，而是优化它们的变分上界$L_{VLB}$，$L_{VLB}$的定义如下:
$$L_{VLB}=E_{q(x_{0:T})}[log\frac{q(X_{1:T}|X_0)}{p_{\Theta }(X_{0:T})}]$$
下面证明$L_{VLB}$是$-log{p}_{\Theta }(X_0)$和$-E_{q(X_0)}log{p}_{\Theta }(X_0)$的上界，即证明$L_{VLB}\ge -log{p}_{\Theta }(X_0)\&L_{VLB}\ge -E_{q(X_0)}log{p}_{\Theta }(X_0)$:
$$\begin{gathered} -log{p}_{\Theta }(X_0)\le -log{p}_{\Theta }(X_0)+D_{KL}(q(X_{1:t}|X_0)||p_{\Theta }(X_{1:T}|X_0)) \\ =-log{p}_{\Theta }(X_0)+E_{X_{1:T}\sim q(X_{1:T}|X_0)}(log\frac{q(X_{1:T}|X_0)}{p_{\Theta }(X_{1:T}|X_0)}) \\ =-log{p}_{\Theta }(X_0)+E_{X_{1:T}\sim q(X_{1:T}|X_0)}(log\frac{q(X_{1:T}|X_0)p_{\Theta }(X_0)}{p_{\Theta }(X_{0:T})}) \\ =-log{p}_{\Theta }(X_0)+E_{X_{1:T}\sim q(X_{1:T}|X_0)}(log\frac{q(X_{1:T}|X_0)}{p_{\Theta }(X_{0:T})}+log(p_{\Theta }(X_0))) \\ =E_{X_{0:T}\sim q(X_{0:T})}(log\frac{q(X_{1:T}|X_0)}{p_{\Theta }(X_{0:T})})=L_{VLB} \end{gathered}$$
$$\begin{gathered} L_{CE}=-\int q(X_0)log{p}_{\Theta }(X_0)d{X}_0=-E_{q(X_0)}log{p}_{\Theta }(X_0) \\ =-E_{q(X_0)}log(\int p_{\Theta }(X_{1:T}|X_0)p_{\Theta }(X_0)d{X}_{1:T}) \\ =-E_{q(X_0)}log(\int p_{\Theta }(X_{0:T})d{X}_{1:T}) \\ =-E_{q(X_0)}log(\int q(X_{1:T}|X_0)\frac{p_{\Theta }(X_{0:T})}{q(X_{1:T}|X_0)}d{X}_{1:T}) \\ =-E_{q(X_0)}(log(E_{q(X_{1:T}|X_0)}\frac{p_{\Theta }(X_{0:T})}{q(X_{1:T}|X_0)})) \\ \le -E_{q(X_0)}(E_{q(X_{1:T}|X_0)}log(\frac{p_{\Theta }(X_{0:T})}{q(X_{1:T}|X_0)})) \\ =E_{q(x_{0:T})}[log\frac{q(X_{1:T}|X_0)}{p_{\Theta }(X_{0:T})}]=L_{VLB} \end{gathered}$$
至此，证明了$L_{VLB}$是$-log{p}_{\Theta }(X_0)$和$-E_{q(X_0)}log{p}_{\Theta }(X_0)$的上界。
下面，对$L_{VLB}$化简：
$$\begin{gathered} L_{VLB}=E_{q(x_{0:T})}[log\frac{q(X_{1:T}|X_0)}{p_{\Theta }(X_{0:T})}] \\ =E_{q(x_{0:T})}[log\frac{{\prod }_{t=1}^{T}q(X_t|X_{t-1})}{p_{\Theta }(X_T){\prod }_{t=1}^T{p}_{\Theta }(X_{t-1}|X_t)}] \\ =E_{q(x_{0:T})}[-log{p}_{\Theta }(X_T)+\sum _{t=1}^{T}log\frac{q(X_t|X_{t-1})}{p_{\Theta }(X_{t-1}|X_t)}] \\ =E_{q(x_{0:T})}[-log{p}_{\Theta }(X_T)+\sum _{t=2}^{T}log\frac{q(X_t|X_{t-1})}{p_{\Theta }(X_{t-1}|X_t)}+log\frac{q(X_1|X_0)}{p_{\Theta }(X_0|X_1)}] \\ =E_{q(X_{0:T})}[-log{p}_{\Theta }(X_T)+\sum _{t=2}^{T}log(\frac{q(X_{t-1}|X_t{X}_0)}{p_{\Theta }(X_{t-1}|X_t)}\ast \frac{q(X_t|X_0)}{q(X_{t-1}|X_0)})+log\frac{q(X_1|X_0)}{p_{\Theta }(X_0|X_1)}] \\ =E_{q(x_{0:T})}[-log{p}_{\Theta }(X_T)+\sum _{t=2}^{T}log\frac{q(X_{t-1}|X_t{X}_0)}{p_{\Theta }(X_{t-1}|X_t)}+\sum _{t=2}^{T}log\frac{q(X_t|X_0)}{q(X_{t-1}|X_0)}+log\frac{q(X_1|X_0)}{p_{\Theta }(X_0|X_1)}] \\ =E_{q(x_{0:T})}[-log{p}_{\Theta }(X_T)+\sum _{t=2}^{T}log\frac{q(X_{t-1}|X_t{X}_0)}{p_{\Theta }(X_{t-1}|X_t)}+log\frac{q(X_T|X_0)}{q(X_1|X_0)}+log\frac{q(X_1|X_0)}{p_{\Theta }(X_0|X_1)}] \\ =E_{q(x_{0:T})}[log\frac{q(X_T|X_0)}{p_{\Theta }(X_T)}+\sum _{t=2}^{T}log\frac{q(X_{t-1}|X_t{X}_0)}{p_{\Theta }(X_{t-1}|X_t)}-log{p}_{\Theta }(X_0|X_1)] \\ =D_{KL}(q(X_T|X_0)||p_{\Theta }(X_T))+\sum _{t=2}^T{D}_{KL}(q(X_{t-1}|X_t{X}_0)||p_{\Theta }(X_{t-1}|X_t))-log{p}_{\Theta }(X_0|X_1) \\ =L_T+L_{T-1}+...+L_0 \\ where:L_T=D_{KL}(q(X_T|X_0)||p_{\Theta }(X_T)) \\ {L}_t=D_{KL}(q(X_t|X_t+1X_0)||p_{\Theta }(X_t|X_{t+1})),1\le t\le T \\ {L}_0=-log{p}_{\Theta }(X_0|X_1) \end{gathered}$$
从$L_t$即可看出，对$p_{\Theta }(X_t|X_{t+1})$的监督就是最小化$p_{\Theta }(X_t|X_{t+1})$和$q(X_t|X_{t+1}{X}_0)$的KL散度。

总结

简单的说，我们的目的是希望学习出一个$p_{\Theta }(X_{t-1}|X_t)$，即能够从噪声图恢复出原图。
为了达到这一个目的，我们使用$q(X_{t-1}|X_t{X}_0)$来监督$p_{\Theta }(X_{t-1}|X_t)$进行训练，$q(X_{t-1}|X_t{X}_0)$是可以用$q(X_t|X_0)$和$q(X_t|X_{t-1})$表示的，即$q(X_{t-1}|X_t{X}_0)$是已知的。

本文是我学习过程中的个人理解，有不对的地方希望大家帮忙指出。希望可以抛砖引玉，欢迎大家在评论区和我交流。

参考

1. What are Diffusion Models?
2. Diffusion Models Beat GANs on Image Synthesis
3. Deep unsupervised learning using nonequilibrium thermodynamics
4. Denoising Diffusion Probabilistic Models