【论文精读】Score-Based Generative Modeling Through Stochastic Differential Equations 通过随机微分方程进行基于分数的生成建模

系列文章目录

【论文精读】DDPM：Denoising Diffusion Probabilistic Models 去噪扩散概率模型

---

---

一、生成模型思路概述

（一）生成模型概述

• 生成模型的目标是希望训练一个神经网络来表征概率分布，从而能够通过其实现采样生成。
• 主流的生成式模型主要可以归纳为两种不同的模式：
  • 隐式(implicit)生成模型：对数据的采样过程进行建模，不从数据分布的概率密度角度出发，而是通过其它方法达到表示概率分布的目的，例如GAN。
  • 显式(explicit)生成模型：通过（近似）最大似然对概率密度进行建模，也被称为基于似然的模型。典型的基于似然的模型包括自回归模型、归一化流动模型、基于能量的模型（EBM） 和变分自编码器（VAE）。

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

（二）扩散模型的三种生成范式

扩散模型（Diffusion Model）是基于似然估计的生成式模型，其首先向数据分布中逐步添加随机噪声到预设的先验分布，然后通过学习其逆过程来重建新的数据样本。通常而言，扩散模型具有三种生成范式：

• Score Matching with Langevin Dynamics (SMLD),
• Denoising Diffusion Probabilistic Model (DDPM)
• Score-based Generative Model (SGM)

二、分数与分数网络

（一）分数函数score

在一般基于似然函数的生成模型中，模型学习的都是数据真实分布的似然$log{p}_{data}(x)$，但是能量模型中为了使得学习的符合概率函数的定义，需要除以一个归一化常数，即对于模型学习的$f(\theta ,x)$。最终的概率函数为$p(x)=\frac{f(\theta ,x)}{z}$，由于z的计算很困难，所以作者提出同时对两边做对数在求导，常数项为0，即$\nabla logp(x)=\nabla logf(\theta ,x)$。

分数函数：通过对分数函数而不是密度函数建模，可以避免难以处理的常数归一化的困难。对于分布为$p(x)$的概率密度函数，其分数的定义为对数概率密度函数对于输入数据的梯度，表达式如下
$$\frac{\partial (logp(x))}{\partial x}={\nabla }_{x}logp(x)$$

（二）分数网络

分数网络：用于对分数建模，通常用$s_{\theta }$表示，输入为D维，输出也为D维。
$$s_{\theta }\approx {\nabla }_{x}logp(x)$$

对分数网络进行训练的目的是使得分数网络的输出逼近真实分布$p_{data}(x)$的分数，因此模型要学习的就是这个分数函数，对应的损失函数为：
$$E[||s_{\theta }(x)-{\nabla }_{x}log{p}_{data}(x)|{|}_2^2]$$

（三）对分数的理解

一般来说，实际场景下的数据往往是多维的。由分数的定义以及从数学的角度出发来看，它应当是一个“矢量场(vector field)”。既然是矢量，那么就有方向，这个方向就是：对于输入数据(样本)来说，其对数概率密度增长最快的方向，其可以用矢量场分布图表示：

三、基于分数的生成式模型

基于分数的生成式模型需要的两个核心要素：分数匹配和朗之万动力

（一）朗之万动力

Langevin dynamics 本指朗之万动力学方程，它是描述物理学中布朗运动(悬浮在液体/气体中的微小颗粒所做的无规则运动)的微分方程，借鉴到这里作为一种生成样本的方法。

朗之万动力可以通过一组迭代的算法从一个分布中进行采样。这个算法依赖于分布的分数。

• 首先从先验分布随机采样一个初始样本；
• 然后利用模型估计出来的分数逐渐将样本向数据分布的高概率密度区域靠近；
• 为保证生成结果的多样性，我们需要采样过程带有随机性。 正好！布朗运动就是带有随机性的，朗之万动力学方程因此也带有随机项。

朗之万动力学的公式表达：

• 假设初始数据满足先验分布$x_0$~$\pi (x)$，就有以下的迭代过程
• 其中$z_i$~$N(0,I)$
• 当分数匹配方法进行训练后，就可以有$s_{\theta }\approx {\nabla }_{x}logp(x)$
• 理论上，当$K\to \infty$且$\epsilon \to 0$时，最终生成的样本就会服从原数据分布。
  

（二）分数匹配

对分数网络进行训练的目的是使得分数网络的输出逼近真实分布的分数${\nabla }_{x}logp(x)$，从而就不需要去直接学习真实分布$p_{data}(x)$。分数匹配就是使得模型输出的分数能够匹配到真实分布的分数，因此可以通过最小化模型和数据分布之间的Fisher散度（简记为F-散度）来训练基于分数的模型，其表达式为：

使用分部积分方法可以将上面的损失函数转化为：（其中${\nabla }_x{s}_{\theta }(x)$表示分数网络的Jacobian矩阵，tr表示对Jacobian矩阵求迹，即主对角元素之和。）

当遇到高维数据(图像特征图的像素数或通道数常都成百上千)时，Jacobian矩阵在计算上十分复杂，需要经历多次(维数多少就需要多少次)反向传播来求偏微分。针对这个问题，原文中提到了两种优化的方法来解决，即去噪分数匹配和切片分数匹配。

1、去噪分数匹配（denoising score matching）

前提条件：噪声非常小
设原始数据为$x$，对原始数据（训练样本）加入噪声后记为$\tilde{x}$（通过预先设定的分布$q_{\sigma }(\tilde{x}|x)$~$N(\tilde{x};x,{\sigma }^{2}I)$来扰动原始数据，方差${\sigma }^2$是提前规定的）。那么加噪声后的数据分布为：

那么此时的分数计算为：
对应的训练目标函数即为：

在去噪分数匹配方法下，只需要从标准高斯分布$N(0,I)$中采样随机高斯噪声$\epsilon$，然后通过重参数化方法$\tilde{x}=x\sigma +\epsilon$就可以得到加入噪声后的样本。加入噪声后的数据分布是满足$q_{\sigma }(\tilde{x}|x)$~$N(\tilde{x};x,{\sigma }^{2}I)$分布的。在这种方法下，训练的网络估计的是$q_{\sigma }(\tilde{x}）$而不是原始数据分布，那么基于这样的情形就要求加入的噪声预定义方差$\sigma$足够小，否则，如果对原始数据扰动过大，模型训练完后生成的样本就会严重偏离原数据分布。

2、切片分数匹配

切片分数匹配的核心思想是用一个服从简单分布的随机变量$v$将网络分数${\nabla }_x{s}_{\theta }(x)$随机映射为一个“Jacobian”向量，即：

则此时的训练损失函数为：

切片分数匹配的损失函数比起正常的网络训练过程来说，是多了一次反向传播的。因为上述的反向传播是为计算训练目标而进行的，而计算出来后还要进行一次反向传播得到的才是模型基于损失而获得的梯度。

（三）基于分数的生成模型的缺陷

数据密度较低的区域，分数估计会不准确，进而郎之万采样也会不准确。

四、噪声条件分数网络（NCSN）

（一）主要思想

核心思想：利用噪声对原始未扰动数据分布中低密度区域进行填充

两难之处：

• 强度过大的噪声会干扰到原始数据的分布进而造成估计分数的误差增大
• 噪声强度小能够获得与原数据分布较为近似的效果，但是却不能够很好地“填充”低概率密度区域

平衡方案：加不同量级的噪声对数据进行扰动，用唯一独立的一个条件分数网络来预测不同噪声下分布的分数，然后采用郎之万采样办法来生成样本。首先去生成噪声加的比较大的情况下分布的分数，然后逐渐减小加噪的量级，最终噪声减小到0，然后再用郎之万采样方法得到符合正确的原始分布的分数。

（二）NCSN的定义

• 噪声：定义一组几何级数序列 { σ i } i = 1 L \{\sigma_i\}_{i=1}^L {σi​}i=1L​，其中${\sigma }_i>0$且$\frac{{\sigma }_1}{{\sigma }_2}=\frac{{\sigma }_2}{{\sigma }_3}=...=\frac{{\sigma }_{L-1}}{{\sigma }_L}>1$ ，这样就可以使得最开始的噪声水平足够大以能够充分“填充"低概率密度区域；而最后的噪声水平足够小以获得对原数据分布良好的近似，避免过度扰动。
• 扰动噪声数据分布：$q_{\sigma }(x)=\int p_{data}(x)N(x|t,{\sigma }^{2}I)dt$
• NCSN的目标是训练一个条件分数网络来估计扰动数据的分布，即$s_{\theta }(x,\sigma )\approx {\nabla }_{x}log{q}_{\sigma }(x)$，一般把$s_{\theta }(x;\sigma )$成为噪声条件分数网络。

NCSN和DSM（denoising score matching）不同，DSM是无条件的分数网络，每次加少量的噪声，而NCSN是以噪声为条件。

（三）NCSN的网络结构

由于生成的图片和原图一样大小，也就是每个像素点都需要由朗之万动力学采样生成，因此模型对于每个像素点都要估计其对应的分数。也就是说，从张量的角度来看，网络的输出(分数)要和输入图像的形状(shape)一致。结合这种输入输出特征图的特征，U-Net恰好可以实现这种任务。

NCSN采用的噪声条件U-Net结构。对U-Net结构进行了如下优化：

• 在其中加入了空洞卷积(dilated/atrous convolution)和以噪声为条件的实例归一化(conditional instance normalization++) 。
• 另外，即使对于同一个像素点，在不同的噪声强度下也要对应估计出不同的分数，于是，模型还需要将噪声强度也作为输入。

（四）NCSN网络的训练

对于训练方法，作者采用了去噪分数匹配。当然，前面提到的切片分数匹配也可以训练 NCSN，只不过需要的计算步骤会多一些，因此训练过程会更慢。

分数匹配的目标函数为：
作者将加入噪声扰动之后的数据分布建模为$q_{\sigma }(\tilde{x}|x)=N(\tilde{x}|x,{\sigma }^{2}I)$，即概率密度函数可以写为：

那么在噪声水平为${\sigma }_i$的损失函数即可表示为：

而 NCSN 使用了多个噪声级别，于是，分别对它们的损失加权求和后再求均值，就得到了联合的损失函数：（$\lambda ({\sigma }_i)={\sigma }_i^2$为噪声权重）

噪声权重的选择：当$\lambda ({\sigma }_i)={\sigma }_i^2$时，有：
由于$q_{\sigma }(\tilde{x}|x)=N(\tilde{x}|x,{\sigma }^{2}I)$，所以$\tilde{x}=x\sigma +\epsilon$，因此$\frac{\tilde{x}-x}{\sigma }=\epsilon$是噪声，也就不会受到${\sigma }_i$的影响，并且当权重$\lambda ({\sigma }_i)={\sigma }_i^2$时，$||{\sigma }_i{s}_{\theta }(\tilde{x},{\sigma }_i)|{|}_2$的量级为1，因此能保证模型在训练时去除噪声条件的影响。

（五）退火的朗之万动力学方法生成样本

• 首先，从一个随机分布中采样得到初始样本$\tilde{x}$；
• 然后设置步长${\alpha }_i=\epsilon \frac{{\sigma }_i^2}{{\sigma }_L^2}$；
• 从最大的噪声级别${\sigma }_i$开始朗之万动力学采样直至最小的噪声级别KaTeX parse error: Undefined control sequence: \sigam at position 1: \̲s̲i̲g̲a̲m̲_L；
• 接着，在每个噪声级别下，一共进行步数为$T$步的迭代朗之万动力学采样
  
• 该噪声级别最后一步采样生成的样本会作为下一个噪声级别的初始样本；
• 最后，待所有噪声级别的朗之万动力学采样过程均完成时，就得到了最终的生成结果。

五、用SDE（随机微分方程）描述扩散模型

引入 SDE 形式来描述扩散模型的本质好处是“将理论分析和代码实现分离开来”，我们可以借助连续性 SDE 的数学工具对它做分析，而实践的时候，则只需要用任意适当的离散化方案对 SDE 进行数值计算。

（一）用SDE描述扩散过程

之前我们用的是有限次数的加噪。现在将噪声量级推广到无穷次，那么就可以得到更加一般的扩散过程，并且这个过程可以用一个SDE来表示，求解也更加方便。

如果$x[t]$是一个连续量（加噪变成无数次），那么扩散过程可以用一个SDE去表达;布朗运动具有增量独立性、增量服从高斯分布、轨迹连续;随机微分方程指的是微分方程中含有随机参数或随机过程或随机初始值或随机边界值，这里的w随机性使得SDE成立)：
$$dx=f(x,t)dt+g(t)dw$$

其中，$f$是漂移系数，$g$是扩散系数，$w$是随机的标准的布朗运动的量。

它可以看成是下述离散形式在$\nabla t\to 0$时的极限：

（二）SDE的逆过程用于样本采样

SDE的逆过程：从先验分布$p_T$不断的通过逆过程获得$p_0$数据分布的样本。

逆向的 SDE 中如果知道了${\nabla }_{x}log{p}_t(x)$ ，那么就可以通过下面的离散化格式完成重建。

（三）通过损失函数进行训练 （得分匹配）

通过一系列理论推导可以得到${\nabla }_{x}log{p}_t(x)$ ，但是它对计算能力的要求过大，而且泛化性也不够好（存在一定的条件才能推导出来）。

因此我们需要用神经网络学一个函数$s_{\theta }(x,t)$，使得它能够直接计算${\nabla }_{x}log{p}_t(x)$ ，这个网络的损失函数即为：（就是“（条件）得分匹配”的损失函数）

参考：
Score Diffusion Model分数扩散模型理论与完整PyTorch代码详细解读
基于分数的生成模型（Score-based Generative Model ）
图像生成别只知道扩散模型，还有基于梯度去噪的分数模型：NCSN
生成扩散模型漫谈：一般框架之SDE篇