【理论推导】扩散模型 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相关的技巧来看待扩散模型

扩散模型

扩散模型是通过向图像多次施加噪声来将图像转化为噪声，该过程称为前向扩散过程 (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,{\overline{z}}_i\sim \mathcal{N}(0,I){\}}_{i=0}^T$为若干独立同分布的随机变量，根据递推公式，有
$$\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\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$

为了网络的训练需要，我们引入另一个性质，即 $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}$$
将 (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}$$

最后，考虑损失函数的设计，假定我们使用含参$\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 }_{\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}}_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}$$
我们希望建模使用的概率模型 ${\mu }_{\theta }$ 具有与 ${\tilde{\mu }}_t$ 相同的形式，使得网络能够更好的拟合，即${\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}$$

算法流程

参考资料

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