扩散模型之DDPM

“What I cannot create, I do not understand.” -- Richard Feynman

近段时间最火的方向无疑是基于文本用AI生成图像，继OpenAI在2021提出的文本转图像模型DALLE之后，越来越多的大公司卷入这个方向，如谷歌在今年相继推出了Imagen和Parti。一些主流的文本转图像模型如DALL·E 2，stable-diffusion和Imagen采用了扩散模型（Diffusion Model）作为图像生成模型，这也引发了对扩散模型的研究热潮。相比GAN来说，扩散模型训练更稳定，而且能够生成更多样的样本，OpenAI的论文Diffusion Models Beat GANs on Image Synthesis也证明了扩散模型能够超越GAN。简单来说，扩散模型包含两个过程：前向扩散过程和反向生成过程，前向扩散过程是对一张图像逐渐添加高斯噪音直至变成随机噪音，而反向生成过程是去噪音过程，我们将从一个随机噪音开始逐渐去噪音直至生成一张图像，这也是我们要求解或者训练的部分。扩散模型与其它主流生成模型的对比如下所示：

目前所采用的扩散模型大都是来自于2020年的工作DDPM: Denoising Diffusion Probabilistic Models，DDPM对之前的扩散模型（具体见Deep Unsupervised Learning using Nonequilibrium Thermodynamics）进行了简化，并通过变分推断（variational inference）来进行建模，这主要是因为扩散模型也是一个隐变量模型（latent variable model），相比VAE这样的隐变量模型，扩散模型的隐变量是和原始数据是同维度的，而且推理过程（即扩散过程）往往是固定的。这篇文章将基于DDPM详细介绍扩散模型的原理，并给出具体的代码实现和分析。

扩散模型原理

扩散模型包括两个过程：前向过程（forward process）和反向过程（reverse process），其中前向过程又称为扩散过程（diffusion process），如下图所示。无论是前向过程还是反向过程都是一个参数化的马尔可夫链（Markov chain），其中反向过程可以用来生成数据，这里我们将通过变分推断来进行建模和求解。

扩散过程

扩散过程是指的对数据逐渐增加高斯噪音直至数据变成随机噪音的过程。对于原始数据${\mathbf{x}}_0\sim q({\mathbf{x}}_0)$\mathbf{x}_0 \sim q(\mathbf{x}_0)，总共包含$T$T步的扩散过程的每一步都是对上一步得到的数据${\mathbf{x}}_{t-1}$\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})$$q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t} \mathbf{x}_{t-1}, \beta_t\mathbf{I}) \\

这里$\{{\beta }_t{\}}_{t=1}^T$\{\beta_t\}^T_{t=1}为每一步所采用的方差，它介于0～1之间。对于扩散模型，我们往往称不同step的方差设定为variance schedule或者noise schedule，通常情况下，越后面的step会采用更大的方差，即满足${\beta }_1<{\beta }_2<\cdots <{\beta }_T$\beta_1 < \beta_2 < \dots < \beta_T。在一个设计好的variance schedule下，的如果扩散步数$T$T足够大，那么最终得到的${\mathbf{x}}_T$\mathbf{x}_{T}就完全丢失了原始数据而变成了一个随机噪音。 扩散过程的每一步都生成一个带噪音的数据${\mathbf{x}}_t$\mathbf{x}_{t}，整个扩散过程也就是一个马尔卡夫链：

$$q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)=\prod _{t=1}^{T}q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})$$q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) = \prod^T_{t=1} q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) \\

另外要指出的是，扩散过程往往是固定的，即采用一个预先定义好的variance schedule，比如DDPM就采用一个线性的variance schedule。

扩散过程的一个重要特性是我们可以直接基于原始数据${\mathbf{x}}_0$\mathbf{x}_{0}来对任意$t$t步的${\mathbf{x}}_t$\mathbf{x}_{t}进行采样：${\mathbf{x}}_t\sim q({\mathbf{x}}_t|{\mathbf{x}}_0)$\mathbf{x}_{t}\sim q(\mathbf{x}_t \vert \mathbf{x}_0)。这里定义${\alpha }_t=1-{\beta }_t$\alpha_t = 1 - \beta_t和${\overline{\alpha }}_t=\prod _{i=1}^t{\alpha }_i$\bar{\alpha}_t = \prod_{i=1}^t \alpha_i，通过重参数技巧（和VAE类似），那么有：

$$\begin{array}{rlr}{\mathbf{x}}_t& =\sqrt{{\alpha }_t}{\mathbf{x}}_{t-1}+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}& \text{ ;where }{ϵ}_{t-1},{ϵ}_{t-2},\cdots \sim \mathcal{N}(\mathbf{0},\mathbf{I})\\ & =\sqrt{{\alpha }_t}(\sqrt{{\alpha }_{t-1}}{\mathbf{x}}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{ϵ}_{t-2})+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{\mathbf{x}}_{t-2}+\sqrt{{\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}}^2+{\sqrt{1-{\alpha }_t}}^2}{\overline{ϵ}}_{t-2}& \text{ ;where }{\overline{ϵ}}_{t-2}\text{ merges two Gaussians (*).}\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{\mathbf{x}}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{ϵ}}_{t-2}\\ & =\dots \\ & =\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ\end{array}$$\begin{aligned} \mathbf{x}_t &= \sqrt{\alpha_t}\mathbf{x}_{t-1} + \sqrt{1 - \alpha_t}\mathbf{\epsilon}_{t-1} & \text{ ;where } \mathbf{\epsilon}_{t-1}, \mathbf{\epsilon}_{t-2}, \dots \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \\ &= \sqrt{\alpha_t}(\sqrt{\alpha_{t-1}}\mathbf{x}_{t-2} + \sqrt{1 - \alpha_{t-1}}\mathbf{\epsilon}_{t-2}) + \sqrt{1 - \alpha_t}\mathbf{\epsilon}_{t-1} \\ &= \sqrt{\alpha_t \alpha_{t-1}} \mathbf{x}_{t-2} + \sqrt{\sqrt {\alpha_t-\alpha_t \alpha_{t-1}}^2+\sqrt{1-\alpha_t }^2} \bar{\mathbf{\epsilon}}_{t-2} & \text{ ;where } \bar{\mathbf{\epsilon}}_{t-2} \text{ merges two Gaussians (*).} \\ &= \sqrt{\alpha_t \alpha_{t-1}} \mathbf{x}_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \bar{\mathbf{\epsilon}}_{t-2} \\ &= \dots \\ &= \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{\epsilon} \end{aligned}\\

上述推到过程利用了两个方差不同的高斯分布$\mathcal{N}(\mathbf{0},{\sigma }_1^2\mathbf{I})$\mathcal{N}(\mathbf{0}, \sigma_1^2\mathbf{I})和$\mathcal{N}(\mathbf{0},{\sigma }_2^2\mathbf{I})$\mathcal{N}(\mathbf{0}, \sigma_2^2\mathbf{I})相加等于一个新的高斯分布$\mathcal{N}(\mathbf{0},({\sigma }_1^2+{\sigma }_2^2)\mathbf{I})$\mathcal{N}(\mathbf{0}, (\sigma_1^2 + \sigma_2^2)\mathbf{I})。反重参数化后，我们得到：

$$q({\mathbf{x}}_t|{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_t;\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0,(1-{\overline{\alpha }}_t)\mathbf{I})$$q(\mathbf{x}_t \vert \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t)\mathbf{I}) \\

扩散过程的这个特性很重要。首先，我们可以看到${\mathbf{x}}_t$\mathbf{x}_{t}其实可以看成是原始数据${\mathbf{x}}_0$\mathbf{x}_{0}和随机噪音$ϵ$\mathbf{\epsilon}的线性组合，其中$\sqrt{{\overline{\alpha }}_t}$\sqrt{\bar \alpha_t}和$\sqrt{1-{\overline{\alpha }}_t}$\sqrt{1 - \bar{\alpha}_t}为组合系数，它们的平方和等于1，我们也可以称两者分别为signal_rate和noise_rate（见https://keras.io/examples/generative/ddim/#diffusion-schedule和Variational Diffusion Models）。更近一步地，我们可以基于${\overline{\alpha }}_t$\bar \alpha_t而不是${\beta }_t$\beta_t来定义noise schedule（见Improved Denoising Diffusion Probabilistic Models所设计的cosine schedule），因为这样处理更直接，比如我们直接将${\overline{\alpha }}_T$\bar \alpha_T设定为一个接近0的值，那么就可以保证最终得到的${\mathbf{x}}_T$\mathbf{x}_{T}近似为一个随机噪音。其次，后面的建模和分析过程将使用这个特性。

反向过程

扩散过程是将数据噪音化，那么反向过程就是一个去噪的过程，如果我们知道反向过程的每一步的真实分布$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$q(\mathbf{x}_{t-1} \vert \mathbf{x}_t)，那么从一个随机噪音${\mathbf{x}}_T\sim \mathcal{N}(\mathbf{0},\mathbf{I})$\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})开始，逐渐去噪就能生成一个真实的样本，所以反向过程也就是生成数据的过程。

估计分布$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$q(\mathbf{x}_{t-1} \vert \mathbf{x}_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};{\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))$$p_\theta(\mathbf{x}_{0:T}) = p(\mathbf{x}_T) \prod^T_{t=1} p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) \quad p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t))\\

这里$p({\mathbf{x}}_T)=\mathcal{N}({\mathbf{x}}_T;\mathbf{0},\mathbf{I})$p(\mathbf{x}_T)= \mathcal{N}(\mathbf{x}_T;\mathbf{0}, \mathbf{I})，而$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t)为参数化的高斯分布，它们的均值和方差由训练的网络${\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t)$\boldsymbol{\mu}_\theta(\mathbf{x}_t, t)和${\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t)$\boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t)给出。实际上，扩散模型就是要得到这些训练好的网络，因为它们构成了最终的生成模型。

虽然分布$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$q(\mathbf{x}_{t-1} \vert \mathbf{x}_t)是不可直接处理的，但是加上条件${\mathbf{x}}_0$\mathbf{x}_0的后验分布$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)却是可处理的，这里有：

$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_{t-1};\tilde{\mathit{\mu }}({\mathbf{x}}_t,{\mathbf{x}}_0),{\tilde{\beta }}_t\mathbf{I})$$q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \color{blue}{\tilde{\boldsymbol{\mu}}}(\mathbf{x}_t, \mathbf{x}_0), \color{red}{\tilde{\beta}_t} \mathbf{I}) \\

下面我们来具体推导这个分布，首先根据贝叶斯公式，我们有：

$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=q({\mathbf{x}}_t|{\mathbf{x}}_{t-1},{\mathbf{x}}_0)\frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}{q({\mathbf{x}}_t|{\mathbf{x}}_0)}$$q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) \frac{ q(\mathbf{x}_{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}_t \vert \mathbf{x}_0) }\\

由于扩散过程的马尔卡夫链特性，我们知道分布$q({\mathbf{x}}_t|{\mathbf{x}}_{t-1},{\mathbf{x}}_0)=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})$q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0)=q(\mathbf{x}_t \vert \mathbf{x}_{t-1})=\mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t} \mathbf{x}_{t-1}, \beta_t\mathbf{I})（这里条件${\mathbf{x}}_0$\mathbf{x}_0是多余的），而由前面得到的扩散过程特性可知：

$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_{t-1};\sqrt{{\overline{\alpha }}_{t-1}}{\mathbf{x}}_0,(1-{\overline{\alpha }}_{t-1})\mathbf{I})，q({\mathbf{x}}_t|{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_t;\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0,(1-{\overline{\alpha }}_t)\mathbf{I})$$q(\mathbf{x}_{t-1} \vert \mathbf{x}_0)=\mathcal{N}(\mathbf{x}_{t-1}; \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0, (1 - \bar{\alpha}_{t-1})\mathbf{I})，q(\mathbf{x}_t \vert \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t)\mathbf{I}) \\

所以，我们有：

$$\begin{array}{rl}q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)& =q({\mathbf{x}}_t|{\mathbf{x}}_{t-1},{\mathbf{x}}_0)\frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}{q({\mathbf{x}}_t|{\mathbf{x}}_0)}\\ & \propto \exp (-\frac{1}{2}(\frac{({\mathbf{x}}_t-\sqrt{{\alpha }_t}{\mathbf{x}}_{t-1}{)}^2}{{\beta }_t}+\frac{({\mathbf{x}}_{t-1}-\sqrt{{\overline{\alpha }}_{t-1}}{\mathbf{x}}_0{)}^2}{1-{\overline{\alpha }}_{t-1}}-\frac{({\mathbf{x}}_t-\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0{)}^2}{1-{\overline{\alpha }}_t}))\\ & =\exp (-\frac{1}{2}(\frac{{\mathbf{x}}_t^2-2\sqrt{{\alpha }_t}{\mathbf{x}}_t{\mathbf{x}}_{t-1}+{\alpha }_t{\mathbf{x}}_{t-1}^2}{{\beta }_t}+\frac{{\mathbf{x}}_{t-1}^2-2\sqrt{{\overline{\alpha }}_{t-1}}{\mathbf{x}}_0{\mathbf{x}}_{t-1}+{\overline{\alpha }}_{t-1}{\mathbf{x}}_0^2}{1-{\overline{\alpha }}_{t-1}}-\frac{({\mathbf{x}}_t-\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0{)}^2}{1-{\overline{\alpha }}_t}))\\ & =\exp (-\frac{1}{2}((\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\overline{\alpha }}_{t-1}}){\mathbf{x}}_{t-1}^2-(\frac{2\sqrt{{\alpha }_t}}{{\beta }_t}{\mathbf{x}}_t+\frac{2\sqrt{{\overline{\alpha }}_{t-1}}}{1-{\overline{\alpha }}_{t-1}}{\mathbf{x}}_0){\mathbf{x}}_{t-1}+C({\mathbf{x}}_t,{\mathbf{x}}_0)))\end{array}$$\begin{aligned} q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) &= q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) \frac{ q(\mathbf{x}_{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}_t \vert \mathbf{x}_0) } \\ &\propto \exp \Big(-\frac{1}{2} \big(\frac{(\mathbf{x}_t - \sqrt{\alpha_t} \mathbf{x}_{t-1})^2}{\beta_t} + \frac{(\mathbf{x}_{t-1} - \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0)^2}{1-\bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}_0)^2}{1-\bar{\alpha}_t} \big) \Big) \\ &= \exp \Big(-\frac{1}{2} \big(\frac{\mathbf{x}_t^2 - 2\sqrt{\alpha_t} \mathbf{x}_t \color{blue}{\mathbf{x}_{t-1}} \color{black}{+ \alpha_t} \color{red}{\mathbf{x}_{t-1}^2} }{\beta_t} + \frac{ \color{red}{\mathbf{x}_{t-1}^2} \color{black}{- 2 \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0} \color{blue}{\mathbf{x}_{t-1}} \color{black}{+ \bar{\alpha}_{t-1} \mathbf{x}_0^2} }{1-\bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}_0)^2}{1-\bar{\alpha}_t} \big) \Big) \\ &= \exp\Big( -\frac{1}{2} \big( \color{red}{(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}})} \mathbf{x}_{t-1}^2 - \color{blue}{(\frac{2\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{2\sqrt{\bar{\alpha}_{t-1}}}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0)} \mathbf{x}_{t-1} \color{black}{ + C(\mathbf{x}_t, \mathbf{x}_0) \big) \Big)} \end{aligned}\\

这里的$C({\mathbf{x}}_t,{\mathbf{x}}_0)$C(\mathbf{x}_t, \mathbf{x}_0)是一个和${\mathbf{x}}_{t-1}$\mathbf{x}_{t-1}无关的部分，所以省略。根据高斯分布的概率密度函数定义和上述结果（配平方），我们可以得到后验分布$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$q(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}, \mathbf{x}_0)的均值和方差：

$$\begin{array}{rl}{\tilde{\beta }}_t& =1/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\overline{\alpha }}_{t-1}})=1/(\frac{{\alpha }_t-{\overline{\alpha }}_t+{\beta }_t}{{\beta }_t(1-{\overline{\alpha }}_{t-1})})=\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}\cdot {\beta }_t\\ {\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)& =(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{\mathbf{x}}_t+\frac{\sqrt{{\overline{\alpha }}_{t-1}}}{1-{\overline{\alpha }}_{t-1}}{\mathbf{x}}_0)/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\overline{\alpha }}_{t-1}})\\ & =(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{\mathbf{x}}_t+\frac{\sqrt{{\overline{\alpha }}_{t-1}}}{1-{\overline{\alpha }}_{t-1}}{\mathbf{x}}_0)\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}\cdot {\beta }_t\\ & =\frac{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{\mathbf{x}}_t+\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}{\mathbf{x}}_0\end{array}$$\begin{aligned} \tilde{\beta}_t &= 1/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}}) = 1/(\frac{\alpha_t - \bar{\alpha}_t + \beta_t}{\beta_t(1 - \bar{\alpha}_{t-1})}) = \color{green}{\frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t} \\ \tilde{\boldsymbol{\mu}}_t (\mathbf{x}_t, \mathbf{x}_0) &= (\frac{\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1} }}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0)/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}}) \\ &= (\frac{\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1} }}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0) \color{green}{\frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t} \\ &= \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t} \mathbf{x}_0\\ \end{aligned}\\

可以看到方差是一个定量（扩散过程参数固定），而均值是一个依赖${\mathbf{x}}_0$\mathbf{x}_0和${\mathbf{x}}_t$\mathbf{x}_t的函数。这个分布将会被用于推导扩散模型的优化目标。

优化目标

上面介绍了扩散模型的扩散过程和反向过程，现在我们来从另外一个角度来看扩散模型：如果我们把中间产生的变量看成隐变量的话，那么扩散模型其实是包含$T$T个隐变量的隐变量模型（latent variable model），它可以看成是一个特殊的Hierarchical VAEs（见Understanding Diffusion Models: A Unified Perspective）：

相比VAE来说，扩散模型的隐变量是和原始数据同维度的，而且encoder（即扩散过程）是固定的。既然扩散模型是隐变量模型，那么我们可以就可以基于变分推断来得到variational lower bound（VLB，又称ELBO）作为最大化优化目标，这里有：

$$\begin{array}{rl}\log p_{\theta }({\mathbf{x}}_0)& =\log \int p_{\theta }({\mathbf{x}}_{0:T})d{\mathbf{x}}_{1:T}\\ & =\log \int \frac{p_{\theta }({\mathbf{x}}_{0:T})q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}d{\mathbf{x}}_{1:T}\\ & \ge {\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[\log \frac{p_{\theta }({\mathbf{x}}_{0:T})}{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}]\end{array}$$\begin{aligned} \log p_\theta(\mathbf{x}_0) &=\log\int p_\theta(\mathbf{x}_{0:T}) d\mathbf{x}_{1:T}\\ &=\log\int \frac{p_\theta(\mathbf{x}_{0:T}) q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} d\mathbf{x}_{1:T}\\ &\geq \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}[\log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}]\\ \end{aligned}\\

这里最后一步是利用了Jensen's inequality（不采用这个不等式的推导见博客What are Diffusion Models?），对于网络训练来说，其训练目标为VLB取负：

$$L=-L_{\text{VLB}}={\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[-\log \frac{p_{\theta }({\mathbf{x}}_{0:T})}{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}]={\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[\log \frac{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})}]$$L=-L_{\text{VLB}}=\mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}[-\log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}]=\mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}[\log \frac{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}{p_\theta(\mathbf{x}_{0:T})}] \\

我们近一步对训练目标进行分解可得：

$\begin{array}{rl}L& ={\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[\log \frac{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})}]\\ & ={\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[\log \frac{\prod _{t=1}^{T}q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})}{p_{\theta }({\mathbf{x}}_T)\prod _{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}]\\ & ={\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=1}^T\log \frac{q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}]\\ & ={\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}+\log \frac{q({\mathbf{x}}_1|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}]\\ & ={\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_t|{\mathbf{x}}_{t-1},{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}+\log \frac{q({\mathbf{x}}_1|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}]& \text{ ;use }q({\mathbf{x}}_t|{\mathbf{x}}_{t-1},{\mathbf{x}}_0)=q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})\\ & ={\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=2}^T\log (\frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}\cdot \frac{q({\mathbf{x}}_t|{\mathbf{x}}_0)}{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)})+\log \frac{q({\mathbf{x}}_1|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}]& \text{ ;use Bayes' Rule }\\ & ={\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_t|{\mathbf{x}}_0)}{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}+\log \frac{q({\mathbf{x}}_1|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}]\\ & ={\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}+\log \frac{q({\mathbf{x}}_T|{\mathbf{x}}_0)}{q({\mathbf{x}}_1|{\mathbf{x}}_0)}+\log \frac{q({\mathbf{x}}_1|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}]\\ & ={\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[\log \frac{q({\mathbf{x}}_T|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_T)}+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}-\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)]\\ & ={\mathbb{E}}_{q({\mathbf{x}}_T|{\mathbf{x}}_0)}[\log \frac{q({\mathbf{x}}_T|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_T)}]+\sum _{t=2}^T{\mathbb{E}}_{q({\mathbf{x}}_t,{\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}[\log \frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}]-{\mathbb{E}}_{q({\mathbf{x}}_1|{\mathbf{x}}_0)}[\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)]\\ & ={\mathbb{E}}_{q({\mathbf{x}}_T|{\mathbf{x}}_0)}[\log \frac{q({\mathbf{x}}_T|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_T)}]+\sum _{t=2}^T{\mathbb{E}}_{q({\mathbf{x}}_t|{\mathbf{x}}_0)}[q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)\log \frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}]-{\mathbb{E}}_{q({\mathbf{x}}_1|{\mathbf{x}}_0)}[\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)]\\ & =\underset{L_T}{\underbrace{D_{\text{KL}}(q({\mathbf{x}}_T|{\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_T))}}+\sum _{t=2}^T\underset{L_{t-1}}{\underbrace{{\mathbb{E}}_{q({\mathbf{x}}_t|{\mathbf{x}}_0)}[D_{\text{KL}}(q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))]}}-\underset{L_0}{\underbrace{{\mathbb{E}}_{q({\mathbf{x}}_1|{\mathbf{x}}_0)}\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}}\end{array}$\begin{aligned} L &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ \log\frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ \log\frac{\prod_{t=1}^T q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{ p_\theta(\mathbf{x}_T) \prod_{t=1}^T p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t) } \Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=1}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} \Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1}, \mathbf{x}_{0})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] & \text{ ;use } q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0)=q(\mathbf{x}_t \vert \mathbf{x}_{t-1})\\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \Big( \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)}\cdot \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1}\vert\mathbf{x}_0)} \Big) + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] & \text{ ;use Bayes' Rule }\\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1} \vert \mathbf{x}_0)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{q(\mathbf{x}_1 \vert \mathbf{x}_0)} + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big]\\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ \log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_T)} + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} - \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1) \Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{T}\vert \mathbf{x}_{0})}\Big[\log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_T)}\Big]+\sum_{t=2}^T \mathbb{E}_{q(\mathbf{x}_{t}, \mathbf{x}_{t-1}\vert \mathbf{x}_{0})}\Big[\log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)}\Big] - \mathbb{E}_{q(\mathbf{x}_{1}\vert \mathbf{x}_{0})}\Big[\log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)\Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{T}\vert \mathbf{x}_{0})}\Big[\log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_T)}\Big]+\sum_{t=2}^T \mathbb{E}_{q(\mathbf{x}_{t}\vert \mathbf{x}_{0})}\Big[q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)\log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)}\Big] - \mathbb{E}_{q(\mathbf{x}_{1}\vert \mathbf{x}_{0})}\Big[\log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)\Big] \\ &= \underbrace{D_\text{KL}(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T))}_{L_T} + \sum_{t=2}^T \underbrace{\mathbb{E}_{q(\mathbf{x}_{t}\vert \mathbf{x}_{0})}\Big[D_\text{KL}(q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t))\Big]}_{L_{t-1}} -\underbrace{\mathbb{E}_{q(\mathbf{x}_{1}\vert \mathbf{x}_{0})}\log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)}_{L_0} \end{aligned}