从VAE到DDPM

$$$$

VAE 回顾

在之前的文章中，我们详细地梳理了一遍 VAE，这里做一个简单回顾。

source:Angus Turner. Diffusion Models as a kind of VAE

在 VAE 中，为了最大化对数似然： $$L(\theta )=\log p_{\theta }(x)=\log ({\int }_z{p}_{\theta }(x|z)p(z)\mathrm{d}z)$$ 我们引入变分后验 $q_{\varphi }(z|x)$： $$\begin{array}{}\text{(1)}& \begin{array}{rl}L(\theta )& =\log p_{\theta }(x)\\ & ={\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\log p_{\theta }(x)]\\ & ={\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\log \frac{p_{\theta }(x,z)}{p_{\theta }(z|x)}]\\ & ={\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\log (\frac{p_{\theta }(x,z)}{p_{\theta }(z|x)}\cdot \frac{q_{\varphi }(z|x)}{q_{\varphi }(z|x)})]\\ & =\underset{\text{ELBO}}{\underbrace{{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\log \frac{p_{\theta }(x,z)}{q_{\varphi }(z|x)}]}}+\underset{\mathrm{KL}(q_{\varphi }(z|x)\Vert p_{\theta }(z|x))}{\underbrace{{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\log \frac{q_{\varphi }(z|x)}{p_{\theta }(z|x)}]}}\\ & \ge \text{ELBO}\end{array}\end{array}$$ 得到证据下界 ELBO，通过最大化 ELBO 来最大化对数似然。进一步地，ELBO 还可以拆写成重构项和 KL 正则项： $$\begin{array}{}\text{(2)}& \begin{array}{rl}\text{ELBO}& ={\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\log \frac{p_{\theta }(x,z)}{q_{\varphi }(z|x)}]\\ & ={\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\log \frac{p_{\theta }(x|z)p(z)}{q_{\varphi }(z|x)}]\\ & =\underset{\text{reconstruction}}{\underbrace{{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\log p_{\theta }(x|z)]}}-\underset{\text{regularization}}{\underbrace{\mathrm{KL}(q_{\varphi }(z|x)\Vert p(z))}}\end{array}\end{array}$$ 为了计算上的方便，实践中常将 $p(z),q_{\varphi }(z|x),p_{\theta }(x|z)$ 都取为正态分布，具体而言，它们分别是： $$\begin{array}{rl}& p(z):=\mathcal{N}(z;0,I)\\ & q_{\varphi }(z|x):=\mathcal{N}(z;{\mu }_{\varphi }(x),\mathrm{diag}({\sigma }_{\varphi }^2(x)))\\ & p_{\theta }(x|z):=\mathcal{N}({\mu }_{\theta }(z),{\sigma }^{2}I)& & \sigma \text{ is a constant}\end{array}$$ 代入 $\text{(2)}$ 式即可得到损失函数： $$\mathcal{L}=\underset{\text{reconstruction}}{\underbrace{\frac{1}{2{\sigma }^2}\Vert x-{\mu }_{\theta }(z){\Vert }^2}}+\underset{\text{KL regularization}}{\underbrace{\frac{1}{2}\sum _{i=1}^d({\mu }_{\varphi }^2(x{)}_i+{\sigma }_{\varphi }^2(x{)}_i-\log {\sigma }_{\varphi }^2(x{)}_i-1)}}$$ 正态分布的假设虽然便于计算，但限制了变分后验的形式，使之不能很好地近似真实后验分布，从而限制了 VAE 的能力。因此，一个自然的改进思路就是用更强的方式去建模变分后验——比如再套一个 VAE？

双层 VAE

把 VAE 中的单个隐变量 $z$ 换成两个隐变量 $z_1,z_2$，形成如下马尔可夫链：

source:Angus Turner. Diffusion Models as a kind of VAE

虽然有两个隐变量，但如果把它们视为一个整体，那证据下界 ELBO 的推导过程与 $\text{(1)}$ 式没有什么本质不同，因此我们只需将 $\text{(2)}$ 式中的 $z$ 换做 $z_1,z_2$ 就得到了双层 VAE 的 ELBO： $$\begin{array}{}\text{(3)}& {\text{ELBO}}_{\text{2-layers}}={\mathbb{E}}_{z_1,z_2\sim q_{\varphi }(z_1,z_2|x)}[\log \frac{p_{\theta }(x,z_1,z_2)}{q_{\varphi }(z_1,z_2|x)}]={\mathbb{E}}_{z_1,z_2\sim q_{\varphi }(z_1,z_2|x)}[\log \frac{p_{\theta }(x|z_1)p_{\theta }(z_1|z_2)p(z_2)}{q_{\varphi }(z_2|z_1)q_{\varphi }(z_1|x)}]\end{array}$$

我们依旧希望从 ELBO 中拆解出重构项和正则项。重构项比较简单，只需要把 $p_{\theta }(x|z_1)$ 拆出来就是了，问题在于剩下的一坨应该如何处理。破局的关键在于利用马尔可夫性质把分母上的 $q_{\varphi }(z_2|z_1)$ 改写作 $q_{\varphi }(z_2|z_1,x)$，然后用贝叶斯公式： $$q_{\varphi }(z_2|z_1)=q_{\varphi }(z_2|z_1,x)=\frac{q_{\varphi }(z_1|z_2,x)q_{\varphi }(z_2|x)}{q_{\varphi }(z_1|x)}$$ 代回 $\text{(3)}$ 式得： $$\begin{array}{rl}{\text{ELBO}}_{\text{2-layers}}& ={\mathbb{E}}_{z_1,z_2\sim q_{\varphi }(z_1,z_2|x)}[\log \frac{p_{\theta }(x|z_1)p_{\theta }(z_1|z_2)p(z_2)}{q_{\varphi }(z_2|z_1)q_{\varphi }(z_1|x)}]\\ & ={\mathbb{E}}_{z_1\sim q_{\varphi }(z_1|x)}[p_{\theta }(x|z_1)]+{\mathbb{E}}_{z_1,z_2\sim q_{\varphi }(z_1,z_2|x)}[\log \frac{p_{\theta }(z_1|z_2)p(z_2)}{\frac{q_{\varphi }(z_1|z_2,x)q_{\varphi }(z_2|x)}{q_{\varphi }(z_1|x)}{q}_{\varphi }(z_1|x)}]\\ & ={\mathbb{E}}_{z_1\sim q_{\varphi }(z_1|x)}[p_{\theta }(x|z_1)]+{\mathbb{E}}_{z_1,z_2\sim q_{\varphi }(z_1,z_2|x)}[\log \frac{p_{\theta }(z_1|z_2)p(z_2)}{q_{\varphi }(z_1|z_2,x)q_{\varphi }(z_2|x)}]\\ & ={\mathbb{E}}_{z_1\sim q_{\varphi }(z_1|x)}[p_{\theta }(x|z_1)]+{\mathbb{E}}_{z_2\sim q_{\varphi }(z_2|x)}[\log \frac{p(z_2)}{q_{\varphi }(z_2|x)}]+{\mathbb{E}}_{z_1,z_2\sim q_{\varphi }(z_1,z_2|x)}[\log \frac{p_{\theta }(z_1|z_2)}{q_{\varphi }(z_1|z_2,x)}]\\ & =\underset{\text{reconstruction}}{\underbrace{{\mathbb{E}}_{z_1\sim q_{\varphi }(z_1|x)}[p_{\theta }(x|z_1)]}}-\underset{\text{regularization}}{\underbrace{\mathrm{KL}(q_{\varphi }(z_2|x)\Vert p(z_2))}}-\underset{\text{matching}}{\underbrace{{\mathbb{E}}_{z_2\sim q_{\varphi }(z_2|x)}[\mathrm{KL}(q_{\varphi }(z_1|z_2,x)\Vert p_{\theta }(z_1|z_2))]}}\end{array}$$ 可以看见，双层 VAE 的 ELBO 由三部分组成：重构项、正则项以及匹配项。对比 VAE 的 ELBO，我们可以认为 $z_1,z_2$ 分担了原本一个隐变量的工作。具体而言，在 VAE 中，隐变量既要重构，又要逼近先验分布，且这两个任务是有点矛盾的；而现在，重构依靠 $z_1$ 完成，逼近先验依靠 $z_2$ 完成，二者由匹配项联系起来，从而增加了模型的灵活性。

DDPM

在双层 VAE 的基础上，我们能再多加几层吗？

source:Angus Turner. Diffusion Models as a kind of VAE

如上图所示，为方便叙述，我们引入两个称呼：

• 称从 $$ 到 $$ 的马尔可夫链为前向过程 (forward process) 或扩散过程 (diffusion process)；
• 称从 $$ 到 $$ 的马尔可夫链为逆向过程 (reverse process) 或去噪过程 (denoising process).

用概率图模型的术语来说，前向过程对应 inference model，逆向过程对应 generative model. 另外，为书写上的方便，下文将 $$ 简写为 $$.

同双层 VAE 一样的道理，把 $$ 整体看作 VAE 中的隐变量，代入 $$ 式就可以得到 DDPM 的 ELBO： $$ 接下来的推导技巧和双层 VAE 如出一辙，即将分母中的 $$ 写作 $$，然后使用贝叶斯公式，即可进行大量的消元： $$ 代回 $$ 式得： $$ 同样出现了重构项、正则项和匹配项。重构项要求 $$ 能够重构 $$，正则项要求 $$ 的后验分布逼近先验分布，而匹配项则建立起相邻两项 $$ 之间的联系。

现在，我们只需要为 $$ 式中出现的所有概率分布设计具体的形式，就可以代入计算了。为了让 KL 散度可解，一个自然的想法就是把它们都设计为正态分布的形式。

前向过程

source:https://cvpr2022-tutorial-diffusion-models.github.io/

首先我们关注前向过程，即从 $$ 到 $$ 的马尔可夫链： $$ DDPM 将 $$ 设计为： $$ 其中 $$ 是事先指定的超参数，代表从 $$ 到 $$ 这一步的方差。直观上理解，如果 $$ 比较小，那么 $$ 均值依旧在 $$ 附近，方差也不大，故 $$ 看起来就是在 $$ 的基础上加了一些噪声。值得注意的是，$$ 不带任何可学习参数，这是 DDPM 与 VAE 不一样的地方。

基于 $$ 式，我们可以推导出 $$ 式中需要的 $$ 和 $$. 首先推导 $$. 为了书写上的方便，做一个变量代换： $$ 那么 $$ 式改写作： $$ 这意味着我们可以用如下方式从 $$ 采样 $$： $$ 类似地，我们可以用如下方式从 $$ 采样 $$：

$$ 合并上面两个式子，从 $$ 直接采样 $$ 写作： $$ 由于两个正态随机变量之和服从均值方差分别相加的正态分布，即： $$ 所以只需采样一个正态随机变量即可： $$ 以此类推，从 $$ 直接采样 $$ 写作： $$ 也就是： $$ 这样就推出了 $$. 进一步地，我们希望无论输入什么，前向过程最后得到的分布都趋近于标准正态分布，即 $$，因此要求： $$ 为满足这个要求，只需 $$，也即 $$ 即可。直观来看，这意味着初期加噪较弱，后期加噪变强。在 DDPM 中，作者取 $$ 为从 $$ 到 $$ 的线性递增序列。

接下来推导 $$，根据贝叶斯公式有： $$ 这意味着 $$ 也是一个正态分布： $$

注意：$$ 时上面的推导没问题，但需要特别考虑 $$ 的情形。当 $$ 时，$$ 其实是一个确定性的分布，即一定取 $$；如果我们合理地补充定义 $$（因为 $$ 代表累乘，第零项设置为 $$ 很合理），会发现 $$，正好符合预期，所以上面 $$ 和 $$ 的表达式对 $$ 都适用。

Tip：推导时不要每一项都打开老老实实地算，直接提取 $$ 的二次项系数和一次项系数即可。

看到这里，不知读者心中是否有疑惑——为什么人为设置后验分布（即前向过程）是合理的？VAE 中 $$ 不是要去拟合真实后验分布吗，现在人为设置好了怎么去拟合啊？私以为，这个问题揭示了 VAE 和 DDPM 出发点的不同。VAE 先定义生成模型 $$，在这个定义下，存在所谓的“真实”后验分布 $$，但是它不可解，所以用 $$ 去近似。DDPM 则是反过来，先定义后验分布（即前向过程），然后根据后验去学习生成模型（即逆向过程）。

逆向过程

source:https://cvpr2022-tutorial-diffusion-models.github.io/

现在我们来关注逆向过程，即从 $$ 到 $$ 的马尔可夫链： $$ 其中 $$ 很容易设计，直接取标准正态分布即可，这也与我们之前设计的 $$ 是匹配的。

对于 $$，考虑到 $$ 式中要最小化它与 $$ 之间的 KL 散度，所以为了计算方便，设计为与之相同的形式，即： $$ 其中 $$ 表示以 $$ 为参数的模型。为了看得更清楚，列表如下：

	$$	$$
表达式	$$	$$
均值	$$	$$
方差	$$	$$

可以看到，$$ 的均值沿用了 $$ 的形式，只不过用模型 $$ 代替了生成过程中我们并不知道的 $$. 对于一个给定的 $$，用 $$ 去近似 $$，本质上就是在用 $$ 去近似 $$，在下一节中我们将优化目标显式写出后可以看得更清楚。

至于方差，DDPM 给出了两个选择 $$ 与 $$. 前者不难理解，就是沿用了 $$ 的方差，但是后者是出自什么考虑呢？这个问题我们暂时放一放，在之后的文章中详细说明。

损失函数

至此，我们已经确定下 $$ 式中出现的所有概率分布的形式，因而可以代入计算了。为避免读者上下翻阅，把所有公式总结于此： $$ 首先看正则项，由于我们设计 $$ 时要求 $$，在 $$ 较大时趋近于标准正态分布，并且 $$ 也设置为标准正态分布，所以正则项可以忽略。

然后看重构项，代入表达式得： $$ 最后看匹配项，根据两个正态分布的 KL 散度计算公式，当取 $$ 时，有： $$ 鉴于 $$ 可以写作 $$，因此重构项与匹配项的格式可以统一起来。综上，总的损失函数为： $$ 也就是 L2 损失或 MSE 损失。这验证了前文我们提到的，对于一个给定的 $$，模型 $$ 的作用是去近似 $$ 的说法。

进一步地，DDPM 对 $$ 做了重参数化。根据 $$ 式有： $$ 因此可以把 $$ 写作相同的形式： $$ 代入 $$ 式得： $$ 其中 $$ 由 $$ 与采样出的 $$ 根据 $$ 式计算，为简便起见没有显式地代入上式。这里的 $$ 就是所谓的“噪声预测模型”，用于近似当前采样出的噪声 $$. 通过实验探索，DDPM 作者发现将模型参数化为 $$ 的效果比 $$ 更好，并且把前面的系数丢掉效果更好。另外，对 $$ 求和可以改作对 $$ 均匀采样，因此损失函数简化为： $$ 最后，注意本文至此的所有推导都建立在给定一个 $$ 的基础上，实际训练时 $$ 是从训练集中采样的，因此最终的损失函数为： $$ 相应算法流程如下：

可见 DDPM 虽然推导有些复杂，但最后得到的算法流程却异常简单，效果也很好，难怪迅速成为了研究的热点。

一些注解

直观上 DDPM 干的事情可以总结为——前向过程对输入图像一步步加噪，使之变成高斯噪声；逆向过程使用模型来预测原图（或预测添加的噪声），进而把带噪图像一步步转换回真实图像。这里容易产生一个误解：既然每一步 $$ 都是去近似 $$，那么岂不是直接一步生成 $$ 就可以了？并不是这样的。注意 $$ 是不断从数据集中采样出来的，由于不同的 $$ 都有可能得到相同的 $$，所以随着训练的进行，$$ 拟合的是这些 $$ 的平均值，是一个模糊的图像。其实通过简单的推导就可以知道：

$$ 因此模型 $$ 拟合的真值是 $$，即 $$ 关于 $$ 的加权平均。同理，$$ 拟合的真值是 $$.

那么，纵观整个生成过程，我们可以把 $$ 理解为大方向。每次我们朝着大方向走一小步，然后重新看看大方向在哪里，再走下一小步。打个比方，我想从成都走到深圳，我知道大致要朝东南 45° 方向走，但是“差之毫厘，谬以千里”，直接走可能一不小心就登陆台湾了，所以我先走一小步到重庆；然后再看地图，大方向变成了东南 50°，于是又走一小步，但是拐弯过猛到了贵阳；没关系，再看地图，大方向变成了东南 30°……这样每走一小步都对大方向做一点修正，最后就能平稳地到达目的地了。

代码实现

Github repo: https://github.com/xyfJASON/Diffusion-Models-Implementations

结果展示

更多内容请查看代码仓库。

关于 clipping

在官方代码^[16]和若干其他实现中，我发现大家普遍喜欢使用 clipping，即对于逆向过程的每一步，在预测 $$ 之后，先算 $$，然后 clip 到 $$ 之间，再算 $$. 这样做为什么合理呢？这是因为 clipping 本质上是对模型误差的人工修正——$$ 是用来估计 $$ 的，本就应该在 $$ 之间，只是出于模型误差而跳脱了这个范围，所以强行把它 clip 回来并不违背理论；另外，clipping 只影响逆向过程，并不需要重新训练模型。

色调偏移问题

早期的实现版本在 MNIST 上 work 得很好，但是在 CelebA-HQ 上训练时出现了色调偏移（color shifting）问题。具体而言，我发现各个 epoch 之间的图片色调会发生明显偏移，比如前一个 epoch 图片都偏红，后一个 epoch 图片都偏蓝，有时候甚至亮/暗得根本看不清人脸，如下图所示：

本以为是模型还没收敛，但是 300 多个 epochs 之后仍然是这样，这就不得不重视起来。一番排查后，发现是我偷懒没有实现 EMA 导致的，特别是原作者把 decay rate 设置为 0.9999，意味着参数更新其实是很慢的。EMA 的本质是对历史权重做了加权平均，可以看作若干历史模型的集成。从这个角度来说，那些色调发生不同偏移的模型互相“抵消”，从而缓解了色调偏移问题。（注意只是缓解，并没有消除！）

后来我读到其实宋飏在论文^[3]里面就提到了这一现象，这也是他引入 EMA 的原因。说到底，色调偏移就是模型还没有收敛到真实分布的一个表现，只不过视觉上给人的冲击比较强烈罢了。

[update 2022.11.27] 虽然 EMA 的 decay rate 设置为 0.9999，但 tensorflow 的官方实现其实是这样的： $$ 随着 num_updates 增加，对应的 decay 序列是 $$，一直到 90000 步左右 decay 才会固定在 0.9999. 这样做能减小初始化的随机权重对整体权重的影响，模型见效更快。

References

1. Ho, Jonathan, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems 33 (2020): 6840-6851. ↩︎
2. Sohl-Dickstein, Jascha, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pp. 2256-2265. PMLR, 2015. ↩︎
3. Song, Yang, and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems 33 (2020): 12438-12448. ↩︎
4. Luo, Calvin. Understanding diffusion models: A unified perspective. arXiv preprint arXiv:2208.11970 (2022). ↩︎
5. Lilian Weng. What are Diffusion Models?. https://lilianweng.github.io/posts/2021-07-11-diffusion-models ↩︎
6. Angus Turner. Diffusion Models as a kind of VAE. https://angusturner.github.io/generative_models/2021/06/29/diffusion-probabilistic-models-I.html ↩︎
7. Denoising Diffusion-based Generative Modeling: Foundations and Applications. https://cvpr2022-tutorial-diffusion-models.github.io ↩︎
8. 苏剑林. (Jul. 06, 2022). 《生成扩散模型漫谈（二）：DDPM = 自回归式VAE 》[Blog post]. Retrieved from https://kexue.fm/archives/9152 ↩︎
9. 苏剑林. (Jul. 19, 2022). 《生成扩散模型漫谈（三）：DDPM = 贝叶斯 + 去噪 》[Blog post]. Retrieved from https://kexue.fm/archives/9164 ↩︎
10. 由浅入深了解Diffusion Model - ewrfcas的文章 - 知乎 https://zhuanlan.zhihu.com/p/525106459 ↩︎
11. 扩散模型之DDPM - 小小将的文章 - 知乎 https://zhuanlan.zhihu.com/p/563661713 ↩︎
12. Probabilistic Diffusion Model概率扩散模型理论与完整PyTorch代码详细解读. https://www.bilibili.com/video/BV1b541197HX ↩︎
13. Diffusion Model：比“GAN”还要牛逼的图像生成模型！https://www.bilibili.com/video/BV1pD4y1179T ↩︎
14. 【炼丹技巧】指数移动平均（EMA）的原理及PyTorch实现 - Nicolas的文章 - 知乎 https://zhuanlan.zhihu.com/p/68748778 ↩︎
15. https://huggingface.co/blog/annotated-diffusion ↩︎
16. https://github.com/lucidrains/denoising-diffusion-pytorch ↩︎
17. https://github.com/hojonathanho/diffusion ↩︎
18. https://github.com/openai/improved-diffusion ↩︎
19. https://github.com/lucidrains/imagen-pytorch ↩︎
20. https://github.com/tqch/ddpm-torch ↩︎
21. https://github.com/abarankab/DDPM ↩︎
22. https://github.com/w86763777/pytorch-ddpm ↩︎