ややプログラム紀行

博士1年のプログラムに関する日記

Hamiltonian Monte Carlo

何となくHamiltonian Monte Carloの挙動を実際に見てみたかったので実装してみた
Hamiltonian Monte Carloについては
arxiv.org
がわかりやすかった*1

Markov Chain Monte Carlo

Hamiltonian Monte Carlo、略してHMCはMarkov Chain Monte Carlo(MCMC)に分子動力学法の知識を援用したものだけど、そもそもMCMCとは何かについて軽く触れておく

MCMCの前2ワードのMarkov chainとはインデックス集合が自然数であるMarkov processのことなのだけど、厳密さを抜きに話すと確率過程 (w_t)_{t \in \mathcal{N}_0}のうち各 t \in \mathcal{N}_0について w_{t+1}の分布が p(w_{t+1}| w_t)に従うようなものを指す

簡単な例としては\mathbb{Z}上のランダムウォーク*2が挙げられる

このMarkov chain (w_t)_{t \in \mathcal{N}_0}の各w_tの分布を\mu_tと書くことにすると、\mu_tは適当な条件のもと t \to \inftyで定常分布 \piに収束することが知られている

そこで、その条件を満たしつつ所望の確率分布\piに収束するようなMarkov chainをうまく構成することができれば、そのMarkov chainを十分なステップ更新することで近似的に\piからサンプリングを行うことができ、この値を用いてMonte Carloを行う手法をMarkov Chain Monte Carloと呼ぶ


これだけでは何だかサンプリングするだけなのになぜこんな遠回りを、、、と思うかもしれないが、MCMCにはベイズ推定などと親和性が高いアルゴリズム的な理由がある

ベイズ推定では確率モデルp(x \mid w)と事前分布\phi(w)に対してX_1,\dots,X_nが観測されたとすると、事後分布は

\begin{align*}
p(w) &= \frac{1}{Z_n(\beta)} \phi(w) \prod_{i=1}^n p(X_i | w)^\beta = \frac{1}{Z_n(\beta)} \exp(-\beta H(w))
\end{align*}
として定式化される(ただし\betaは逆温度で、H(w) = -\sum_{i=1}^n \log p(X_i | w) - \frac{1}{\beta} \log \phi(w)かつZ_n(\beta)は正規化項)

ベイズ推定では一般的に、正規化項 Z_n(\beta)を解析的に求めることが難しく、それゆえ事後分布から直接サンプリングを試みることは出来ないわけだが、MCMCでは正規化項が計算できなくとも、 H(w)さえ計算できれば事後分布から近似的にサンプリングを行うことができる

Metropolis Algorithm

上ではベイズ推定を例に出したが、より一般に\frac{1}{Z}\exp(-\beta H(w))の形をしている分布に対し、正規化項Zを計算せずにサンプリングを行えることをMCMCの中でも一番有名なMetropolis法を用いて見ていく

適当な初期状態 w(1)からはじめ、t-1ステップ後にw(t)が得られている時、まず提案分布 r(w' | w(t))によって次の遷移先の候補w'をサンプリングする

ここで r(w' | w(t))は基本的には何でも良いが、対称的、すなわちr(w_1|w_2) = r(w_2|w_1)である必要がある*3

こうして得られた遷移先の候補 w'に対して、

\begin{align*}
P := \min \{1, \exp(-\beta \Delta H)\} \quad (\text{ただし}\Delta H := H(w') - H(w))
\end{align*}
と定義し、確率Pw'に遷移(すなわち w(t+1) = w')し、確率1-Pwにとどまる(すなわち w(t+1) = w)ことを繰り返すことで、十分大きな tに対してw(t)が近似的に \frac{1}{Z}\exp(-\beta H(w))に従い、正規化項Zを計算することなくサンプリングを行うことができる


上で適当な条件のもとMarkov chainは \piに収束すると述べたが、その具体的な条件とは雑に言って

  •  \piが定常分布*4
  • 任意の状態から任意の状態に到達可能
  • Markov chainが周期的でない*5

が満たされれば良く、MCMCでは特に\piが定常分布であることの十分条件である詳細釣り合い条件

\begin{align*}
p(w_b | w_a)\pi(w_a) = p(w_a | w_b)\pi(w_b)
\end{align*}
が満たされているかどうかを確かめることが多く*6、上のMetropolis法も詳細釣り合い条件を満たしていることが示せる


下のgifはMetropolis法を用いて混合正規分布

\begin{align*}
&0.3 \times \mathcal{N}((-1,-1)^\top, 1_2) + 0.7 \times \mathcal{N}((2,2)^\top, 1_2) \\
&= 0.3 \times \frac{1}{2\pi}e^{-\frac{(x_1 + 1)^2 + (x_2 + 1)^2}{2}} + 0.7 \times \frac{1}{2\pi}e^{-\frac{(x_1 - 2)^2 + (x_2 - 2)^2}{2}}
\end{align*}
からサンプリングを行ったもの(提案分布は標準正規分布を用いている)

大きい丸は遷移先の候補で、それが採用されると小さい丸になるようにプロットしている
https://github.com/TongariCorn/MCMC/blob/master/metro.gif?raw=true
このgifを見ると、結構な割合で候補が棄却されており何となく効率が悪そうに見えるのと、混合正規分布の2つの山の谷にあたる部分を飛び越えるのに苦労している、すなわち各ステップの相関が高いことがわかる

Hamiltonian Monte Carlo

HMCもMetropolis法の一種であり、同様の手続きを踏むことになるが、名前の通りHamiltonianの概念が登場する

HMCでは座標 w \in \mathbb{R}^dと運動量 p \in \mathbb{R}^dに対して

\begin{align*}
\mathcal{H}(w,p) = \frac{1}{2}p^\top p + \beta H(w)
\end{align*}
として定義されるHamiltonianを用いて以下のような遷移を行う:
1. 状態(w,p)から次の遷移を行う場合、まず pを新たに標準正規分布 \mathcal{N}(0,1_d)からサンプリングし置き換える
2. 1で得られた(w,p)正準方程式に従って T秒後まで更新する:
\begin{align*}
\frac{dw_i}{dt} &= \frac{\partial H}{\partial p_i} \\
\frac{dp_i}{dt} &= -\frac{\partial H}{\partial w_i}\end{align*}

3. 2で得られた遷移先の候補 (w^*, p^*)を、確率
\begin{align*}
P := \min\{1, \exp(-\Delta \mathcal{H})\} \quad (\text{ただし}\Delta \mathcal{H} := \mathcal{H}(w^*, p^*) - \mathcal{H}(w, p))
\end{align*}
で採用し、確率1-Pで棄却する

この手続きを繰り返して得られる(w,p)\frac{1}{Z}\exp(-\mathcal{H}(w,p))に従い、\mathcal{H}(w,p)の形からwpは独立であるから、特に w \frac{1}{Z_w}\exp(-\beta H(w))に従う*7

HMCが\frac{1}{Z}\exp(-\mathcal{H}(w,p))を定常分布に持つことは詳細釣り合い条件が満たされることを確かめることによって示されるが、Metropolis法のように直接詳細釣り合い条件を確認するわけではない

というのも、上のステップ1は運動量のみを更新するため定常分布には影響を与えず、ゆえに2,3の手続きが詳細釣り合い条件を満たすことを確認すれば良く、これは簡単な計算によって示される

Leapfrog Method

上の正準方程式の数値解法としては、leapfrog法と呼ばれる手法を用いることが推奨されている:

\begin{align*}
p_i(t + \epsilon/2) &= p_i(t) - \frac{\epsilon \beta}{2} \frac{\partial H}{\partial w_i}(w(t)) \\
w_i(t + \epsilon) &= w_i(t) + \epsilon p_i(t + \epsilon/2) \\
p_i(t + \epsilon) &= p_i(t + \epsilon/2) - \frac{\epsilon \beta}{2} \frac{\partial H}{\partial w_i}(w(t + \epsilon))\end{align*}
というのも、leapfrog法は更新を通してHamiltonianを不変に保つことが知られており*8、更新が安定するとともに遷移先の候補の採用率Pがほぼ1になるという性質を備えているからである

Langevin Equation

HMCでは本来、正準方程式を時刻tに依存しない、時間T分だけ進めるわけだが、ここでleapfrog法を各ステップごとに1イテレーションだけ実行した場合、Langevin Monte Carlo (LMC)と呼ばれる、これまた別のMCMCと等価になることが知られている

LMCはLangevin方程式と呼ばれる確率微分方程式*9

\begin{align*}
\frac{d w}{dt} = -\beta \nabla H(w) + \frac{d B}{dt}
\end{align*}
を数値的に解くというもので、1ステップのleapfrog法を見てみると、 w(t+\epsilon)に含まれる p(t+\epsilon/2)を展開することで
\begin{align*}
w_i(t+\epsilon) = w_i(t) - \frac{\epsilon^2\beta}{2}\frac{\partial H}{\partial w_i}(w(t)) + \epsilon p_i(t)\end{align*}

となり、p_i正規分布からサンプリングされることから上の確率微分方程式のEuler-Maruyama近似になっていることがわかる*10

Langevin方程式は粘性が大きく加速度による影響を無視できるとした上での、粒子がポテンシャル \beta Hとホワイトノイズとして定式化される他の粒子との衝突を受ける時の運動を記述したものであり、少なくとも連続的な確率微分方程式の場合はポテンシャルが適当な正則性を満たせば*11 t \to \infty\frac{1}{Z}\exp(-\beta H(w))に収束することが知られている

下のgifはMetropolis法と同じ混合正規分布を、leapfrog法を1イテレーションのみ更新するHMCでサンプリングした様子
https://github.com/TongariCorn/MCMC/blob/master/langevin.gif?raw=true
ノイズが乗った勾配降下法のような挙動をしており、また、leapfrog法はHamiltonianを保存するため棄却率が大幅に減ったためMetropolis法よりは効率が良さそうだが、やはりランダムウォークのような動きをするため点同士の相関が高い


最後に、各ステップにleapfrog法を10イテレーション回した場合の様子*12
https://github.com/TongariCorn/MCMC/blob/master/hmc.gif?raw=true
推進力、または突破力のようなものが加わり、各ステップの点の位置が大きく離れるようになった

*1:そもそもはHybrid Monte Carloという名前で提案された手法だけど、ここで引用した論文の著者はあえてHamiltonian Monte Carloと呼んでいる

*2:適当な x \in \mathbb{Z}から始めて、各ステップ1/2の確率で-11/2の確率で+1に移動するようなもの

*3:何でも良いと言ってもMarkov chainが収束するための条件は満たしている必要があり、また対称的でなくても遷移候補の棄却確率を少し煩雑にすることで対応できる

*4:X_t\piに従うときX_{t+1}もまた\piに従うこと

*5:\mathrm{gcd}(\{n \in \mathbb{N}_0 \mid p^n(x,x) > 0\} = 1が任意の状態xについて成り立つこと

*6:本来は2つ目と3つ目の条件も確認するべきなんだろうけれど、基本的に無視されていると思われる

*7:Z_wは正規化項

*8:なので天体の軌道計算などにも使われるみたいな話を何かの授業で聞いたことある気がする

*9:Bブラウン運動

*10:正確には\epsilon^2/2を新しいステップ幅にとる必要があるので、時間軸がスケール変換される気がする

*11:例えばポテンシャルの裾が重すぎると粒子が遠くに行きがちになり収束しない

*12:このgifではleapfrog法による途中経過もプロットしている