参考: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时，我们通过最小化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散度，而无需直接计算数据得分。

分数匹配(flow 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})$的样本。