HMC:哈密顿蒙特卡洛方法

时间:2024-03-06 20:08:39

哈密顿蒙特卡洛方法

1. 简介

哈密顿蒙特卡洛方法(Hamiltonian Monte Carlo)这一算法最早出现在由Dunne等人在1987年发表的一篇论文中,最初他们称之为混合蒙特卡洛(Hybird Monte Carlo),是因为该方法将分子动力学理论和马尔科夫链蒙特卡洛方法(Markov Chain Monte Carlo)混合在了一起。但Hamiltonian MC这个称呼更能具体的描述所引入的Hamiltonian动力学,因而更广为人知。HMC在统计问题上的应用始于(Neal,1996a)用于神经网络模型的参数估计。在Neal看来,这是一个被统计学家低估的算法,时至今日,对于高维的参数分布估计问题,HMC仍然是最好的方法,如贝叶斯深度学习中使用HMC推断网络参数。
掌握HMC方法需要把握三个关键:

  • Metropolis-Hasting算法
  • 哈密顿动力学基础
  • 蛙跳法积分

2. 准备

2.1 Metropolis-Hasting算法

本文假设读者拥有一定的相关统计知识,更详细的马尔可夫链蒙特卡洛内容等可以参考

  1. 统计学习方法 李航
  2. 随机过程基础
  3. 博客:基础马尔可夫链概念,MCMC算法。

简单来说,MH算法是MCMC算法的一种易于实现的改良版本。在回顾MH算法之前,这里对Markov Chain相关概念也进行简要介绍。马尔可夫链是在参数空间中通过序列地应用被称作马尔可夫转移的随机映射而产生一系列连续点集。一个任意的马尔可夫链可以简单地在参数空间中漫游,它本身在计算期望是没有特殊作用的,然而当它漫游次数足够,接近并保持为目标分布,这就意味着该链是通过目标分布生成的一个样本集合。而且后续的漫游产生的样本集合也任然保持在这个目标分布当中。这样的样本属于目标分的典型集(Typical set,可以类似理解为大部分的样本所在的集合中,如正态分布的\(\mu\pm2\sigma\)区间)。因此,理论上无论从参数空间的什么位置开始,相应的马尔可夫链最终都会进入并且跨越目标分布的典型集。通过这些样本集进行期望计算 \(E(\widehat{f})\),以及估计整个参数空间的期望 \(E(f)\)

图 1 绿色的部分代表马尔可夫转移密度,它决定了新点产生的概率

MCMC方法定义了在给定目标分布的马尔可夫转移情况下量化典型集的一般策略。然而如何构建这样的转移本身就是一个不小的问题。幸运的是,对于任意给定的目标分布,都有各种自动构建合适的转移的程序,其中最重要的就是MH算法。
MH算法分两步:建议(proposal)与 修正(correction)。这里的建议是指给初始状态施加一个任意的随机扰动,而修正指的是拒绝任何偏离目标分布的典型集太远的的建议。更正式的说,假设 \(Q(q^*|q)\)为定义的每一个建议的概率密度,接受给定的建议的的概率则是:

\[a(q^*|q)= \min(1, \frac{Q(q|q^*)\pi(q^*)}{Q(q^*|q)\pi(q)}) \]

\(\pi\)表示目标分布,\(q\)是感兴趣的变量,一般来说,采用高斯分布作为建议,这一算法至今仍然被广泛使用。即:

\[Q(q^*|q)=\mathcal{N}(q^*|q,\Sigma)) \]

采用高斯扰动的好处是建议机制在初始点和建议点是对称等价的,因此它的接受率可以简化成;

\[a(q^*|q)= \min(1, \frac{\pi(q^*)}{\pi(q)}) \]

随机游走的MH算法实现简单,非常符合直觉,建议点的转移趋向于概率大的邻域(若是从后续提到的正则分布中的能量的角度来说,建议点的转移则是优先能量低的点)。这样的 建议 然后修正 的组合程序优先选择了落在高概率质量的点,并集中于典型集之内。但随着目标分布的维度和复杂性的增加,这一算法的性能大打折扣。尽管可以通过调整建议分布的尺度来获得更高的接受率,但会导致马尔可夫链移动缓慢,甚至是失败的探索,得到高度偏置的估计量。如果需要扩展到高维概率分布中去,那么就需要一个更好的方法来探索典型集,特别是需要利用好典型集本身的集合性质,HMC就是这样的方法。

2.2 哈密顿动力学简介

哈密顿动力学(Hamiltonian Dynamics)是拉格朗日力学的勒让德变换。它可以从一个直观且经典的物理解释来简单理解。拉格朗日运动方程是基于广义坐标和广义速度的,而哈密顿运动方程则是基于广义坐标和广义动量的。想像一个冰球在一个光滑表面上滑动,这个系统的能量状态包括冰球的势能 \(U\) 和动能 \(K\)。冰球的势能 \(U(q)\) 与当前的位置 \(q\) 有关,它的动能 \(K(q)\) 与动量 \(p\) 有关。当冰球遇到上坡会减速,动能减少,势能增加;当它遇到下坡会加速,动能增加,势能减少。在非物理的MCMC应用中,HMC中的位置对应于感兴趣的变量。势能等于这些变量的负对数概率密度。每一个位置变量都会对应一个动量变量,动量由人为引入。这一物理先验的引入,最大的意义在于在MCMC方法中(更确切的说MH算法),通过动量项避免了采样点完全的随机游走。

2.2.1 哈密顿方程

哈密顿动力学作用于一个 \(d\)维位置向量 \(q\)和一个 \(d\) 维动量向量 \(p\) ,因此整个状态空间是2维的。这个系统由一个成为哈密顿量的函数 \(H=(q,p)\)
它的运动方程描述如下:

\[\begin{align} &\frac{dq_{i}}{dt}= \frac{∂H}{∂p_{i}} \tag{2.1}\\ &\frac{dp_{i}}{dt}=-\frac{∂H}{∂q_{i}} \notag \\ \end{align} \]

对于任何的时间间隔 \(s\),都定义了一个从任意时间状态 \(t\)\(t+s\) 的映射 \(T_{s}\)。作为替代,可以使用向量\(q\)和向量\(p\)构造一个新的二维状态向量\(z=(p,q)\),此时哈密顿方程可以写为

\[\frac{dz}{dt}=J∇H(z) \]

\(∇H\)\(H\) 的梯度(即 \([∇H]_{k}=∂H/∂z_{k}\)),并且

\[J=\left[\begin{array}{rl} 0_{d \times d} & I_{d \times d} \\ -I_{d \times d} & 0_{d \times d} \end{array}\right]\]

是一个由单位矩阵和零矩阵构建的 \(2d×2d\) 的矩阵。通常情况下,HMC当中可以将哈密顿方程写成如下形式 $$H(q, p)=U(q)+K(p) \tag{2.2}$$ 这里的 \(U(q)\)被称为势能,并且被定义为我们想要采样的参数 \(q\)对应分布的负对数概率密度加上一个常数。\(K(p)\)被称作动能,通常被定义为

\[K(p)=p^{T}M^{-1}p/2 \tag{2.3} \]

这里 \(M\) 是一个对称对角正定的“质量矩阵”,通常为单位矩阵的标量倍数。这种形式的的 \(K(p)\) 对应于一个均值为0,协方差矩阵为\(M\)的高斯分布的对数概率密度。 \(H\)\(K\)在这种形势下,方程\((2.1)\)可以写成

\[\begin{align} \frac{d q_{i}}{d t} &=\left[M^{-1} p\right]_{i} \tag{2.4}\\ \frac{d p_{i}}{d t} &=-\frac{\partial U}{\partial q_{i}} \notag \end{align} \]

2.2.2 哈密顿动力学的性质

在构造MCMC更新时,哈密顿动力学的几个性质是至关重要的,这些性质的在相关文献中有详细介绍:

  1. 可逆性:保持期望的分布不变
  2. 守恒性
  3. 体积保持性:在进行MH更新时,不需要考虑积分空间体积变化,概率质量得以保证。
  4. 辛性

2.3 离散哈密顿方程——蛙跳法(Leapfrog Method)

为了在计算机上模拟哈密顿动力学方程,它必须使用若干离散的小步长 \(\varepsilon\)的时间来求和近似。从时刻0开始计算 \(\varepsilon,2\varepsilon,3\varepsilon...\)。通常假设哈密顿量的形式为\((2.2)\),虽然\(K(p)\)可以时任意形式,但为了简单起见,假设动能形如\((2.3)\)。而且,\(M\)是由对角元素 \(m_{1},...,m_{d}\) 构成的对角矩阵,所以

\[K(p)=\sum^{d}_{i=1}\frac{p^2_{i}}{2m_{i}} \tag{2.5} \]

从方程 \((2.4)\)利用欧拉法可以很简单的得到:

\[\begin{aligned} &p_{i}(t+\varepsilon)=p_{i}(t)+\varepsilon \frac{d p_{i}}{d t}(t)=p_{i}(t)-\varepsilon \frac{\partial U}{\partial q_{i}}(q(t)) \\ &q_{i}(t+\varepsilon)=q_{i}(t)+\varepsilon \frac{d q_{i}}{d t}(t)=q_{i}(t)+\varepsilon \frac{p_{i}(t)}{m_{i}} \end{aligned} \]

但是欧拉法容易发散,在分子动力学中的运动方程模拟主要还是使用蛙跳法,可以在精度和收敛方便做得更好,且蛙跳法积分在时间上是可逆的,这对于马尔可夫链达到细致平衡条件更有利。

\[\begin{aligned} p_{i}(t+\varepsilon / 2) &=p_{i}(t)-(\varepsilon / 2) \frac{\partial U}{\partial q_{i}}(q(t)) \\ q_{i}(t+\varepsilon) &=q_{i}(t)+\varepsilon \frac{p_{i}(t+\varepsilon / 2)}{m_{i}} \\ p_{i}(t+\varepsilon) &=p_{i}(t+\varepsilon / 2)-(\varepsilon / 2) \frac{\partial U}{\partial q_{i}}(q(t+\varepsilon)) \end{aligned} \]

它先更新半个时间步长获取动量新值,然后用动量的新值对位置变量进行一个完整的时间步长 \(\varepsilon\)的更新,最后再用新的位置变量对动量变量的新值进行另一个半步更新获得完整的动量更新。这种更新方法每次只更新动能或者位置,但因为\((2.6)\)是剪切变换,因此蛙跳法准确的保留了积分空间体积。

图 2 $H(q,p)=q^2/2+p^2/2$时,不同的积分算法和参数设置对精度的影响。模拟20个时间步

3. HMC算法

推荐阅读:

3.1 概率和哈密顿量:正则分布

准备好以上知识之后,HMC算法便可以这样理解:将目标分布的概率密度函数转化为哈密顿动力学中对应的势能函数,然后引入动量变量 \(p\)来配合位置变量 \(q\)(感兴趣的参数)模拟该变量一定时间长度的运动学轨迹并到达该次模拟的终点(即建议点),然后对哈密顿动力学发现的建议点进行Metropolis-Hasting更新,之后在每个迭代中重新采样动量,重复执行直至模拟出一个马尔可夫链。
实际上,在任何MCMC方法中我们想要采样的分布都可以通过统计力学中的吉布斯正则分布(Canonical Distribution)与能量函数联系起来。给定某个物理系统的某个状态 \(x\)的能量函数 \(E(x)\),在统计力学中它的正则分布有概率且概率密度函数为: $$P(x)=\frac{1}{Z}\exp(\frac{-E(x)}{T}) \tag{3.1}$$ 这里的 \(T\)代表系统的温度,\(Z\)代表归一化常数。反过来看,如果我们对一些概率密度函数 \(P(x)\)感兴趣,可以假设 \(T=1\)\(E(x)=-\log{P(x)}-\log{Z}\)来得到一个正则分布,其中 \(Z\) 可以是任意值,因为后续的计算中会约去。
HMC中哈密顿量是位置 \(q\)和动量\(p\)联合的能量函数,因此 $$P(q,p)=\frac{1}{Z}\exp(\frac{-H(q,p)}{T})$$如果 \(H(q,p)=U(q)+K(p)\), 那么联合概率分布为:

\[P(q,p)=\frac{1}{Z}\exp(\frac{-U(q)}{T})\exp(\frac{-K(q)}{T}) \tag{3.2} \]

并且注意到 \(p\)\(q\)是相互独立的且均存在一个能量分别为 \(U(q)\), \(K(p)\)的独立的正则分布。在贝叶斯统计学中,参数的后验分布往往是关注的焦点。这些参数扮演着HMC中的位置变量 \(q\)。我们可以将后验分布的正则分布 \((T=1)\)用势能函数表达为: $$U(q)=-\log{[L(q|D)\pi(q)]} \tag{3.3}$$这里的 \(\pi(q)\)表示 \(q\)的先验概率密度函数 \(L(q|D)\)表示在给定数据 \(D\)情况下的似然函数。
至此,我们对HMC有了整体的了解。HMC方法只能用于 \(R^d\)上的连续分布中取样,这些分布的概率密度函数可以被评估且处处不为零。我们还必须能够计算该概率密度函数的对数的偏导数。HMC从公式 \((3.2)\)定义的正则分布中抽取 \(q\)\(p\),其中 \(q\)由势能函数\(U(q)\)指定。我们可以选择动量 \(p\)的分布,正如之前在公式 \((2.3),(2.5)\)所描述一样。

3.2 HMC的两个步骤

HMC算法的每个迭代都有两个步骤:第一步动量会改变,第二步可能同时改变位置和动量。在第一步中,动量变量是从给定的高斯分布中随机抽取的,与位置变量的当前值无关。对于公式\((2.5)\)中每个动量变量 \(p_i\)它的方差为 \(m_i\)。在第二步中,利用哈密顿动力学模拟并建议一个新的位置变量,然后进行一次Metropolis更新。从当前状态 \((q,p)\)开始,利用leapfrog方法(或其他可以保持体积的可逆方法)进行 \(L\) 次步长为 \(\varepsilon\)的动力学模拟。这两个参数需要调节以获得最好的模拟效果,且较为敏感,一旦设置不合理,不能保证MCMC方法所需要满足的各种性质将不能获得正确的结果。在模拟完该次动力学轨迹后,在建议位置将动量变量置负 (保证Metropolis建议的对称性,实践中是不需要的,因为二次动能 \(K(p)=K(-p)\)),然后给出该次模拟最终的建议状态 \((q^*,p^*)\)。这一状态被作为马尔可夫链的下一个状态的概率,即接受率为:

\[\min[1,\exp(-H(q^*,p^*)+H(q,p))] = \min[1,\exp(-U(q^*)+U(q)-K(q^*)+K(p))] \]

如果状态不被接受,那下一个状态与当前状态一样(拒绝的点也会参与蒙特卡罗积分,只是没有发生转移,但也是马氏链的一个状态)。
值得注意的是,第一步的动量重采样对实施HMC是至关重要的,因为哈密顿量 \(H\)是不变或基本不变的,否则模拟的运动方程将只在同一概率测度的超表面(可以理解为能级)移动,并不能近似模拟出目标分布的典型集。通过不断的改变每次模拟的动量,就像是在无摩擦光滑地面上给冰球随便踹一脚,冰球就有机会遍历该地面的不同势能的位置。这种解释不够严谨但很直观。另外第二步中的Metropolis更新在一定程度上补偿了初始位置和建议位置的点因模拟而产生的能量差异,没有被精确模拟的运动轨迹也有一定的概率接受,这也是HMC算法在超参数合理情况下接受率高的另一个原因。

4. Python简单实例

在熟悉了HMC的理论之后,我们可以将其用于从已知概率密度复杂分布中抽取样本。但很多情况下,我们面对的都是模型参数估计问题。在贝叶斯数据分析中,感兴趣的往往是参数难以计算的后验分布,应该这里给出一个简单的例子来说明如何实现HMC并估计模型中的未知参数。
已知模型为 \(f(x)=ax^{b} + 450\)。 模型参数 \(q=(a,b)\), 其中 \(a=10\), \(b=0.2\)。假设我们现在只观测到区间 \(x\in [0,4]\)区间上的若干带噪音的数据 \(\mathcal{D}=(\boldsymbol{x_{obs}},\boldsymbol{y_{obs}})\)(图3), 并不清楚该非线性模型中参数的具体数据,希望通过这些观测值去推理模型参数的大致的取值分布。

图 3 观测数据与真实数据

首先,我们需要构建公式\((3.3)\)所描绘的势能函数将需要解决的问题与已知数据连接起来。观测值与模型输出之间的残差为: $$\boldsymbol{\varepsilon}=\boldsymbol{y_{obs}}-f(\boldsymbol{x_{obs}},q)$$ 假设残差\(\varepsilon\)为同方差高斯噪音,即 \(\boldsymbol{\varepsilon} \sim \mathcal{N}(0,I)\),因为各观测值相互独立,模型参数 \(q\)的似然函数便可以写为:

\[L(q|\boldsymbol{\varepsilon})=P(\boldsymbol{\varepsilon}|q)=\Pi^{n}_{i=0}P(\varepsilon_i|q)=\Pi^n_{i=0}\frac{1}{\sqrt{2\pi}}\exp(-\frac{\varepsilon_i^2}{2}) \]

两边取负对数:

\[U(q)=-\log{L(q|\boldsymbol{\varepsilon})}=\sum^n_{i=0}\frac{1}{2}\log{2\pi}+\frac{1}{2}\varepsilon_i^2 \]

因为这里使用的是极大似然估计,没有关于参数的先验信息(无信息先验),该式即为本算例的势能函数。利用构造的势能函数带入HMC流程中,模拟出目标参数的典型集即可求解模型的参数分布。图4,图5给出了最后的估计结果。算例代码

图 4 HMC采样的结果。当积分步长和积分次数选择得当,该系统运动方程将被精确模拟,因此接受率很高,且模拟的马氏链可以很快接近并很好的跨越参数联合分布的典型集。从图中可以看出各状态之间的相关性是比较低的,相邻两个状态并不是只停留在很小的邻域之内

图 5 参数a,b的边缘分布。因为HMC更好的模拟的参数的典型集样本,所以对于参数的期望和方差分析也更为合理

总结

HMC算法毫无疑问是优秀的,它给高维复杂分布的蒙特卡洛积分提供了优秀的解决方案。但是HMC对于leapfrog积分方法中的\(L\)\(\varepsilon\)太过敏感,这对于HMC在很多实际应用中的实施面临困难,同时HMC也有可能面对陷入“局部”地区的困境。HMC依赖于势能函数的一阶导数,多次的积分运算模拟使得在单次蒙特卡洛模拟中耗时较多。但相比于MH采样,整体是还是大大提升了MCMC方法在较复杂问题上的效率和精度,接受率大大提高,穿越并停留在目标分布的典型集的能力大大提升,因此也是优秀的不确定性量化方法。学者们还不遗余力的发展了诸如NUTS, SGHMC等一系列改进的算法。其中NUTS算法凭借其自适应的调节\(L\)\(\varepsilon\)的能力和有效性在多个库中作为默认采样算法。