[翻译+笔记] 什么是Diffusion模型?

---

原博客地址为https://lilianweng.github.io/posts/2021-07-11-diffusion-models/

该笔记是[翻译+笔记]GAN和[翻译+笔记]VAE的续集.

注意: 这个笔记中score-based模型的部分暂时略去, 请参见我的另一个笔记, 其中有更详细的解释:https://blog.csdn.net/wjpwjpwjp0831/article/details/141647920

---

之前的生成模型或多或少都有一些缺陷, 例如, GAN的训练不够稳定, VAE的训练实际上要依靠一些代理损失, 并不直接作用在网络上, 而流模型需要用特定的结构来实现取逆的操作.

Diffusion模型是基于非平衡热力学的, 其定义了一个Markov过程来给原始数据逐步加噪, 然后学习从噪声中恢复样本. 和VAE不同, Diffusion模型可以学习数据的高维表示, 甚至和原始数据同一个维度(pixel-based的).

几种模型对比图如下:

1. 基础的扩散模型

1.1 前向加噪过程

假定原始数据服从分布${\mathbf{x}}_0\sim q(\mathbf{x})$, 现在我们逐步对其加噪, 加入的是高斯噪声. 对于每一步加噪, 我们希望将分布$q$逐渐向高斯过程靠近, 也即让$q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}$. 在每一步, 我们假定高斯分布的均值与过去的值${\mathbf{x}}_{t-1}$有关, 而协方差为固定值(对角阵):

$$q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})$$

其中${\beta }_t$为一小正数, 在0~1之间.

注意我们假定加噪过程为Markov过程, 因此当前状态${\mathbf{x}}_t$只假定与上一状态${\mathbf{x}}_{t-1}$有关.

因此, 当前时刻${\mathbf{x}}_t$是由前一时刻${\mathbf{x}}_{t-1}$决定的正态分布, 其均值为$\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1}$, 方差为${\beta }_t\mathbf{I}$. 为了表示${\mathbf{x}}_t$, 这里我们使用一下重参数化(Re-parametrization)技巧. 重参数化是说, 如果我们从一个高斯分布中取样, 也等效于从标准高斯分布中取样, 只不过是加上均值, 以及乘以标准差. 这是因为一个$x\sim \mathcal{N}(\mu ,{\sigma }^2)$的高斯分布可以等价于$\mu +\sigma ϵ$, 其中$ϵ\sim \mathcal{N}(0,1)$.

由高斯分布性质立即得到.

因此, ${\mathbf{x}}_t$可以表示为:

$${\mathbf{x}}_t=\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1}+\sqrt{{\beta }_t}{ϵ}_{t-1}$$

其中${ϵ}_{t-1}\sim \mathcal{N}(0,I)$. 为了表示方便, 令$\sqrt{1-{\beta }_t}=\sqrt{{\alpha }_t}$, 将上式递归展开, 我们有:

$${\mathbf{x}}_t=\sqrt{{\alpha }_t}(\sqrt{{\alpha }_{t-1}}{\mathbf{x}}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{ϵ}_{t-2})+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}$$

我们注意到后面两项可以合并为一个新的高斯分布, 其均值为0, 方差为$1-{\alpha }_t{\alpha }_{t-1}$, 按此规律展开, 我们得到:

$${\mathbf{x}}_t=\sqrt{{\Pi }_i{\alpha }_i}{\mathbf{x}}_0+\sqrt{1-{\Pi }_i{\alpha }_i}ϵ,ϵ\sim \mathcal{N}(0,I)$$

所以, 我们可以直接从${\mathbf{x}}_0$得到${\mathbf{x}}_t$的分布:

$$q({\mathbf{x}}_t|{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_t;\sqrt{{\Pi }_i{\alpha }_i}{\mathbf{x}}_0,\sqrt{1-{\Pi }_i{\alpha }_i}\mathbf{I})$$

令${\bar{\alpha }}_t={\Pi }_{i=1}^t$, 则$q({\mathbf{x}}_t|{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,\sqrt{1-{\bar{\alpha }}_t}\mathbf{I})$

所以, 随着时间的增加, ${\mathbf{x}}_t$会越来越趋向于标准正态分布. 以上就是加噪的过程.

1.2 反向去噪过程

假如我们知道$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$的真实分布, 我们就可以采样${\mathbf{x}}_t$从而一步步得到${\mathbf{x}}_0$. 和VAE的境况相似, 我们很难准确知道这个分布, 又只能自己估计一个模型$p_{\theta }$来近似这个分布.

当然我们最后的目标是得到${\mathbf{x}}_0$, 就希望$p_{\theta }({\mathbf{x}}_0)$最大化, 也就是我们估计的分布, 要尽量符合真实的样本分布, 即在有意义的样本处最大. 有:

$$p_{\theta }({\mathbf{x}}_0)=p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)p_{\theta }({\mathbf{x}}_1|{\mathbf{x}}_2)\dots p_{\theta }({\mathbf{x}}_{T-1}|{\mathbf{x}}_T)p_{\theta }({\mathbf{x}}_T)$$

注意, 如果加噪的步长取得足够小, 实际上每一步去噪的过程$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$也符合高斯分布, 相应地我们假定我们估计的$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$也符合高斯分布, 其均值和方差与${\mathbf{x}}_t$有关:

$$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mu }_{\theta }({\mathbf{x}}_t,t),{\Sigma }_{\theta }({\mathbf{x}}_t,t))$$

下面让我们来看一下我们能不能写出$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$的分布形式. 注意到在前向加噪过程中我们得到了任意时间步加噪样本和初始样本的关系, 即$q({\mathbf{x}}_t|{\mathbf{x}}_0)$. 受此启发, 便于用到之前的结论, 我们将初始样本考虑进去: $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$.

根据Bays公式, 有:

$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\frac{q({\mathbf{x}}_t,{\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}{q({\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)}$$

这样就巧妙地转换成了三个前向加噪的过程, 而前向加噪我们已经知道其高斯分布的具体形式了, 因此有:(过程有点复杂, 此处手写)

我们知道形如$a{x}^2+bx+c$整理成完全平方式的形式为$a(x+\frac{b}{2a}{)}^2+(c-\frac{b^2}{4a})$, 因此, $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$这个新分布的方差就是$1/a$, 均值就是$-b/2a$, 假定$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_{t-1};\tilde{\mu }({\mathbf{x}}_t,{\mathbf{x}}_0),\widetilde{{\beta }_t}\mathbf{I})$, 则

因此, 理论上, 如果我们可以预测出每一个去噪过程$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$的均值, 那么就可以知道分布, 也就能够采样了.

1.3 训练

我们先来证明我们按照上述方法估计$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$的合理性. 我们希望最大化我们估计的分布中真实样本处的值, 即$p_{\theta }({\mathbf{x}}_0)$. 作为损失函数来讲, 理论上的优化目标是最小化$-\log q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$. 我们有:

令$L_{VLB}={\mathbb{E}}_q[\dots ]$, 即等号右边的项, 得到

$$L_{VLB}\ge -{\mathbb{E}}_{q_{{\mathbf{x}}_0}}{p}_{\theta }({\mathbf{x}}_0)$$

因此让我们估计的分布$p$和真实的$q$相似, 就等价于降低$-\log q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$的值. 因此这么做是合理的.

当然现在的期望我们很难计算, 我们接下来可以整理成已知的高斯分布组合的形式, 从而可以得到闭式解:

看完理论上的合理性后, 我们再来看实际中到底怎么去算损失. 注意我们在估计$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$的时候推导出了其均值和方差. 因此我们的目标就是用网络预测出它的均值, 从而对该分布有一个良好的近似. 也就是说, 对我们近似的分布$p_{\theta }$, 预测均值为:

$${\mu }_{\theta }({\mathbf{x}}_t,t)=\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_t)$$

${ϵ}_t$是未知的加噪, 因此我们需要从${\mathbf{x}}_t$将其预测出来, 即:

相当于对该噪声项做了一个重参数化.

因此损失函数就定义为预测的均值和实际均值之差:

实际上在DDPM的论文中, 作者发现不需要前面的方差系数项:

下图为DDPM的训练和推理过程:

1.4 ${\beta }_t$的参数化

在传统的Diffusion模型中, 在加噪过程中, ${\beta }_t$是渐进线性变化的, 例如从0.0001到0.02, 但是, 这样生成的效果和其他生成式模型来说还是有差距.

因此有人就提出了一系列的改进方式, 来达到更低的NLL(负对数似然, 也就是我们要的优化目标). 其中一个就是提出一个cosine的变化方式, 对于变化函数的选择, 可以是任意的, 只要在训练过程中是近似线性的下降, 以及在末端($t=0,T$)有细微的变化即可, 即

从下图中可以看到, 采用cosine方式可以让${\beta }_t$的变化更加平缓.

1.5 反向去噪过程方差${\Sigma }_{\theta }$的参数化

根据前面的推导, 我们估计的分布$p_{\theta }$的方差为:

$${\Sigma }_{\theta }({\mathbf{x}}_t,t)=\frac{1-{\bar{\alpha }}_{t-1}}{1-\overline{{\alpha }_t}}{\beta }_t\cdot \mathbf{I}$$

因此该方差为定值, 是不含训练参数的, 因为实验证明如果是将方差设定为可学习的对角矩阵, 则会导致训练很不稳定.

后来有人想训练这个方差, 并将方差定为${\beta }_t$和${\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-\overline{{\alpha }_t}}{\beta }_t$的线性组合, 至于组合的系数多少, 则是可学习的:

$${\Sigma }_{\theta }({\mathbf{x}}_t,t)=\exp (\mathbf{v}\log {\beta }_t+(1-\mathbf{v})\log {\tilde{\beta }}_t)$$

然而, 传统DDPM的损失函数中没有该方差项(是均值的差距), 为了包含进去, 将之前推导的$L_{VLB}$纳入损失函数, 并赋予一个小权重. 还记得$L_{VLB}$最终被推导成了若干相邻时间步去噪分布的KL散度的组合, 是有解析解的, 并且也与分布的方差有关. 因此, 在计算$L_{VLB}$时, 对均值进行stop gradient操作, 只训练其中的方差项. 实际上这样还是比较难收敛, 因此最终采用了time smoothing的trick来平滑梯度噪声带来的不稳定.

2. 加速Diffusion模型

2.1 去噪过程的加速

经典的Diffusion模型的生成过程是很慢的, 对DDPM来说, 需要接近20小时来采样5w张分辨率仅有32x32的图片. 因此, 要加速Diffusion模型, 有两大思路: 减少采样步骤, 或者做蒸馏.

一个简单的方式就是直接减少采样步数, 每$ceil(T/S)$步采样一次, 这样一共采样$S$次.

还有一种方式, 就是DDIM的方式. 我们重新考虑去噪过程$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$.

我们先从加噪过程推导:

对于前面去噪过程的推导, 我们知道$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_{t-1};\tilde{\mu }({\mathbf{x}}_t,{\mathbf{x}}_0),{\tilde{\beta }}_t\mathbf{I})$, 其中

对比刚刚推导的式子, 我们令${\sigma }_t^2=\eta {\tilde{\beta }}_t$, 那么$\eta$就控制了方差的大小, 也就是, 去噪过程的不确定性. 因此, 如果$\eta =0$, 说明去噪过程是完全确定的, 其可以确定性地将噪声映射回原始样本, 所以叫做DDIM(Implicit Model).

这样的话, 在去噪过程中, 我们就没必要遵循整个Markov链了, 而是跳跃性地执行. 因此去噪过程变为:

当步长比较大(即前面提到的$S$比较小)的时候, DDIM的表现会比DDPM好很多.

2.2 潜在变量空间(latent variable space)

前面说的都是直接在原始数据分布, 即pixel-wise的加噪与去噪过程, 这会导致比较高的数据维度和冗余信息的计算. 当然就想能不能将数据样本映射到更低的变量空间来减少计算量, 因此Latent Diffusion Model就是干这个事情的.

因此, 加噪和去噪过程是在latent space进行的, 首先对图像用VAE的encoder进行编码, 然后得到去噪的latent后再用VAE的decoder进行解码, 就得到了原始图像.

在去噪过程中, 采用了一个以时间步为条件的UNet, 并加入了交叉注意力机制, 这是用于融合模态. 如下式所示:

整个模型结构如下图: