4. 基于分数的生成模型(Score-based generative models)

通过前面的学习,我们发现扩散模型可以从不同的角度进行解释。 其中一个等价的解释是基于分数的生成模型,前面章节虽然简单介绍了下, 但没有详细说明,本章我们详细讨论下基于分数的生成模型。 基于分数的生成模型是由宋旸 [] 等人在2019年提出 [1] 的, 后来他们又提出了基于随机微分方程的更一般的形式 [2], 本章我们一起讨论学习一下。

4.1. 基于分数的生成模型

在前面 DDPM 和 DDIM 的章节中,已经探讨了 DDPM 的降噪过程,可以看做是沿着分数(梯度) \(\nabla\log p(x_t)\) 前进。然而宋旸等人提出基于分数的论文 [] 相关工作并不是建立在 DDPM 的基础上, 所以论文里不是从DDPM的马尔科夫链式结构讨论和导出的,而是直接从分数匹配估计算法推导。

首先看下基于分数的生成模型的核心思想。 假设我们有一些服从某个分布的观察数据,比如大量的图片数据, 但是这些数据背后真实的概率分布是未知的, 我们暂且用符号 \(p_{\text{data}}(x)\) 表示这些观测数据的真实分布。 我们希望能生成新的数据,也就是想从分布 \(p_{\text{data}}(x)\) 中采样新的样本,比如采样生成新的图片。 但分布 \(p_{\text{data}}(x)\) 是未知的,怎么办呢? 此时有一种利用分数(score)进行采样的方法,虽然我们不知道 \(p_{\text{data}}(x)\) 的具体形式, 但如果我们能得到它的分数(基于数据变量 \(x\) 的一阶偏导) \(\nabla_{x} \log p_{\text{data}}(x)\), 那么就可以利用它的分数 \(\nabla_{x}\log p_{\text{data}}(x)\)\(p_{\text{data}}(x)\) 中随机采样。 这样基于分数的采样方法有很多,可以从中选择一个合适的采样方法。 所以基于分数的生成模型,宏观来看就两步:

  1. 利用分数匹配方法估计出数据分布 \(p_{\text{data}}(x)\) 的近似分数 \(s_{\theta}(x) \approx \nabla_{x} \log p_{\text{data}}(x)\)

  2. 使用某个基于分数 \(s_{\theta}(x)\) 的采样算法,随机采样近似数据分布 \(p_{\text{data}}(x)\) 的样本。

4.1.1. 分数匹配算法(Score Matching)

分数匹配(Score Matching),简单来说它就是一种概率密度估计的方法。 该方法在多元高斯和独立分量分析模型上得到了证明 [3]。 在概率统计学中,很多概率密度函数可以写成如下的形式,比如无向图中的概率密度,通过贝叶斯公式推导出的后验概率分布等等。

(4.1.51)\[p(x;\theta) = \frac{q(x;\theta)}{Z}\]

其中 \(\theta\) 是未知参数, \(x\) 是我们感兴趣的目标随机变量,可以是一个多元向量。 \(Z\) 作为分母,它起到归一化的作用,使得式子 \(q(x;\theta)/Z\) 的积分为 \(1\) ,进而满足概率密度的约束。 多数情况下,尤其是使用了贝叶斯定理的情况,\(Z\) 本身就是分子的积分

(4.1.52)\[Z = \int q(x;\theta) dx\]

然而很多情况下这个积分是难以计算的,甚至无法计算的,这就导致我们无法得到概率密度 \(p(x;\theta)\) 的准确形式, 虽然用最大似然估计能帮我们估计出未知参数 \(\theta\),但没办法得到 \(p(x;\theta)\) 的完整形式, 也就不能从 \(p(x;\theta)\) 进行样本生成等等工作。 而分数匹配就是帮我们估计出 \(p(x;\theta)\) 的一个近似表示,我们用这个近似表示来代替原本的 \(p(x;\theta)\)

假设我们有变量 \(x\) 的一些观测样本,我们用 \(p_{\text{data}}(x)\) 表示这些数据的真实分布 ,然后定义一个近似的参数化分布(训练的模型)表示为 \(p_{\theta}(x)\), 分数匹配的目标函数是

(4.1.53)\[ \begin{align}\begin{aligned}J(\theta) &= \frac{1}{2} \int p_{\text{data}}(x) \left\lVert \nabla \log p_{\text{data}}(x) - \nabla \log p_{\theta}(x) \right\rVert^2 d x\\&=\frac{1}{2} \EE[p_{\text{data}}(x)]{ \left\lVert \nabla \log p_{\text{data}}(x) - \nabla \log p_{\theta}(x) \right\rVert^2 }\end{aligned}\end{align} \]

为了符号简洁,分别定义 \(s_{\theta}(x)=\nabla \log p_{\theta}(x)\)\(s_{data}(x)=\nabla \log p_{\text{data}}(x)\), 注意这里的梯度都是对变量 \(x\) 的 ,不是参数 \(\theta\),别搞混了。

(4.1.54)\[ \begin{align}\begin{aligned}J(\theta) &= \frac{1}{2} \int p_{\text{data}}(x) \left\lVert s_{data}(x) - s_{\theta}(x) \right\rVert^2_2 d x\\&= \frac{1}{2} \EE[p_{\text{data}}(x)]{ \left\lVert s_{data}(x) - s_{\theta}(x) \right\rVert^2_2 }\end{aligned}\end{align} \]

通过极小化目标函数求得参数化分布的未知参数

(4.1.55)\[\hat{\theta} = \operatorname*{\arg \min}_{\theta} \ J(\theta)\]

单纯的看这个目标函数似乎很简单,他就是一个均方误差而已,它的目标是令近似分布的对数概率密度梯度和真实数据分布的对数概率密度的梯度尽量相近, 这其实建立在:如果两个分布的对数密度梯度是一样的,那么他们的概率密度也是一样的。 这个目标函数形式很简单,但是有个问题,就是 \(s_{data}(x)\) 是不知道的,是无法计算的。

然后一顿操作猛如虎,得到了目标函数的一个等价表示

(4.1.56)\[ \begin{align}\begin{aligned}J(\theta) &= \int p_{\text{data}}(x) \left [ tr( \nabla_{x} s_{\theta}(x)) + \frac{1}{2} \left\lVert s_{\theta}(x) \right\rVert^2_2 \right ] d x\\&= \EE[p_{\text{data}}(x)] { tr( \nabla_{x} s_{\theta}(x)) + \frac{1}{2} \left\lVert s_{\theta}(x) \right\rVert^2_2 }\end{aligned}\end{align} \]

其中 \(\nabla_{x} s_{\theta}(x)\)\(s_{\theta}(x)\) 的一阶偏导, 是 \(\log p_{\theta}(x)\) 的二阶偏导,因为是二阶偏导,所以他是一个方阵, 也被称为 Hessian 矩阵, \(tr(\cdot)\) 表示对角线元素。\(s_{\theta}(x)\) 是我们的近似分布(拟合模型)的 一阶偏导,它是可以计算的。在这个等价表示中,没有了数据分布的梯度,只剩下了可以计算的拟合模型的梯度。 从 公式(4.1.55)公式(4.1.56) 的推导过程这里不再赘述, 详细过程可以参考论文 []

然而事情还没有完,二阶偏导 \(tr( \nabla_{x} s_{\theta}(x))\) 虽然可以计算, 但是计算成本时非常高的,尤其是在变量 \(x\) 很高维或者神经网络层次很深的时候, 通常是无法接受的。针对这个情况,一般有两种解决方法, 分层分数匹配(Sliced score matching)和降噪分数匹配(Denoising score matching)。

分层分数匹配(Sliced score matching)

分层匹配使用一个随机投影矩阵可以 近似 计算 \(tr( \nabla_{x} s_{\theta}(x))\) ,改变后的目标函数为

(4.1.57)\[ \mathbb{E}_{p_v} \EE[p_{\text{data}}(x)] { v^T \nabla_{x} s_{\theta}(x) v + \frac{1}{2} \left\lVert s_{\theta}(x) \right\rVert^2_2 }\]

其中 \(p_v\) 是一个简单的随机向量即可,比如多元正态分布。 其中 \(v^T \nabla_{x} s_{\theta}(x) v\) 可以直接利用正向模式的自动微分计算, 但是仍然要四倍的计算量。

降噪分数匹配(Denoising score matching)

另一种解决分数匹配的方法是降噪分数匹配(Denoising score matching),它是分数匹配算法的一个变种, 它可以完全避开 \(tr( \nabla_{x} s_{\theta}(x))\) 的计算。 首先在观测数据 \(x\) 上添加一些预先设定好的噪声数据,得到了新的数据 \(\tilde{x}\) ,这相当于构建了一条件概率分布 \(q_{\sigma}(\tilde{x}|x)\), 根据边际化方法,边缘分布 \(q_{\sigma}(\tilde{x})\) 的计算方法为

(4.1.58)\[q_{\sigma}(\tilde{x}) \triangleq \int q_{\sigma}(\tilde{x}|x) p_{\text{data}}(x) d x\]

然后把分数匹配算法应用在这个加噪后的数据分布上,

(4.1.59)\[ \begin{align}\begin{aligned}& \frac{1}{2} \mathbb{E}_{q_{\sigma}(\tilde{x}|x) } \left [ \left\lVert s_{\theta}(\tilde{x}) -\nabla_{\tilde{x}} \log q_{\sigma}(\tilde{x}) \right\rVert^2_2 \right ]\\=& \frac{1}{2} \mathbb{E}_{q_{\sigma}(\tilde{x}|x) p_{\text{data}}(x) } \left [ \left\lVert s_{\theta}(\tilde{x}) -\nabla_{\tilde{x}} \log q_{\sigma}(\tilde{x}|x) \right\rVert^2_2 \right ]\end{aligned}\end{align} \]

这么做的一个前提是,如果添加的噪声足够小,那么 \(q_{\sigma}(\tilde{x}) \approx p_{\text{data}}(x)\) 成立,此时有 \(\nabla_{\tilde{x}} \log q_{\sigma}(\tilde{x}) \approx \nabla_{x} \log p_{\text{data}}(x)\) 成立,这时我们可以用分数匹配算法估计出 \(q_{\sigma}(\tilde{x})\) 的分数 \(\nabla_{\tilde{x}} \log q_{\sigma}(\tilde{x})\) ,并用它近似表示原数据分布 \(p_{\text{data}}(x)\) 的分数。

4.1.2. 基于分数的生成模型面临的困难

理想是美好的,现实是残酷的。虽然我们想到了这种基于分数的采样模型, 然而在实际应用时需要面临一些困难,论文中主要提出了两类困难:

  1. 分数匹配算法估计分数的准确性问题。

  2. 基于分数的采样算法的准确性问题。

4.1.2.1. 分数估计不准的问题

论文中一共提出两个原因会导致分数估计不准确,一个是流体假设问题,一个是低密度区域样本不足的问题。

流形假设问题

首先看流体假设问题,论文指出了流形假设的存在,即现实世界中的数据往往集中在嵌入高维空间的低维流形上。 这句话对数学不好的同学来说跟天书一样,说人话就是,你观测的数据 \(x\)\(n\) 维, 你以为它真的就是 \(n\) 维吗?现实是:多数情况下,数据中含有信息的真正维度往往小于 \(n\) 维。 比如一张 \(4096 \times 2160\) 的 4K 高清图像,它的像素点总数是 \(4096 \times 2160=8847360\) ,你以为这 \(880w\) 个像素点都是独立有意义的么?未必,我把它压缩到 \(1920 \times 1080=1080P\) 后,图像的关键信息可能并没有丢失。也就是说原来 4K 图像中很多像素点其实可有可无,并没有包含有意义的信息。

这里我换个容易理解的说法,从小学数学说起。 在学线性方程组时,我们知道方程组的解有无解、唯一解、无穷解等等情况。 假设有一个包含 \(n\) 个方程和 \(m\) 个未知参数的方程组, 有 \(m\) 个参数就有 \(m\) 对应的系数。 但是你的方程与方程之间,系数与系数之间可能存在着线性相关, 即某个方程可以通过其它方程线性变化得到,某个系数可以通过其它系数线性变换得到, 这样的方程称为无效方程,这样的系数对应的参数称为无效参数或者自由参数。 假设方程组中有意义的方程为 \(\tilde{n}\) 个,即方程组中独立无关的方程有 \(\tilde{n}\) 个; 相互独立无关的系数有 \(\tilde{m}\) 个。 有两种情况会导致方程组有无穷解,无穷解意为着不能确定未知参数的值。

  1. \(\tilde{n}<\tilde{m}\) 时,意为着有意义方程数量不足以确定全部参数的值,此时就方程组会有无穷解。 这种情况,就相当于在机器学习场景中,有意义的观测样本(方程)数量少于观测变量(系数)的数量。

  2. \(\tilde{n}<n\) 时,此时意为着有部分系数之间不是相互独立的,这些系数对应的参数就成了自由参数,它的值可以是任意的,此时就方程组也会有无穷解。 这种情况,就相当于在机器学习场景中,你的特征变量(观测变量)不是相互独立的,存在线性相关。

进入到线性代数,线性方程组可以用形如 \(A\theta=b\) 的矩阵表达, 其中 \(A\) 是系数矩阵,对应着观测变量的观测值矩阵,或者说特征矩阵,\(\theta\) 是未知参数组成的向量, \(b\) 是方程组等号右侧值组成的向量,对应着机器学习场景下的 Label 值向量。 矩阵 \(A\) 行的数量对应着观测样本的数量,列数对应着观测变量的维数(或者是个数), 列数和 \(\theta\) 向量长度(未知参数数量)是一样的。

如果矩阵 \(A\) 列满秩,意味着观测变量(系数)之间是相互独立的,每一个观测变量都含有有意义的信息; 反之,如果观测变量(系数)之间是部分存在线性关系(或者一一映射),则矩阵 \(A\) 列不满秩, 矩阵 \(A\) 的列秩就等于有意义观测变量的数量。

回到流形假设问题,流行假设描述的是 类似列不满秩的情况 ,就是说我们观测到的数据维度是 \(m\) ,但很可能这 \(m\) 个变量(或者说维度)并不是都有意义,其中某些变量(维度)可能没有意义, 可以由其它变量组合变换的到,这样的变量本身是不包含任何有价值信息的。 反映到分数估计算法上,这些变量的分数(偏导数)理论上有任意解,不稳定,呈现出来的现象就是训练时 Loss 不稳定不收敛,会反复横跳 ,具体可以看原论文里的实验和图解。

../_images/score_001.png

图 4.1.1 左图流体假设导致的LOSS不收敛,右图添加了噪声后的 LOSS情况 (图片来自 []

低密度的问题

另一个影响分数估计准确性的是概率密度的低密度问题,这个问题其实很简单, 本质上就是观测样本不足的问题。 我们知道,对于一个非均匀分布的概率分布来说,其概率密度是不均匀的,概率密度低的地方意为着产生概率低, 那理论上获得这些区域的观测样本的概率就小。也就说,大多数情况下,对于低概率密度的区域, 我们观测样本覆盖是不足的,这时自然分数估计算法这些区域的拟合和学习就不足, 显然它对这些区域的预测也将不准。

4.1.2.2. 郎之万动力采样不准的问题

郎之万动力采样(Langevin dynamics sample)算法,是一种利用分数从目标概率分布采样的方法, 假设目标概率分布是 \(p(x)\), 它的分数表示为 \(\nabla_{x} \log p(x)\) ,给定一个固定的步进 \(\epsilon >0\) ,以及一个初始样本 \(\tilde{x} \sim \pi(x)\)\(\pi(x)\) 可以认为是一个先验分布, 郎之万动力采样通过如下迭代方程得到一个目标分布 \(p(x)\) 的采样。

(4.1.60)\[\tilde{x}_{t} = \tilde{x}_{t-1} + \frac{\epsilon}{2} \nabla_{x} \log p( \tilde{x}_{t-1}) + \sqrt{\epsilon} z_t , \quad z_t \sim \mathcal{0,\textit{I}}\]

当满足 \(\epsilon \rightarrow 0\)\(T \rightarrow \infty\) 和一些正则性条件时, \(\tilde{x}_{T}\) 就服从分布 \(p(x)\), 此时 \(\tilde{x}_{T}\) 可以看成目标分布 \(p(x)\) 的一个采样。 当然如果 \(\epsilon > 0\) 并且 \(T < \infty\) 采样就不是很准,但实际应用中可以忽略这些误差,勉强能用, 我们要做的就是让 \(\epsilon\) 尽量小,让 \(T\) 尽量大。

郎之万动力采样的优势就是它只需要有目标分布的分数就行了,不需要知道目标分布 \(p(x)\) 具体形式。 仔细观察下这个公式,这不就是一个梯度迭代法么,让 \(x_t\) 逐步沿着 \(p(x)\) 的梯度向着 \(p(x)\) 概率密度最大的点前进, 迭代过程中加入一个随机高斯噪声 \(z_t\),使其具备随机性。万变不离其宗,搞来搞去就这么点事!!!

小技巧

既然就是个梯度迭代,那我加入二阶梯度是不是采样能更快点?

然而郎之万动力采样存在一个显著的缺陷, 当数据分布是一个很复杂的分布时,比如存在低密度区域,并且低密度区域把整个概率密度空间分割成多个区域时, 郎之万动力采样法无法在合理的时间能得到正确的采样,最典型的例子就是高斯混合分布, 原论文就是用两个分量的混合高斯分布做实验和举例, 论证了在这种情况下郎之万动力采样法无法健康工作。 当然如果 \(T\) 足够大还是可以合理采样的, 但实际应用中,没办法令 \(T\) 足够大,那样的话迭代次数太多,效率太慢了。

4.1.3. 通过加噪的方法估计分布的近似分数

通过上一节,我们知道了基于分数的生成模型,面临两个困难, 1)高维空间有效性问题,有意义信息的维度可能远远小于观测数据的维度,导致分数估计算法不收敛。 2)数据分布低密度区域问题,低密度区域一方面因为样本不足令分数匹配算法估计不准,另一方面使郎之万采样算法无法在可接受的步数内得到有效采样。 那怎么解决呢?你猜对了,通过对数据添加高斯噪声解决。

为什么添加噪声能解决上述问题呢?思考下,原数据分数是 \(p_{data}(x)\), 再它基础上添加一些高斯噪声,这个过程对应条件概率分布 \(q_\sigma(\tilde{x}|x)\), 添加噪声后得到新的数据 \(\tilde{x}\) 的边缘分布是 \(q_\sigma(\tilde{x}) = \int q_\sigma(\tilde{x}|x) p_{data}(x) dx\), 添加噪声就相当于改变了原来的数据分布,这种改变影响两个方面:

  1. 会破坏掉原来数据变量 \(x\) 各个维度(分量)之间的相关性,使得各个维度(分量)间相关性减弱,相当于 \(x\) 变得满秩了。

  2. 会改变 \(p_{data}(x)\) 的密度,各分量之间添加的噪声是同等权重的,低密度区域会密度变高,相对来说整个密度空间变得均匀了。

添加的噪声越多,对上述两点的影响越大,对问题的改善就越好。这种加噪声的方法又完美契合了上面介绍的 分数匹配两种方法之一的降噪分数匹配,就是这么的巧合。 但是,但是,这里有冲突!!!回看一下降噪分数匹配算法,它成立的条件是添加的噪声不能太多啊, 不然 \(q_\sigma(\tilde{x})\) 就离 \(p_{data}(x)\) 太远了, 就不能用 \(q_\sigma(\tilde{x})\) 的分数近似 \(p_{data}(x)\) 的分数了。 而这里有需要足够的大噪声去改善上述两个问题,噪声不够大,这两个问题解决不彻底啊!!!

那么怎么解决?既然有的人贪心要的多,有的人佛系要的少,重口难调。 那么就按需分配,设置不同强度等级噪声,各种强度等级的噪声都安排上,总有一款适合你。

\(\{\sigma_i\}_{i=1}^L\) 是一个满足 \(\frac{\sigma_1}{\sigma_2}=\cdots=\frac{\sigma_{L-1}}{\sigma_{L}}\) 的正几何序列。 令条件分布 \(q_{\sigma_i}(\tilde{x}|x)\) 是一个高斯条件分布, 它表示加噪过程,则有

(4.1.61)\[q_{\sigma_i}(\tilde{x}|x) \sim \mathcal{N}( \tilde{x}|x, \sigma_i^2 \textit{I} )\]

\(q_{\sigma_i}(\tilde{x})\) 表示加噪扰动后的数据边缘分布, 则有

(4.1.62)\[q_{\sigma_i}(\tilde{x}) \triangleq \int p_{data}(x) q_{\sigma_i}(\tilde{x}|x) dt\]

和之前一样,可以通过采样法得到 \(\tilde{x}\) 的样本

(4.1.63)\[\tilde{x} = x + \sigma^2_i \epsilon , \quad \epsilon \sim \mathcal{0,\textit{I}}\]

相比于 DDPM ,这里没有定义一个 \(\alpha\) 参数,事实上 \(\sigma^2_i\) 起到了类似的作用。 \(\sigma^2_i\) 从大到小变化, \(\sigma^2_0\) 足够大用来解决上面的问题, \(\sigma^2_L\) 足够小用来提高分数匹配估计的效果。 直观地,高噪声有助于分数函数的估计,但也会导致样本损坏; 而较低的噪声给出了干净的样本,但使得分函数更难估计。

接下来利用上述的分层分数匹配(Sliced score matching)或者降噪分数匹配(Denoising score matching) 学习一个分数估计模型,它能在不同噪声强度下估计出扰动加噪后数据的分数。 在这里分层分数匹配和降噪分数匹配两种方法都可以使用,没有特别的限制。 条件高斯分布 \(q_{\sigma_i}(\tilde{x}| x)\) 的分数为

(4.1.64)\[\nabla_{\tilde{x}} \log q_{\sigma_i}(\tilde{x}| x)= - \frac{\tilde{x}-x}{\sigma_i^2}\]

按照降噪分数匹配的方法,单一噪声强度的目标函数为

(4.1.65)\[ \begin{align}\begin{aligned}\ell(\theta;\sigma_i) & \triangleq \frac{1}{2} \mathbb{E}_{p_{data}(x)} \mathbb{E}_{\tilde{x} \sim \mathcal{x,\sigma_i^2 \textit{I}}} \left [ \left\lVert s_{\theta}(\tilde{x},\sigma_i) -\nabla_{\tilde{x}} \log q_{\sigma}(\tilde{x}|x) \right\rVert^2_2 \right ]\\& = \frac{1}{2} \mathbb{E}_{p_{data}(x)} \mathbb{E}_{\tilde{x} \sim \mathcal{x,\sigma_i^2 \textit{I}}} \left [ \left\lVert s_{\theta}(\tilde{x},\sigma_i) + \frac{\tilde{x}-x}{\sigma_i^2} \right\rVert^2_2 \right ]\end{aligned}\end{align} \]

\(s_{\theta}(\tilde{x},\sigma_i)\) 代表神经网络模型,它的输入是 \(\tilde{x}\)\(\sigma_i\) ,输出是预测的分数。最后把所有噪声等级结合在一起

(4.1.66)\[\mathcal{L}(\theta;\{\sigma_i\}_{i=1}^L ) \triangleq \frac{1}{L} \sum_{i=1}^L \lambda(\sigma_i) \ell(\theta;\sigma_i)\]

这就是最终的目标函数了,其中 \(\lambda(\sigma_i)\) 是一个与 \(\sigma_i\) 有关的权重参数,它的作用是可以对不同噪声等级 \(\sigma_i\) 设置不同的重要性权重。 目标函数中包含一个方差项 \(\sigma_i^2\) ,它是平方,当 \(\sigma_i\) 取不同值时波动比较大,可以认为当 \(\sigma_i\) 取不同值时 公式(4.1.65) 的数量级不同,所以这里可以令 \(\lambda(\sigma_i)=\sigma_i^2\)

(4.1.67)\[ \begin{align}\begin{aligned} \lambda(\sigma_i) \ell(\theta;\sigma_i) &= \sigma_i^2 \frac{1}{2} \mathbb{E}_{p_{data}(x)} \mathbb{E}_{\tilde{x} \sim \mathcal{x,\sigma_i^2 \textit{I}}} \left [ \left\lVert s_{\theta}(\tilde{x},\sigma_i) + \frac{\tilde{x}-x}{\sigma_i^2} \right\rVert^2_2 \right ]\\ &= \frac{1}{2} \mathbb{E}_{p_{data}(x)} \mathbb{E}_{\tilde{x} \sim \mathcal{x,\sigma_i^2 \textit{I}}} \left [ \left\lVert \sigma_i s_{\theta}(\tilde{x},\sigma_i) + \frac{\tilde{x}-x}{\sigma_i} \right\rVert^2_2 \right ]\end{aligned}\end{align} \]

此时有 \(\frac{\tilde{x}-x}{\sigma_i} \sim \mathcal{N}(0,\textit{I})\) ,这样一来不同的噪声等级都有相同的数量级(量纲)。 原论文中把这个模型称为噪声条件分数网络(Noise Conditional Score Networks,NSCN)。

4.1.4. 基于分数的改进采样算法

通过上一步的加噪分数匹配方法,训练出了一个神经网络模型,这个神经网络模型可以预测不同等级的加噪数据的分数 \(s_{\theta}(\tilde{x},\sigma_i)\), 接下来就是如何利用这个分数预测模型来近似采样生成原始数据分布 \(p_{data}(x)\) 的样本。 基于分数的采样方法其实有多种,作者这里采样的郎之万动力采样法(Langevin dynamic sample), 然而前文我们讨论过,郎之万动力采样法存在着不足,对于那些存在低密度区域分割成多个高密度区域的复杂分布,需要较多的采样步骤才能得到相对可靠的采样结果, 无法在一个可接受的步骤内得到较好的采样结果,针对这个问题,作者提出了一个改进的郎之万动力采样法,称为退火郎之万动力采样法 (annealed Langevin dynamics),算法的伪代码如图 4.1.2 所示。

../_images/score_annealed_Langevin_dynamics.png

图 4.1.2 退火郎之万动力采样法伪代码过程 (图片来自 []

算法的过程其实不复杂,首先初始化超参数 \(\{\sigma_i\}_{i=1}^L\) ,然后初始 \(\tilde{x}_0\),它可以是一个均匀分布的随机采样,也可以是高斯分布的随机采样。 接下来就是两层循环,外层循环是噪声等级的循环,从较大噪声等级的 \(\sigma_1 \approx=1\) ,到较小的噪声等级 \(\sigma_L \approx 0\)。 内循环是一个郎之万动力采样过程,负责对每个噪声等级下的 \(q_{sigma_i}(\tilde{x})\) 进行采样。 这样一直到最后一步 \(i=L\) 时,噪声等级足够小了,\(\sigma_L \approx 0\), 此时 \(q_{sigma_L}(\tilde{x}) \approx p_{data}(x)\) ,最后的得到的采样就近似是原数据分布 \(p_{data}(x)\) 的采样。 因为是两层循环,内循环是原始的郎之万动力采样,外循环是一个噪声等级逐步降低的过程, 两层循环加起来构成了退火郎之万动力采样。

4.1.5. 改进的分数生成模型

原始的分数模型,作者在 \(32 \times 32\) 的低分辨率图像上做了实验,得到了不错的效果。 但是在高分辨率上效果变差了很多。因此作者紧接着又发布了一篇论文,重点在如何改进分数模型,使其在更高分辨率的图像上也能有较好的效果。 作者进行了大量的分析和试验后,提出了了5项改进,这里不再赘述分析和试验过程的细节,只列出这5项改进内容, 对细节感兴趣的同学可以阅读原论文 [4]

  1. 初始噪声等级 \(\sigma_1\) 的最佳设定,它的大小影响图像生成的多样性(这其实很符合直觉认知), \(\sigma_1\) 越大多样性越好,\(\sigma_1\) 的最佳选择是训练数据集里样本对之间的最大的欧式距离。

  2. \(\{\sigma_i\}_{i=1}^L\) 设置成公比为 \(\gamma\) 的几何级数,即 \(\gamma=\sigma_{i-1}/\sigma_i\),此外 \(\gamma\) 满足约束 \(\Phi(\sqrt{2 D}(\gamma-1)+3\gamma)-\Phi(\sqrt{2 D}(\gamma-1)-3\gamma)\approx 0.5\), 其中 \(\Phi\) 是标准高斯分布的累积分布函数,\(D\) 是观测变量 \(x\) 的维数。

  3. 神经网络模型由原来的的 \(s_{\theta}(\tilde{x},\sigma_i)\) 改成 \(s_{\theta}(\tilde{x})/\sigma_i\) , 把噪声等级 \(\sigma_i\) 从模型输入中去掉, \(s_{\theta}(\tilde{x})\) 称为无条件分数网络(unconditional score network)。

  4. 在计算预算允许的范围内选择 \(T\) ,然后选择一个 \(\epsilon\) ,使下列方程最大接近 \(1\), 其中 \(x_{T} \sim \mathcal{N}(0,s^{2}_T \textit{I})\)

    (4.1.68)\[\frac{s^2_T}{\sigma_i^2} = \left( 1-\frac{\epsilon}{\sigma_L^2} \right)^{2T} \left( \gamma^2-\frac{2\epsilon}{\sigma_L^2-\sigma_L^2 \left( 1-\frac{\epsilon}{\sigma_L^2} \right)^2 } \right) + \frac{2\epsilon}{\sigma_L^2-\sigma_L^2 \left( 1-\frac{\epsilon}{\sigma_L^2} \right)^2 }\]
  5. 训练模型时加入指数滑动平均(exponential moving average,EMA)技术, \(\theta_{ema} \leftarrow m\theta_{ema} + (1 − m)\theta_t\) ,得到一份 EMA 的模型参数副本, 其中 \(m = 0.999\) 是动量参数。 在推理时(生成图像)时,用EMA的模型参数 \(\theta_{ema}\)。注意,在应用 EMA 时,实际上是有两套模型参数的, 模型训练过程中仍然按照原来的优化器(比如 Adam)去更新参数,只是在每一轮更新结束后,额外更新和保存一份 EMA 版本的参数 \(\theta_{ema}\)。 训练完成后输出保存的是两套模型参数,在推理时可以选择用原始模型参数还是使用 EMA 版本的参数。EMA 版本的参数理论上能缓解原始参数过拟合的情况, 但这也不是总成立的。

经过这五项改进后模型在生成图像的多样性和质量上都得到大幅提升,可以生成较高质量的 \(256 \times 256\) 尺寸的图像。

4.2. 随机微分方程

我们提出了一个随机微分方程(SDE),通过缓慢注入噪声将复杂的数据分布平滑地转换为已知的先验分布,以及相应的逆时间SDE,通过缓慢去除噪声将先验分布转换回数据分布。

我们表明,该框架封装了以前在基于分数的生成建模和扩散概率建模中的方法,允许新的采样过程和新的建模能力 我们还推导了一个等效的神经ODE,它从与SDE相同的分布中进行采样,但还能够进行精确的似然计算,并提高采样效率。 此外,我们提供了一种新的方法来解决基于分数的模型的反问题,正如类条件生成、图像修复和着色实验所证明的那样

4.2.1. 微分方程

在研究客观现象时,常常遇到这样一类数学问题,即其中某个变量和其他变量之间的函数依赖关系是未知的, 但是这个未知的函数关系以及它的某些阶的导数(或微分)连同自变量都由一个已知的方程联系在一起, 这样的方程称为微分方程(Differential Equation)。简单来说, 微分方程指的是:含有未知函数及其导数的方程

微分方程是联系自变量 \(x\),关于自变量 \(x\) 的未知函数 \(f\) 和 它的某些阶导数 \(\frac{df}{dx},\frac{d^2f}{dx^2},\dots,\frac{d^nf}{dx^n}\) 的关系式:

(4.2.10)\[F\left(x,f,\frac{df}{dx},\frac{d^2f}{dx^2},\dots,\frac{d^nf}{dx^n} \right ) = 0\]

如果未知函数是一元的,对应的微分方程称为常微分方程(Ordinary Differential Equation, ODE); 如果未知函数是多元的,对应的微分方程称为偏微分方程(Partial Differential Equations, PDE)。 方程中出现的最高阶导数的阶数称为这个微分方程的阶(order)。

接下是关于微分方程的解,微分方程解的数量是不固定的,有些微分方程有无穷多解,有的微分方程无解,有的微分方程则仅有有限个解。 特解 指的是满足微分方程的某一个解; 通解 指的是满足微分方程的一组解。

4.2.2. 随机微分方程

上节介绍的微分方程,其对象为可导函数,可以称之为一般微分方程。 然而我们经常需要面对一些 随机过程函数,即随时间变化的随机变量, 此时一般微分方程就不再适用了。在处理随机过程时,就需要有特殊的处理方法。 在所有随机过程中,扩散过程(diffusion process)是一种最基本的、常见的随机过程, 从名字也能看出,扩散过程和本章的主题扩散模型显然是相关的。

扩散过程

在概率论和统计学中,扩散过程(diffusion process)是一种随机过程(random process), 它是一种连续时间上马尔科夫过程(Markov process)。 扩散过程本质上是随机的,因此用于模拟许多现实生活中的随机系统, 布朗运动、反射布朗运动和奥恩斯坦-乌伦贝克过程是扩散过程的例子。 它被大量用于统计物理学、统计分析、信息论、数据科学、神经网络、金融和市场营销。

所谓的马尔科夫过程是指,未来的状态至于当前时刻的状态有关,与过去的状态无关。 用条件独立的表述为:在当前时刻的条件下,未来时刻与过去时刻是独立无关的。 扩散过程(diffusion process)可以用如下特殊的微分方程去描述

(4.2.11)\[d X_t = a(X_t,t)dt + b(X_t,t) d W_t\]

其中 \(t \in [0,T]\) 表示连续的时间, \(a(X_t,t)\) 漂移系数(drift coefficient), \(b(X_t,t)\) 噪声系数(noise coefficient),又名扩散系数(diffusion coefficient)。 \(W_t\) 表示一个 标准温拿过程(standard Wiener process),又名标准布朗运动(standard Brownian motion)。 如果 \(b(X_t,t)\)\(X\) 无关,则称为加性(additive)噪声。 如果 \(b(X_t,t)\)\(X\) 相关,那么噪声是乘性(multiplicative)的。

这个方程也被称为随机微分方程(stochastic differential equation,SDE), 如果去掉噪声项 \(b(X_t,t) d W_t\),就称为常微分方程(ordinary differential equation,ODE)。 随机微分方程多用于对一些多样化现象进行建模,比如不停变动的股票价格,部分物理现象如热扰动等。

从直觉上理解,扩散过程就是一个随机变量会随着时间的变化而变化,当然它还要满足马尔科夫性。 公式(4.2.11)\(d X_t\) 就是对 \(X_t\) 的微分, 可以理解成,在极小的时间变化中 \(\Delta t \rightarrow 0\)\(X_t\) 的变化量。 直接看 公式(4.2.11) 不是很容易理解,可以写成它的积分形式

(4.2.12)\[X_t = X_0 + \int_0^t a(X_s,t)dt + \int_0^t b(X_s,s) d W_s\]

\(X_0\) 表示在初始时刻 \(t=0\) 时,变量 \(X\) 的状态。 \(X_t\) 表示在任意某个时刻 \(t\) 时,变量 \(X\) 的状态。 显然,从公式可以看出 \(X_t\) 可以由 \(X_0\) 加上漂移项和噪声项的在时间 \(t\) 上的积分得到。

4.2.3. 基于随机微分方程的生成模型

回到图像生成模型的主题来,图像生成扩散模型可以用上述随机微分方程来描述。 我们的目标是构建一个,建立在连续时间 \(t \in [0,T]\) 上的扩散过程 \(\{X_t\}_{t=0}^T\)\(X_0 \sim p_0\) 表示初始时刻的随机变量,也就是真实图像背后的概率分布, 同时,我们有一些满足 i.i.d 的真实图像的样本 \(x_0\)。 换句话说 \(X_0 \sim p_0\) 是数据样本的概率分布。 定义一个扩散过程

(4.2.13)\[d x = f(x,t)dt + g(t) dw\]

其中 \(x\) 是一个向量,表示多个独立随机变量的值组成的向量。 这里我们对噪声系数进行了简化, 简化为关于时间的函数 \(g(t)\), 只与时间 \(t\) 相关,而与 \(x\) 无关,因此 \(g(t)\) 是一个标量。

这个SDE(公式(4.2.13))表达一个扩散过程,相当于 DDPM 中的前向扩散过程, 只不过在 DDPM 中,时间 \(t\) 是离散的,而在SDE中,时间是连续的。 因此 SDE 可以看做是DDPM在连续时间上的扩展

定义 \(s\) 表示早于 \(t\) 的时刻,即 \(0 \le s \lt t \le T\), 从时刻 \(s\) 到时刻 \(t\) ,变量转换的条件概率表示为 \(p_{st}(x_t|x_s)\), 也可以称 \(p_{st}(x_t|x_s)\)\(x_s\)\(x_t\) 的转换核(transition kernel)。

定义 \(X_T \sim p_T\) 为最终时刻的概率分布,它是一个无结构的先验分布(unstructured prior distribution ), 并且它不包含 \(p_0\) 的任何信息,这里把它定义成了一个标准高斯分布。

和DDPM一样,这个前向扩散过程不需要模型去学习,我们关注的是它的逆过程, 利用逆过程从数据分布 \(p_0\) 中去采样新的样本。 这里直接给出 SDE 的逆过程方程,我们不需要关注它的推导过程, SDE的逆过程称为 逆时间SDE (reverse-time SDE)

(4.2.14)\[d x = \left [ f(x,t) -g(t)^2 \Delta_x \log p_t(x) \right ] dt + g(t) d \bar{w}\]

4.3. 参考文献