Langevin Dynamics

在 Score-Based Generative Models 一文中，我们用到了 Langevin dynamics 进行采样，但没有解释其原理。另外，在 Elucidated Diffusion Models 一文中，我们提到常用的扩散 SDE 和 PF ODE 其实可以推广为一个由 PF ODE 与 Langevin diffusion SDE 组合而成的广义 SDE. 由此可见，Langevin dynamics 在扩散模型的研究中有着举足轻重的地位，因此本文探究其背后的原理。

物理背景

Langevin dynamics 脱胎于经典的布朗运动模型：将花粉放入水杯之中，则花粉会因受到水分子的无规律撞击而做布朗运动，如图所示：

本节将通过分析该模型下的牛顿第二定律、功能关系与 Langevin 方程，推导出 Langevin diffusion SDE，最后借助 Fokker-Planck 方程找到该 SDE 的平稳分布，为下一节基于 Langevin dynamics 的采样方法做准备。

牛顿第二定律

设花粉的质量为 $m$，在 $t$ 时刻位移为 $\mathbf{x}=\mathbf{x}(t)$，受到的瞬时合力为 $\mathbf{F}=\mathbf{F}(\mathbf{x}(t),t)$，则根据牛顿第二定律，有： $$\begin{array}{}\text{(1)}& \mathbf{F}(\mathbf{x},t)=m\frac{{\mathrm{d}}^2\mathbf{x}(t)}{\mathrm{d}{t}^2}\end{array}$$

功能关系

根据动能定理，合外力对物体做的功等于物体动能变化量，因此在一个微小时间 $\mathrm{\Delta }t$ 内，有： $$⟨\mathbf{F}(\mathbf{x},t),\mathrm{\Delta }\mathbf{x}⟩=\mathrm{\Delta }E(\mathbf{x})$$ 其中 $\mathrm{\Delta }\mathbf{x}$ 表示这一段时间内花粉的位移，$E(\mathbf{x})$ 表示花粉具有的动能。取 $\mathrm{\Delta }t\to 0$ 得： $$⟨\mathbf{F}(\mathbf{x},t),\mathrm{d}\mathbf{x}⟩=\mathrm{d}E(\mathbf{x})⟹\mathbf{F}(\mathbf{x},t)=\mathrm{\nabla }E(\mathbf{x})$$ 现在把整个水杯看作一个系统，则系统不受外力作用且只有保守力做功，因此整个系统机械能守恒，即： $$E(\mathbf{x})+U(\mathbf{x})\equiv \text{const},\mathrm{\forall }\mathbf{x}$$ 其中 $U(\mathbf{x})$ 表示系统的势能。这意味着 $\mathrm{\nabla }E(\mathbf{x})+\mathrm{\nabla }U(\mathbf{x})\equiv 0$，因此： $$\begin{array}{}\text{(2)}& \mathbf{F}(\mathbf{x},t)=\mathrm{\nabla }E(\mathbf{x})=-\mathrm{\nabla }U(\mathbf{x})\end{array}$$

Langevin 方程

花粉的运动还可以用 Langevin 方程描述： $$ Langevin 方程将花粉受到的力分为两部分：

• $$ 表示与速度大小正相关、方向相反的阻尼力，其中 $$ 为阻尼系数，由 Stokes' Law 给出；
• $$​ 表示粒子受周围水分子无规律撞击的随机涨落力，是一个零均值正态随机变量。

Langevin SDE

联立 $$ 式得： $$

为了将上式写作 SDE 的形式，引入标准维纳过程 $$ 代替 $$，那么有 $$，代入得：

$$ 注意到 $$ 是一个物理常数，简便起见不妨抛开其实际物理意义，直接令 $$ 得： $$ 进一步地，为了后续形式上的简便，将参数 $$ 换成 $$ 得： $$ 称上式为 Langevin diffusion SDE. 下图模拟了单独一个粒子依 Langevin diffusion SDE 运动的动态：

https://odie2630463.github.io/2020/09/28/langevin/

Langevin SDE 的平稳分布

考虑一堆依 Langevin diffusion SDE 运动的粒子，初始时它们服从某种概率分布，那么随运动进行它们的概率分布也将随时间变化。我们感兴趣的是，当运动足够长时间后，粒子的概率分布将收敛到怎样的平稳分布。

设 $$ 表示 $$ 时刻 $$ 处的概率密度， 则计算 $$ 式的 Fokker-Planck 方程，有： $$ Fokker-Planck 方程描述了概率分布随时间是如何变化的。为求解平稳分布（即 $$ 不随时间 $$ 变化），令上式为 0 得： $$ 可以验证，如下形式的玻尔兹曼分布是上式的解（事实上是唯一的解）： $$ 其中 $$ 为归一化系数。下图模拟了大量粒子依 Langevin diffusion SDE 运动的动态，并可视化了概率分布的变化：

source:https://odie2630463.github.io/2020/09/28/langevin/

Langevin MCMC

假设我们希望从分布 $$ 中采样，根据上一节的理论，只需要找到合适的 $$，使得平稳分布 $$，那么按照 Langevin dynamics 的过程模拟即可。这一思想与 MCMC 是一致的。事实上，考虑到 SDE 就是连续时间、连续状态的马尔可夫过程，所以人们常将这类采样方法称作 Langevin MCMC.

ULA

令 $$，代入 $$ 式得： $$ 因此我们要找的 $$ 就是 $$ 差一个常数的若干倍。两边求梯度： $$ 这就出现了熟悉的 score function. 代入 $$ 式得： $$

这就是用于采样的 Langevin diffusion SDE 形式。

实际操作时需要将 $$ 式进行 Euler-Maruyama 离散化： $$ 其中 $$ 为离散化的步长。称该采样算法为 ULA (Unadjusted Langevin Algorithm).

联系梯度下降

注意到，如果 $$ 式中没有随机项，那么 ULA 就变成了梯度下降（梯度上升）： $$ 其中梯度下降的“损失函数”就是负对数概率密度 $$. 从这个角度看，ULA 可以直观理解为梯度下降+随机扰动：

• 梯度下降项使得采样点不断向 $$ 增大的方向移动，即向高密度区域聚集；
• 随机扰动项使得优化算法变成了采样算法——如果只有梯度下降项，那么数据点最终会收敛到 $$ 极大的地方，但采样应该允许取到 $$​ 较小的区域，因此加入随机扰动项实现这一点。

下图展示了梯度下降与 ULA 的区别：

source:https://zhuanlan.zhihu.com/p/699598518

MALA

在 ULA 中，我们不假思索地直接将 $$ 式离散化为了 $$ 式，但离散化后我们真的能得到想要的分布吗？具体而言，离散化后的 Langevin dynamics 形成了离散时间的马尔可夫链，在合适的条件下，它是不可约正常返的，因此存在极限分布 $$ 且极限分布就是平稳分布。不过，由于离散误差，$$ 与我们想要的 $$ 之间存在一定的差异，一些理论工作正是研究不同步长 $$ 下 $$ 与 $$ 之间的差异。

为了改进 ULA 的误差，受到 Metropolis-Hastings 采样的启发，Roberts & Tweedie 提出了 Metropolis Adjusted Langevin Agorithm (MALA)：

1. 提议 $$；

2. 计算接受率： $$ 其中 $$ 表示取最小值操作；

3. 接受或拒绝提议。

这个链接中有非常直观的 MALA 算法采样过程的可视化。

References

1. Paweł A. Pierzchlewicz. Langevin Monte Carlo: Sampling using Langevin Dynamics. https://perceptron.blog/langevin-dynamics/ ↩︎
2. Chan, Stanley H. Tutorial on Diffusion Models for Imaging and Vision. arXiv preprint arXiv:2403.18103 (2024). ↩︎
3. 朗之万动力学原理简介 - Openverse的文章 - 知乎. https://zhuanlan.zhihu.com/p/699598518 ↩︎
4. Yuansi Chen. Lecture 3 – Langevin algorithms. https://metaphor.ethz.ch/x/2024/fs/401-4634-24L/lec/Lecture03.pdf ↩︎
5. Kirill Neklyudov. Langevin Dynamics for sampling and global optimization. https://docs.google.com/presentation/d/1_yekoTv_CHRgz6vsT57RMDESHjlnbGQvq8tYCxKLyW0 ↩︎
6. Alain Durmus, Eric Moulines. The Langevin MCMC: Theory and Methods. https://www.icts.res.in/sites/default/files/paap-2019-08-08-Eric%20Moulines.pdf ↩︎
7. Roberts, Gareth O., and Richard L. Tweedie. Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli (1996): 341-363. ↩︎
8. Roy Friedman. A Simplified Overview of Langevin Dynamics. https://friedmanroy.github.io/blog/2022/Langevin/ ↩︎