DDPM( Denoising Diffusion Probabilistic Model )

Denoising Diffusion Probabilistic Models

Jonathan Ho, Ajay Jain, Pieter Abbeel
NeurIPS 2020

1 Background

Diffusion 模型为隐变量模型，${\mathbf{x}}_1,...,{\mathbf{x}}_T$ 为与原始数据 ${\mathbf{x}}_0\sim q({\mathbf{x}}_0)$ 维度一致的隐变量，所有隐变量之间满足马尔科夫性。已知 $p({\mathbf{x}}_T)\sim \mathcal{N}({\mathbf{x}}_T;0,\mathbf{I})$，计算联合概率 $p_{\theta }({\mathbf{x}}_{0:T})$ 称为 逆过程。
$$p_{\theta }({\mathbf{x}}_{0:T})=p({\mathbf{x}}_T)\prod _{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t),p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))\text{\hspace{1em}(1)}$$

其中，${\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t)$ 和 ${\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t)$ 为神经网络。
估计后验 $q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)$ 称为 扩散过程。
$$q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)=\prod _{t=1}^{T}q({\mathbf{x}}_t|{\mathbf{x}}_{t-1}),q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})\text{\hspace{1em}(2)}$$

扩散过程可以看做，根据超参数 ${\beta }_t$，向数据中添加高斯噪声的过程。
下面介绍扩散过程的一个性质，定义 ${\alpha }_t=1-{\beta }_t,{\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$，根据式 2 可得，
$$\begin{array}{rl}{\mathbf{x}}_t& =\sqrt{{\alpha }_t}{\mathbf{x}}_{t-1}+\sqrt{{\beta }_t}{\mathbf{ϵ}}_t\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{\mathbf{x}}_{t-2}+\sqrt{{\alpha }_t{\beta }_{t-1}}{\mathbf{ϵ}}_{t-1}+\sqrt{{\beta }_t}{\mathbf{ϵ}}_t\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}{\alpha }_{t-2}}{\mathbf{x}}_{t-3}+\sqrt{{\alpha }_t{\alpha }_{t-1}{\beta }_{t-2}}{\mathbf{ϵ}}_{t-2}+\sqrt{{\alpha }_t{\beta }_{t-1}}{\mathbf{ϵ}}_{t-1}+\sqrt{{\beta }_t}{\mathbf{ϵ}}_t\\ & =...\\ & =\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{{\alpha }_t{\alpha }_{t-1}...{\alpha }_2{\beta }_1}{\mathbf{ϵ}}_1+...+\sqrt{{\alpha }_t{\beta }_{t-1}}{\mathbf{ϵ}}_{t-1}+\sqrt{{\beta }_t}{\mathbf{ϵ}}_t\end{array}\text{\hspace{1em}(3)}$$

其中，所有 ${\mathbf{ϵ}}_t\sim \mathcal{N}(0,\mathbf{I})$ 均为随机噪声，根据高斯噪声的叠加性可得，
$$\begin{array}{rl}q({\mathbf{x}}_t|{\mathbf{x}}_0)& =\mathcal{N}({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,({\alpha }_t{\alpha }_{t-1}...{\alpha }_2{\beta }_1+...+{\alpha }_t{\beta }_{t-1}+{\beta }_t)\mathbf{I})\\ & =\mathcal{N}({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,(1-{\bar{\alpha }}_t)\mathbf{I})\end{array}\text{\hspace{1em}(4)}$$

根据这条性质，我们可以从 ${\mathbf{x}}_0$ 直接采样 ${\mathbf{x}}_{t-1}$，无需中间步骤。

我们的目的为训练神经网络，使逆过程得到的 $p_{\theta }({\mathbf{x}}_0)$ 尽可能与真实分布 $q({\mathbf{x}}_0)$ 接近，即最小化KL散度。
$$\begin{array}{rl}{D}_{KL}(q({\mathbf{x}}_0)||p_{\theta }({\mathbf{x}}_0))& =\underset{constant}{\underbrace{\int q({\mathbf{x}}_0)\log q({\mathbf{x}}_0)d{\mathbf{x}}_0}}-\int q({\mathbf{x}}_0)\log p_{\theta }({\mathbf{x}}_0)d{\mathbf{x}}_0\\ & =const-\int q({\mathbf{x}}_0)\log p_{\theta }({\mathbf{x}}_{0:T})d{\mathbf{x}}_{0:T}\\ & =const-\int q({\mathbf{x}}_0)\log \frac{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{p}_{\theta }({\mathbf{x}}_{0:T})d{\mathbf{x}}_{0:T}\\ & \le const-\int q({\mathbf{x}}_{0:T})\log \frac{p_{\theta }({\mathbf{x}}_{0:T})}{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}d{\mathbf{x}}_{0:T}\\ & =const+{\mathbb{E}}_q[-\log \frac{p_{\theta }({\mathbf{x}}_{0:T})}{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}]\\ & =const+\underset{L}{\underbrace{{\mathbb{E}}_q[-\log p({\mathbf{x}}_T)-\sum _{t\ge 1}\log \frac{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}{q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})}]}}\end{array}\text{\hspace{1em}(5)}$$

由于 $q({\mathbf{x}}_0)$ 为真实分布，所以第一项为常数，训练目标变为最小化 $L$。
$$\begin{array}{rl}L& ={\mathbb{E}}_q[-\log p({\mathbf{x}}_T)-\sum _{t>1}\log \frac{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}{q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})}-\log \frac{p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}{q({\mathbf{x}}_1|{\mathbf{x}}_0)}]\\ & ={\mathbb{E}}_q[-\log p({\mathbf{x}}_T)-\sum _{t>1}\log \frac{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}{q({\mathbf{x}}_t|{\mathbf{x}}_{t-1},{\mathbf{x}}_0)}-\log \frac{p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}{q({\mathbf{x}}_1|{\mathbf{x}}_0)}]\\ & ={\mathbb{E}}_q[-\log p({\mathbf{x}}_T)-\sum _{t>1}\log \frac{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}\frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}{q({\mathbf{x}}_t|{\mathbf{x}}_0)}-\log \frac{p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}{q({\mathbf{x}}_1|{\mathbf{x}}_0)}]\\ & ={\mathbb{E}}_q[-\log p({\mathbf{x}}_T)-\sum _{t>1}\log \frac{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}-\sum _{t>1}\log \frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}{q({\mathbf{x}}_t|{\mathbf{x}}_0)}-\log \frac{p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}{q({\mathbf{x}}_1|{\mathbf{x}}_0)}]\\ & ={\mathbb{E}}_q[-\log \frac{p({\mathbf{x}}_T)}{q({\mathbf{x}}_T|{\mathbf{x}}_0)}-\sum _{t>1}\log \frac{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}-\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)]\\ & ={\mathbb{E}}_q[\underset{L_T}{\underbrace{D_{KL}(q({\mathbf{x}}_T|{\mathbf{x}}_0)||p({\mathbf{x}}_T))}}+\sum _{t>1}\underset{L_{t-1}}{\underbrace{D_{KL}(q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)||p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))}}\underset{L_0}{\underbrace{-\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}}]\end{array}\text{\hspace{1em}(6)}$$

2 Diffusion models and denoising autoencoders

2.1 Forward process and $L_T$

由于 $q({\mathbf{x}}_T|{\mathbf{x}}_0)$ 和 $p({\mathbf{x}}_T)$ 均为正态分布，所以 $L_T=0$。

2.2 Reverse process and $L_{1:T-1}$

$L_{t-1}$ 中，$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_{t-1};{\tilde{\mathbf{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0),{\tilde{\mathbf{\beta }}}_t\mathbf{I})$，证明如下图所示。图源：54、Diffusion Model扩散模型理论与完整PyTorch代码详细解读。

本文作者将 $p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$ 中的 ${\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t)$ 替换为 ${\sigma }_t^2$，${\sigma }_t^2$ 取 ${\beta }_t$ 或 ${\tilde{\beta }}_t$，取得了相似的效果。
由于 $D_{KL}(\mathcal{N}({\mu }_1,{\sigma }_1^2)||\mathcal{N}({\mu }_2,{\sigma }_2^2))=\frac{({\mu }_1-{\mu }_2{)}^2+{\sigma }_1^2}{2{\sigma }_2^2}+\log \frac{{\sigma }_2}{{\sigma }_1}-\frac{1}{2}$，所以，
$$\begin{array}{rl}{L}_{t-1}& ={\mathbb{E}}_q[\frac{1}{2{\sigma }_t^2}\Vert {\tilde{\mathbf{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)-{\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t){\Vert }^2]+C\\ & ={\mathbb{E}}_q[\frac{1}{2{\sigma }_t^2}\Vert \frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-\overline{{\alpha }_t}}{\mathbf{x}}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-\overline{{\alpha }_t}}{\mathbf{x}}_0-{\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t){\Vert }^2]+C\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}[\frac{1}{2{\sigma }_t^2}\Vert \frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t({\mathbf{x}}_0,\mathbf{ϵ})-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}\mathbf{ϵ})-{\mathbf{\mu }}_{\theta }({\mathbf{x}}_t({\mathbf{x}}_0,\mathbf{ϵ}),t){\Vert }^2]+C\end{array}\text{\hspace{1em}(7)}$$

其中，${\mathbf{x}}_t({\mathbf{x}}_0,\mathbf{ϵ})=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ}$，$\mathbf{ϵ}$ 为由 ${\mathbf{x}}_0$ 采样 ${\mathbf{x}}_t$ 时引入的高斯噪声，C 为常数。
根据式 7 ，可以将神经网络 ${\mathbf{\mu }}_{\theta }$ 的形式变为，
$${\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t)=\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t))\text{\hspace{1em}(8)}$$

损失 $L_{t-1}$ 就变为，
$${\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}[\frac{{\beta }_t^2}{2{\sigma }_t^2{\alpha }_t(1-{\bar{\alpha }}_t)}\Vert \mathbf{ϵ}-{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t({\mathbf{x}}_0,\mathbf{ϵ}),t){\Vert }^2]+C\text{\hspace{1em}(9)}$$

2.3 Reverse process and $L_0$

原文中作者根据图像的性质做出了一些变化，本文在此进行忽略。
根据式 1 可知，
$$p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)=\mathcal{N}({\mathbf{x}}_0;{\mathbf{\mu }}_{\theta }({\mathbf{x}}_1,1),{\sigma }_1^2\mathbf{I})\text{\hspace{1em}(10)}$$

所以，
$$L_0=\frac{1}{2{\sigma }_1^2}{\mathbb{E}}_{{\mathbf{x}}_0,{\overline{\mathbf{ϵ}}}_1}[\Vert {\mathbf{x}}_0-{\mathbf{\mu }}_{\theta }({\mathbf{x}}_1({\mathbf{x}}_0,{\overline{\mathbf{ϵ}}}_1),1){\Vert }^2]+C^{\prime }\text{\hspace{1em}(11)}$$

其中，$C^{\prime }$ 为常数。

2.4 Simplified training objective

作者发现，将损失 $L$ 简化为 $L_{simple}$，能提高模型的效果。
$$L_{simple}(\theta )={\mathbb{E}}_{t,{\mathbf{x}}_0,\mathbf{ϵ}}[\Vert \mathbf{ϵ}-{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t({\mathbf{x}}_0,\mathbf{ϵ}),t){\Vert }^2]$$

训练的伪代码如下图所示。

3 Sampling

采样的过程就是采样 ${\mathbf{x}}_T$ 逆扩散生成 ${\mathbf{x}}_0$ 的过程，伪代码如下图所示。