【AI绘图学习笔记】Probabilistic Diffusion Model概率扩散模型

参考：Probabilistic Diffusion Model概率扩散模型理论与完整PyTorch代码详细解读
由浅入深了解Diffusion Model

---

回顾VAE

在上节我们了解了VAE的原理，总体来说可以分为两个过程，一个是使用$q(z|x)$给Encoder进行前向学习的过程，另一个是使用$p(x|z)$给Decoder进行逆向推理的过程。

多层VAE模型和单层VAE实际上是类似的，只不过在单层基础上，我们将其视为一个马尔科夫链，因此每一个概率只与前一个概率有关，

上面是我们利用詹森不等式得到的置信下界。

上面是概率的链式法则（马尔科夫链下的），我们把这个代入上面的最大似然，可以得到下界可以写为这种形式：

这就是多层VAE的目标函数。

---

Diffusion Model

之所以要先介绍VAE，是因为实际上多层VAE的过程和Diffusion Model很像。

Diffusion Model的原理是先对目标$x_0$逐步正向加噪得到最终的分布$x_T$,然后再用逆向推理的过程逐步去噪,由$x_T$得到$x_0$。

上图就是Diffusion Model的可视化过程，总的来说就是加噪的扩散过程（熵增过程）$q(x_t|x_{t-1})$，也就是上图第一行，我们可以看到随着加噪图像逐渐变得无序。
而对于一张给定的噪声图片$x_T$，Diffusion Model在学习了去噪的逆向推理过程$p(x_{t-1}|x_t)$后，就可以产生一个新的图片，也就是上图的第二行，（从T时刻到0时刻）可以看到我们生产了新的样本，它和原来我们用于加噪的训练图片大致是相似的。
第三行是漂移量，从中我们可以看到上一时刻和下一时刻的图像像素运动方向。

---

扩散过程forward

1.给定初始数据分布$x_0\sim q(x)$，可以不断地向原始分布中添加高斯噪声（仿射变换），该噪声的标准差是以固定值${\beta }_t$确定的，均值（期望）以固定值${\beta }_t$和当前$t$时刻的数据$x_t$决定。这个过程是一个马尔科夫链。

2.随着$t$的不断增大，最终的数据分布$x_T$就变成了一个各向独立的高斯分布

关于$x_t$服从的高斯分布的算法，其实就是我们上节讲过的重采样技巧，我们从正态分布中取样一个$z$，然后计算$x=\sigma z+\mu$得到$x$的采样值。利用重采样技巧对$x_{t-1}$迭代就能求得$x_t$的高斯分布。并且$q(x_t|x_0)$满足马尔科夫链。

注意：${\beta }_t\in (0,1)$，并且随着时间推移会越来越大。

3.任意时刻的$q(x_t)$推导也可以完全基于$x_0$和${\beta }_t$来计算，而不需要迭代，下面是计算过程：

这里我们需要用到参数重整化技巧，我们令${\alpha }_t=1-{\beta }_t，{\overline{\alpha }}_t=\prod _{t=1}^T{\alpha }_i$，接下来将上式$q(x_t|x_{t-1})$代入重采样技巧$x_t=\sigma z_{t-1}+\mu =\sqrt{\beta }{z}_{t-1}+\sqrt{1-\beta }{x}_{t-1}$得到：

$x_t=\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{z}_{t-1}{z}_{t-1},z_{t-2}均...\sim N(0,I)$

我们将上式中的$x_{t-1}$用关于$x_{t-2}$的式子替换：
$$\begin{gathered} x_t=\sqrt{{\alpha }_t}(\sqrt{{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{z}_{t-2})+\sqrt{1-{\alpha }_t}{z}_{t-1} \\ =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}{z}_{t-2}+\sqrt{1-{\alpha }_t}{z}_{t-1} \end{gathered}$$

给出一个基本结论(独立分布可加性）：两个正态分布$X\sim N({\mu }_1,{\sigma }_1^2)和Y\sim N({\mu }_2,{\sigma }_2^2)$得到的叠加分布$aX+bY$的均值为$a{\mu }_1+b{\mu }_2$，方差为$a^2{\sigma }_1^2+b^2{\sigma }_2^2$，所以$\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}{z}_{t-2}+\sqrt{1-{\alpha }_t}{z}_{t-1}$由于$z\sim N(0,I)$，所以对应的叠加分布均值就是$\mu =0+0=0$，方差${\sigma }^2=a^2+b^2={\alpha }_t-{\alpha }_t{\alpha }_{t-1}+1-{\alpha }_t=1-{\alpha }_t{\alpha }_{t-1},$因此代入重采样公式即为$\sigma z+\mu =\sqrt{1-{\alpha }_t{\alpha }_{t-1}}z$

$$\begin{gathered} x_t=\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{z}}_{t-2}({\overline{z}}_{t-2}是混合高斯，但仍是标准正态分布) \\ =... \\ =\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}z \end{gathered}$$

结论：
$x_t=\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}z(3)$

因此任意时刻的$q(x_t)$都可以基于$x_0$和${\beta }_t$来计算而不需要迭代：

$q(x_t|x_0)=N(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)I)$，这样我们就能以这个高斯分布采样出$x_t$而不需要t次迭代了。

所以$\alpha$更像是类似学习率的参数，因为$\beta$会越来越大，所以$\alpha$会越来越小，因此足够多的时刻之后$\sqrt{\alpha }\to 0,\sqrt{1-\alpha }\to 1$，因此当到达 t 时刻，$q(x_t|x_0)=N(x_t;0,I)$会收敛于一个标准的正态分布，这样我们就能求出这个最大的时刻 $t$ 了。从学习率上表现的性质就是如果用$x_t$预测$x_0$，前面阶段的生成会很快显示出图像的底样，而越往后越慢，越是需要生成更多的细节。

（上图表示了${\bar{\alpha }}_t$和扩散步骤的关系)

从这里我们可以看到Diffusion Model 和VAE的一些区别，
首先在VAE中参数是通过前向逆向过程预测出来的，但diffusion里是给出一个固定参数后进行训练。
其次在VAE中采样的隐变量$z$和$x$是有一定关联的，而diffusion中的$x_t$最终是一个标准正态分布，和$x$无关了。
此外在VAE中$x$和$z$的维度不一定一样，而diffusion中$x_0....x_t$的维度始终是一样的。

---

逆扩散过程reverse

如果说前向过程是一个加噪的过程，那么逆向过程就说去噪推断的过程，如果我们能够逐步得到逆转后的分布$q(x_{t-1}|x_t)$，就可以从标准正态分布$x_t\sim N(0,I)$还原出原图分布$x_0$，文献¹中证明了如果$q(x_{t-1}|x_t)$满足高斯分布且${\beta }_t$足够小，那么$q(x_{t-1}|x_t)$仍然是高斯分布，如果逐步地对$x_t...x_0$进行拟合来找到其服从的高斯分布参数，实在难以计算。所以我们需要构建一个参数分布来去做估计，逆向扩散过程仍然是一个马尔科夫链。

我们使用深度学习模型（参数为$\theta$,目前主流是U-Net+attention的结构）去预测这样的一个逆向的分布$p_{\theta }$。（类似VAE）:

虽然我们无法得到逆转后的分布$q(x_{t-1}|x_t)$,但是如果我们知道$x_0$，可以通过以下的公式进行计算:

$q(x_{t-1}|x_t,x_0)=N(x_{t-1};\tilde{\mu }(x_t,x_0),{\tilde{\beta }}_{t}I)(6)$

推理如下：

我们用贝叶斯公式推理下(7-1):
$q(a|b,c)=\frac{q(a,b,c)}{q(b,c)}$
其中链式法则$q(a,b,c)=q(b|a,c)q(a|c)q(c)$，$q(b,c)=q(b|c)q(c)$
因此代入原式$=\frac{q(b|a,c)q(a|c)q(c)}{q(b|c)q(c)}=\frac{q(b|a,c)q(a|c)}{q(b|c)}$，由于是马尔科夫链，上式还等价于$q(b|c)\frac{q(a|c)}{q(b|c)}$

我们可以看到在(7-1)式子中，我们使用贝叶斯公式，将这个逆向过程转化为了前向过程，(7-2)为其对应的高斯分布概率密度函数。我们将其展开得到式子(7-3),其中与$x_{t-1}$无关的项例如只包含了$x_t\&x_0$的项都被归入了$C(x_t,x_0)$

（高斯分布的概率密度函数是$f(x)=\frac{1}{\sqrt{2\pi }\sigma }{e}^{-\frac{(x-\mu {)}^2}{2{\sigma }^2}}$)

我们之前说过$q(x_{t-1}|x_t)$仍然是高斯分布，也就意味着上式可以被整理为高斯分布的式子，对应的一般高斯概率密度函数的指数部分为$exp(-\frac{(x-\mu {)}^2}{2{\sigma }^2})=exp(-\frac{1}{2}(\frac{1}{{\sigma }^2}{x}^2-\frac{2\mu }{{\sigma }^2}x+\frac{{\mu }^2}{{\sigma }^2}))$，和(7-3)整理出来的形式是对应的，因此：

用前向时讲到的$\beta$代替方差${\sigma }^2$，得到：

由于扩散过程讲过的推断3，我们得知任意时刻的$x_t$可以由$x_0$和$\beta$表示。因此:

$x_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\bar{z}}_t)$($x_t$式子变形）

将其代入到（8-2）得到

其中的高斯分布${\bar{z}}_t$为深度模型所预测的噪声（用于去噪），可看做为$z_{\theta }(x_t,t),$得到：

我们经过了一连串的计算，先得到$q(x_{t-1}|x_t)$的高斯分布概率密度函数，然后将其整理为高斯分布指数部分的一般形式，由此得到对应的均值和方差，其中方差$\beta$是定值，而均值$\mu$再表示为由$x_t,t$为参数的函数形式

这样一来,DDPM的每一步的推断可以总结为：

1) 每个时间步通过$x_t$和$t$来预测高斯噪声$z_{\theta }(x_t,t)$,随后根据(9)得到均值${\mu }_{\theta }(x_t,t)$

2) 得到方差${\Sigma }_{\theta }(x_t,t)$,DDPM中使用untrained ${\Sigma }_{\theta }(x_t,t)={\tilde{\beta }}_t$（也就是不训练方差作为固定参数），且认为${\tilde{\beta }}_t={\beta }_t$和${\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}\cdot {\beta }_t$结果近似

3) 根据(5-2)得到$q(x_{t-1}|x_t)$,利用重参数技巧得到$x_{t-1}$

重复上述步骤逐步去噪，直到计算出$x_0$,去噪过程完毕。

---

Diffusion训练

讲完了前向加噪的扩散过程和逆向去噪的推断过程（虽然文字上来看原理很简单，但是公式好繁杂）。现在我们讲讲如何训练diffusion model以得到靠谱的参数${\mu }_{\theta }(x_t,t)和{\Sigma }_{\theta }(x_t,t)$，方法还是最大对数似然（此处用的最小化负对数似然）。

由于整个Diffusion模型和VAE很相似，训练过程也是，由于KL散度恒大于0，因此我们在负对数似然上加上一个KL散度就构成了它的上界（和VAE最大对数似然的时候正好相反，那时是减去一个KL散度是下界)：

利用詹森不等式，我们就能得到（这块看的不太细，记住结论就好)：

我们进一步对$L_{VLB}$进行推导，可以得到熵与多个KL散度的累加²，其中分母是扩散过程，分子是逆扩散过程：

（上式【从$L_{VLB}$开始为第一行】第四行到第五行又是应用了贝叶斯公式，先逆向马尔科夫链补上了一个$x_0$再应用了和之前前向中讲到的一模一样的贝叶斯公式)
第六行的第三项${\sum }_{t=2}^{T}log\frac{q(x_t|x_0)}{q(x_{t-1}|x_0)}$可化简,最后与第四项以及第一项可合并，最终得到了第七行的式子

最后一行将其简化为了含有KL散度的式子，其中$L_T$不含参（q分布不含参，$x_T$逆向过程最终为纯高斯噪声)相当于常量可以直接忽略，$L_{t-1}$是逆扩散过程的KL散度，最后考虑的还是$L_{t-1}和L_0$。

并且由于$q和p_{\theta }$其实都是高斯分布，并且$q$是关于参数${\beta }_t$的高斯分布，且${\beta }_t$是untrained的固定参数(忘记的,点此回去)，$p_{\theta }$的高斯分布均值是${\mu }_{\theta }$,方差是${\Sigma }_{\theta }$且${\Sigma }_{\theta }$也是untrained的固定参数。因此可训练的参数$\theta$只在$p_{\theta }$中。

上面刚才推导的式子也可写为：

给出一个结论：对于两个高斯分布p，q而言，它们的KL散度可等价为
$D_{KL}(p,q)=log\frac{{\sigma }_2}{{\sigma }_1}+\frac{{\sigma }_1^2+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }_2^2}-\frac{1}{2}$
让我们把$L_{t-1}$的式子用上式表示出来，得到:

(此处$L_{t-1}$应该是笔误，实为$L_t$，$L_{VLB}$给出的是$t=2$开始，在(14-3)中已经被改为$t=1$开始到$T-1$)

然后将${\tilde{\mu }}_t$用(8-2)替换，${\mu }_{\theta }$用(9)替换，$x_t$用(3)替换，得到：

从(16)可以看出，diffusion训练的核心就是取学习高斯噪声${\bar{z}}_t,z_{\theta }$之间的均方误差MSE。论文中作者说我们可以将(16)式子中前面的这个系数给直接丢掉，这样训练会更稳定。

最后论文给出的式子，我们将${\bar{z}}_t$替换为$ϵ$，$z_{\theta }$替换为${ϵ}_{\theta }$,DDPM将loss进一步简化为：

训练过程可以看做：

1）获取输入$x_0$，从1…T随机采样一个$t$
2) 从标准高斯分布采样一个噪声$ϵ\sim N(0,I)$
3) 最小化loss函数

---

总结

最后我们给出DDPM提供的训练/测试（采样）流程图

在训练过程中，我们要输入$x_0$进行随机采样时刻$t$并采样噪声$ϵ$，然后对loss函数进行梯度下降直到拟合。而在测试采样过程中，我们则利用马尔科夫链逆向逐步去噪计算$x_T$直到$x_0$作为最后的生成结果。

---

加速Diffusion采样和方差的选择(DDIM)

通过遵循反向扩散过程的马尔可夫链从DDPM生成样品非常慢，因为高质量生成需要的$T$最多要走一千步或是几千步。“例如，从 DDPM 采样大小为 20 × 50 的 32k 图像大约需要 32 小时，但从 Nvidia 2080 Ti GPU 上的 GAN 采样不到一分钟。”
这就导致diffusion的前向过程非常缓慢。在denoising diffusion implicit model (DDIM)中提出了一种牺牲多样性来换取更快推断的手段。

一种简单的方法是运行一个跨步的采样，通过每隔$\lceil T/S\rceil$步进行采样更新，总共采样$S$步，这样就能有效减少采样数量。

另一种方法是重写$q_{\sigma }(x_{t-1}|x_t,x_0)$的标准差${\sigma }_t$：
根据(3)我们可知：

（第二步利用了独立高斯分布可加性)

最终得到的(18)将方差${\sigma }_t^2$迎入到了均值中，当${\sigma }_t^2={\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}$时，(18)等价于(6)。我们给定一个$\eta$作为控制采样随机性的超参数${\sigma }_t^2=\eta {\tilde{\beta }}_t$（作用类似于调整方差的大小),当$\eta =1$的时候就等价于DDPM，当$\eta =0$的时候是DDIM

（上图是不同设置的扩散模型在 CIFAR10 和 CelebA 数据集上的 FID 得分，包括了DDIM($\eta =0$)以及DDIM($\hat{\sigma }$)($\eta =1$））

根据上表可知，数据量较小时DDIM的训练更快，而数据量较大时使用DDPM并采用更大的方差效果更好

与DDPM相比，DDIM能够：

1.使用更少的步骤生成更高质量的样本。
2.具有“一致性”属性，因为生成过程是确定性的，这意味着以同一潜在变量为条件的多个样本应该具有类似的高级特征。
3.由于一致性，DDIM 可以在潜在变量中执行语义上有意义的插值。

---

1. Feller, William. “On the theory of stochastic processes, with particular reference to applications.” Proceedings of the [First] Berkeley Symposium on Mathematical Statistics and Probability. University of California Press, 1949. ↩︎

2. Feller, William. “On the theory of stochastic processes, with particular reference to applications.” Proceedings of the [First] Berkeley Symposium on Mathematical Statistics and Probability. University of California Press, 1949. ↩︎