DDPM

原论文标题：Denoising diffusion probabilistic models

背景

动机

这篇文章主要的动机，就是在Diffusion Probabilistic Models¹基础上，将对样本的预测转换为对噪声的预测，并简化训练目标，让其能生成高质量图像。

贡献

• 改进了DDPM：简化训练目标，使训练更加稳定，生成的图像质量更高。
• 建立等价性：揭示了特定参数下的DPM，其训练目标等价于多噪声级别的去噪分数匹配 (denoising score matching over multiple noise levels)，其采样过程等价于退火朗之万动力学 (annealed Langevin dynamics）。（其实就是说和Yang Song的NCSN模型²等价）

方法

回顾DPM

记初始样本为随机向量$x_0\sim q(x_0)$，DPM定义了如下的forward/diffusion process：

$$q\left({\mathbf{x}}_{1:T}|{\mathbf{x}}_0\right):=\prod _{t=1}^{T}q\left({\mathbf{x}}_t|{\mathbf{x}}_{t-1}\right),q\left({\mathbf{x}}_t|{\mathbf{x}}_{t-1}\right):=\mathcal{N}\left({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I}\right)\text{\hspace{1em}(2)}$$

其中${\beta }_1,\dots ,{\beta }_T$是不同方差的高斯噪声。由于$q$未知，那么反向转换$q(x_{t-1}|x_t)$是intractable的。因此DPM设计$p_{\theta }$去近似$q$，因此定义了如下的reverse process（从$p\left({\mathbf{x}}_T\right)=\mathcal{N}\left({\mathbf{x}}_T;0,\mathbf{I}\right)$开始的、具有高斯转换的马尔科夫链）：

$$p_{\theta }\left({\mathbf{x}}_{0:T}\right):=p\left({\mathbf{x}}_T\right)\prod _{t=1}^T{p}_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right),p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right):=\mathcal{N}\left({\mathbf{x}}_{t-1};{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t,t\right),{\mathbf{\Sigma }}_{\theta }\left({\mathbf{x}}_t,t\right)\right)\text{\hspace{1em}(1)}$$

这么定义forward/reverse process的背后原因之一：对于任一复杂的分布，理论上都能找到一个复杂函数将其映射到高斯分布上；那么生成样本时，可以在这个高斯分布上采样，然后通过这个复杂函数的逆操作来得到对应样本。DPM定义的markov chain就是这样的复杂函数的一种选择。

根据公式(2)，不难推导出：

$$q\left({\mathbf{x}}_t|{\mathbf{x}}_0\right)=\mathcal{N}\left({\mathbf{x}}_t;{\sqrt{\bar{\alpha }}}_t{\mathbf{x}}_0,\left(1-{\bar{\alpha }}_t\right)\mathbf{I}\right)\text{\hspace{1em}(4)}$$

其中${\alpha }_t:=1-{\beta }_t,{\bar{\alpha }}_t:={\prod }_{s=1}^t{\alpha }_s$，这里使用reparameterization trick可以得到：${\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_t,{ϵ}_t\sim \mathcal{N}\left(0,\mathbf{I}\right)$（后文会用到此结论，公式4的推导可参考《生成扩散模型漫谈（一）：DDPM = 拆楼 + 建楼》）。

求解$\theta$就是优化下面这个variational bound(变分下界)：

$$\begin{array}{rl}\mathbb{E}\left[-\log p_{\theta }\left({\mathbf{x}}_0\right)\right]& \le {\mathbb{E}}_q\left[-\log \frac{p_{\theta }\left({\mathbf{x}}_{0:T}\right)}{q\left({\mathbf{x}}_{1:T}|{\mathbf{x}}_0\right)}\right]\\ & ={\mathbb{E}}_q\left[-\log p\left({\mathbf{x}}_T\right)-\sum _{t\ge 1}\log \frac{p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)}{q\left({\mathbf{x}}_t|{\mathbf{x}}_{t-1}\right)}\right]\\ & =:L\end{array}\text{\hspace{1em}(3)}$$

公式(3)推导

其中的中间公式：

$$\underset{\text{evidence}}{\underbrace{E_q\left[\log p_{\theta }\left({\mathbf{x}}_0\right)\right]}}=\underset{\text{ELBO}}{\underbrace{E_q\left[\log \frac{p_{\theta }\left({\mathbf{x}}_{T:1},{\mathbf{x}}_0\right)}{q\left({\mathbf{x}}_{T:1}|{\mathbf{x}}_0\right)}\right]}}+KL\left(q\left({\mathbf{x}}_{T:1}|{\mathbf{x}}_0\right)\Vert p_{\theta }\left({\mathbf{x}}_{T:1}|{\mathbf{x}}_0\right)\right)$$

第一项在VAE里称为“证据”（表述数据集的统计信息），第二项称为evidence lower bound (ELBO)。在这篇文章里似乎也称为variational lower bound (VLB)。

公式(3)可展开简化为：

$$\begin{array}{rl}L& ={\mathbb{E}}_q\left[-\log \frac{p_{\theta }\left({\mathbf{x}}_{0:T}\right)}{q\left({\mathbf{x}}_{1:T}|{\mathbf{x}}_0\right)}\right]\\ & ={\mathbb{E}}_q\left[-\log p\left({\mathbf{x}}_T\right)-\sum _{t\ge 1}\log \frac{p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)}{q\left({\mathbf{x}}_t|{\mathbf{x}}_{t-1}\right)}\right]\\ & ={\mathbb{E}}_q\left[-\log p\left({\mathbf{x}}_T\right)-\sum _{t>1}\log \frac{p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)}{q\left({\mathbf{x}}_t|{\mathbf{x}}_{t-1}\right)}-\log \frac{p_{\theta }\left({\mathbf{x}}_0|{\mathbf{x}}_1\right)}{q\left({\mathbf{x}}_1|{\mathbf{x}}_0\right)}\right]\\ & ={\mathbb{E}}_q\left[-\log p\left({\mathbf{x}}_T\right)-\sum _{t>1}\log \frac{p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)}{q\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0\right)}\cdot \frac{q\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_0\right)}{q\left({\mathbf{x}}_t|{\mathbf{x}}_0\right)}-\log \frac{p_{\theta }\left({\mathbf{x}}_0|{\mathbf{x}}_1\right)}{q\left({\mathbf{x}}_1|{\mathbf{x}}_0\right)}\right]\\ & ={\mathbb{E}}_q\left[-\log \frac{p\left({\mathbf{x}}_T\right)}{q\left({\mathbf{x}}_T|{\mathbf{x}}_0\right)}-\sum _{t>1}\log \frac{p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)}{q\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0\right)}-\log p_{\theta }\left({\mathbf{x}}_0|{\mathbf{x}}_1\right)\right]\\ & ={\mathbb{E}}_q\left[D_{\mathrm{KL}}\left(q\left({\mathbf{x}}_T|{\mathbf{x}}_0\right)\Vert p\left({\mathbf{x}}_T\right)\right)+\sum _{t>1}{D}_{\mathrm{KL}}\left(q\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0\right)\Vert p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)\right)-\log p_{\theta }\left({\mathbf{x}}_0|{\mathbf{x}}_1\right)\right]\end{array}$$

因此得到公式(5)：

$${\mathbb{E}}_q[\underset{L_T}{\underbrace{D_{\mathrm{KL}}\left(q\left({\mathbf{x}}_T|{\mathbf{x}}_0\right)\Vert p\left({\mathbf{x}}_T\right)\right)}}+\sum _{t>1}\underset{L_{t-1}}{\underbrace{D_{\mathrm{KL}}\left(q\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0\right)\Vert p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)\right)}}\underset{L_0}{\underbrace{-\log p_{\theta }\left({\mathbf{x}}_0|{\mathbf{x}}_1\right)}}]\text{\hspace{1em}(5)}$$

训练时，每次迭代将优化公式(5)的随机某一项$L_t$。

虽然得到了公式(5)，但是需要考察每一项是否可处理（tractable，是否能得到解析式或被神经网络建模）。

$L_0$： 其中$p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)$由被$\theta$参数化的模型所给出，因此$L_0$可处理。

$L_T$： 由公式(4)可得$q\left({\mathbf{x}}_T|{\mathbf{x}}_0\right)=\mathcal{N}\left({\mathbf{x}}_T;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,\left(1-{\bar{\alpha }}_T\right)\mathbf{I}\right)$，而$p\left({\mathbf{x}}_T\right)=\mathcal{N}\left({\mathbf{x}}_T;0,\mathbf{I}\right)$。因此也是可处理。

$L_{t-1}$： $p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$由模型$\theta$计算，而$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,x_0)$可由贝叶斯公式导出：

$$\begin{array}{rl}q\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0\right)& =\mathcal{N}\left({\mathbf{x}}_{t-1};{\tilde{\mathbf{\mu }}}_t\left({\mathbf{x}}_t,{\mathbf{x}}_0\right),{\tilde{\beta }}_t\mathbf{I}\right),\\ \text{ where }{\tilde{\mathbf{\mu }}}_t\left({\mathbf{x}}_t,{\mathbf{x}}_0\right)& :=\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{\mathbf{x}}_0+\frac{\sqrt{{\alpha }_t}\left(1-{\bar{\alpha }}_{t-1}\right)}{1-{\bar{\alpha }}_t}{\mathbf{x}}_t\text{ and }{\tilde{\beta }}_t:=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t\end{array}\text{\hspace{1em}(6)(7)}$$

公式(6)推导

因此$L_{t-1}$也是可处理的，那么公式(5)所定义的损失函数是可以优化的。

DPM简化→DDPM

DPM的实现方式具有很大的自由度，论文作者们进行了如下的简化。

$L_T=D_{\mathrm{KL}}\left(q\left({\mathbf{x}}_T|{\mathbf{x}}_0\right)\Vert p\left({\mathbf{x}}_T\right)\right)$：

作者将${\beta }_t$固定为常量，因此$L_T$没可学习参数，也成为了常量，训练中可忽略。

$L_{1:T-1}$：

主要是讨论$p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)=\mathcal{N}\left({\mathbf{x}}_{t-1};{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t,t\right),{\mathbf{\Sigma }}_{\theta }\left({\mathbf{x}}_t,t\right)\right)\text{ for }1<t\le T$中的${\mu }_{\theta }$和${\Sigma }_{\theta }$的选取。

作者设定${\Sigma }_{\theta }({\mathbf{x}}_t,t)={\sigma }_t^2$，作者的实验表明，令${\sigma }_t^2={\beta }_t$以及${\sigma }_t^2={\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t$具有相似的效果。并且这两种选择对应于${\mathbf{x}}_0$为极端的两种分布（正态分布和狄拉克分布）下的最优选择（分别reverse process中具有coordinatewise unit variance的数据样本的熵上界和下界）。实验中，作者选择了第一种（第二种不就是公式(7)$q$分布的方差吗？为啥reverse process得到的数据样本分布会是下界呢？留坑，有时间了探索）。

对于${\mu }_{\theta }$的选择：将$p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)=\mathcal{N}\left({\mathbf{x}}_{t-1};{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t,t\right),{\sigma }_t^2\mathbf{I}\right)$代入公式(5)中的$L_{t-1}$可得到：

$$L_{t-1}={\mathbb{E}}_q\left[\frac{1}{2{\sigma }_t^2}{\Vert {\tilde{\mathbf{\mu }}}_t\left({\mathbf{x}}_t,{\mathbf{x}}_0\right)-{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t,t\right)\Vert }^2\right]+C\text{\hspace{1em}(8)}$$

公式(8)推导

由公式(8)可见，${\mu }_{\theta }$的选择最直观的条件是要能预测后验均值${\tilde{\mu }}_t$。那么将公式(4)重参数后的表达式以及公式(7)代入公式(8)可得：

$$\begin{array}{rl}{L}_{t-1}-C& ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}\left[\frac{1}{2{\sigma }_t^2}{\Vert {\tilde{\mathbf{\mu }}}_t\left({\mathbf{x}}_t\left({\mathbf{x}}_0,\mathbf{ϵ}\right),\frac{1}{\sqrt{{\bar{\alpha }}_t}}\left({\mathbf{x}}_t\left({\mathbf{x}}_0,\mathbf{ϵ}\right)-\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ}\right)\right)-{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t\left({\mathbf{x}}_0,\mathbf{ϵ}\right),t\right)\Vert }^2\right]\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}\left[\frac{1}{2{\sigma }_t^2}{\Vert \frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t\left({\mathbf{x}}_0,\mathbf{ϵ}\right)-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}\mathbf{ϵ}\right)-{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t\left({\mathbf{x}}_0,\mathbf{ϵ}\right),t\right)\Vert }^2\right]\end{array}\text{\hspace{1em}(9)(10)}$$

公式(10)表明${\mu }_{\theta }$应预测$\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t\left({\mathbf{x}}_0,\mathbf{ϵ}\right)-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}\mathbf{ϵ}\right)$，因此${\mu }_{\theta }$可按照${\tilde{\mu }}_t$的形式来参数化：

$${\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t,t\right)={\tilde{\mathbf{\mu }}}_t\left({\mathbf{x}}_t,\frac{1}{\sqrt{{\bar{\alpha }}_t}}\left({\mathbf{x}}_t-\sqrt{1-{\bar{\alpha }}_t}{\mathbf{ϵ}}_{\theta }\left({\mathbf{x}}_t\right)\right)\right)=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }\left({\mathbf{x}}_t,t\right)\right)\text{\hspace{1em}(11)}$$

我的理解是$p_{\theta }$和$q$均值都固定了，那么这两个分布靠近的话，近似为期望相等，那么他们的期望应该有类似的表达形式。

按照上述定义的$p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)=\mathcal{N}\left({\mathbf{x}}_{t-1};{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t,t\right),{\sigma }_t^2\mathbf{I}\right)$，reverse process中，采样$x_{t-1}$就等价于计算：

$${\mathbf{x}}_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }\left({\mathbf{x}}_t,t\right)\right)+{\sigma }_t\mathbf{z},\text{ where }\mathbf{z}\sim \mathcal{N}(0,\mathbf{I})\text{\hspace{1em}(11-1)}$$

将公式(11)代入到公式(10)，可进一步得到简化后的训练目标：

$${\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}\left[\frac{{\beta }_t^2}{2{\sigma }_t^2{\alpha }_t\left(1-{\bar{\alpha }}_t\right)}{\Vert \mathbf{ϵ}-{\mathbf{ϵ}}_{\theta }\left(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ},t\right)\Vert }^2\right]\text{\hspace{1em}(12)}$$

可证³公式(12)和Yang Song提出的NCSN中的损失（原文公式(2)）等价VDM等价性证明 ，所以DDPM的训练目标等价于多噪声级别的去噪分数匹配，其采样过程等价于退火朗之万动力学。

综上，公式(12)将对$\tilde{\mu }$-prediction最终转换为了$ϵ$-prediction。
这样做的好处可能有：由于${\mu }_{\theta }$的值域范围大，很容易导致训练不稳定，而相对来说${ϵ}_{\theta }$的训练更稳定，实验结果也证明了这一点。

$L_0=-\log p_{\theta }\left({\mathbf{x}}_0|{\mathbf{x}}_1\right)$：

$x_1$是连续数据，而$x_0$是离散数据（从$\{0,1,\dots ,255\}$缩放到$[-1,1]$），因此不能按照公式(11-1)来从$x_1$得到$x_0$。 作者设计了如下的离散化方式：

$$\begin{array}{rl}{p}_{\theta }\left({\mathbf{x}}_0\mid {\mathbf{x}}_1\right)& =\prod _{i=1}^D{\int }_{{\delta }_{-}\left(x_0^i\right)}^{{\delta }_{+}\left(x_0^i\right)}\mathcal{N}\left(x;{\mu }_{\theta }^i\left({\mathbf{x}}_1,1\right),{\sigma }_1^2\right)dx\\ {\delta }_{+}(x)& =\{\begin{array}{ll}\infty & \text{ if }x=1\\ x+\frac{1}{255}& \text{ if }x<1\end{array}{\delta }_{-}(x)=\{\begin{array}{ll}-\infty & \text{ if }x=-1\\ x-\frac{1}{255}& \text{ if }x>-1\end{array}\end{array}\text{\hspace{1em}(13)}$$

其中$D$为数据维度，$i$代表维度索引。计算$L_0$时，按照公式(13)来计算似然。

最终简化 作者实验中发现去除公式12中范数前的权重会达到更好的效果：

$$L_{\text{simple }}(\theta ):={\mathbb{E}}_{t,{\mathbf{x}}_0,\mathbf{ϵ}}\left[{\Vert \mathbf{ϵ}-{\mathbf{ϵ}}_{\theta }\left(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ},t\right)\Vert }^2\right]$$

作者给出的理论解释是：去除权重后，将低噪声下（低${\beta }_t$值）的权重变小了，模型更关注更难（高${\beta }_t$值）任务的学习，从而得到更好的学习效果。

思考：

• $T$越大越好吗？（扩散充分性，单步学习难度，离散过程近似连续过程，计算效率）
• ${\beta }_t$变化规律有可能对结果有影响吗？

参考读物

• 扩散模型(Diffusion Model)详解：直观理解、数学原理、PyTorch 实现
• 降噪扩散概率模型（Denoising diffusion probabilistic model,DDPM）

Footnotes

1. Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256–2265, 2015. ↩

2. Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, pages 11895–11907, 2019. ↩

3. Calvin Luo. Understanding diffusion models: a unified perspective. 2022. arXiv:2208.11970. ↩