diffusion model原理和算法伪代码

Diffusion model

奠基性的工作：

1. Ho,(2020),Denoising diffusion peobabilistic models
2. Sohi,(2015), Deep unsupervised learning using nonequilibruim thermodynamics

前置数学知识

1. 条件概率的一般形式
   $$P(B,C|A)=P(B|A)P(C|A,B)$$

2. 基于马尔可夫假设的条件概率

   假设马尔可夫链关系$A\to B\to C$，有
   $$P(A,B,C)=P(C|B)P(B|A)P(A)$$

3. 高斯分布的KL散度

   对于两个单一变量的高斯分布p和q而言，他们的KL散度满足
   $$KL(\mathcal{N}({\mu }_1,{\sigma }_1^2),\mathcal{N}({\mu }_2,{\sigma }_2^2))=\log \frac{{\sigma }_2}{{\sigma }_1}-\frac{1}{2}+\frac{{\sigma }_1^2+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }_2^2}$$
   推导详见CSDN博客

4. 参数重整化

   若希望从高斯分布$\mathcal{N}(\mu ,{\sigma }^2)$中采样，可以先从标准分布$\mathcal{N}(0,1)$得到$z$，得到$\sigma \cdot z+\mu$。

   这样就可以将$\sigma$和$\mu$也作为仿射网络的一部分，而不是不可导的环境参数。

   这个技巧在VAE和Diffusion model中大量被使用。

VAE和多层VAE回顾

1. 单层VAE的原理

$$\begin{gathered} x\to z,q_{\varphi }(z|x) \\ z\to x,p_{\theta }(x|z) \end{gathered}$$

此时$x$的边缘概率分布可以改写为关于z的期望式
$$\begin{array}{rl}p(x)& ={\int }_z{p}_{\theta }(x|z)p(z)\text{d}z\\ & ={\int }_z{q}_{\varphi }(z|x)\frac{p_{\theta }(x|z)p(z)}{q_{\varphi }(z|x)}\text{d}z\\ & ={\mathbb{E}}_{z\sim q_{\varphi }(z|x)}\frac{p_{\theta }(x|z)p(z)}{q_{\varphi }(z|x)}\end{array}$$
此时的Evidence存在一个lower bound（ELBO）
$$\log p(x)=\log {\mathbb{E}}_{z\sim q_{\varphi }(z|x)}\frac{p_{\theta }(x|z)p(z)}{q_{\varphi }(z|x)}\ge {\mathbb{E}}_{z\sim q_{\varphi }(z|x)}\log \left[\frac{p_{\theta }(x|z)p(z)}{q_{\varphi }(z|x)}\right]$$
在训练中,我们需要最大化对数似然，即Evidence，可以通过最小化lower bound实现，而这个lower bound可以分为两部分：

1. ${\mathbb{E}}_{z\sim q_{\varphi }(z|x)}{p}_{\theta }(x|z)$，可以通过神经网络实现预测
2. $-{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}\log \frac{q_{\varphi }(x|z)}{p(z)}$，即两个分布的KL的散度，一般可假设$z$服从高斯分布，而$q_{\varphi }(x|z)$也逼近高斯分布，而两个高斯分布的KL散度存在公式

所以，单层VAE的损失函数是可优化的。

2. 多层VAE的原理

基于同样的原理，
$$\begin{array}{rl}p(x)& ={\int }_{z_1}{\int }_{z_2}{p}_{\theta }(x,z_1,z_2)\text{d}{z}_1\text{d}{z}_2\\ & ={\int }_{z_1}{\int }_{z_2}{q}_{\varphi }(z_1,z_2|x)\frac{p_{\theta }(x,z_1,z_2)}{q_{\varphi }(z_1,z_2|x)}\text{d}{z}_1\text{d}{z}_2\\ & ={\mathbb{E}}_{z1,z_2\sim q_{\varphi }(z_1,z_2|x)}\frac{p_{\theta }(x,z_1,z_2)}{q_{\varphi }(z_1,z_2|x)}\end{array}$$
得到
$$\log p(x)\ge {\mathbb{E}}_{z1,z_2\sim q_{\varphi }(z_1,z_2|x)}\log \frac{p_{\theta }(x,z_1,z_2)}{q_{\varphi }(z_1,z_2|x)}$$
如果上述过程满足马尔科夫假设，即
$$\begin{gathered} p_{\theta }(x,z_1,z_2)=p(x|z_1)p(z_1|z_2)p(z_2) \\ q(z_1,z_2|x)=q(z_1|x)q(z_2|z_1) \end{gathered}$$
(6)式能够被进一步简化为
$$\mathcal{L}(\theta ,\varphi )={\mathbb{E}}_{q(z_1,z_2|x)}\left[\log p(x|z_1)-\log q(z_1|x)+\log p(z_1|z_2)-\log q(z_2|z_1)+\log p(z_2)\right]$$

Diffusion model

从右往左，从目标分布到噪声分布称为扩散过程，而我们希望学习到从左往右的逆扩散过程。上图中的第一行从左往右是扩散过程，第二行从右往左是逆扩散过程，而第三行是前两者的差值，称为偏移量。

扩散过程（Diffusion Process）

1. 给定初始数据分布${\mathbf{x}}_0\sim q(\mathbf{x})$，不断向分布中添加高斯噪声，噪声的标准差是以${\beta }_t$确定的，均值是以固定值${\beta }_t$和当前时刻的数据${\mathbf{x}}_{\mathbf{t}}$决定的，所以该过程并没有需要学习的参数，而且是一个马尔科夫链过程。

2. 随着$t$的不断增大，最终数据分布$x_T$变成了一个各项独立的高斯分布
   $$q({\mathbf{x}}_{\mathbf{t}}|{\mathbf{x}}_{\mathbf{t}-1})=\mathcal{N}({\mathbf{x}}_{\mathbf{t}};\sqrt{1-{\beta }_t}{\mathbf{x}}_{\mathbf{t}-1},{\beta }_{\mathbf{t}}\mathbf{I})$$

   $$q({\mathbf{x}}_{1:\mathbf{T}}|{\mathbf{x}}_{\mathbf{o}})=\prod _{t=1}^{T}q({\mathbf{x}}_{\mathbf{t}}|{\mathbf{x}}_{\mathbf{t}-1})$$

   这充分体现了参数重整化的技巧。

3. 任意时刻的$q({\mathbf{x}}_{\mathbf{t}})$推导也可以完全基于${\mathbf{x}}_0$和${\beta }_t$计算得到闭式解，而不需要做迭代。（令${\alpha }_t=1-{\beta }_t$）

   两个正态分布$X\sim \mathcal{N}({\mu }_1,{\sigma }_1)$和$Y\sim \mathcal{N}({\mu }_2,{\sigma }_2)$叠加后的分布$aX+bY$服从分布$\mathcal{N}(a{\mu }_1+b{\mu }_2,a^2{\sigma }_1^2+b^2{\sigma }_2^2)$。

   对于第$t$步的分布$x_t$等于上一步的分布$x_{t-1}$加上高斯噪声$z_{t-1}$，即
   $$\begin{array}{rl}{\mathbf{x}}_t& =\sqrt{{\alpha }_t}{\mathbf{x}}_{t-1}+\sqrt{1-{\alpha }_t}{\mathbf{z}}_{t-1};\text{where}{\mathbf{z}}_{\mathbf{t}-1},{\mathbf{z}}_{\mathbf{t}-2},...\sim \mathcal{N}(0,\mathbf{I})\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{\mathbf{x}}_{t-2}+\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}{\mathbf{z}}_{t-2}+\sqrt{1-{\alpha }_t}{\mathbf{z}}_{\mathbf{t}-1}\end{array}$$
   这里借助参数重整化的技巧，将红色部分的两个高斯分布合并为新的高斯分布，整理如下所示
   $$\begin{array}{rl}{\mathbf{x}}_t& =\sqrt{{\alpha }_t{\alpha }_{t-1}}{\mathbf{x}}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\bar{\mathbf{z}}}_{t-2}\end{array}$$
   其中，${\bar{\mathbf{z}}}_{t-2}\sim \mathcal{N}(0,\mathbf{I})$

   重复上面的步骤，最终可以得到${\mathbf{z}}_t$的闭式解
   $${\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{z};\text{where}{\bar{\alpha }}_t=\prod _{i=1}^T{\alpha }_i$$
   此时，作者认为${\mathbf{x}}_t\sim \mathcal{N}({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,\sqrt{1-{\bar{\alpha }}_t}\mathbf{I})$，（此处应该是认为${\mathbf{x}}_0$是完全已知的，方差为零），最终当上述分布趋近于$\mathcal{N}(0,\mathbf{I})$的时候，可认为模型已经基本完成扩散过程。因此，作者给出了一种${\beta }_t$的设置经验，${\beta }_1<{\beta }_2<\cdot \cdot \cdot <{\beta }_T$，即随着扩散深度的加深，逐步扩大$\beta$。

   逆扩散过程（Reverse Process）

   逆过程是从高斯分布中恢复原始数据，当${\beta }_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};{\mu }_{\theta }({\mathbf{x}}_t,t),\sum _{\theta }({\mathbf{x}}_t,t))$$
   逆扩散过程可以被总结为如下形式
   $$p_{\theta }({\mathbf{x}}_{0:T})=p({\mathbf{x}}_T)\prod _{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$$
   此处通过使用网络估计参数$\theta$以实现逆扩散过程。

   后验的扩散条件概率$q(x_{t-1}|x_t,x_0)$存在闭式解

   根据条件概率的贝叶斯公式
   $$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)q({\mathbf{x}}_t|{\mathbf{x}}_0)=q({\mathbf{x}}_t|{\mathbf{x}}_{t-1},{\mathbf{x}}_0)q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)$$
   得到
   $$\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 \left(-\frac{1}{2}\left(\frac{({\mathbf{x}}_t-\sqrt{{\alpha }_t}{\mathbf{x}}_{t-1}{)}^2}{1-{\alpha }_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}\right)\right)\\ & =\exp \left(-\frac{1}{2}\left(a{\mathbf{x}}_{t-1}^2+b{\mathbf{x}}_{t-1}+c({\mathbf{x}}_t,{\mathbf{x}}_0)\right)\right)\end{array}$$
   可见，上述分布的核心可以用一个二次函数来描述，那对应的中轴线应该是
   $$\begin{array}{rl}{\tilde{\mu }}_{\mathbf{t}}({\mathbf{x}}_t,{\mathbf{x}}_0)& =-\frac{b}{2a}\\ & =\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 }}{\mathbf{x}}_0\end{array}$$
   容易从扩散过程的表达式（式11）得到${\mathbf{x}}_0$的表达式
   $${\mathbf{x}}_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}\left({\mathbf{x}}_t-\sqrt{1-{\bar{\alpha }}_t}\mathbf{z}\right)$$
   带入得到
   $$\begin{array}{rl}{\tilde{\mu }}_{\mathbf{t}}({\mathbf{x}}_t,{\mathbf{x}}_0)& =\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 }}\frac{1}{\sqrt{{\bar{\alpha }}_t}}\left({\mathbf{x}}_t-\sqrt{1-{\bar{\alpha }}_t}\mathbf{z}\right)\\ & =\frac{1}{{\sqrt{\alpha }}_t}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{z}}_t\right)\end{array}$$
   这就是${\mathbf{x}}_{t-1}$分布的均值表达式，即给定${\mathbf{x}}_0$的条件下，后验条件高斯分布的均值计算只与${\mathbf{x}}_t$和${\mathbf{z}}_t$有关。

目标-数据分布的似然函数

我们在待优化的目标数据分布的似然函数（负）上加一个非负的KL散度，构成负对数似然的上界，通过最小化上界，负对数似然自然取得最小。
$$\begin{array}{rl}-\log p_{\theta }({\mathbf{x}}_0)& =-\log p_{\theta }({\mathbf{x}}_0)+D_{KL}(q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)||p_{\theta }({\mathbf{x}}_{1:T}|{\mathbf{x}}_0))\\ & =-\log p_{\theta }({\mathbf{x}}_0)+{\mathbb{E}}_{{\mathbf{x}}_{1:T}\sim q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}\left[\log \frac{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})/p_{\theta }({\mathbf{x}}_0)}\right]\\ & =-\log p_{\theta }({\mathbf{x}}_0)+{\mathbb{E}}_{{\mathbf{x}}_{1:T}\sim q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}\left[\log \frac{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})}+\log p_{\theta }({\mathbf{x}}_0)\right]\\ & ={\mathbb{E}}_{{\mathbf{x}}_{1:T}\sim q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}\left[\log \frac{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})}\right]=L_{VLB}\end{array}$$
我们也可以继续对$L_{VLB}$进行展开，过程比较繁琐，建议查看论文，最终的形式如下
$$L_{VLB}={\mathbb{E}}_q\left[D_{KL}\left(q({\mathbf{x}}_T|{\mathbf{x}}_0)||p_{\theta }({\mathbf{x}}_T)\right)+\sum _{t=1}^T{D}_{KL}(q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)||p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))\right]$$
其中第一项是不含待优化参数的，仅仅需要优化第二项即可。而且作者将$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$的方差设置为与$\beta$有关的常数，可训练参数仅存在其均值中。

我们已经知道$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$服从高斯分布，并给出了其均值的表达式，而且知道$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$也服从高斯分布，其方差设置为常数，仅需优化均值即可。使用文章开头给出的两个单一变量的高斯分布的KL散度表达式，两个分布的方差均为常数，最终的损失函数可以写作两个分布的均值的关系：
$$L_{t-1}={\mathbb{E}}_q\left[\frac{1}{2{\sigma }_t^2}||{\tilde{\mu }}_t({\mathbf{x}}_t,{\mathbf{x}}_0)-{\mu }_{\theta }({\mathbf{x}}_t,t)|{|}^2\right]+C$$
我们可以将已经得到的${\mu }_t$的表达式，进行简化得到最终的损失函数：
$$L_{\text{simple}}(\theta ):={\mathbb{E}}_{t,{\mathbf{x}}_0,ϵ}\left[||ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t)|{|}^2\right]$$
这里，${ϵ}_{\theta }$就是可学习的网络，输入${\mathbf{x}}_0$和高斯噪声$ϵ$以及时刻$t$。

Diffusion Probabilistic Model的算法代码

Training

1. repeat
2. $x_0\sim q(x_0)$
3. $t\sim \text{Uniform}(\{1,2,..,T\})$
4. $ϵ\sim \mathcal{N}(0,\mathbf{I})$
5. Take gradient descent step on
6. ${\nabla }_{\theta }||ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t)|{|}^2$
7. until converged

Sampling

优化好网络${ϵ}_{\theta }$之后，可以从$x_T$逐步获得$x_0$

1. $x_T\sim \mathcal{N}(0,\mathbf{I})$
2. for $t=T,T-1,...,1$ do
3. $\mathbf{z}\sim \mathcal{N}(0,\mathbf{I})$ if $t>1$ else $\mathbf{z}=0$
4. ${\mathbf{x}}_{t-1}={\mu }_{\theta }({\mathbf{x}}_t,t)+{\sigma }_t\mathbf{z}=\frac{1}{\sqrt{{\alpha }_t}}\left(\mathbf{x}-\frac{1-{\alpha }_t}{\sqrt{1-{\alpha }_t}}{ϵ}_{\theta }({\mathbf{x}}_t,t\right)+{\sigma }_t\mathbf{z}$
5. end for
6. return $x_0$