扩散模型Diffusion model | 基于分数的生成模型

扩散模型Diffusion model | 基于分数的生成模型

论文参考：Score-Based Generative Modeling through Stochastic Differential Equations

在 DDIM 中，其作者对 DDPM 的采样及前项传播公式进行了修改。从而实现了采样的确定性以及速度提升。但是此过程也可以与随机微分方程（Stochastic Differential Equation，SDE）和常微分方程ODE相联系

生成模型概述

现有的图像生成技术在很大程度上可以根据它们如何表示概率分布分为两类。

• 基于似然的模型，通过（近似）最大似然直接学习分布的概率密度（或质量）函数。典型的基于似然的模型包括自回归模型、归一化流动模型、基于能量的模型（EBM） 和变分自编码器 （VAE）。
• 隐式生成模型，其中概率分布由其采样过程的模型隐式表示。如生成对抗网络（GAN），其中通过用神经网络转换随机高斯向量来合成来自数据分布的新样本。

基于似然的模型要么需要对模型架构进行严格限制，以确保似然计算具有可处理的归一化常数，要么必须依赖代理目标来逼近最大似然训练。而隐式生成模型通常需要对抗性训练，这是出了名的不稳定，可能导致模式崩溃。

基于分数的生成模型

此模型的关键思想是对对数概率密度函数的梯度进行建模，该函数通常称为评分函数。这种基于分数的模型不需要具有可处理的归一化常数，并且可以通过得分匹配直接学习。

评分函数可以具体定义为：$\frac{\partial \left(\log p\left(\mathbf{x}\right)\right)}{\partial \mathbf{x}}={\nabla }_{\mathbf{x}}\log p\left(\mathbf{x}\right)$，那么可以利用神经网络构建分数模型$s_{\theta }(x)$来拟合${\nabla }_x\log p(x)$。从而可以通过最小化模型和数据分布之间的Fisher散度（简记为F-散度）来训练基于分数的模型，其表达式为
$${\mathbb{E}}_{p\left(\mathbf{x}\right)}\left[{\Vert s_{\theta }\left(\mathbf{x}\right)-{\nabla }_{\mathbf{x}}\log p\left(\mathbf{x}\right)\Vert }_2^2\right]={\int }_{\mathbf{x}\sim \mathcal{X}}p\left(\mathbf{x}\right){\Vert s_{\theta }\left(\mathbf{x}\right)-{\nabla }_{\mathbf{x}}\log p\left(\mathbf{x}\right)\Vert }_2^{2}d\mathbf{x}$$

对于输入数据(样本)来说，其对数概率密度增长最快的方向，其可以用矢量场分布图表示。因此，在采样过程中沿着分数的方向走，就能够走到数据分布的高概率密度区域，最终生成的样本就会符合原数据分布。

Langevin 动力学

Langevin dynamics 提供了一个 MCMC 过程来使用其评分函数${\nabla }_{\mathbf{x}}\log p(\mathbf{x})$从分布中采样$p(\mathbf{x})$。具体来说，它从任意先验分布初始化${\mathbf{x}}_0\sim \pi (\mathbf{x})$，然后迭代以下内容：
$${\mathbf{x}}_{i+1}\leftarrow {\mathbf{x}}_i+ϵ{\nabla }_{\mathbf{x}}\log p(\mathbf{x})+\sqrt{2ϵ}{\mathbf{z}}_i,i=0,1,\cdots ,K$$

其中，${\mathbf{z}}_i\sim \mathcal{N}(0,I)$为标准高斯分布。理论上，当$ϵ\to 0$且$K\to \infty$时，最终生成的样本$x_T$将会服从原数据分布$p(x)$。这个模型就叫基于分数的生成模型SMLD。

基于分数的生成模型的缺陷与改进

在低密度区域，估计的分数函数是不准确的，因为在这些区域可用于计算分数匹配目标的数据点很少。一种解决方案是用噪声扰动数据点，并在嘈杂的数据点上训练基于分数的模型。当噪声幅度足够大时，它可以填充低数据密度区域，以提高估计分数的准确性。

但噪声的选择上依然存在一定程度上的问题，强度太小起不到效果，而强度太大会破坏数据的原始分布。因此可以同时使用多个尺度${\sigma }_1<{\sigma }_2<\cdots <{\sigma }_L$的同性高斯噪声扰动加入到$p(x)$中，即可获得噪声扰动分布
$$p_{{\sigma }_i}(\mathbf{x})=\int p(\mathbf{y})\mathcal{N}(\mathbf{x};\mathbf{y},{\sigma }_i^{2}I)\mathrm{d}\mathbf{y}.$$

之后就可以用噪声条件评分模型$s_{\theta }(x,t)$来估计经过噪声扰动过的分布分数，
$$\sum _{i=1}^L\lambda (i){\mathbb{E}}_{p_{{\sigma }_i}(\mathbf{x})}[\Vert {\nabla }_{\mathbf{x}}\log p_{{\sigma }_i}(\mathbf{x})-{\mathbf{s}}_{\theta }(\mathbf{x},i){\Vert }_2^2]$$
采样方式也要做出相应的变化，我们对于不同的噪声强度$i=L,L-1,\cdots ,1$做Langevin采样，噪声尺度${\sigma }_t$随时间逐渐减小(退火)，同时上一个scale的结果作为这一次的初始化。

与SDM的联系

前向过程

DDPM的前向传播过程被划分为了固定的$T$步，我们可以将它们理解为一个在时间上连续的变换过程：
$$dx=f(x,t)dt+g(t)d\mathbf{w}$$

其中，$t$代表时间步，$w$为标准维纳过程（也称为布朗运动，可视为随机项），$f(x,t)$和$g(t)$分别为漂移系数和扩散系数。和 DDIM 相比，DDIM的前向过程可以看作 ODE。

SDE 与 ODE 的主要差别在于 SDE 在传播过程中，添加了随机项 $g(t)d\mathbf{w}$，而引入SDE形式来描述扩散模型的本质好处是“将理论分析和代码实现分离开来”

上式也可以看做下述离散形式在$\Delta t\to 0$时的极限:

$$x_{t+\Delta t}-x_t=f(x,t)\Delta t+g(t)\sqrt{\Delta t}\mathbf{\epsilon },\mathbf{\epsilon }\sim \mathcal{N}(0,\mathbf{I})$$
上式可以用概率描述为：
$$p(x_{t+\Delta t}|x_t)=N(x_{t+\Delta t};x_t+f_t(x_t)\Delta t,g_t^2\Delta tI)\propto \exp \left(-\frac{||x_{t+\Delta t}-x_t-f_t(x_t)\Delta t|{|}^2}{2g_t^2\Delta t}\right)$$

反向过程

使用贝叶斯公式可以得到
$$\begin{array}{rl}p({\mathbf{x}}_t|{\mathbf{x}}_{t+\Delta t})& =\frac{p({\mathbf{x}}_{t+\Delta t}|{\mathbf{x}}_t)p({\mathbf{x}}_t)}{p({\mathbf{x}}_{t+\Delta t})}=p({\mathbf{x}}_{t+\Delta t}|{\mathbf{x}}_t)\exp (\log p({\mathbf{x}}_t)-\log p({\mathbf{x}}_{t+\Delta t}))\\ & \propto \exp \left(-\frac{\Vert {\mathbf{x}}_{t+\Delta t}-{\mathbf{x}}_t-{\mathbf{f}}_t({\mathbf{x}}_t)\Delta t{\Vert }^2}{2{\mathbf{g}}_t^2\Delta t}+\log p({\mathbf{x}}_t)-\log p({\mathbf{x}}_{t+\Delta t})\right)\end{array}$$

将$\log p({\mathbf{x}}_{t+\Delta t})$进行泰勒展开：
$$\log p({\mathbf{x}}_{t+\Delta t})\approx \log p({\mathbf{x}}_t)+({\mathbf{x}}_{t+\Delta t}-{\mathbf{x}}_t)\cdot {\nabla }_{{\mathbf{x}}_t}\log p({\mathbf{x}}_t)+\Delta t\frac{\partial }{\partial t}\log p({\mathbf{x}}_t)$$

再代回原式化简即可得到反向传播过程对应SDE为:
$$dx=[f(x,t)-g^2(t){\nabla }_x\log p_t(x)]dt+g(t)d\bar{\mathbf{w}}$$

那么此时我们只需要得到${\nabla }_x\log p_t(x)$，就可以进行推理生成的反向传递过程了。

我们利用分数匹配模型去$s_{\theta }(x(t),t)$拟合${\nabla }_x\log p_t(x)$，对${\mathbf{x}}_t$积分即可得到最终的损失函数：

$$\mathcal{L}={\mathbb{E}}_{t\sim U[0,T]}\left[\lambda (t)\int p_t(x)\Vert s_{\theta }(x,t)-{\nabla }_x\log p_t(x){\Vert }^{2}dx\right]$$

其中$\mathcal{U}(0,T)$表示在时间间隔内的均匀分布$[0,T]$，而一般令$\lambda (t)\propto 1/{\mathbb{E}}_{x\sim p_t(x)}^2[{\nabla }_x\log p_t(x\mid x_0)]$，只是起到调节Loss权重的作用，简单起见可以直接去掉它，这不会影响最优解的结果。

如果$x_T\sim \pi$，在理想情况下我们可以得到$p_{\theta }\approx p_0$，而当$\lambda (t)=g^2(t)$时，F-散度的加权组合和$p_0$和$p_{\theta }$之间的KL散度之间有一个重要的联系：
$$\mathrm{KL}(p_0(\mathbf{x})\Vert p_{\theta }(\mathbf{x}))\le \frac{T}{2}{\mathbb{E}}_{t\in \mathcal{U}(0,T)}{\mathbb{E}}_{p_t(\mathbf{x})}[\lambda (t)\Vert {\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})-{\mathbf{s}}_{\theta }(\mathbf{x},t){\Vert }_2^2]+\mathrm{KL}(p_T\Vert \pi )$$

由等式可知，在此条件下最小化KL散度和最大化模型训练的似然之间具有等价性，可以基于此似然加权函数训练基于分数的生成模型，以实现较高的似然率。

具体实例

上述$dx=f(x,t)dt+g(t)d\mathbf{w}$给出了一般化的SDE的定义，实际上$f(x,t)$和$g(t)$的不同取值也会得到不同的生成模型。SMLD 和 DDPM 中使用的噪声扰动可以看作是两个不同 SDE 的离散化。

VE-SDE在$t\to \infty$时总是产生方差爆炸的过程。相反，VP-SDE产生具有有界方差的过程。

VE-SDE（Variance Exploding SDE）

当使用总共N个噪声尺度时，SMLD的每个扰动核$p_{{\sigma }_i}\left(x\mid x_0\right)$对应于$x_i$的分布，即可写成一下马尔可夫链形式：
$${\mathbf{x}}_i={\mathbf{x}}_{i-1}+\sqrt{{\sigma }_i^2-{\sigma }_{i-1}^2}{\mathbf{z}}_{i-1},i=1,\cdots ,N$$
在上述表达中，我们假设初始随机变量${\mathbf{z}}_{i-1}$符合均值为$0$、协方差矩阵为单位矩阵$\mathbf{I}$的多元正态分布$\mathcal{N}(0,\mathbf{I})$。为了简化符号，引入了参数${\sigma }_0=0$。在$N\to \infty$的极限下，随机变量序列 { σ i } i = 1 N \{\sigma_i\}_{i=1}^N {σi​}i=1N​变成了连续函数$\sigma (t)$，${\mathbf{z}}_i$变成了连续随机过程$\mathbf{z}(t)$。马尔可夫链 { x i } i = 1 N \{\mathbf{x}_i\}_{i=1}^N {xi​}i=1N​也演变为了连续随机过程 { x ( t ) } t = 0 1 \{\mathbf{x}(t)\}_{t=0}^1 {x(t)}t=01​，其中我们使用连续时间变量$t\in [0,1]$进行索引，而不再使用整数$i$。由此可以将上式写为：

$$\mathbf{x}(t+\Delta t)=\mathbf{x}(t)+\sqrt{{\sigma }^2(t+\Delta t)-{\sigma }^2(t)}\mathbf{z}(t)\approx \mathbf{x}(t)+\sqrt{\frac{\mathrm{d}\left[{\sigma }^2(t)\right]}{\mathrm{d}t}\Delta t}\mathbf{z}(t),$$

当$\Delta t\ll 1$时，近似等式成立。在$\Delta t\to 0$的极限下，它收敛于以下SDE：
$$\mathrm{d}\mathbf{x}=\sqrt{\frac{\mathrm{d}\left[{\sigma }^2(t)\right]}{\mathrm{d}t}}\mathrm{d}\mathbf{w}$$

VP-SDE（Variance Preserving SDE）

同样对于 DDPM 的扰动核 { p α i ( x ∣ x 0 ) } i = 1 N \text{}\{p_{\alpha_i}(\mathbf{x}\mid\mathbf{x}_0)\}_{i=1}^N\text{} {pαi​​(x∣x0​)}i=1N​，离散马尔可夫链是
$${\mathbf{x}}_i=\sqrt{1-{\beta }_i}{\mathbf{x}}_{i-1}+\sqrt{{\beta }_i}{\mathbf{z}}_{i-1},i=1,\cdots ,N$$

同样简化时间变量到$t\in [0,1]$的范围，上式可改写为：
$${\mathbf{x}}_i=\sqrt{1-\frac{{\bar{\beta }}_i}{N}}{\mathbf{x}}_{i-1}+\sqrt{\frac{{\bar{\beta }}_i}{N}}{\mathbf{z}}_{i-1},i=1,\cdots ,N.$$

令$\Delta t=1/N$，上式可以改写为离散形式下的：
$$\begin{array}{rl}\mathbf{x}(t+\Delta t)& =\sqrt{1-\beta (t+\Delta t)\Delta t}\mathbf{x}(t)+\sqrt{\beta (t+\Delta t)\Delta t}\mathbf{z}(t)\\ & \approx \mathbf{x}(t)-\frac{1}{2}\beta (t+\Delta t)\Delta t\mathbf{x}(t)+\sqrt{\beta (t+\Delta t)\Delta t}\mathbf{z}(t)\\ & \approx \mathbf{x}(t)-\frac{1}{2}\beta (t)\Delta t\mathbf{x}(t)+\sqrt{\beta (t)\Delta t}\mathbf{z}(t),\end{array}$$

当$\Delta t\ll 1$时，近似等式成立。在$\Delta t\to 0$的极限下，它收敛于以下SDE：
$$\mathrm{d}\mathbf{x}=-\frac{1}{2}\beta (t)\mathbf{x}\mathrm{d}t+\sqrt{\beta (t)}\mathrm{d}\mathbf{w}$$

整合联系

DDPM与SMLD的联系

DDPM中我们希望找到一个${ϵ}_{\theta }(x_t,t)$去近似$\frac{x_t-{\bar{\alpha }}_t{x}_0}{{\bar{\beta }}_t}$，而SMLD中则是希望找到一个$s_{\theta }(x,t)$来估计${\nabla }_x\log p(x)$

而在这两个模型中，${\mathbf{x}}_t$的条件概率都可以写成$p({\mathbf{x}}_t|{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_t;{\bar{\alpha }}_t{\mathbf{x}}_0,{\bar{\beta }}_t^2\mathbf{I})$（DDPM要求${\bar{\alpha }}_t^2+{\bar{\beta }}_t^2=1$；而SMLD要求${\bar{\alpha }}_t=1$）

于是可以得到${\nabla }_{{\mathbf{x}}_t}\log p({\mathbf{x}}_t|{\mathbf{x}}_0)=-\frac{{\mathbf{x}}_t-{\bar{\alpha }}_t{\mathbf{x}}_0}{{\bar{\beta }}_t^2}=-\frac{\epsilon }{{\bar{\beta }}_t}$，令$s_{\theta }(x_t,t)=-\frac{{ϵ}_{\theta }(x_t,t)}{{\bar{\beta }}_t}$，此时SMLD的损失函数可以写为：

$$\frac{1}{{\bar{\beta }}_t^2}{\mathbb{E}}_{{\mathbf{x}}_0\sim \tilde{p}({\mathbf{x}}_0),\mathbf{\epsilon }\sim \mathcal{N}(0,\mathbf{I})}\left[{\Vert {\mathbf{\epsilon }}_{\theta }({\bar{\alpha }}_t{\mathbf{x}}_0+{\bar{\beta }}_t\mathbf{\epsilon },t)-\mathbf{\epsilon }\Vert }^2\right]$$

忽略系数后就是DDPM的损失函数，这说明二者实质上是相通的，也即VE和VP-SDE是相通的。由此即将DDPM、SMLD和SDE联系了起来。

SDE与ODE的演变

从数学上也可以证明存在不同方差的前向过程，它们产生的边缘分布$p(x_t)$是一样的，从而可以推出对于所有扩散过程，都存在一个相应的确定性过程，其轨迹与SDE具有相同的边缘概率密度，${\sigma }_t=0$此时SDE退化为常微分方程ODE：
$$\mathrm{d}\mathbf{x}=[\mathbf{f}(\mathbf{x},t)-\frac{1}{2}g(t{)}^2{\nabla }_{\mathbf{x}}\log p_t(\mathbf{x})]\mathrm{d}t$$

这个ODE称为概率流ODE（Probability flow ODE），此时扩散过程是确定过程的变换，因而可以基于这个式子做精确的似然计算、在隐空间编码并编辑等任务。同时，对于反向传播$x_T$到$x_0$的过程也可以采用一些ODE算法去加速。

总结

扩散模型通过定义一个不断加噪声的前向过程来将图片逐步变为高斯噪声，再通过定义了一个逆向过程将高斯噪声逐步去噪变为清晰图片以得到采样。

在采样过程中，根据是否添加额外的噪声，可以将扩散模型分为两类：

• 一类是扩散随机微分方程模型（Diffusion SDE） ，其反向过程可以建模为$dx=[f(x,t)-g^2(t)s_{\theta }(x(t),t)]dt+g(t)d\bar{\mathbf{w}}$
• 另一类是扩散常微分方程（Diffusion ODE），其反向过程可以建模为$dx=[f(x,t)-\frac{1}{2}g(t{)}^2{s}_{\theta }(x(t),t)]\mathrm{d}t$