前一篇概率采样文章讨论了基本采样、重要性采样和马尔可夫链蒙特卡洛:当目标分布已经写出,却无法直接完成积分时,用样本近似期望。还有一类模型更进一步,只给出每个状态的相对权重,连归一化常数都不便计算。高斯、伯努利等标准分布有成熟的直接采样算法;复杂模型则需要先解决怎样在未归一化密度上移动的问题。

能量模型(Energy-Based Model,EBM)把这种相对权重写成能量。低能量状态比高能量状态更可能出现,模型不必显式构造一个方向固定的生成过程。代价是似然中的配分函数通常无法精确计算,训练和生成都会转化为采样问题。

本文从能量函数定义概率分布,推导最大似然梯度中的数据期望与模型期望,再讨论什么条件下样本均值是无偏估计。后半部分用郎之万动力学(Langevin dynamics)构造只依赖能量梯度的马尔可夫链,并说明离散步长、链的相关性和混合速度会怎样影响实际结果。

能量模型

设随机变量 x\mathbf{x} 的状态空间为 X\mathcal{X},能量函数 Eθ(x)E_\theta(\mathbf{x}) 为每个状态分配一个实数。对应的概率密度为

pθ(x)=exp ⁣[Eθ(x)]Zθ,p_\theta(\mathbf{x}) = \frac{\exp\!\left[-E_\theta(\mathbf{x})\right]}{Z_\theta},

其中

Zθ=Xexp ⁣[Eθ(x)]dxZ_\theta = \int_{\mathcal{X}} \exp\!\left[-E_\theta(\mathbf{x})\right] \,\mathrm{d}\mathbf{x}

称为配分函数(partition function)。离散变量把积分换成求和。只要 ZθZ_\theta 有限,归一化后的 pθp_\theta 就是合法分布。

这个定义只关心能量差。若给所有状态的能量同时加上与 x\mathbf{x} 无关的常数 cc,分子和配分函数都会乘上 ece^{-c},概率分布保持不变。因此,能量的绝对零点没有概率意义;模型学习的是状态之间的相对偏好。

分类器的 logits 也可以从能量角度理解。令类别 yy 对应的能量为 Eθ(x,y)=fθ(x)yE_\theta(\mathbf{x},y)=-f_\theta(\mathbf{x})_y,条件分布就是

pθ(yx)=exp[Eθ(x,y)]yexp[Eθ(x,y)].p_\theta(y\mid\mathbf{x}) = \frac{\exp[-E_\theta(\mathbf{x},y)]} {\sum_{y'}\exp[-E_\theta(\mathbf{x},y')]}.

类别数有限时,分母可以直接求和,交叉熵训练因而很方便。若要为一张图像的所有可能像素组合建立联合分布,状态空间几乎无法枚举,配分函数便成为主要困难。

能量函数

能量函数可以是二次型,也可以由神经网络参数化。最简单的一维高斯分布可以写成

E(x)=(xμ)22σ2.E(x) = \frac{(x-\mu)^2}{2\sigma^2}.

代入能量模型后,配分函数会补上高斯密度中的归一化系数。双稳态能量

E(x)=14(x21)2E(x) = \frac{1}{4}(x^2-1)^2

则在 x=1x=-1x=1x=1 附近各有一个低能量区域,中间的 x=0x=0 是能垒。这样的例子能直接显示采样与优化的区别:优化只需找到一个低能量点,采样却要按每个区域所占的概率质量反复访问两个模式。

能量函数并非任意实值函数都能定义概率分布。在无界连续空间中,如果 Eθ(x)E_\theta(\mathbf{x}) 没有随 x\lVert\mathbf{x}\rVert 增大而充分上升,exp[Eθ(x)]\exp[-E_\theta(\mathbf{x})] 的积分可能发散。神经网络输出有限并不自动保证配分函数存在,模型结构、输入域与训练约束需要共同处理这个问题。

能量还决定了分布的分数函数(score):

xlogpθ(x)=xEθ(x).\nabla_{\mathbf{x}}\log p_\theta(\mathbf{x}) = -\nabla_{\mathbf{x}}E_\theta(\mathbf{x}).

logZθ\log Z_\theta 与当前状态 x\mathbf{x} 无关,求输入梯度时会消失。这个性质很重要:即使不知道归一化常数,仍然可以计算概率密度增长最快的局部方向。

无偏估计

很多概率计算最终都可以写成某个函数在目标分布下的期望:

μ=Exp[f(x)].\mu = \mathbb{E}_{\mathbf{x}\sim p}[f(\mathbf{x})].

若能得到 NN 个独立同分布的精确样本 x1,,xNp\mathbf{x}_1,\ldots,\mathbf{x}_N\sim p,样本均值

μ^N=1Ni=1Nf(xi)\hat\mu_N = \frac{1}{N}\sum_{i=1}^{N}f(\mathbf{x}_i)

满足 E[μ^N]=μ\mathbb{E}[\hat\mu_N]=\mu,因此是无偏估计。其方差为 Varp[f]/N\operatorname{Var}_p[f]/N,样本数增加时,估计会按蒙特卡洛方法的典型速度收敛。

“用样本求平均”本身不保证无偏。马尔可夫链中的相邻状态通常相关;若链从任意初值出发,有限步后的边缘分布也未必已经等于目标分布。相关性主要降低有效样本量,未达到平稳分布则还会引入初始化偏差。丢弃一段预热样本可以减弱初值影响,但不能证明剩余样本已经精确来自目标分布。

重要性采样也有相似区别。已知目标密度的归一化形式时,普通重要性采样可以构造无偏期望估计;能量模型通常只知道未归一化权重,此时常用自归一化重要性采样。权重再归一化形成的比值估计在有限样本下通常有偏,虽然在适当条件下仍是一致估计。讨论“无偏”时,必须同时说明样本如何获得、链是否平稳,以及估计式是否包含随机比值。

最大似然

设训练数据来自未知分布 pdatap_{\mathrm{data}}。单个样本的对数似然为

logpθ(x)=Eθ(x)logZθ.\log p_\theta(\mathbf{x}) = -E_\theta(\mathbf{x})-\log Z_\theta.

对参数求梯度,配分函数一项可以改写成模型分布下的期望:

θlogpθ(x)=θEθ(x)+Expθ[θEθ(x)].\nabla_\theta \log p_\theta(\mathbf{x}) = -\nabla_\theta E_\theta(\mathbf{x}) + \mathbb{E}_{\mathbf{x}'\sim p_\theta} \left[ \nabla_\theta E_\theta(\mathbf{x}') \right].

对数据分布再取期望,最大化平均对数似然等价于沿下面的方向更新:

θEpdata[logpθ(x)]=Epdata[θEθ(x)]+Epθ[θEθ(x)].\nabla_\theta \mathbb{E}_{p_{\mathrm{data}}} [\log p_\theta(\mathbf{x})] = - \mathbb{E}_{p_{\mathrm{data}}} [\nabla_\theta E_\theta(\mathbf{x})] + \mathbb{E}_{p_\theta} [\nabla_\theta E_\theta(\mathbf{x})].

第一项通常称为正相(positive phase)。它降低真实数据附近的能量,只需从训练集抽取小批量即可估计。第二项称为负相(negative phase)。它提高模型当前偏爱状态的能量,防止模型把整个空间都压成低能量。负相要求从不断变化的 pθp_\theta 中采样,这正是最大似然训练的计算瓶颈。

直接计算 ZθZ_\theta 需要遍历状态空间;用数值积分只能处理很低的维度;朴素重要性采样在目标分布与提议分布重叠很小时会产生极不稳定的权重。神经能量模型因此常用马尔可夫链近似负相期望。近似链若只停留在少数模式,训练梯度就只能纠正这些区域,其他被模型错误赋予高概率的状态可能长期没有进入负样本。

郎之万动力学

郎之万动力学把沿能量下降的确定性运动与布朗噪声结合起来。采用一种常见的尺度约定,连续时间随机微分方程为

dXt=xE(Xt)dt+2dWt,\mathrm{d}\mathbf{X}_t = -\nabla_{\mathbf{x}}E(\mathbf{X}_t)\,\mathrm{d}t + \sqrt{2}\,\mathrm{d}\mathbf{W}_t,

其中 Wt\mathbf{W}_t 是标准布朗运动。梯度项把状态推向低能量区域,噪声项让状态能够离开局部低谷。二者达到平衡时,这个扩散过程的平稳密度与 exp[E(x)]\exp[-E(\mathbf{x})] 成正比。

若只保留梯度项,方程退化为梯度流,轨迹最终会停在某个局部极小值,无法按概率质量产生样本。若只保留噪声,状态会做没有目标的随机游走。郎之万动力学利用能量梯度决定局部方向,再用噪声保留分布所要求的不确定性。

引入温度 TT 后,目标分布写成 pT(x)exp[E(x)/T]p_T(\mathbf{x})\propto\exp[-E(\mathbf{x})/T],对应动力学可以写成

dXt=E(Xt)dt+2TdWt.\mathrm{d}\mathbf{X}_t = -\nabla E(\mathbf{X}_t)\,\mathrm{d}t + \sqrt{2T}\,\mathrm{d}\mathbf{W}_t.

温度较低时,样本更集中在低能量区域,也更难越过能垒;温度较高时,链更容易探索,但目标分布同时变得更平坦。温度不是只影响运行速度的调参项,它会改变采样分布本身。

未校正郎之万算法

计算机只能使用离散时间。用 Euler–Maruyama 方法离散化上面的随机微分方程,可以得到未校正郎之万算法(Unadjusted Langevin Algorithm,ULA):

xk+1=xkεxE(xk)+2εξk,ξkN(0,I).\mathbf{x}_{k+1} = \mathbf{x}_k - \varepsilon\nabla_{\mathbf{x}}E(\mathbf{x}_k) + \sqrt{2\varepsilon}\,\boldsymbol{\xi}_k, \qquad \boldsymbol{\xi}_k\sim\mathcal{N}(\mathbf{0},\mathbf{I}).

也有文献把步长写成 η/2\eta/2,相应地把噪声尺度写成 η\sqrt{\eta}。两种形式只是时间尺度约定不同,使用公式时需要让漂移项与噪声项保持配套。

下面的 PyTorch 函数实现了这一更新。energy 接收一批状态并为每个状态返回一个标量能量:

import math
import torch


def ula_sample(energy, initial, step_size, steps):
    x = initial.detach().clone()

    for _ in range(steps):
        x.requires_grad_(True)
        total_energy = energy(x).sum()
        (energy_grad,) = torch.autograd.grad(total_energy, x)

        with torch.no_grad():
            noise = torch.randn_like(x)
            x = x - step_size * energy_grad
            x = x + math.sqrt(2.0 * step_size) * noise

    return x

步长 ε\varepsilon 过大时,离散轨迹可能越过能量曲面的细节,甚至数值发散;步长过小时,单步近似更准确,但有限计算预算内移动距离很短。更关键的是,固定非零步长的 ULA 通常不以精确目标分布为平稳分布,它带有离散化偏差。让步长逐渐减小可以在一定条件下恢复渐近正确性,却会降低后期探索速度。

Metropolis 调整郎之万算法(Metropolis-Adjusted Langevin Algorithm,MALA)把郎之万更新作为提议,再用接受—拒绝步骤修正离散化误差。这样可以保持精确的目标分布,但每步需要计算提议概率与能量差,拒绝也会造成重复状态。高维神经能量模型训练常直接使用有限步 ULA,因为目标是得到有用的负样本,而非完成严格的无偏积分;这项工程选择同时意味着训练梯度是近似的。

对比散度

若每次负相采样都从随机噪声开始,并把链运行到充分混合,成本通常难以接受。对比散度(Contrastive Divergence,CD)从真实数据附近启动短链,只观察模型经过少数步后会把数据推向哪里。它计算便宜,也能在一些模型上得到有用参数,但短链分布不等于 pθp_\theta,所得更新不是精确最大似然梯度。

持续性对比散度(Persistent Contrastive Divergence,PCD)保存上一轮的链状态,参数更新后继续推进,而不是每次重新初始化。若参数变化相对缓慢,持久链有机会跟随模型分布。现代实现还常维护一个重放缓冲区,在历史负样本和随机重启之间混合,以减少所有链长期困在同一批模式的风险。

无论使用 CD、PCD 还是重放缓冲区,训练循环都包含同一组动作:对数据样本计算正相,对当前负样本运行若干步郎之万更新,再用正负相的能量梯度更新参数。实现时通常要在采样阶段对状态求梯度,但不保留跨越所有采样步的参数计算图,否则内存会随链长迅速增长。采样得到的状态会先 detach,再用于负相损失。

采样诊断

单看样本是否“像数据”不足以判断采样器是否正确。链可能生成局部清晰的结果,却只覆盖目标分布中的少数模式。对于低维问题,可以比较直方图、核密度估计和已知目标密度;高维问题则需要同时检查多条独立链、能量轨迹、不同初始化的结果和任务相关统计量。

自相关函数描述同一条链间隔若干步后仍保留多少依赖。高度自相关意味着保存一万个状态也可能只相当于少量独立样本。有效样本量用自相关时间折算样本数,能比原始链长更准确地反映蒙特卡洛误差。多链诊断还可以发现不同初值长期停留在不同模式的情况,但任何有限诊断都不能证明高维链已经遍历完整状态空间。

预热、抽稀和多链运行解决的问题不同。预热用于减弱初值偏差;抽稀通过丢弃中间状态减少存储和相邻相关性,却不会凭空增加单位计算量得到的信息;多链可以暴露部分混合问题,也会增加并行计算成本。若能保存完整链并在分析阶段处理自相关,机械地每隔固定步数抽稀通常没有必要。

适用边界

郎之万采样要求能量对状态连续可微,因此不能直接更新离散符号。离散数据可以改用 Gibbs 采样、Metropolis–Hastings,或者先映射到连续表示再采样;每种处理都会改变模型或引入新的近似。带有边界约束的连续变量也需要投影、反射或变量变换,不能让高斯噪声把状态随意推到定义域外。

高维多峰分布是更普遍的困难。局部梯度能指向附近的低能量区域,却不会告诉采样器远处还有另一个模式。能垒较高时,普通郎之万链可能在训练时间内从未跨越模式。退火、并行温度链、哈密顿蒙特卡洛和学习式提议分布可以改善特定问题,但没有一种局部采样器能自动消除所有混合瓶颈。

能量模型的吸引力在于,它允许直接规定哪些状态应当兼容,并且采样时不需要知道配分函数。相同的配分函数又构成最大似然训练的核心障碍。郎之万动力学利用 xlogp=xE\nabla_{\mathbf{x}}\log p=-\nabla_{\mathbf{x}}E 绕开归一化常数,把低能量方向与随机探索结合起来;有限步离散链仍然带有初始化、相关性、混合和步长误差。理解这些误差来自哪里,才能判断样本均值何时接近目标期望,何时只是模型局部能量地形的一次短暂巡游。

参考文献

  1. Y. LeCun, S. Chopra, R. Hadsell, M. Ranzato, and F. Huang. A Tutorial on Energy-Based Learning. Predicting Structured Data, 2006
  2. M. Welling and Y. W. Teh. Bayesian Learning via Stochastic Gradient Langevin Dynamics. ICML, 2011
  3. Y. Du and I. Mordatch. Implicit Generation and Modeling with Energy Based Models. NeurIPS, 2019
  4. E. Nijkamp, M. Hill, S.-C. Zhu, and Y. N. Wu. Learning Non-Convergent Non-Persistent Short-Run MCMC Toward Energy-Based Model. NeurIPS, 2019