生成模型 (2.1):Energy-based Model

生成模型 #生成模型 #基于能量的模型 约 23 分钟 · 8033 字

引言

在前面的章节中,我们介绍了在变分推断的框架下的生成模型,包括VAE和DDPM两大类。在变分的视角下,生成模型实际上是在学习一个带参的变分分布 来近似难解的后验分布 ,并通过优化经验下界ELBO来最大化对数似然函数 。

从这一篇文章开始,我们将站在一个全新的视角来看待生成模型:基于分数 (score-based) 的视角。这里的分数实际上是一个梯度场 (gradient field),指向数据分布中高概率的方向。这种方法的中心思想在于:我们可以通过建模分数来进行生成式建模,因为我们可以利用梯度场将随机采样的样本推向数据分布 中高概率的方向,从而避免计算归一化常数中难解的积分项。因此,这类方法并不是直接学习数据分布的形式,而是学习如何将先验分布推向数据分布。

基于能量的模型 (Energy-based Models, EBMs) 是利用分数进行生成式建模的最简单的方法。能量是建模概率分布的一种方法,概率越高的地方能量越低。EBMs往往基于朗之万动力学 (Langevin dynamics) 进行采样,根据能量地图的梯度方向(即分数)将样本推向高概率区域。

一、基于能量的模型

令 表示一个数据点。EBMs通过一个能量函数 来定义一个概率分布:

(1)

其中,归一化项 用于保证概率分布的归一性:

(2)

从能量函数的定义我们可以看到,概率越低的数据点能量越高,反之亦然。

EBMs可以通过最大化对数似然来进行训练:

(3)

其中,第一项用于保证位于 的支撑集上的数据点具有较低的能量,第二项是正则化项。

正如前面提到的,当数据维度 较大时,计算第二项的积分是非常困难的。在变分方法中,我们使用可解的ELBO来近似难解的对数似然。在基于分数的方法中,我们通过 score matching 来直接避免计算这一项。

在下面的内容中,我们将首先介绍 score function 的形式化定义,然后介绍 score matching 是如何避免计算难解的归一化项的。最后,我们介绍利用 score function 进行采样的朗之万采样法。

1.1. 什么是分数?

对于定义在 上的概率分布 ,score function 是其对数密度的梯度:

(4)

因此,分数函数组成了 上的一个向量场,其中每一个向量都指向 中的高概率区域。

一个很自然且重要的问题是:为什么我们不直接建模概率密度,而是建模密度的梯度(即分数)?这是因为建模分数具有下面几个理论上和实践上的优势。

1. 不需要归一化常数。由于难解的归一化常数 往往与数据点 无关,因此其梯度为0,我们也就不需要计算这个难解的积分。

比如说在EBMs中,分数的定义为:

(5)

2. 分数函数完整刻画了概率密度。由于分数是对数密度的梯度,因此我们可以利用下面的公式,从分数重建概率密度:

(6)

其中, 是某个固定的参考数据点。

因此,建模分数和建模密度具有相同的表达能力,而分数往往是可解的。

1.2. 利用 score matching 训练EBMs

我们希望利用 中的样本来训练一个神经网络 用于近似分数函数 :

(7)

score matching 方法通过最小化真实分数和预测分数的均方误差 (MSE) 来拟合这个向量场:

(8)

然而,由于真实分数 是未知的,上面的损失函数是难解的。幸运的是,Hyvarinen找到了一个与上式等价且可解的损失函数,仅依赖于 和数据 ,不需要真实分数。下面的定理说明了这一点。

Theorem 1 (Hyvarinen’s Tractable Form of SM). 我们有下面等式成立:

(9)

其中, 是一个与 无关的常数,且

(10)

定理1的证明放在附录中。

我们可以从直观上理解这个可解的损失 :

  1. 第二项 是范数项,使得在 的高密度区域中,预测分数尽可能接近0,即梯度很小,让这些高密度区域成为驻点。
  2. 第一项 是散度项,鼓励 的散度为负,确保驻点是吸引子 (attractive sinks)。因为散度越小,表示向量场在该点向内汇聚,保证了这些高密度区域周围的梯度方向会指向驻点。

二、朗之万动力学采样

我们一般使用朗之万动力学 (Langevin dynamics) 来对EBMs进行采样。在下面的内容中,我们将首先介绍离散时间中的朗之万采样,并将其逼近至连续时间的场景中,使用随机微分方程来求解。最后,我们简要介绍一下朗之万动力学背后的物理直觉,以及其为什么适用于对EBMs进行采样。

2.1. 离散时间朗之万动力学

离散时间朗之万动力学 (Discrete-Time Langevin Dynamics) 可以又下面的迭代方程描述:

(11)

其中, 是从一个先验分布(一般是高斯分布)中采样得到的数据点。 是步长。 是高斯噪声,用于增加随机性,帮助模型跳出局部最小值。

2.2. 连续时间朗之万动力学

随着采样步长 ,上面的更新方程会收敛与下面的朗之万SDE:

(12)

其中, 表示一个标准布朗运动 (Standard Brownian Motion),也称为维纳过程 (Wiener Process),用于添加随机性。布朗运动的增量满足:。

这里的 是用于确保目标分布 是该SDE对应的平稳分布,这是由福克-普朗克方程 (Fokker-Planck Equation) 推导得到的,后面的内容中我们将进一步深入。

2.3. 为什么使用朗之万采样

理解朗之万采样的一种自然的方式是从物理世界的角度去理解。EBMs通过能量方程 定义了一个势能场,用于规范其中粒子的行为。根据牛顿力学,一个粒子在势能场中的运动轨迹由下面的常微分方程描述:

(13)

上面的常微分方程将粒子确定性地推向梯度下降的方向,直到去到能量方程的某个局部最小值,导致无法对整个数据分布进行全面地探索,可能会错过更优的全局最小值。

为了克服这个缺陷,朗之万动力学引入了一个由布朗运动所描述的随机扰动,得到了下面的SDE:

(14)

随机扰动可以帮助粒子跨越能量势垒以逃离局部最小值,使得粒子的轨迹形成一个随机过程,且其平稳分布将收敛止玻尔兹曼分布:

(15)

从这个角度来看,EBMs实际上学习了一个向量场,并使用朗之万采样将数据点推向高概率的区域。

然而,朗之万采样在高维空间中会面临一些问题。比如说,它对于一些超参数(如采样步长 、迭代步数等)非常敏感。这种缺陷的根源在于朗之万采样的“混合时间” (mixing time) 过长:在高维复杂概率分布中往往存在多个孤立的高概率区域,而朗之万采样需要非常长的迭代步数来在这些区域之间迁移。

附录

Proof on Theorem 1.

定理1提供了 score matching 训练目标的一个可解形式:

(16)

其中, 是一个与 无关的常数。

下面我们证明这一点。

首先我们展开 的形式:

(17)

这里我们关注第三项,即交叉相乘的项。由于:

(18)

代入得:

(19)

其中, 是分数函数的第 个分量:

(20)

下面我们将使用一个分部积分公式:

Lemma (Integration by Parts). 令 是以原点为中心,半径为 的封闭超球体,其边界 是光滑超球面。设函数 是在 上的可微实函数,则对 都有下面等式成立:

(21)

其中, 是边界 上的外法向单位向量, 是 上的面积微元。

令 ,并假设当 时,。

代入得:

(22)

因此,我们有:

(23)

其中, 是与 无关的常数。

得证。