DDPM = 自回归式 VAE

• DDPM 本质上已经不是传统的扩散模型了，它更多的是一个变分自编码器 VAE，实际上 DDPM 的原论文中也是将它按照 VAE 的思路进行推导的。所以，下面就从 VAE 的角度来重新介绍 DDPM

多步突破

• 在传统的 VAE 中，编码过程和生成过程都是一步到位的。这样做就只涉及到三个分布：编码分布 $p(z|x)$、生成分布 $q(x|z)$ 以及先验分布 $q(z)$，它的好处是形式比较简单，$x$ 与 $z$ 之间的映射关系也比较确定，因此可以同时得到编码模型和生成模型，实现隐变量编辑等需求；但是它的缺点也很明显，因为我们建模概率分布的能力有限，三个分布都只能建模为正态分布，这限制了模型的表达能力，最终通常得到偏模糊的生成结果
• 为了突破这个限制，DDPM 将编码过程和生成过程分解为 $T$ 步，每一步编码过程 $p({\mathbf{x}}_t|{\mathbf{x}}_{t-1})$ 和生成过程 $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$ 仅仅负责建模一个微小变化，它们依然建模为正态分布。对于微小变化来说，可以用正态分布足够近似地建模，类似于曲线在小范围内可以用直线近似，多步分解就有点像用分段线性函数拟合复杂曲线，因此理论上可以突破传统单步 VAE 的拟合能力限制。编码和生成公式为
  $$\begin{array}{rl}{\mathbf{x}}_t& ={\alpha }_t{\mathbf{x}}_{t-1}+{\beta }_t{\mathbf{\epsilon }}_t,{\mathbf{\epsilon }}_t\sim \mathcal{N}(0,\mathbf{I})\\ {\mathbf{x}}_{t-1}& =\frac{1}{{\alpha }_t}\left({\mathbf{x}}_t-{\beta }_t{\mathbf{ϵ}}_{\mathbf{\theta }}({\mathbf{x}}_t,t)\right)+{\sigma }_t\mathbf{z}\\ & =\mathbf{\mu }({\mathbf{x}}_t)+{\sigma }_t\mathbf{z},\mathbf{z}\sim \mathcal{N}(0,\mathbf{I})\end{array}$$可以将编码过程和生成过程写为下式：
  $$\begin{gathered} p({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;{\alpha }_t{\mathbf{x}}_{t-1},{\beta }_t^2\mathbf{I}) \\ q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};\mathbf{\mu }({\mathbf{x}}_t),{\sigma }_t^2\mathbf{I}) \end{gathered}$$可以看到，在编码时，传统 VAE 的均值方差都是用神经网络学习出来的，而 DDPM 放弃了模型的编码能力，最终只得到一个纯粹的生成模型；在生成时，DDPM 则将生成过程建模成均值向量可学习的正态分布 $\mathcal{N}({\mathbf{x}}_{t-1};\mathbf{\mu }({\mathbf{x}}_t),{\sigma }_t^2\mathbf{I})$。整个模型拥有可训练参数的就只有 $\mathbf{\mu }({\mathbf{x}}_t)$

联合散度

• 理解 VAE 最简洁的理论途径，就是将其理解为在最小化联合分布的 KL 散度，对于 DDPM 也是如此。生成和编码的联合分布分别为
  $$\begin{array}{rl}& p({\mathbf{x}}_0,{\mathbf{x}}_1,{\mathbf{x}}_2,\cdots ,{\mathbf{x}}_T)=p({\mathbf{x}}_T|{\mathbf{x}}_{T-1})\cdots p({\mathbf{x}}_2|{\mathbf{x}}_1)p({\mathbf{x}}_1|{\mathbf{x}}_0)\tilde{p}({\mathbf{x}}_0)\\ & q({\mathbf{x}}_0,{\mathbf{x}}_1,{\mathbf{x}}_2,\cdots ,{\mathbf{x}}_T)=q({\mathbf{x}}_0|{\mathbf{x}}_1)\cdots q({\mathbf{x}}_{T-2}|{\mathbf{x}}_{T-1})q({\mathbf{x}}_{T-1}|{\mathbf{x}}_T)q({\mathbf{x}}_T)\end{array}$$DDPM 的目的就是最小化两个联合分布之间的 KL 散度
  $$KL(p\Vert q)=\int p({\mathbf{x}}_T|{\mathbf{x}}_{T-1})\cdots p({\mathbf{x}}_1|{\mathbf{x}}_0)\tilde{p}({\mathbf{x}}_0)\log \frac{p({\mathbf{x}}_T|{\mathbf{x}}_{T-1})\cdots p({\mathbf{x}}_1|{\mathbf{x}}_0)\tilde{p}({\mathbf{x}}_0)}{q({\mathbf{x}}_0|{\mathbf{x}}_1)\cdots q({\mathbf{x}}_{T-1}|{\mathbf{x}}_T)q({\mathbf{x}}_T)}d{\mathbf{x}}_{0}d{\mathbf{x}}_1\cdots d{\mathbf{x}}_T$$

分而治之

• 下面对上述 KL 散度进行化简。由于目前分布 $p$ 不含任何的可训练参数，因此关于 $p$ 的积分就只是贡献一个可以忽略的常数，所以 KL 散度等价于
  $$\begin{array}{rl}& -\int p({\mathbf{x}}_T|{\mathbf{x}}_{T-1})\cdots p({\mathbf{x}}_1|{\mathbf{x}}_0)\tilde{p}({\mathbf{x}}_0)\log q({\mathbf{x}}_0|{\mathbf{x}}_1)\cdots q({\mathbf{x}}_{T-1}|{\mathbf{x}}_T)q({\mathbf{x}}_T)d{\mathbf{x}}_{0}d{\mathbf{x}}_1\cdots d{\mathbf{x}}_T\\ =& -\int p({\mathbf{x}}_T|{\mathbf{x}}_{T-1})\cdots p({\mathbf{x}}_1|{\mathbf{x}}_0)\tilde{p}({\mathbf{x}}_0)\left[\log q({\mathbf{x}}_T)+\sum _{t=1}^T\log q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)\right]d{\mathbf{x}}_{0}d{\mathbf{x}}_1\cdots d{\mathbf{x}}_T\end{array}$$由于先验分布 $q({\mathbf{x}}_T)$ 一般都取标准正态分布，也是没有参数的，所以这一项也只是贡献一个常数。因此需要计算的就是每一项
  $$\begin{array}{rl}& -\int p({\mathbf{x}}_T|{\mathbf{x}}_{T-1})\cdots p({\mathbf{x}}_1|{\mathbf{x}}_0)\tilde{p}({\mathbf{x}}_0)\log q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)d{\mathbf{x}}_{0}d{\mathbf{x}}_1\cdots d{\mathbf{x}}_T\\ =& -\int p({\mathbf{x}}_t|{\mathbf{x}}_{t-1})\cdots p({\mathbf{x}}_1|{\mathbf{x}}_0)\tilde{p}({\mathbf{x}}_0)\log q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)d{\mathbf{x}}_{0}d{\mathbf{x}}_1\cdots d{\mathbf{x}}_t\\ =& -\int p({\mathbf{x}}_t|{\mathbf{x}}_{t-1})p({\mathbf{x}}_{t-1},{\mathbf{x}}_{t-2},...,{\mathbf{x}}_1|{\mathbf{x}}_0)\tilde{p}({\mathbf{x}}_0)\log q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)d{\mathbf{x}}_{0}d{\mathbf{x}}_{t-1}d{\mathbf{x}}_t\\ =& -\int p({\mathbf{x}}_t|{\mathbf{x}}_{t-1})p({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)\tilde{p}({\mathbf{x}}_0)\log q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)d{\mathbf{x}}_{0}d{\mathbf{x}}_{t-1}d{\mathbf{x}}_t\\ =& \int p({\mathbf{x}}_t|{\mathbf{x}}_{t-1})p({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)\left[{\mathbb{E}}_{{\mathbf{x}}_0\sim \tilde{p}({\mathbf{x}}_0)}-\log q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)\right]d{\mathbf{x}}_{t-1}d{\mathbf{x}}_t\end{array}$$其中第一个等号是因为 $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$ 至多依赖到 ${\mathbf{x}}_t$，因此 $t+1$ 到 $T$ 的分布可以直接积分为 1；第三个等号则是因为 $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$ 也不依赖于 ${\mathbf{x}}_1,\cdots ,{\mathbf{x}}_{t-2}$，所以关于它们的积分我们也可以事先算出。此外，$p({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)$ 也可以直接写为 $p({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_{t-1};{\bar{\alpha }}_{t-1}{\mathbf{x}}_0,{\bar{\beta }}_{t-1}^2\mathbf{I})$
• 此外，由于拥有可训练参数的就只有 $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$，因此损失函数等价于
  $${\mathbb{E}}_{{\mathbf{x}}_0\sim \tilde{p}({\mathbf{x}}_0)}-\log q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$$

场景再现

• 正态分布为 $\mathcal{N}(\mathrm{x}\mid \mathbf{\mu },\mathbf{\Sigma })=\frac{1}{(2\pi {)}^{D/2}}\frac{1}{|\mathbf{\Sigma }{|}^{1/2}}\exp \left\{-\frac{1}{2}(\mathrm{x}-\mathbf{\mu }{)}^{\mathrm{T}}{\mathbf{\Sigma }}^{-1}(\mathrm{x}-\mathbf{\mu })\right\}$，因此除去优化无关的常数，有
  $$-\log q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)\sim \frac{1}{2{\sigma }_t^2}{\Vert {\mathbf{x}}_{t-1}-\mathbf{\mu }({\mathbf{x}}_t)\Vert }^2$$由 ${\mathbf{x}}_{t-1}=\frac{1}{{\alpha }_t}\left({\mathbf{x}}_t-{\beta }_t{\mathbf{\epsilon }}_t\right)$ 启发，将 $\mathbf{\mu }({\mathbf{x}}_t)$ 写为 $\mathbf{\mu }({\mathbf{x}}_t)=\frac{1}{{\alpha }_t}\left({\mathbf{x}}_t-{\beta }_t{\mathbf{ϵ}}_{\mathbf{\theta }}({\mathbf{x}}_t,t)\right)$ 的形式，可得
  $$\begin{array}{rl}-\log q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)& \sim \frac{1}{{\sigma }_t^2}{\Vert {\mathbf{x}}_{t-1}-\mathbf{\mu }({\mathbf{x}}_t)\Vert }^2\\ & =\frac{{\beta }_t^2}{{\alpha }_t^2{\sigma }_t^2}{\Vert {\mathbf{\epsilon }}_t-{\mathbf{ϵ}}_{\mathbf{\theta }}({\mathbf{x}}_t,t)\Vert }^2\\ & =\frac{{\beta }_t^2}{{\alpha }_t^2{\sigma }_t^2}{\Vert {\mathbf{\epsilon }}_t-{\mathbf{ϵ}}_{\mathbf{\theta }}({\bar{\alpha }}_t{\mathbf{x}}_0+{\alpha }_t{\bar{\beta }}_{t-1}{\bar{\mathbf{\epsilon }}}_{t-1}+{\beta }_t{\mathbf{\epsilon }}_t,t)\Vert }^2\end{array}$$再按照 “降低方差” 一节做换元，可得
  $$\begin{array}{rl}-\log q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)& \sim \frac{{\beta }_t^4}{{\bar{\beta }}_t^2{\alpha }_t^2{\sigma }_t^2}{\Vert \mathbf{\epsilon }-\frac{{\bar{\beta }}_t}{{\beta }_t}{\mathbf{ϵ}}_{\mathbf{\theta }}({\bar{\alpha }}_t{\mathbf{x}}_0+{\bar{\beta }}_t\mathbf{\epsilon },t)\Vert }^2\end{array}$$
• 现在损失函数可以写为
  $$\frac{{\beta }_t^4}{{\bar{\beta }}_t^2{\alpha }_t^2{\sigma }_t^2}{\mathbb{E}}_{\mathbf{\epsilon }\sim \mathcal{N}(0,\mathbf{I}),{\mathbf{x}}_0\sim \tilde{p}({\mathbf{x}}_0)}\left[{\Vert \mathbf{\epsilon }-\frac{{\bar{\beta }}_t}{{\beta }_t}{\mathbf{ϵ}}_{\mathbf{\theta }}({\bar{\alpha }}_t{\mathbf{x}}_0+{\bar{\beta }}_t\mathbf{\epsilon },t)\Vert }^2\right]$$这就得到了 DDPM 的训练目标了（原论文通过实验发现，去掉上式前面的系数后实际效果更好些）

超参设置

• 从最小化两个联合分布的 KL 散度出发，我们可以进一步理解为什么要设计适当的 ${\alpha }_t$ 使得 ${\bar{\alpha }}_T\approx 0$. 前面说了，$q({\mathbf{x}}_T)$ 一般都取标准正态分布 $\mathcal{N}({\mathbf{x}}_T;0,\mathbf{I})$。而我们的学习目标是最小化两个联合分布的 KL 散度，即希望 $p=q$，那么它们的边缘分布自然也相等，所以我们也希望
  $$q({\mathbf{x}}_T)=\int p({\mathbf{x}}_T|{\mathbf{x}}_0)\tilde{p}({\mathbf{x}}_0)d{\mathbf{x}}_0$$由于数据分布 $\tilde{p}({\mathbf{x}}_0)$ 是任意的，所以要使上式恒成立，只能让 $p({\mathbf{x}}_T|{\mathbf{x}}_0)=q({\mathbf{x}}_T)$，即退化为与 ${\mathbf{x}}_0$ 无关的标准正态分布

References

• 苏剑林. (Jul. 06, 2022). 《生成扩散模型漫谈（二）：DDPM = 自回归式 VAE 》[Blog post]. Retrieved from https://kexue.fm/archives/9152