引言
在前面的章节中，我们介绍了在变分推断的框架下的生成模型，包括VAE和DDPM两大类。在变分的视角下，生成模型实际上是在学习一个带参的变分分布 $q_\phi$ 来近似难解的后验分布 $p(x\mid z)$，并通过优化经验下界ELBO来最大化对数似然函数 $\log p(x)$。
从这一篇文章开始，我们将站在一个全新的视角来看待生成模型：基于分数 (score-based) 的视角。这里的分数实际上是一个梯度场 (gradient field)，指向数据分布中高概率的方向。这种方法的中心思想在于：我们可以通过建模分数来进行生成式建模，因为我们可以利用梯度场将随机采样的样本推向数据分布 $p_{data}$ 中高概率的方向，从而避免计算归一化常数中难解的积分项。因此，这类方法并不是直接学习数据分布的形式，而是学习如何将先验分布推向数据分布。
基于能量的模型 (Energy-based Models, EBMs) 是利用分数进行生成式建模的最简单的方法。能量是建模概率分布的一种方法，概率越高的地方能量越低。EBMs往往基于朗之万动力学 (Langevin dynamics) 进行采样，根据能量地图的梯度方向（即分数）将样本推向高概率区域。
一、基于能量的模型
令 $x\in\mathbb{R}^D$ 表示一个数据点。EBMs通过一个能量函数 $E_\phi(x)$ 来定义一个概率分布：
\[\begin{equation} p_\phi(x):=\frac{\exp\left(-E_\phi(x)\right)}{Z_\phi} \end{equation}\]其中，归一化项 $Z_\phi$ 用于保证概率分布的归一性：
\[\begin{equation} Z_\phi=\int_{\mathbb{R}^D}\exp\left(-E_\phi(x)\right)\mathrm{d}x \end{equation}\]从能量函数的定义我们可以看到，概率越低的数据点能量越高，反之亦然。
EBMs可以通过最大化对数似然来进行训练：
\[\begin{equation} \begin{aligned} \mathcal{L}_{MLE}(\phi) &=\mathbb{E}_{p_{data}(x)}\left[\log\frac{\exp\left(-E_\phi(x)\right)}{Z_\phi}\right]\\ &=-\mathbb{E}_{p_{data}(x)}\left[E_\phi(x)\right]-\log\int\exp\left(-E_\phi(x)\right)\mathrm{d}x \end{aligned} \end{equation}\]其中，第一项用于保证位于 $p_{data}$ 的支撑集上的数据点具有较低的能量，第二项是正则化项。
正如前面提到的，当数据维度 $D$ 较大时，计算第二项的积分是非常困难的。在变分方法中，我们使用可解的ELBO来近似难解的对数似然。在基于分数的方法中，我们通过 score matching 来直接避免计算这一项。
在下面的内容中，我们将首先介绍 score function 的形式化定义，然后介绍 score matching 是如何避免计算难解的归一化项的。最后，我们介绍利用 score function 进行采样的朗之万采样法。
1.1. 什么是分数？
对于定义在 $\mathbb{R}^D$ 上的概率分布 $p(x)$，score function 是其对数密度的梯度：
\[\begin{equation} s(x):=\nabla_x\log p(x) \end{equation}\]因此，分数函数组成了 $\mathbb{R}^D$ 上的一个向量场，其中每一个向量都指向 $p(x)$ 中的高概率区域。
一个很自然且重要的问题是：为什么我们不直接建模概率密度，而是建模密度的梯度（即分数）？这是因为建模分数具有下面几个理论上和实践上的优势。
1. 不需要归一化常数。由于难解的归一化常数 $Z$ 往往与数据点 $x$ 无关，因此其梯度为0，我们也就不需要计算这个难解的积分。
比如说在EBMs中，分数的定义为：
\[\begin{equation} \begin{aligned} s(x) &:=\nabla_x\log p(x)\\ &=\nabla_x\left[\log \frac{\exp\left(-E_\phi(x)\right)}{Z_\phi}\right]\\ &=-\nabla_xE_\phi(x) \end{aligned} \end{equation}\]2. 分数函数完整刻画了概率密度。由于分数是对数密度的梯度，因此我们可以利用下面的公式，从分数重建概率密度：
\[\begin{equation} \log p(x)=\log p(x_0)+\int_0^1 s(x_0+t(x-x_0))^T(x-x_0)\mathrm{d}t \end{equation}\]其中，$x_0$ 是某个固定的参考数据点。
因此，建模分数和建模密度具有相同的表达能力，而分数往往是可解的。
1.2. 利用 score matching 训练EBMs
我们希望利用 $p_{data}$ 中的样本来训练一个神经网络 $s_\phi(x)$ 用于近似分数函数 $s(x)=\nabla_x\log p_{data}(x)$：
\[\begin{equation} s_\phi(x)\approx s(x) \end{equation}\]score matching 方法通过最小化真实分数和预测分数的均方误差 (MSE) 来拟合这个向量场：
\[\begin{equation} \mathcal{L}_{SM}(\phi):=\frac{1}{2}\mathbb{E}_{x\sim p_{data}}\left[ \left\|s_\phi(x)- s(x)\right\|_2^2 \right] \end{equation}\]然而，由于真实分数 $s(x)$ 是未知的，上面的损失函数是难解的。幸运的是，Hyvarinen找到了一个与上式等价且可解的损失函数，仅依赖于 $s_\phi$ 和数据 $x$，不需要真实分数。下面的定理说明了这一点。
Theorem 1 (Hyvarinen’s Tractable Form of SM). 我们有下面等式成立：
\[\begin{equation} \mathcal{L}_{SM}(\phi)=\widetilde{\mathcal{L}}_{SM}(\phi)+C \end{equation}\]其中，$C$ 是一个与 $\phi$ 无关的常数，且
\[\begin{equation} \widetilde{\mathcal{L}}_{SM}(\phi):=\mathbb{E}_{x\sim p_{data}}\left[ \text{Tr}\left(\nabla_xs_\phi(x) \right) + \frac{1}{2}\left\|s_\phi(x)\right\|_2^2 \right] \end{equation}\]定理1的证明放在附录中。
我们可以从直观上理解这个可解的损失 $\widetilde{\mathcal{L}}_{SM}(\phi)$：
- 第二项 $\left|s_\phi(x)\right|2^2$ 是范数项，使得在 $p{data}$ 的高密度区域中，预测分数尽可能接近0，即梯度很小，让这些高密度区域成为驻点。
- 第一项 $\text{Tr}\left(\nabla_xs_\phi(x) \right)=\text{div}(s_\phi)$ 是散度项，鼓励 $s_\phi$ 的散度为负，确保驻点是吸引子 (attractive sinks)。因为散度越小，表示向量场在该点向内汇聚，保证了这些高密度区域周围的梯度方向会指向驻点。
二、朗之万动力学采样
我们一般使用朗之万动力学 (Langevin dynamics) 来对EBMs进行采样。在下面的内容中，我们将首先介绍离散时间中的朗之万采样，并将其逼近至连续时间的场景中，使用随机微分方程来求解。最后，我们简要介绍一下朗之万动力学背后的物理直觉，以及其为什么适用于对EBMs进行采样。
2.1. 离散时间朗之万动力学
离散时间朗之万动力学 (Discrete-Time Langevin Dynamics) 可以又下面的迭代方程描述：
\[\begin{equation} \begin{aligned} x_{n+1} &=x_n-\eta\nabla_xE_\phi(x_n)+\sqrt{2\eta}\epsilon_n\\ &=x_n+\eta s_\phi(x)+\sqrt{2\eta}\epsilon_n\\ \end{aligned} \end{equation}\]其中，$x_0\sim p_{prior}$ 是从一个先验分布（一般是高斯分布）中采样得到的数据点。$\eta\gt0$ 是步长。$\epsilon_n\sim\mathcal{N}(0,1)$ 是高斯噪声，用于增加随机性，帮助模型跳出局部最小值。
2.2. 连续时间朗之万动力学
随着采样步长 $\eta\rightarrow 0$，上面的更新方程会收敛与下面的朗之万SDE：
\[\begin{equation} \mathrm{d}x(t)=s_\phi(x(t))\mathrm{d}t+\sqrt{2}\mathrm{d}w(t) \end{equation}\]其中，$w(t)$ 表示一个标准布朗运动 (Standard Brownian Motion)，也称为维纳过程 (Wiener Process)，用于添加随机性。布朗运动的增量满足：$w(t+\eta)-w(t)\sim\mathcal{N}(0,\eta)$。
这里的 $\sqrt{2}$ 是用于确保目标分布 $p_\phi$ 是该SDE对应的平稳分布，这是由福克-普朗克方程 (Fokker-Planck Equation) 推导得到的，后面的内容中我们将进一步深入。
2.3. 为什么使用朗之万采样
理解朗之万采样的一种自然的方式是从物理世界的角度去理解。EBMs通过能量方程 $E_\phi(x)$ 定义了一个势能场，用于规范其中粒子的行为。根据牛顿力学，一个粒子在势能场中的运动轨迹由下面的常微分方程描述：
\[\begin{equation} \mathrm{d}x(t)= -\nabla_xE_\phi(x(t)) \mathrm{d}t \end{equation}\]上面的常微分方程将粒子确定性地推向梯度下降的方向，直到去到能量方程的某个局部最小值，导致无法对整个数据分布进行全面地探索，可能会错过更优的全局最小值。
为了克服这个缺陷，朗之万动力学引入了一个由布朗运动所描述的随机扰动，得到了下面的SDE：
\[\begin{equation} \mathrm{d}x(t)=-\nabla_xE_\phi(x(t))\mathrm{d}t+\sqrt{2}\mathrm{d}w(t) \end{equation}\]随机扰动可以帮助粒子跨越能量势垒以逃离局部最小值，使得粒子的轨迹形成一个随机过程，且其平稳分布将收敛止玻尔兹曼分布：
\[\begin{equation} p_\phi(x)\propto e^{-E_\phi(x)} \end{equation}\]从这个角度来看，EBMs实际上学习了一个向量场，并使用朗之万采样将数据点推向高概率的区域。
然而，朗之万采样在高维空间中会面临一些问题。比如说，它对于一些超参数（如采样步长 $\eta$、迭代步数等）非常敏感。这种缺陷的根源在于朗之万采样的“混合时间” (mixing time) 过长：在高维复杂概率分布中往往存在多个孤立的高概率区域，而朗之万采样需要非常长的迭代步数来在这些区域之间迁移。
附录
Proof on Theorem 1.
定理1提供了 score matching 训练目标的一个可解形式：
\[\begin{equation} \begin{aligned} \mathcal{L}_{SM}(\phi) &:=\frac{1}{2}\mathbb{E}_{x\sim p_{data}}\left[ \left\|s_\phi(x)- s(x)\right\|_2^2 \right]\\ &=\widetilde{\mathcal{L}}_{SM}(\phi)+C\\ &=\mathbb{E}_{x\sim p_{data}}\left[ \text{Tr}\left(\nabla_xs_\phi(x) \right) + \frac{1}{2}\left\|s_\phi(x)\right\|_2^2 \right]+C \end{aligned} \end{equation}\]其中，$C$ 是一个与 $\phi$ 无关的常数。
下面我们证明这一点。
首先我们展开 $\mathcal{L}_{SM}(\phi)$ 的形式：
\[\begin{equation} \begin{aligned} \mathcal{L}_{SM}(\phi) &:=\frac{1}{2}\mathbb{E}_{x\sim p_{data}}\left[ \left\|s_\phi(x)- s(x)\right\|_2^2 \right]\\ &=\frac{1}{2}\mathbb{E}_{x\sim p_{data}}\left[ \left\|s_\phi(x)\right\|_2^2-2\left\langle s_\phi(x),s(x) \right\rangle + \left\|s(x)\right\|_2^2 \right]\\ &=\frac{1}{2}\mathbb{E}_{x\sim p_{data}}\left[ \left\|s_\phi(x)\right\|_2^2 \right]+\frac{1}{2}\mathbb{E}_{x\sim p_{data}}\left[ \left\|s(x)\right\|_2^2 \right]-\mathbb{E}_{x\sim p_{data}}\left[ \left\langle s_\phi(x),s(x) \right\rangle \right] \end{aligned} \end{equation}\]这里我们关注第三项，即交叉相乘的项。由于：
\[\begin{equation} s(x)=\nabla_x\log p_{data}(x)=\frac{\nabla_x p_{data}(x)}{p_{data}(x)} \end{equation}\]代入得：
\[\begin{equation} \begin{aligned} \mathbb{E}_{x\sim p_{data}}\left[ \left\langle s_\phi(x),s(x) \right\rangle \right] &=\int s_\phi(x)^T s(x)p_{data}(x)\mathrm{d}x\\ &=\int s_\phi(x)^T \nabla_x p_{data}(x)\mathrm{d}x\\ &=\sum_{i=1}^D \int s_\phi^{(i)}(x) \partial_{x_i} p_{data}(x)\mathrm{d}x\\ \end{aligned} \end{equation}\]其中，$s_\phi^{(i)}(x)$ 是分数函数的第 $i$ 个分量：
\[\begin{equation} s_\phi(x)=\left( s_\phi^{(1)}(x),s_\phi^{(2)}(x),\cdots,s_\phi^{(D)}(x) \right) \end{equation}\]下面我们将使用一个分部积分公式：
Lemma (Integration by Parts). 令 $\mathbb{B}(0,R)\subset\mathbb{R}^D$ 是以原点为中心，半径为 $R>0$ 的封闭超球体，其边界 $\partial \mathbb{B}(0,R)$ 是光滑超球面。设函数 $u,v$ 是在 $\mathbb{B}(0,R)$ 上的可微实函数，则对 $i=1,\dots,D$ 都有下面等式成立：
\[\begin{equation} \int_{\mathbb{B}(0,R)}u\partial_{x_i}v\mathrm{d}x=-\int_{\mathbb{B}(0,R)}v\partial_{x_i}u\mathrm{d}x+\int_{\partial\mathbb{B}(0,R)}uv\nu_i\mathrm{d}S \end{equation}\]其中，$\nu=(\nu_1,\dots,\nu_D)$ 是边界 $\partial \mathbb{B}(0,R)$ 上的外法向单位向量，$\mathrm{d}S$ 是 $\partial \mathbb{B}(0,R)$ 上的面积微元。
令 $u=s_\phi^{(i)},v=p_{data}$，并假设当 $R\rightarrow \infty$ 时，$\lvert s_\phi^{(i)}(x)p_{data}(x)\rvert\rightarrow 0$。
代入得：
\[\begin{equation} \begin{aligned} \mathbb{E}_{x\sim p_{data}}\left[ \left\langle s_\phi(x),s(x) \right\rangle \right] &=\sum_{i=1}^D \int s_\phi^{(i)}(x) \partial_{x_i} p_{data}(x)\mathrm{d}x\\ &=-\sum_{i=1}^D \int p_{data} \partial_{x_i} s_\phi^{(i)}\mathrm{d}x+\sum_{i=1}^D \int s_\phi^{(i)} p_{data} \nu_i\mathrm{d}S\\ &=-\sum_{i=1}^D \int p_{data} \partial_{x_i} s_\phi^{(i)}\mathrm{d}x\\ &=- \int p_{data} \sum_{i=1}^D\partial_{x_i} s_\phi^{(i)}\mathrm{d}x\\ &=- \int p_{data}(x) \text{Tr}(\nabla_xs_\phi(x)) \mathrm{d}x\\ &=-\mathbb{E}_{x\sim p_{data}}\left[ \text{Tr}(\nabla_xs_\phi(x)) \right] \end{aligned} \end{equation}\]因此，我们有：
\[\begin{equation} \begin{aligned} \mathcal{L}_{SM}(\phi) &=\frac{1}{2}\mathbb{E}_{x\sim p_{data}}\left[ \left\|s_\phi(x)\right\|_2^2 \right]+\frac{1}{2}\mathbb{E}_{x\sim p_{data}}\left[ \left\|s(x)\right\|_2^2 \right]-\mathbb{E}_{x\sim p_{data}}\left[ \left\langle s_\phi(x),s(x) \right\rangle \right]\\ &=\mathbb{E}_{x\sim p_{data}}\left[ \text{Tr}\left(\nabla_xs_\phi(x) \right) + \frac{1}{2}\left\|s_\phi(x)\right\|_2^2 \right]+\frac{1}{2}\mathbb{E}_{x\sim p_{data}}\left[ \left\|s(x)\right\|_2^2 \right]\\ &=\mathbb{E}_{x\sim p_{data}}\left[ \text{Tr}\left(\nabla_xs_\phi(x) \right) + \frac{1}{2}\left\|s_\phi(x)\right\|_2^2 \right]+C \end{aligned} \end{equation}\]其中，$C=\frac{1}{2}\mathbb{E}{x\sim p{data}}\left[ \left|s(x)\right|_{2}^{2} \right]$ 是与 $\phi$ 无关的常数。
得证。
