参考:https://yang-song.net/blog/2021/score/
https://yang-song.net/blog/2019/ssm/
原论文:https://arxiv.org/pdf/1907.05600
https://arxiv.org/pdf/2011.13456
…

前言

首先，无论是学习还是回顾，都非常建议先阅读这篇博文，作者为Yang Song，是score系列模型及其相关工作的核心人物。该文章从传统的likelihood-based models和implicit generative models讲起，通过它们的不足引出score functions,score-based models,score matching,Langevin dynamics sampling,annealed Langevin sampling,SDE,Probability flow ODE等一些score流派的核心概念，串起了score领域的两篇核心论文（即这里的“原论文”），讲解非常深入浅出，思路清晰，逻辑通透，很好的建立起了该领域的整体框架。本文是对该博文的再度精简，仅作备忘和记忆锚点之用。

引入

生成式模型的终极目的，就是通过神经网络$p_\theta(x)$拟合真实数据的分布$p_(x)$。这里$p_\theta(x)$是PDF，数值上为正数，而神经网络本身倾向于输出实数域上的值，因此需要通过下列公式进行转换：

$p_\theta(x)=\displaystyle\frac{e^{-f_\theta(x)}}{Z_\theta}$

这个公式的形式借鉴了玻尔兹曼能量分布。这里$f_\theta(x)$是R上的神经网络原始输出，可以看作是一个粒子（对应图像空间中的一个样本点）的能量大小，能量越大，越不稳定，因此出现的概率越小，反之亦然，这就是通过$e^{-f_\theta(x)}$来建模概率的原因。至于分母${Z_\theta}$，是为了满足$\int p_\theta(\mathbf{x}) \, dx = 1。$这一限制条件而添加的取决于 θ 的归一化常数，即$Z_\theta=\int e^{-f_\theta(x)}dx$。

我们只需要

$\max_{\theta} \sum_{i=1}^{N} \log p_{\theta}(\mathbf{x}_i).$

就可以得到最佳的模型以及相应的预测分布$p_\theta(x)$

然而，由于${Z_\theta}$需要在x的真实分布上积分，而我们不可能拿到所有可能的x，那么也就无法计算${Z_\theta}$。先前的likelihood-based模型和implicit generative模型都是通过改进架构或近似的方式来对这个${Z_\theta}$进行求解。然而，所有这些处理都在一定程度上限制了模型的灵活性和可扩展性。

那么，有没有什么方法可以完全绕过${Z_\theta}$的计算，直接学习一个模型来拟合数据分布呢？

Score-based models

答案是肯定的，这就是score-based models的核心思想。

通过对得分函数而非密度函数进行建模，我们可以避免处理难以处理的归一化常数的难题。分布 p(x) 的得分函数定义为

$score(\mathbf{x}) = \nabla_{\mathbf{x}} \log p(\mathbf{x})$

而得分函数的模型称为基于得分的模型。我们将其记为$s_\theta(\mathbf{x})$，基于分数的模型经过学习，使得

$\mathbf{s}_\theta(\mathbf{x}) \approx \nabla_{\mathbf{x}} \log p(\mathbf{x})$

成立，并且可以在不考虑归一化常数的情况下进行参数化。我们可以通过以下方式使用基于能量的模型来参数化基于分数的模型：

$\mathbf{s}_\theta(\mathbf{x}) = \nabla_{\mathbf{x}} \log p_\theta(\mathbf{x}) = - \nabla_{\mathbf{x}} f_\theta(\mathbf{x}) - \nabla_{\mathbf{x}} \log Z_\theta = - \nabla_{\mathbf{x}} f_\theta(\mathbf{x}). \quad$

基于分数的模型$s_\theta(\mathbf{x})$与归一化常数$Z_\theta$无关,这大大扩展了我们可以方便使用的模型族，因为我们不需要任何特殊的架构来处理归一化常数。

此外，score-based models还与逆问题求解有着天然的联系：

$\nabla_{\mathbf{x}} \log p(\mathbf{x}\mid \mathbf{y}) = \nabla_{\mathbf{x}} \log p(\mathbf{x}) + \nabla_{\mathbf{x}} \log p(\mathbf{y}\mid \mathbf{x}). \quad (15)$

我们可以训练一个模型来估计无条件数据分布的得分函数，即 $s_\theta(\mathbf{x}) \approx \nabla_{\mathbf{x}} \log p(\mathbf{x})$。这将使我们能够通过公式 (15) 轻松地从已知的正向过程 $p(\mathbf{y}\mid \mathbf{x})$ 计算后验得分函数 $\nabla_{\mathbf{x}} \log p(\mathbf{x}\mid \mathbf{y})$，并使用朗之万型采样对其进行采样。

在训练score-based models时，我们通过最小化Fisher散度来使得模型的得分函数与数据分布的得分函数接近。Fisher散度定义为：

$\mathbb{E}_{p(\mathbf{x})} \left[ \left\| \nabla_{\mathbf{x}} \log p(\mathbf{x}) - \nabla_{\mathbf{x}} \log q(\mathbf{x}) \right\|_2^2 \right]. \quad$

也就是

$\mathbb{E}_{p(\mathbf{x})} \left[ \left\| \nabla_{\mathbf{x}} \log p(\mathbf{x}) - \mathbf{s}_\theta(\mathbf{x}) \right\|_2^2 \right]. \quad$

直观地说，Fisher 散度比较的是真实数据得分与基于得分的模型之间的平方距离。然而，直接计算此散度是不可行的，因为它需要访问未知的数据得分$\nabla_{\mathbf{x}} \log p(\mathbf{x})$。幸运的是，我们可以使用分数匹配方法来最小化Fisher散度，而无需直接计算数据得分。

分数匹配(score matching)方法无需了解真实数据得分即可最小化 Fisher 散度。该方法通过数学推导，将Fisher散度转化为一个仅依赖于模型得分函数的表达式，从而使得我们可以直接优化模型参数$\theta$:

$\frac{1}{2} \, \mathbb{E}_{p_{\text{data}}} \left[ \left( \nabla_{\mathbf{x}} \log p_{\text{data}}(\mathbf{x}) - \nabla_{\mathbf{x}} \log p_\theta(\mathbf{x}) \right)^2 \right] = \mathbb{E}_{p_{\text{data}}} \left[ \operatorname{tr} \left( \nabla_{\mathbf{x}}^2 \log p_\theta(\mathbf{x}) \right) + \frac{1}{2} \left\| \nabla_{\mathbf{x}} \log p_\theta(\mathbf{x}) \right\|_2^2 \right] + \text{const}.$

对于高维数据，直接计算hessian矩阵$\nabla_{\mathbf{x}}^2 \log p_\theta(\mathbf{x})$的迹会导致计算复杂度太高，作者提出了Sliced score matching方法，通过将数据投影到随机方向上来近似计算Fisher散度，从而大大降低了计算复杂度:

$\frac{1}{2} \, \mathbb{E}_{p_{\text{data}}} \left[ \left( \nabla_{\mathbf{x}} \log p_{\text{data}}(\mathbf{x}) - \nabla_{\mathbf{x}} \log p_\theta(\mathbf{x}) \right)^2 \right] = \frac{1}{NM} \sum_{i=1}^{N} \sum_{j=1}^{M} \left[ \mathbf{v}_{ij}^{\top} \nabla_{\mathbf{x}}^2 \log p_\theta(\mathbf{x}_i) \, \mathbf{v}_{ij} + \frac{1}{2} \left( \mathbf{v}_{ij}^{\top} \nabla_{\mathbf{x}} \log p_\theta(\mathbf{x}_i) \right)^2 \right].$

Langevin dynamics sampling

训练好score-based model后，我们可以使用Langevin dynamics来从模型中采样。Langevin 动力学提供了一种 MCMC 程序，仅使用其得分函数$\nabla_{\mathbf{x}} \log p(\mathbf{x})$进行采样。具体来说，它从任意先验分布$\mathbf{x}_0 \sim \pi(\mathbf{x})$初始化链，然后迭代执行以下操作：

$\mathbf{x}_{i+1} \leftarrow \mathbf{x}_i + \epsilon \nabla_{\mathbf{x}} \log p(\mathbf{x}) + \sqrt{2\epsilon} \, \mathbf{z}_i, \quad i = 0, 1, \cdots, K, \quad$

其中$\mathbf{z}_i \sim \mathcal{N}(0, \mathbf{I})$，当$\epsilon \to 0$和$K \to \infty$时，在某些正则性条件下，$\mathbf{x}_K$的分布将收敛到$p(\mathbf{x})$。实际上，当 ϵ 足够小且 K 足够大时，误差可以忽略不计。

可以看到，该方法实际上类似于梯度下降（实际上是“能量下降”），让样本点逐渐往数据分布的高密度区域移动，但又加入了噪声项$\sqrt{2\epsilon} \, \mathbf{z}_i$，以确保样本点不会全部掉进某个局部最小值，保障了多样性和随机性。

问题

Score-based models虽然避免了归一化常数的计算，但在低密度区域，由于可用于计算得分匹配目标的数据点较少，估计的得分函数不够准确。这种行为会导致结果不理想。在使用朗之万动力学进行采样时，如果数据位于高维空间，初始样本极有可能位于低密度区域。因此，不准确的基于得分的模型会从一开始就破坏朗之万动力学的运行，使其无法生成能够代表数据的高质量样本。

多尺度噪声扰动

为了缓解这个问题，作者用噪声扰动数据点，并基于这些带噪声的数据点训练基于分数的模型。当噪声幅度足够大时，它可以填充低数据密度区域，从而提高估计分数的准确性。

较大的噪声显然可以覆盖更多低密度区域，从而获得更好的分数估计，但它会过度破坏数据，使其与原始分布显著不同。另一方面，较小的噪声对原始数据分布的破坏较小，但对低密度区域的覆盖效果却不如我们预期。

为了兼顾两者的优点，我们同时使用多种尺度的噪声扰动。对于不同的噪声水平$\sigma_1 < \sigma_2 < \cdots < \sigma_L$，我们可以先得到不同的噪声扰动分布$p_{\sigma_i}(x)$，然后训练一个训练一个基于噪声条件评分的模型$s_\theta(x, i)$(也称为噪声条件评分网络，或 NCSN)来估计每个噪声扰动分布的score function $\nabla_x \log p_{\sigma_i}(x)$。

训练的目标函数是所有噪声尺度下 Fisher 散度的加权和。依然可以通过分数匹配方法来优化目标函数。

在训练好基于噪声条件得分的模型$s_\theta(x, i)$后，我们可以通过依次运行$i = L, L-1, \cdots, 1$的annealed Langevin dynamics来采样生成样本。噪声尺度随时间逐渐减小，因此被称为退火朗之万动力学。对于每个噪声水平$\sigma_i$，我们从任意先验分布$\mathbf{x}_0 \sim \pi(\mathbf{x})$初始化链，然后迭代执行以下操作：

$\mathbf{x}_{k+1} \leftarrow \mathbf{x}_k + \epsilon \, s_\theta(\mathbf{x}_k, i) + \sqrt{2\epsilon} \, \mathbf{z}_k, \quad k = 0, 1, \cdots, K_i, \quad$

通过依次运行$i = L, L-1, \cdots, 1$的annealed Langevin dynamics，我们可以从$p_{\sigma_L}(\mathbf{x})$逐渐过渡到$p_{\sigma_1}(\mathbf{x})$，最终得到接近原始数据分布$p(\mathbf{x})$的样本。

SDE

在前面，我们通过添加多个噪声尺度改善了模型的表现。而通过将噪声尺度的数量推广到无穷大，我们不仅可以获得更高质量的样本，还可以进行精确的对数似然计算，以及可控的逆问题求解生成。当噪声尺度数量趋于无穷大时，我们实际上是在用不断增加的噪声水平扰动数据分布。在这种情况下，噪声扰动过程是一个连续时间随机过程。而许多随机过程（特别是扩散过程）都是随机微分方程（SDE）的解。一般来说，SDE 具有以下形式：

$d\mathbf{x} = \mathbf{f}(\mathbf{x}, t)\, dt + g(t)\, d\mathbf{w}, \quad (8)$

其中，$\mathbf{f}(\cdot, t): \mathbb{R}^d \to \mathbb{R}^d$ 是一个向量值函数，称为漂移系数；$g(t) \in \mathbb{R}$ 是一个实值函数，称为扩散系数；$\mathbf{w}$ 表示标准布朗运动；$d\mathbf{w}$ 可以视为无穷小白噪声。随机微分方程的解是一组连续的随机变量 ${\mathbf{x}(t)}_{t \in [0, T]}$。这些随机变量沿着随机轨迹运动，时间索引 $t$ 从起始时间 $0$ 增长到结束时间 $T$。令 $p_t(\mathbf{x})$ 表示 $\mathbf{x}(t)$ 的（边缘）概率密度函数。这里，$t \in [0, T]$ 类似于噪声尺度数量有限时的 $i = 1, 2, \cdots, L$，$p_t(\mathbf{x})$ 类似于 $p_{\sigma_i}(\mathbf{x})$。显然，$p_0(\mathbf{x}) = p(\mathbf{x})$ 是数据分布，因为在 $t = 0$ 处没有对数据施加扰动。在用随机过程对 $p(\mathbf{x})$ 进行足够长时间的扰动后 $T$，$p_T(\mathbf{x})$ 将接近于一个易于处理的噪声分布 $\pi(\mathbf{x})$，称为先验分布。我们注意到，在有限噪声尺度的情况下，$p_T(\mathbf{x})$ 类似于 $p_{\sigma_L}(\mathbf{x})$，这对应于对数据施加最大噪声扰动 $\sigma_L$。

通俗地讲，可以将图像空间中的每个样本点（即图片）看作该空间中的粒子，这些粒子从初始的真实分布开始，根据SDE所描述运动过程在时间上连续的运动，因此所有这些粒子的分布，即$p_t(\mathbf{x})$也在时间上连续地演化，直到最后时刻T，这些粒子的分布完全变成某种先验分布（如高斯分布），即$p_T(\mathbf{x})=\pi(\mathbf{x})$。$\mathbf{f}(\mathbf{x}, t)$这一项可以看作是粒子的速度场$v(\mathbf{x}, t)$，它告诉我们在时间t，粒子在位置$\mathbf{x}$处的运动方向和速度大小，该速度乘以一个小的时间dt构成了粒子的稳定漂移趋势；$g(t) \, d\mathbf{w}$这一项可以看作是粒子在时间t受到的随机扰动，扰动的强度由$g(t)$控制，扰动的方向和大小由$d\mathbf{w}$（即无穷小白噪声）决定，由于$\mathbf{w}$是一个随机采样的噪声项，因此这一项构成了整个SDE的随机性。也就是说，粒子在确定的漂流趋势的基础上，还会受到随机扰动的影响，最终他们的分布接近于完全散乱的状态。注意，尽管粒子受到的总噪声强度是随时间增加，但这并不意味着g(t)是单调递增的函数，因为总的累积噪声实际上是$\text{variance} \sim \int_0^t g^2(s)\,ds$。

(8) 中的 SDE 是手工设计的，类似于我们在有限噪声尺度的情况下手工设计 $\sigma_1 < \sigma_2 < \cdots < \sigma_L$ 的方式。添加噪声扰动的方法有很多种，SDE 的选择也并非唯一。例如，以下 SDE：

$d\mathbf{x} = e^{t} \, d\mathbf{w} \quad (9)$

用均值为零、方差呈指数增长的高斯噪声扰动数据，这类似于当 $\sigma_1 < \sigma_2 < \cdots < \sigma_L$ 为等比数列时用
$\mathcal{N}(0, \sigma_1^2 \mathbf{I}), \mathcal{N}(0, \sigma_2^2 \mathbf{I}), \cdots, \mathcal{N}(0, \sigma_L^2 \mathbf{I})$ 扰动数据。因此，SDE 应被视为模型的一部分，就像 ${\sigma_1, \sigma_2, \cdots, \sigma_L}$ 一样。

对于无限数量的噪声尺度，我们可以类似地使用反向随机微分方程来逆转扰动过程以生成样本。任何随机微分方程都有对应的反向随机微分方程。其封闭形式由下式给出

$d\mathbf{x} = \left[\mathbf{f}(\mathbf{x}, t) - g^2(t)\nabla_{\mathbf{x}} \log p_t(\mathbf{x})\right] dt + g(t)\, d\mathbf{w}. \quad (10)$

这里 $dt$ 表示一个负的无穷小时间步长，因为随机微分方程 (10) 需要反向求解（从 $t = T$ 到 $t = 0$）。为了计算反向随机微分方程，我们需要估计 $\nabla_{\mathbf{x}} \log p_t(\mathbf{x})$，它恰好是 $p_t(\mathbf{x})$ 的得分函数。

求解逆随机微分方程需要我们知道最终分布 $p_T(\mathbf{x})$ 和得分函数 $\nabla_{\mathbf{x}} \log p_t(\mathbf{x})$。根据设计，终端分布 $p(\mathbf{x})$ 非常接近于完全可处理的先验分布 $\pi(\mathbf{x})$。为了估计 $\nabla_{\mathbf{x}} \log p(\mathbf{x})$，我们训练一个时间相关的基于得分的模型 $s_\theta(\mathbf{x}, t)$，使得 $s_\theta(\mathbf{x}, t)\approx \nabla_{\mathbf{x}} \log p_t(\mathbf{x})$。这类似于用于有限噪声尺度的噪声条件基于得分的模型 $s_\theta(\mathbf{x}, i)$，该模型经过训练后满足 $s_\theta(\mathbf{x}, i)\approx \nabla_{\mathbf{x}} \log p_{\sigma_i}(\mathbf{x})$。

我们对 $s_\theta(\mathbf{x}, t)$ 的训练目标是 Fisher 散度的连续加权组合，由下式给出：

$\mathbb{E}_{t \sim \mathcal{U}(0, T)} \mathbb{E}_{p_t(\mathbf{x})} \left[ \lambda(t) \left\| \nabla_{\mathbf{x}} \log p_t(\mathbf{x}) - s_\theta(\mathbf{x}, t) \right\|_2^2 \right], \quad (11)$

其中 $\mathcal{U}(0, T)$ 表示在时间区间 $[0, T]$ 上的均匀分布，$\lambda: \mathbb{R} \to \mathbb{R}_{>0}$ 是一个正的加权函数。通常我们使用 $\lambda(t) \propto \frac{1}{\mathbb{E}\left[\left| \nabla_{\mathbf{x}(t)} \log p(\mathbf{x}(t)\mid \mathbf{x}(0)) \right|_2^2\right]}$ 来平衡不同时间上的得分匹配损失的幅度。

与之前一样，我们利用分数匹配方法可以有效地优化 Fisher 散度的加权组合。一旦我们的基于分数的模型 $s_\theta(\mathbf{x}, t)$ 训练到最优状态，我们就可以将其代入 (10) 中的反向 SDE 表达式，从而获得估计的反向 SDE。

$d\mathbf{x} = \left[\mathbf{f}(\mathbf{x}, t) - g^2(t)s_\theta(\mathbf{x}, t)\right]dt + g(t)d\mathbf{w}. \quad (12)$

我们可以从 $\mathbf{x}(T) \sim \pi$ 开始，求解上述反向随机微分方程得到样本 $\mathbf{x}(0)$。我们将以此方式得到的 $\mathbf{x}(0)$ 的分布记为 $p_\theta$。当基于分数的模型 $s_\theta(\mathbf{x}, t)$ 训练良好时，我们有 $p_\theta \approx p_0$，此时 $\mathbf{x}(0)$ 是数据分布 $p_0$ 的一个近似样本。

通过使用数值随机微分方程 (SDE) 求解器求解估计的反向 SDE，我们可以模拟用于样本生成的反向随机过程。数值SDE求解器包括欧拉-丸山 (Euler-Maruyama) 方法，自适应步长 SDE 求解器，一种类似于Actor-Critic架构的预测-校正采样器等，它们的核心思想是通过极小的步长将连续过程离散化，迭代求解。

Probability flow ODE

SDE很好，它可以精确模拟连续时间的随机扩散过程，但它的随机性也带来了一些问题：首先，SDE 的求解比较困难；其次，SDE的求解过程包含大量的随机采样，这会增加计算成本和采样时间；最后，由于SDE是非确定性的，因此无法从中精确计算出likelihood。

经证明，实际上可以可以在不改变其边缘分布 ${p_t(\mathbf{x})}_{t \in [0, T]}$ 的情况下，将任何随机微分方程（SDE）转化为常微分方程（ODE）。因此，通过求解该 ODE，我们可以从与反向 SDE 相同的分布中进行采样。SDE 对应的 ODE 被称为概率流 ODE：

$d\mathbf{x} = \left[\mathbf{f}(\mathbf{x}, t) - \frac{1}{2} g^2(t)\nabla_{\mathbf{x}} \log p_t(\mathbf{x})\right] dt. \quad (14)$

对于图像空间的粒子群而言，虽然 ODE 的轨迹明显比 SDE 的轨迹更平滑，但它们将相同的数据分布转换为相同的先验分布，反之亦然，并且共享相同的边缘分布集 ${p_t(\mathbf{x})}_{t \in [0, T]}$。换句话说，通过求解概率流 ODE 得到的轨迹与 SDE 的轨迹具有相同的边缘分布。

如何理解SDE与ODE的关系？首先，可以看到，ODE中不含SDE的随机项$g(t) \, d\mathbf{w}$，而ODE的速度场变成了$\mathbf{f}(\mathbf{x}, t) - \frac{1}{2} g^2(t)\nabla_{\mathbf{x}} \log p_t(\mathbf{x})$，相比于SDE中的速度场$\mathbf{f}(\mathbf{x}, t)$，ODE多了一个包含score的修正项，这是为了模拟SDE的随机结果，在每个时刻t，ODE的分布要保持与SDE的$p_t(\mathbf{x})$一致，因此ODE的速度场需要包含一个修正项来补偿SDE中随机项的影响。其次，在SDE的随机性与ODE的确定性方面，同一粒子放到SDE里跑很多遍，每一遍的轨迹都不一样，放到ODE里跑很多遍，每一遍的轨迹是一样的。但是，如果有大量粒子同时放到SDE和ODE里跑，在任意时刻t，这个时刻的“粒子云/概率云”（即分布）的形状是一样的，尽管这两个粒子云中同一位置的粒子在开始时刻并不是同一个位置对应的粒子。ODE有点类似于“先射箭，后画靶”，通过精心设计的速度，让粒子在每个时刻都往自己应该去的地方移动，在这个过程中就需要用到移动的预期结果$p_t$。但是真实分布在t时刻的$p_t$不知道，怎么办？答案是用神经网络预测出$\nabla_{\mathbf{x}} \log p_t(\mathbf{x})$，然后再得出相应的ODE。得到了想要的ODE，比原来SDE的优势在哪里？1，ODE是Deterministic sampling，同一个初始点 → 唯一结果，因此采样更稳定（无随机 variance）；2，这是一个Continuous normalizing flow，可以用$\log p(x_0) = \log p(x_T) - \int \nabla \cdot v(x,t)\,dt$计算精确 likelihood；3，可以用高阶 ODE solver；4，flow matching 本质就是直接学习这个probability flow ODE的速度场$v(x,t)$，因此ODE架起了与flow方法的桥梁。

流派的统一

基于评分的生成模型和扩散概率模型都可以看作是由评分函数确定的随机微分方程的离散化，基于分数的生成模型（包含多个噪声扰动）和扩散概率模型是同一模型族的不同视角。这项工作从微分方程的数学视角，统一了生成式模型长期以来的两大流派：score-based models和 diffusion probability models。