diffusion model原理和算法伪代码_diffusion loss

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的期望式
p ( x ) = ∫ z p θ ( x ∣ z ) p ( z ) d z = ∫ z q ϕ ( z ∣ x ) p θ ( x ∣ z ) p ( z ) q ϕ ( z ∣ x ) d z = E z ∼ q ϕ ( z ∣ x ) p θ ( x ∣ z ) p ( z ) q ϕ ( z ∣ x )

$$\begin{aligned} p(x)&=\int_zp_\theta(x|z)p(z)\text{d}z\\ &=\int_zq_\phi(z|x)\frac{p_\theta(x|z)p(z)}{q_\phi(z|x)}\text{d}z\\ &=\mathbb{E}_{z\sim q_\phi(z|x)}\frac{p_\theta(x|z)p(z)}{q_\phi(z|x)} \end{aligned}$$ p(x)​=∫z​pθ​(x∣z)p(z)dz=∫z​qϕ​(z∣x)qϕ​(z∣x)pθ​(x∣z)p(z)​dz=Ez∼qϕ​(z∣x)​qϕ​(z∣x)pθ​(x∣z)p(z)​​
此时的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的原理

基于同样的原理，
p ( x ) = ∫ z 1 ∫ z 2 p θ ( x , z 1 , z 2 ) d z 1 d z 2 = ∫ z 1 ∫ z 2 q ϕ ( z 1 , z 2 ∣ x ) p θ ( x , z 1 , z 2 ) q ϕ ( z 1 , z 2 ∣ x ) d z 1 d z 2 = E z 1 , z 2 ∼ q ϕ ( z 1 , z 2 ∣ x ) p θ ( x , z 1 , z 2 ) q ϕ ( z 1 , z 2 ∣ x )

$$\begin{aligned} 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_\phi(z_1,z_2|x)\frac{p_\theta(x,z_1,z_2)}{q_\phi(z_1,z_2|x)}\text{d}z_1\text{d}z_2\\ &=\mathbb{E}_{z1,z_2\sim q_\phi(z_1,z_2|x)}\frac{p_\theta(x,z_1,z_2)}{q_\phi(z_1,z_2|x)} \end{aligned}$$ p(x)​=∫z1​​∫z2​​pθ​(x,z1​,z2​)dz1​dz2​=∫z1​​∫z2​​qϕ​(z1​,z2​∣x)qϕ​(z1​,z2​∣x)pθ​(x,z1​,z2​)​dz1​dz2​=Ez1,z2​∼qϕ​(z1​,z2​∣x)​qϕ​(z1​,z2​∣x)pθ​(x,z1​,z2​)​​
得到
$$\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}$，即
   x t = α t x t − 1 + 1 − α t z t − 1 ; where  z t − 1 , z t − 2 , . . . ∼ N ( 0 , I ) = α t α t − 1 x t − 2 + α t − α t α t − 1 z t − 2 + 1 − α t z t − 1

   $$\begin{aligned} \bold{x}_t&=\sqrt{\alpha_t}\bold{x}_{t-1}+\sqrt{1-\alpha_t}\bold{z}_{t-1}\qquad ;\text{where} \ \bold{z_{t-1}},\bold{z_{t-2}},...\sim \mathcal{N}(\bold{0},\bold{I})\\ &=\sqrt{\alpha_t\alpha_{t-1}}\bold{x}_{t-2}+{\color{red} \sqrt{\alpha_t-\alpha_t\alpha_{t-1}}\bold{z}_{t-2}+\sqrt{1-\alpha_t}\bold{z_{t-1}}} \end{aligned}$$ xt​​=αt​ ​xt−1​+1−αt​ ​zt−1​;where zt−1​,zt−2​,...∼N(0,I)=αt​αt−1​ ​xt−2​+αt​−αt​αt−1​ ​zt−2​+1−αt​ ​zt−1​​
   这里借助参数重整化的技巧，将红色部分的两个高斯分布合并为新的高斯分布，整理如下所示
   $$\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)$$
   得到
   q ( x t − 1 ∣ x t , x 0 ) = q ( x t ∣ x t − 1 , x 0 ) q ( x t − 1 ∣ x 0 ) q ( x t ∣ x 0 ) ∝ exp ⁡ ( − 1 2 ( ( x t − α t x t − 1 ) 2 1 − α t + ( x t − 1 − α ˉ t − 1 x 0 ) 2 1 − α ˉ t − 1 − ( x t − α ˉ t x 0 ) 2 1 − α ˉ t ) ) = exp ⁡ ( − 1 2 ( a x t − 1 2 + b x t − 1 + c ( x t , x 0 ) ) )

   $$\begin{aligned} q(\bold{x}_{t-1}|\bold{x}_t,\bold{x}_0)&=q(\bold{x}_{t}|\bold{x}_{t-1},\bold{x}_0)\frac{q(\bold{x}_{t-1}|\bold{x}_0)}{q(\bold{x}_t|\bold{x}_0)}\\ &\propto \exp \left(-\frac{1}{2}\left(\frac{(\bold{x}_t-\sqrt{\alpha_t}\bold{x}_{t-1})^2}{1-\alpha_t}+\frac{(\bold{x}_{t-1}-\sqrt{\bar{\alpha}_{t-1}}\bold{x}_0)^2}{1-\bar{\alpha}_{t-1}}-\frac{(\bold{x}_{t}-\sqrt{\bar{\alpha}_{t}}\bold{x}_0)^2}{1-\bar{\alpha}_{t}}\right)\right)\\ &=\exp\left(-\frac{1}{2}\left({\color{blue} a}\bold{x}_{t-1}^2+{\color{red} b}\bold{x}_{t-1}+c(\bold{x}_t,\bold{x}_0)\right)\right) \end{aligned}$$ q(xt−1​∣xt​,x0​)​=q(xt​∣xt−1​,x0​)q(xt​∣x0​)q(xt−1​∣x0​)​∝exp(−21​(1−αt​(xt​−αt​ ​xt−1​)2​+1−αˉt−1​(xt−1​−αˉt−1​ ​x0​)2​−1−αˉt​(xt​−αˉt​ ​x0​)2​))=exp(−21​(axt−12​+bxt−1​+c(xt​,x0​)))​
   可见，上述分布的核心可以用一个二次函数来描述，那对应的中轴线应该是
   $$\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散度，构成负对数似然的上界，通过最小化上界，负对数似然自然取得最小。
− log ⁡ p θ ( x 0 ) = − log ⁡ p θ ( x 0 ) + D K L ( q ( x 1 : T ∣ x 0 ) ∣ ∣ p θ ( x 1 : T ∣ x 0 ) ) = − log ⁡ p θ ( x 0 ) + E x 1 : T ∼ q ( x 1 : T ∣ x 0 ) [ log ⁡ q ( x 1 : T ∣ x 0 ) p θ ( x 0 : T ) / p θ ( x 0 ) ] = − log ⁡ p θ ( x 0 ) + E x 1 : T ∼ q ( x 1 : T ∣ x 0 ) [ log ⁡ q ( x 1 : T ∣ x 0 ) p θ ( x 0 : T ) + log ⁡ p θ ( x 0 ) ] = E x 1 : T ∼ q ( x 1 : T ∣ x 0 ) [ log ⁡ q ( x 1 : T ∣ x 0 ) p θ ( x 0 : T ) ] = L V L B

$$\begin{aligned} -\log p_\theta(\bold{x}_0) &=-\log p_\theta(\bold{x}_0) + D_{KL}(q(\bold{x}_{1:T}|\bold{x}_0)||p_\theta(\bold{x}_{1:T}|\bold{x}_0))\\ &=-\log p_\theta(\bold{x}_0) +\mathbb{E}_{\bold{x}_{1:T}\sim q(\bold{x}_{1:T}|\bold{x}_0)}\left[\log \frac{q(\bold{x}_{1:T}|\bold{x}_0)}{p_\theta(\bold{x}_{0:T})/p_\theta (\bold{x}_0)}\right]\\ &=-\log p_\theta(\bold{x}_0) +\mathbb{E}_{\bold{x}_{1:T}\sim q(\bold{x}_{1:T}|\bold{x}_0)}\left[\log \frac{q(\bold{x}_{1:T}|\bold{x}_0)}{p_\theta(\bold{x}_{0:T})}+\log p_\theta (\bold{x}_0)\right]\\ &=\mathbb{E}_{\bold{x}_{1:T}\sim q(\bold{x}_{1:T}|\bold{x}_0)}\left[\log \frac{q(\bold{x}_{1:T}|\bold{x}_0)}{p_\theta(\bold{x}_{0:T})}\right]\qquad \color{blue}{=L_{VLB}} \end{aligned}$$ −logpθ​(x0​)​=−logpθ​(x0​)+DKL​(q(x1:T​∣x0​)∣∣pθ​(x1:T​∣x0​))=−logpθ​(x0​)+Ex1:T​∼q(x1:T​∣x0​)​[logpθ​(x0:T​)/pθ​(x0​)q(x1:T​∣x0​)​]=−logpθ​(x0​)+Ex1:T​∼q(x1:T​∣x0​)​[logpθ​(x0:T​)q(x1:T​∣x0​)​+logpθ​(x0​)]=Ex1:T​∼q(x1:T​∣x0​)​[logpθ​(x0:T​)q(x1:T​∣x0​)​]=LVLB​​
我们也可以继续对$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 ∼ Uniform ( { 1 , 2 , . . , T } ) t\sim \text{Uniform}(\{1,2,..,T\}) t∼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$