本文整理了 Ho, Jonathan, Ajay Jain, and Pieter Abbeel. “Denoising diffusion probabilistic models.” Advances in Neural Information Processing Systems 33 (2020): 6840-6851.
本文也整理参考了大佬deep_thoughts的视频（https://www.bilibili.com/video/BV1b541197HX?spm_id_from=333.337.search-card.all.click&vd_source=45da6ec16c326593a0c40e7da0ae6d21）
1 前置知识
1.1 KL散度
给定两个概率分布p和q，两者的KL散度为
\[ D_{KL}(p||q)=E_p\big[ log\space p(x)-log\space q(x) \big] \] 当两个概率分布为离散时
\[ D_{KL}(p||q)=\sum_{i=1}^{N}p(x_i)\big[ log\space p(x_i)-log\space q(x_i) \big] \]
易得两个一维高斯分布的p和q的KL散度为
\[D_{KL}(p||q)=log\frac{\sigma_2}{\sigma_1}+\frac{\sigma_1^2+(\mu_1-\mu_2)^2}{2\sigma^2_2}-\frac{1}{2}\]
2. VAE
VAE，也就是变分自编码器，与Diffusion model扩散模型有相似之处。
2.1 单层VAE
单层变分自编码器，假设x由隐变量z生成，隐变量的分布\(p(z)\)符合正态分布
\[p(x)=\int_zp_{\theta}(x|z)p(z)dz\]
加入后验概率分布\(q_{\phi}(z|x)\)，我们同样假设该后验条件分布符合正态分布
\[p(x)=\int_z q_{\phi}(z|x)\frac{ p_{\theta}(x|z)p(z)}{q_{\phi}(z|x)}dz\] \[log(p(x))=log \space E_{z\sim q_{\phi}}\big[ \frac{ p_{\theta}(x|z)p(z)}{q_{\phi}(z|x)}\big]\] 从Jensen’s不等式，可以得到
\[ log(p(x))\geq \space E_{z\sim q_{\phi}}\big[ log\frac{ p_{\theta}(x|z)p(z)}{q_{\phi}(z|x)}\big]\\ =E_{z\sim q_{\phi}}\big[ log \space p_{\theta}(x|z)\big]+E_{z\sim q_{\phi}}\big[ log\space p(z)-log \space q_{\phi}(z|x)\big]\\ =E_{z\sim q_{\phi}}\big[ log \space p_{\theta}(x|z)\big]-D_{KL}(q_{\phi}(z|x)||p(z)) \]
同样的，上面的不等式也可以看作ELBO（证据下限）
\[ log(p(x))\geq \space E_{z\sim q_{\phi}}\big[ log\frac{ p_{\theta}(x,z)}{q_{\phi}(z|x)}\big]\\ =E_{z\sim q_{\phi}}\big[ log \space p_{\theta}(x,z)\big]-E_{z\sim q_{\phi}}\big[log \space q_{\phi}(z|x)\big]\\ =ELBO(\phi,\theta) \]
2.2 多层VAE
多层隐变量有\(z_1\), \(z_2\)两个，类似单层，我们同样可以推导到
\[p(x)=\int_{z_1}\int_{z_2}p_{\theta}(x,z_1,z_2)dz_1dz_2 \]
\[p(x)=\int_{z_1}\int_{z_2}q_{\phi}(z_1,z_2|x)\frac{p_{\theta}(x,z_1,z_2)}{q_{\phi}(z_1,z_2|x)}dz_1 dz_2\] \[p(z)=E_{z_1,z_2\sim q_{\phi}(z_1,z_2|x)}\big[ \frac{p_{\theta}(x,z_1,z_2)}{q_{\phi}(z_1,z_2|x)} \big]\]
\[log\space p(z)\geq E_{z_1,z_2\sim q_{\phi}(z_1,z_2|x)}\big[ log\frac{p_{\theta}(x,z_1,z_2)}{q_{\phi}(z_1,z_2|x)} \big]\] 通过马尔可夫链下的贝叶斯法则，我们可以知道
\[p(x,z_1,z_2)=p(x|z_1)p(z_1|z_2)p(z_2)\] \[q(z_1,z_2|x)=q(z_1|x)q(z_2|z_1)\]
通过上面的两个性质，不等式可以改为
\[log\space p(z)\geq E_{z_1,z_2\sim q_{\phi}(z_1,z_2|x)}\big[ log\space p(x|z_1)-log\space q(z_1|x)+log \space p(z_1|z_2)-log\space q(z_2|z_1)+log\space p(z_2) \big]\]
同样，看作ELBO，上式也可以表达为
\[log\space p(z)\geq E_{z_1,z_2\sim q_{\phi}(z_1,z_2|x)}\big[ log\frac{p_{\theta}(x,z_1,z_2)}{q_{\phi}(z_1,z_2|x)} \big]=ELBO(\phi,\theta)\]
3 Diffusion Model
Fig.1 Diffision Process
上面就是Diffusion Model的流程图，分为正向扩散过程\(q(x_{t}|x_{t-1})\)与逆向扩散过程\(p_\theta(x_{t-1}|x_t)\)
3.1 正向扩散过程
假设初始数据的分布为\(x_0\sim q(x)\)，扩散过程就是不断向这个分布添加确定的高斯噪声，扩散过程是不含参数的，噪声的方差为\(\beta_t\in(0,1)\)，均值由\(\beta_t\)与\(x_t\)决定。由\(t=1\)到\(t=T\)的过程是应该马尔可夫过程，所以可以得到
\[q(x_t|x_{t-1})=N(x_t;\sqrt{1-\beta_t}x_{t-1},\beta_tI)\] \[q(x_{1:T}|x_{0})=\prod_{t=1}^Tq(x_t|x_{t-1})\]
这个马尔可夫过程中，任意时刻t的分布\(q(x_t|x_0)\)可以表示为:
\[x_t=\sqrt{\alpha_t}x_{t-1}+\sqrt{1-\alpha_t}z_{t-1}=\sqrt{\alpha_t\alpha_{t-1}}x_{t-2}+\sqrt{1-\alpha_t\alpha_{t-1}}\bar{z_{t-2}}...=\sqrt{\bar{\alpha_t}}x_0+\sqrt{1-\bar{\alpha_t}}z_t\]
Where \(z_{t-1},z_{t-2},...\sim N(0,I)\), \(\bar{z_{t-2}}\)是\(z_{t-1}\)与\(z_{t-2}\)的叠加分布, \(\alpha_t=1-\beta_t\) and \(\bar{\alpha_t}=\prod_{i=1}^T\alpha_i\)
上式中的\(z\)使用了参数重整化的技巧，\(z\)是由\(z_{t-1}\)到\(z_0\)叠加而成的。比如说，当\(X\sim N(\mu_1,\sigma_1)\), \(Y\sim N(\mu_2,\sigma_2)\)，两个分布叠加\(aX+bY\)之后的均值为\(a\mu_1+b\mu_2\)，方差为\(a^2\sigma_1^2+b^2\sigma_2^2\)，所以\(\sqrt{1-\alpha_t\alpha_{t-1}}\bar{z_{t-2}}=\sqrt{\alpha_t-\alpha_t\alpha_{t-1}}z_{t-2}+\sqrt{1-\alpha_t}z_{t-1}\)
所以，我们迭代得到
\[q(x_t|x_0)=N(x_t|\sqrt{\bar{\alpha_t}}x_0,(1-\bar{\alpha_t})I)\]
在Diffusion model的原论文中，作者给出了一个假设，就是当样本中的噪声越来越多时，我们通常可以使用更大的更新，也就是说，我们可以假设\(\beta_1<\beta_2<...<\beta_T\)
我们可以看到Diffusion model与VAE的区别，扩散过程对应于VAE中的\(x\to z\)，但不同的是，VAE中使用后验网络，是有参的，Diffusion Model这个扩散过程是不含参数的，我们不用去学习参数，比如\(\beta\)。Diffusion model中\(X_T\)基本与\(X_0\)无关了，由于大量的噪声，\(X_T\)趋向与高斯噪声。在扩散过程中，Diffusion model数据的维度与大小不改变，VAE中会改变。
3.2 逆扩散过程
逆扩散过程就是从\(X_T\)中恢复原始数据的过程，它是一个马尔可夫过程，具体恢复步骤如下面公式所示
\[p_{\theta}(x_{0:T})=p_\theta(x_T)\prod_{t=1}^Tp_{\theta}(X_{t-1}|x_t)\] 我们实际上是不能直接去定义并且跟踪一个逆向过程\(q(x_{t-1}|x_t)\)的，可以预见的是，直接去拟合\(q(x_{t-1}|x_t)\)是非常困难的，因为尽管它可以看作一个高斯分布，但参数的估计需要整个数据集去拟合参数，非常耗费时间与空间。所以diffusion model使用构建了一个参数估计\(p_\theta(x_{t-1}|x_t)\)，也就是说，我们关心的上式中的\(p_{\theta}(x_{t-1}|x_t)\)，
\[p_\theta(x_{t-1}|x_t)=N(x_{t-1};\mu_{\theta}(x_t,t),\Sigma_\theta(x_t,t))\]
3.3 后验扩散条件概率分布
我们引入一个后验的扩散条件概率分布，\(q(x_{t-1}|x_0,x_t)\)，由于马尔可夫过程，\(q(x_{t-1}|x_0,x_t)\)。它是可以用公式表达的
\[q(x_{t-1}|x_0,x_t)=N(x_{t-1};\widetilde{\mu}(x_t,x_0),\widetilde{\beta_t}I)\]
使用贝叶斯定理，我们可以得到
\[ q(x_{t-1}|x_t,x_0)=q(x_t|x_{t-1},x_0)\frac{x_{t-1}|x_0}{q(x_t|x_0)}\\ \propto exp(-\frac{1}{2}\big(\frac{(x_t-\sqrt{\alpha_t}x_{t-1})^2}{\beta_t}+\frac{(x_{t-1}-\bar{\alpha_{t-1}}x_0)^2}{1-\bar{\alpha_{t-1}}}-\frac{(x_t-\bar{\alpha_{t}}x_0)^2}{1-\bar{\alpha_{t}}}\big))\\ =exp(-\frac{1}{2}\big( (\frac{\alpha_t}{\beta_t}+\frac{1}{1-\bar{\alpha_{t-1}}})x_{t-1}^2-(\frac{2\sqrt{\alpha_{t}}}{\beta_t}x_t+\frac{2\bar{\alpha_{t-1}}}{1-\bar{\alpha_{t-1}}})x_{t-1}+C(x_t,x_0) \big)) \]
我们可以省略\(C(x_t,x_0)\)因为它不涉及\(x_{t-1}\)
我们可以很容易的得到后验扩散条件概率分布的方差与均值
\[\widetilde{\beta_t}=\frac{1}{\frac{\alpha_t}{\beta_t}+\frac{1}{1-\bar{\alpha_{t-1}}}}\]
\[\widetilde{\mu_t}(x_t,x_0)=\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{1-\bar{\alpha_{t}}}x_t+\frac{\sqrt{\bar \alpha_{t-1}}\beta_t}{1-\bar{\alpha_{t}}}x_0\]
在公式（19），我们已经得到了\(x_0\)与\(x_t\)之间的关系
\[x_0=\frac{1}{\sqrt{\bar\alpha_t}}(x_t-\sqrt{1-\bar{\alpha_{t}}}z_t)\]
将\(x_0\)代入\(q(x_{t-1}|x_t,x_0)\)的均值中，可以得到
\[\widetilde{\mu_t}(x_t,x_0)=\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{1-\bar{\alpha_{t}}}x_t+\frac{\sqrt{\bar \alpha_{t-1}}\beta_t}{1-\bar{\alpha_{t}}}x_0=\frac{1}{\sqrt{\alpha_t}}(x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}z_t)\]
3.4 目标数据分布的似然函数
我们需要寻找一个优化函数loss function，在这里，我们加入KL散度寻找负对数似然函数的一个上界，这个上界就是Diffusion model的loss function。
\[ -log\space p_{\theta}(x_0)\leq -log\space p_{\theta}(x_0)+D_{KL}(q(x_{1:T}|x_0)||p_{\theta}(x_{1:T}|x_0))\\ =-log\space p_{\theta}(x_0)+E_{x_{1:T}\sim q(x_{1:T}|x_0)}\big[ log\frac{q(x_{1:T}|x_0)p_\theta(x_0)}{p_{\theta}(x_{0:T})} \big]\\ =-log\space p_{\theta}(x_0)+E_{x_{1:T}\sim q(x_{1:T}|x_0)}\big[ log\frac{q(x_{1:T}|x_0)}{p_{\theta}(x_{0:T})}+log \space p_\theta(x_0) \big]\\ =E_{x_{1:T}\sim q(x_{1:T}|x_0)}\big[ log\frac{q(x_{1:T}|x_0)}{p_{\theta}(x_{0:T})}\big] \]
定义\(L_{VLB}=E_{ q(x_{0:T})}\big[ log\frac{q(x_{1:T}|x_0)}{p_{\theta}(x_{0:T})}\big]\geq -E_{q(x_0)}log\space p_\theta(x_0)\)，这就是交叉熵\(-E_{q(x_0)}log\space p_\theta(x_0)\)的一个上界
\(L_{VLB}\)还是可以简化的
\[ L_{VLB}=E_{ q(x_{0:T})}\big[ log\frac{q(x_{1:T}|x_0)}{p_{\theta}(x_{0:T})}\big]\\ =E_{ q(x_{0:T})}\big[ log\frac{\prod_{t=1}^Tq(x_t|x_{t-1})}{p_\theta(x_T)\prod_{t=1}^Tp_{\theta}(X_{t-1}|x_t)}\big]\\ =E_{ q(x_{0:T})}\big[ \sum_{t=1}^Tlog\space\frac{ q(x_t|x_{t-1})}{p_{\theta}(X_{t-1}|x_t)}-log \space p_\theta(x_T)\big]\\ =E_{ q(x_{0:T})}\big[ \sum_{t=2}^Tlog\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)}-log \space p_\theta(x_T)\big]\\ =E_{ q(x_{0:T})}\big[ \sum_{t=2}^Tlog(\frac{q(x_t|x_{t-1},x_0)}{p_{\theta}(X_{t-1}|x_t)}\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)}-log \space p_\theta(x_T)\big]\\ =E_{ q(x_{0:T})}\big[ \sum_{t=2}^Tlog\frac{q(x_t|x_{t-1},x_0)}{p_{\theta}(X_{t-1}|x_t)}+\sum_{t=2}^Tlog\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)}-log \space p_\theta(x_T)\big]\\ =E_{ q(x_{0:T})}\big[ \sum_{t=2}^Tlog\frac{q(x_t|x_{t-1},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)}-log \space p_\theta(x_T)\big]\\ =E_{ q(x_{0:T})}\big[ \sum_{t=2}^Tlog\frac{q(x_t|x_{t-1},x_0)}{p_{\theta}(X_{t-1}|x_t)}+log\frac{q(x_T|x_0)}{p_\theta(x_T)}-log \space p_\theta(x_0|x_1)\big]\\ =E_{ q(x_{0:T})}\big[\sum_{t=2}^T D_{KL}(q(x_t|x_{t-1},x_0)||p_{\theta}(X_{t-1}|x_t))+ D_{KL}(q(x_T|x_0)||p_\theta(x_T)- log \space p_\theta(x_0|x_1) \big] \]
上式（30）中，第一项是一个后验扩散的条件概率分布与逆扩散过程的条件概率分布的KL散度，第二项是从0到T的扩散过程与\(x_T\)的逆扩散过程分布的KL散度，是不含参数的，完全可以省略，第三项可以和第一项一起考虑，它就是第一项\(t=1\)的特殊形式。
对于公式（30）的第一项，根据公式（22）与（23）我们可以发现，两者都是单一变量的高斯分布，并且方差都是不含参数的，所以由公式（3），我们可以进一步化简，需要注意的是，我们只需要关心\(\frac{(\mu_1-\mu_2)^2}{2\sigma_2^2}\)，因为其他项最后都是常数C。
\[E_{ q(x_{0:T})}\big[\sum_{t=2}^T D_{KL}(q(x_t|x_{t-1},x_0)||p_{\theta}(X_{t-1}|x_t))\big]\\ =E_{ q(x_{0:T})}\big[ \frac{1}{2\sigma_t^2} (\widetilde{\mu}(x_t,x_0)-\mu_{\theta}(x_t,t))^2 \big]+C\\ =E_{ x_{0},\epsilon}\big[ \frac{1}{2\sigma_t^2} (\widetilde{\mu}(x_t(x_0,\epsilon),\frac{1}{\sqrt{\bar\alpha_t}}(x_t(x_0,\epsilon)-\sqrt{1-\bar{\alpha_{t}}}\epsilon))-\mu_{\theta}(x_t(x_0,\epsilon),t))^2 \big]+C \]
在式（31）中，我们将原本式（27）中的\(z_t\)改为了\(\epsilon\)，这是因为\(z_t\)在前面的假设中是t时刻的单位高斯分布，我们在这边就用\(\epsilon\)替代。式（31）进一步化简，根据公式（28）
\[E_{ x_{0},\epsilon}\big[\sum_{t=2}^T D_{KL}(q(x_t|x_{t-1},x_0)||p_{\theta}(X_{t-1}|x_t))\big]\\ =E_{ x_{0},\epsilon}\big[ \frac{1}{2\sigma_2^2} \big(\frac{1}{\sqrt{\bar\alpha_t}}(x_t(x_0,\epsilon)-\sqrt{1-\bar{\alpha_{t}}}\epsilon)-\mu_{\theta}(x_t(x_0,\epsilon),t)\big)^2 \big]+C \]
下面就是上神经网络的部分了，对于Diffusion model，神经网络的输入是\(x_t\)与\(t\)，但有很多输出目标可以选择。在deep_thoughts大佬的文章中，他总结了三种输出目标。
神经网络的输出等于正向扩散过程中的后验分布均值\(\widetilde{\mu}(x_t,x_0)\)
神经网络的输出直接等于\(x_0\)，但是这样子的话效果会很差，需要通过马尔可夫过程高斯条件迭代才能获得最终的生成样本
神经网络的输出等于随机变量\(\epsilon_\theta\)
我们可以发现在个loss function中\(\mu_{\theta}(x_t(x_0,\epsilon),t)\)需要基于\(x_t\)去预测\(\frac{1}{\sqrt{\bar\alpha_t}}(x_t(x_0,\epsilon)-\sqrt{1-\bar{\alpha_{t}}}\epsilon)\)，\(x_t\)是已知的，唯一的参数是\(\epsilon\)，所以，在\(\mu_{\theta}\)中，使用了\(\epsilon_\theta\)。新的逆扩散过程的条件分布的均值可以表示为
\[\mu_\theta(x_t,t)=\widetilde{\mu}(x_t,\frac{1}{\sqrt{\bar\alpha_t}}(x_t-\sqrt{1-\bar{\alpha_{t}}}\epsilon_\theta(x_t)))=\frac{1}{\sqrt{\alpha_t}}(x_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha_{t}}}}\epsilon_\theta(x_t,t))\]
于是，loss function可以表示为
\[E_{ x_{0},\epsilon}\big[ \frac{1}{2\sigma_t^2\alpha_t(1-\bar\alpha_t)} \big( \epsilon-\epsilon_\theta(\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon,t) \big)^2 \big]\] DDPM的作者发现将系数去掉效果更好，所以，最终的loss function 为
\[E_{t,x_0,\epsilon}\big[ \big( \epsilon-\epsilon_\theta(\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon,t) \big)^2 \big]\]
3.5 算法步骤
Fig.2 Diffision Algorithm
在DDPM原论文中，算法分为两个部分，训练部分与采样部分。在训练步骤之后，得到了噪声系数\(\epsilon_\theta\)，接着迭代采样，\(x_{t-1}\)可以表示为
\[x_{t-1} =\frac{1}{\sqrt\alpha_t}(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha_t}}}\epsilon_\theta(x_t,t))+\sigma_tz\]
Where \(z\sim N(0,I)\)