【理论推导】扩散模型 Diffusion Model

VAE 与 多层 VAE

回顾之前的文章 【理论推导】变分自动编码器 Variational AutoEncoder(VAE)，有结论
$$\log p(x)={\mathbb{E}}_{z\sim q(z|x)}[\log \frac{p(x,z)}{q(z|x)}]+\text{KL}(q||p)\ge {\mathbb{E}}_{z\sim q(z|x)}[\log \frac{p(x,z)}{q(z|x)}]$$
该不等式的另一种推导方式如下所示
$$\log p(x)=\log {\mathbb{E}}_{z\sim q(z|x)}[\frac{p(x,z)}{q(z|x)}]\ge {\mathbb{E}}_{z\sim q(z|x)}[\log \frac{p(x,z)}{q(z|x)}]$$
其中不等号由 Jensen 不等式给出

将单层 VAE 扩展到多层 VAE，如下所示

$$\begin{array}{rl}\log p(x)& =\log {\int }_{z_1}{\int }_{z_2}p(x,z_1,z_2)d{z}_{1}d{z}_2\\ & =\log {\int }_{z_1}{\int }_{z_2}q(z_1,z_2|x)\frac{p(x,z_1,z_2)}{q(z_1,z_2|x)}d{z}_{1}d{z}_2\\ & =\log {\mathbb{E}}_{z_1,z_2\sim q(z_1,z_2|x)}[\frac{p(x,z_1,z_2)}{q(z_1,z_2|x)}]\\ & \ge {\mathbb{E}}_{z_1,z_2\sim q(z_1,z_2|x)}[\log \frac{p(x,z_1,z_2)}{q(z_1,z_2|x)}]\\ & \stackrel{(i)}{=}{\mathbb{E}}_{z_1,z_2\sim q(z_1,z_2|x)}[\log \frac{p(x|z_1)p(z_1|z_2)p(z_2)}{q(z_1|x)q(z_2|z_1)}]\end{array}$$
其中 (i) 处要求变量之间满足Markov假设，如果我们将多层 VAE 扩展到更多层，可以得到与扩散模型相近的图示形式，因此我们可以借助VAE相关的技巧来看待扩散模型

DDPM

扩散模型是通过向图像多次施加噪声来将图像转化为噪声，该过程称为前向扩散过程 (forward diffusion process)，而从某个先验噪声分布中采样一个噪声图作为初值，通过不断去噪来生成图像的过程称为是扩散的逆过程，可以类比于使用 Langevin Dynamics 进行图像生成的思路。

扩散过程

假定 $x_0\sim q(x)$ 是采样自真实数据分布 $q$ 的样本，我们向其添加 $T$ 步的高斯噪声，公式如下
$$q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)$$
其中 ${\beta }_t\in [0,1]$，整个过程服从Markov假设，因此有 $q(x_{1:T}|x_0)={\prod }_{t=1}^{T}q(x_t|x_{t-1})$，当 $T\to \infty$，$x_T$ 服从高斯分布

如果我们希望快速得到 $x_t$，可以不通过递推式而是求一个通项的表达形式。假定 ${\alpha }_t=1-{\beta }_t$，$\overline{{\alpha }_t}={\prod }_{i=1}^t{\alpha }_i$， { z i , z ‾ i ∼ N ( 0 , I ) } i = 0 T \{z_i, \overline{z}_i \sim \mathcal N(0,I)\}_{i=0}^T {zi​,zi​∼N(0,I)}i=0T​为若干独立同分布的随机变量，根据递推公式，有
$$\begin{array}{rl}{x}_t& =\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{z}_{t-1}\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-1}+\sqrt{{\alpha }_t}\sqrt{1-{\alpha }_{t-1}}{z}_{t-2}+\sqrt{1-{\alpha }_t}{z}_{t-1}\\ & \stackrel{(i)}{=}\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-1}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{z}}_{t-2}\\ & =...\\ & =\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}{\overline{z}}_0\end{array}$$
其中，等式 (i) 为两个高斯分布的线性叠加仍为一个高斯分布，即对于 $A\sim \mathcal{N}({\mu }_a,{\sigma }_a^2)$，$B\sim \mathcal{N}({\mu }_b,{\sigma }_b^2)$，线性叠加 $mA+nB\sim \mathcal{N}(m{\mu }_a+n{\mu }_b,m^2{\sigma }_a^2+n^2{\sigma }_b^2)$。因此，有
$$\begin{array}{r}{x}_t|x_0\sim \mathcal{N}(\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)I)\end{array}$$
对于扩散过程，我们希望加噪的强度从小到大，即 ${\beta }_1<{\beta }_2<...<{\beta }_{T-1}<{\beta }_T$，有 $1>{\overline{\alpha }}_1>...>{\overline{\alpha }}_T>0$

逆扩散过程/采样过程

我们希望从 $x_T$ 中恢复出 $x_0$，为此需要建模条件概率 $q(x_{t-1}|x_t)$，注意到，根据Bayes公式，有
$$q(x_{t-1}|x_t)=q(x_t|x_{t-1})\frac{q(x_{t-1})}{q(x_t)}$$
我们无法得到真实的 $q(x_{t-1})$，因此采用条件概率分布来逼近 $q(x_{t-1}|x_t)$，有
$$q(x_{t-1}|x_t)\approx q(x_{t-1}|x_t,x_0)=q(x_t|x_{t-1},x_0)\frac{q(x_{t-1}|x_0)}{q(x_t|x_0)}$$
其中 $q(x_{t-1}|x_t,x_0)$是可以计算出来的。使用贝叶斯公式，有
$$\begin{array}{rl}q(x_{t-1}|x_t,x_0)& =q(x_t|x_{t-1},x_0)\frac{q(x_{t-1}|x_0)}{q(x_t|x_0)}\\ & \propto \exp \left(-\frac{1}{2}\left(\frac{(x_t-\sqrt{{\alpha }_t}{x}_{t-1}{)}^2}{{\beta }_t}+\frac{(x_{t-1}-\sqrt{{\overline{\alpha }}_{t-1}}{x}_0{)}^2}{1-{\overline{\alpha }}_{t-1}}-\frac{(x_t-\sqrt{{\overline{\alpha }}_t}{x}_0{)}^2}{1-{\overline{\alpha }}_t}\right)\right)\\ & =\exp \left(-\frac{1}{2}\left((\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\overline{\alpha }}_t})x_{t-1}^2-(\frac{2\sqrt{{\alpha }_t}}{{\beta }_t}{x}_t+\frac{2\sqrt{{\overline{\alpha }}_t}}{1-{\overline{\alpha }}_t}{x}_0)x_{t-1}+C(x_0,x_t)\right)\right)\end{array}$$
对比高斯分布的形式，可以得到条件概率分布 $x_{t-1}|x_t,x_0$ 服从均值，方差为如下形式的高斯分布
$$\begin{array}{rl}\mu & =(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{x}_t+\frac{\sqrt{{\overline{\alpha }}_t}}{1-{\overline{\alpha }}_t}{x}_0)/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\overline{\alpha }}_{t-1}})=\frac{{\sqrt{\alpha }}_t(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{x}_t+\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}{x}_0\\ {\sigma }^2& ={\tilde{\beta }}_t=\frac{1}{\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\overline{\alpha }}_{t-1}}}=\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}{\beta }_t\end{array}$$
我们使用神经网络来拟合${\overline{z}}_0$，即 ${ϵ}_{\theta }(x_t,t)\approx {\overline{z}}_0$，注意到，我们通过 ${\overline{z}}_0$的加噪方式得到的 $x_t$，因此，神经网络本质上是拟合的添加的噪声。将 (1) 式代入到其中，消掉 $x_0$，得到
$$\begin{array}{rl}\mu & ={\tilde{\mu }}_t=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{\overline{z}}_0\right)\end{array}$$
因此，有
$$\begin{array}{r}{x}_{t-1}|x_t\sim \mathcal{N}\left(\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{\overline{z}}_0\right),\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}{\beta }_t\right)\end{array}$$

损失函数

考虑损失函数的设计，假定我们使用含参$\theta$ 的概率模型 $p_{\theta }$ 去拟合真实数据分布 $q$，根据 KL 散度的性质，有
$$\begin{array}{rl}-\log p_{\theta }(x_0)& \le -\log p_{\theta }(x_0)+\text{KL}(q(x_{1:T}|x_0)||p_{\theta }(x_{1:T}|x_0))\\ & =-\log p_{\theta }(x_0)+{\mathbb{E}}_{q(x_{1:T}|x_0)}[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})/p_{\theta }(x_0)}]\\ & ={\mathbb{E}}_{q(x_{1:T}|x_0)}[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}]\end{array}$$
对左右两边求期望，有
$${\mathbb{E}}_{q(x_0)}[-\log p_{\theta }(x_0)]\le {\mathbb{E}}_{q(x_{0:T}|x_0)}[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}]\stackrel{△}{=}{L}_{\text{VLB}}$$
对 $L_{\text{VLB}}$ 进行化简，有
$$\begin{array}{rl}{L}_{\text{VLB}}& ={\mathbb{E}}_{q(x_{0:T}|x_0)}[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}]\\ & ={\mathbb{E}}_{q(x_{0:T}|x_0)}[-\log p(x_T)+\sum _{i=1}^T\log \frac{q(x_t|x_{t-1})}{p_{\theta }(x_{t-1}|x_t)}]\\ & \stackrel{(i)}{=}{\mathbb{E}}_{q(x_{0:T}|x_0)}[-\log p(x_T)+\sum _{i=2}^T\log (\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}\frac{q(x_t|x_0)}{q(x_{t-1}|x_0)})+\log \frac{q(x_1|x_0)}{p_{\theta }(x_0|x_1)}]\\ & ={\mathbb{E}}_{q(x_{0:T}|x_0)}[-\log p(x_T)+\sum _{i=2}^T\log \frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}+\log \frac{q(x_T|x_0)}{q(x_1|x_0)}+\log \frac{q(x_1|x_0)}{p_{\theta }(x_0|x_1)}]\\ & ={\mathbb{E}}_{q(x_{0:T}|x_0)}[\sum _{i=2}^T\log \frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}+\log \frac{q(x_T|x_0)}{p(x_T)}-\log p_{\theta }(x_0|x_1)]\\ & =\text{KL}(q(x_T|x_0)||p(x_T))+\sum _{t=2}^T\text{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))-\log p_{\theta }(x_0|x_1)\end{array}$$
其中，等号 (i) 处推导如下所示
$$q(x_t|x_{t-1})=q(x_t|x_{t-1},x_0)=\frac{q(x_{t-1}|x_t,x_0)q(x_t|x_0)}{q(x_{t-1}|x_0)}$$
我们固定方差${\beta }_t$为一超参数，因此对于公式(5)中的第一项是无参的常量，可以忽略；对于最后一项，作者提出简化掉来训练会更好。因为 $p_{\theta }$ 是我们拟合分布使用的模型，所以我们可以假定 $p_{\theta }(x_{t-1}|x_t)=\mathcal{N}({\mu }_{\theta }(x_t,t),{\sigma }_t^{2}I))$，因此该分布仅均值部分与输入有关，可以得到如下式子
$$\begin{array}{rl}\text{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))& ={\mathbb{E}}_q[\frac{1}{2{\sigma }_t^2}||{\tilde{\mu }}_t(x_t,x_0)-{\mu }_{\theta }(x_t,t)|{|}_2^2]+C\end{array}$$
设 $ϵ\sim \mathcal{N}(0,I)$，使用公式 (1) $x_0$ 与 $ϵ$ 替换掉里面的 $x_t$，同时使用 (4) 式替换掉其中的 ${\tilde{\mu }}_t$
$$\begin{array}{rl}\text{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))& ={\mathbb{E}}_q[\frac{1}{2{\sigma }_t^2}||\frac{1}{\sqrt{{\alpha }_t}}(x_t(x_0,ϵ)-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}ϵ)-{\mu }_{\theta }(x_t,t)|{|}_2^2]+C\end{array}$$
其中公式(4)给出了${\mu }_{\theta }$ 与 ${ϵ}_{\theta }$ 满足如下关系
$${\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}(x_t(x_0,ϵ)-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_{\theta }(x_t,t))$$
因此，有
$$\begin{array}{rl}\text{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))& ={\mathbb{E}}_{x_0,ϵ}[\frac{{\beta }_t^2}{2{\sigma }_t^2{\alpha }_t(1-{\overline{\alpha }}_t)}||ϵ-{ϵ}_{\theta }(\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ,t)|{|}_2^2]+C\end{array}$$
损失函数即为
$$\begin{array}{rl}\mathcal{L}(\theta )& ={\mathbb{E}}_{t,x_0,ϵ}[||ϵ-{ϵ}_{\theta }(\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ,t)|{|}_2^2]\end{array}$$

算法流程

DDIM

DDPM的核心问题在于采样需要迭代足够多的次数，而且理论推导中的概率分布是 $q(x_{t-1}|x_t)$，因此每次迭代的下标变化为1，如果我们希望下标变化可以不局限为1，例如如果我们支持计算 $q(x_s|x_t)(s<t)$ 那么我们可以任意设置从 $t:T\to 0$ 迭代的次数，这需要我们在采样时突破下式中的 $q(x_t|x_{t-1})$
$$q(x_{t-1}|x_t)\approx q(x_{t-1}|x_t,x_0)=q(x_t|x_{t-1},x_0)\frac{q(x_{t-1}|x_0)}{q(x_t|x_0)}=q(x_t|x_{t-1})\frac{q(x_{t-1}|x_0)}{q(x_t|x_0)}$$
如果没有 $q(x_t|x_{t-1},x_0)$，但仍可以通过下式来求解
$$\int p(x_{t-1}|x_t,x_0)p(x_t|x_0)d{x}_t=p(x_{t-1}|x_0)$$
其中 $p(x_t|x_0),p(x_{t-1}|x_0)$ 都为高斯分布，可以假定 $p(x_{t-1}|x_t,x_0)$ 也为高斯分布，其均值为 $x_t$ 与 $x_0$ 的线性组合

更一般地，我们考虑任意两个下标 $x_s,x_t(s<t)$，假定 $x_s=m_{s|t}{x}_t+n_{s|t}{x}_0+{\sigma }_{s|t}{\epsilon }_1$，我们已知 $q(x_s|x_0)$，$q(x_t|x_0)$，联立可得
$$\{\begin{array}{c}{x}_s=m_{s|t}{x}_t+n_{s|t}{x}_0+{\sigma }_{s|t}{\epsilon }_1\\ x_t=\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}{\epsilon }_2\\ x_s=\sqrt{{\overline{\alpha }}_s}{x}_0+\sqrt{1-{\overline{\alpha }}_s}{\epsilon }_3\end{array}$$
可得关于 $m_{s|t}$ 与 $n_{s|t}$ 的联立表达式
$$\{\begin{array}{c}{m}_{s|t}\sqrt{{\overline{\alpha }}_t}+n_{s|t}=\sqrt{{\overline{\alpha }}_s}\\ m_{s|t}^2(1-{\overline{\alpha }}_t)+{\sigma }_{s|t}^2=1-{\overline{\alpha }}_s\end{array}$$
解得
$$\{\begin{array}{c}{m}_{s|t}=\sqrt{\frac{1-{\overline{\alpha }}_s-{\sigma }_{s|t}^2}{1-{\overline{\alpha }}_t}}\\ n_{s|t}=\sqrt{{\overline{\alpha }}_s}-\sqrt{\frac{{\overline{\alpha }}_t}{1-{\overline{\alpha }}_t}(1-{\overline{\alpha }}_s-{\sigma }_{s|t}^2)}\end{array}$$
带回到原式，可得任意两下标 $s<t$ 的采样公式
$$x_s=\sqrt{{\overline{\alpha }}_s}{x}_0+\sqrt{1-{\overline{\alpha }}_s-{\sigma }_{s|t}^2}\frac{x_t-\sqrt{{\overline{\alpha }}_t}{x}_0}{\sqrt{1-{\overline{\alpha }}_t}}+{\sigma }_{s|t}\epsilon$$
注意到，DDIM并没有使用 $q(x_t|x_s)$，因此 DDIM 相比于 DDPM 具有更加泛化的形式，这里的 $x_0$ 是使用 $x_t$、${ϵ}_{\theta }(x_t,t)$ 和公式 (1) 给出的估计值，而 $\frac{x_t-\sqrt{{\overline{\alpha }}_t}{x}_0}{\sqrt{1-{\overline{\alpha }}_t}}$ 对应 ${ϵ}_{\theta }(x_t,t)$，即
$$x_s=\sqrt{\frac{{\overline{\alpha }}_s}{{\overline{\alpha }}_t}}(x_t-\sqrt{1-{\overline{\alpha }}_t}{ϵ}_{\theta }(x_t,t))+\sqrt{1-{\overline{\alpha }}_s-{\sigma }_{s|t}^2}{ϵ}_{\theta }(x_t,t)+{\sigma }_{s|t}\epsilon$$

参考资料

Denoising Diffusion Probabilistic Models
Denoising Diffusion Implicit Models
Probabilistic Diffusion Model概率扩散模型理论
扩散模型 Diffusion Model