朗之万动力学与 SGLD

仅凭得分函数,通过带噪声的梯度上升从未归一化密度中采样。

朗之万动力学是一条马尔可夫链,它仅凭目标密度 p(x)p(x)得分 xlogp(x)\nabla_x \log p(x) 便能从中采样。每一步沿对数密度做梯度上升,并注入高斯噪声:

xt+1=xt+η2xlogp(xt)+η  ξt,ξtN(0,I).x_{t+1} = x_t + \frac{\eta}{2}\,\nabla_x \log p(x_t) + \sqrt{\eta}\;\xi_t, \qquad \xi_t \sim \mathcal{N}(0, I).

逐项来看:xtx_t 是第 tt 步的样本;η\eta 是步长(学习率);xlogp(xt)\nabla_x \log p(x_t) 是当前位置的得分,指向对数密度上升最快的方向;ξt\xi_t 是每步独立采样的标准高斯噪声,按 η\sqrt{\eta} 缩放后注入。漂移项前的系数是 η/2\eta/2、噪声项前的系数是 η\sqrt{\eta},这两个量并非任意——正是这个特定比例使该链的平稳分布恰为 pp 而非 pp 的某个幂或形变。

当步长 η\eta 足够小时,该链在分布意义下收敛到 pp。漂移项把样本拉向高概率区域,而噪声防止迭代塌缩到单一模态,使其能遍历整个分布。

直觉

logp\log p 的纯梯度上升会停在某个模态——用于寻找极大值有用,用于采样则无效。所加噪声把这种”上升”转为探索:行走者在概率高处徘徊,却仍四处游走,按各区域的质量比例造访它们。

可以把它想成一颗滚进势能面谷底的小球,谷越深代表概率越高,但小球同时被持续的随机踢动(温度)所搅动。没有踢动,它会卡在最近的谷底;踢动太强,它会无视势阱、近乎均匀地乱走。朗之万动力学把踢动的强度恰好调到使小球停留在每个谷的时间正比于该谷的概率质量——它在最深的阱里待得最久,却偶尔越过势垒去拜访较浅的阱。

数百个行走者从均匀分布出发,按上式漂移加噪,逐渐聚到双峰目标分布上:

目标分布 p(x)样本直方图

数百个粒子从均匀分布出发,按朗之万动力学沿对数密度的梯度漂移、并叠加随机噪声,逐渐聚到两个峰上。步长越大越快,但太大会越过细节、采样失真。只需要梯度、无需归一化常数。

注意行走者并非全部冲向较高的那个峰:两个峰都按其各自的概率质量留住了一部分行走者。这正是采样与优化的分野——优化只要那个全局最优,采样要的是整条分布。

关键的实用之处在于,得分只通过 xlogp\nabla_x \log p 依赖 pp,其中任何归一化常数都被消去。对于能量模型 p(x)eEθ(x)p(x)\propto e^{-E_\theta(x)}

xlogp(x)=xEθ(x),\nabla_x \log p(x) = -\nabla_x E_\theta(x),

因此朗之万动力学能从模型中采样而无需计算不可解的配分函数 ZθZ_\theta。原因在于 ZθZ_\theta 是不依赖 xx 的常数,对 xx 求梯度即归零;采样只看能量的形状(即不同位置之间的能量差),而不看它的绝对标定。这使朗之万动力学成为 EBM 与基于得分的生成模型的标准采样器:只要你能对一个未归一化的能量求导,就能从它对应的分布里抽样。

举一个具体的数:设一维高斯 p(x)=N(0,1)p(x)=\mathcal{N}(0,1),其得分为 xlogp(x)=x\nabla_x\log p(x)=-x。代入更新式得 xt+1=xtη2xt+ηξt=(1η2)xt+ηξtx_{t+1}=x_t-\tfrac{\eta}{2}x_t+\sqrt{\eta}\,\xi_t=(1-\tfrac{\eta}{2})x_t+\sqrt{\eta}\,\xi_t。这是一个一阶自回归过程:漂移把样本按比例往原点收缩,噪声再把方差补回去;在 η\eta 小的极限下,其平稳方差正好趋于 11,复现出目标高斯。这就是上面那条抽象更新式在最简单情形下的样子。

深入

上述离散更新只在 η0\eta\to 0 的连续极限下严格采到 pp。它是如下朗之万随机微分方程的欧拉–丸山离散:

dx=12xlogp(x)dt+dWt,dx = \tfrac{1}{2}\nabla_x \log p(x)\,dt + dW_t,

其中 WtW_t 为布朗运动。这条 SDE 的平稳分布恰为 pp,可由其福克–普朗克方程验证。有限步长 η\eta 会引入 O(η)O(\eta) 量级的离散偏差,使链收敛到一个略微走样的分布;Metropolis 调整朗之万算法(MALA)在每步后加一道接受/拒绝判定来消去这一偏差,代价是需要可求值的(未归一化)密度。

随机梯度朗之万动力学(SGLD)则面向另一种规模问题:当得分本身是对大数据集 NN 个样本求和(如贝叶斯后验 logp(θ)+ilogp(dataiθ)\log p(\theta)+\sum_i \log p(\mathrm{data}_i\mid\theta))时,用小批量估计取代全数据梯度,从而可扩展至大型数据集。当步长 ηt0\eta_t \to 0 按适当的调度(满足 tηt=\sum_t\eta_t=\inftytηt2<\sum_t\eta_t^2<\infty)衰减时,小批量噪声相对注入的朗之万噪声变得可忽略,链便从随机优化平滑过渡到后验采样。SGLD 因而能以约等于随机梯度下降的代价产生贝叶斯后验样本——早期大步长像 SGD 一样快速找到高概率区,后期小步长把迭代变成真正的后验采样器。

朗之万采样为从能量先验中抽取候选重构提供了内循环。在 Cryo-ET 重构里,密度 p(x)p(x) 是一个后验:能量先验(哪些三维结构合理)与数据项(重构需与所测倾斜投影一致)相加给出 logp-\log p,朗之万动力学便从中抽样。CryoWGEN-II 正是直接、迭代地用 Langevin/SGLD 采这个玻尔兹曼后验,得到的不是单一答案而是一族重构——每个样本是一个与数据相容的体,它们之间的离散度刻画了缺失楔形所留下的不确定性。这把生成建模的能量视角与变分视角联系起来:变分推断用一个可解分布去逼近后验,朗之万采样则直接从后验抽样,二者是逼近同一目标的两条路径。

← 推断与采样