Hamiltonian Monte Carloについて学んだ

Hamiltonian Monte Carlo(HMC)について学んだのでメモ。

前提/対象/自身の学習前の知識
参考文献
Metropolis法
Metropolis-Hastings法
Hamiltonian Monte-Carlo
このあと

前提/対象/自身の学習前の知識

MCMCについてなんとなく知っている
- ランダムウォークをすることで確率分布からのサンプリングやら積分の近似ができる、程度の理解
物理に関してはよく知らない
- ニュートンの運動方程式、力学的エネルギー保存則ぐらいなら知ってる
MCMCの主な用途はBayes
- 共役事前分布でない事前分布を使ったときもどうやって事後分布/予測分布を得るのかきちんと知っておきたい。
- 解りにくかった理由の一つに参考にするものによって記号がいろいろ違ったのがあるので、この記事では位置は全て $\theta$ とした。

参考文献

ゼロからできるMCMC

Metropolis法

ある(確率密度)関数 $f(\theta)$ を近似するのに、 $\theta$ の $k$ 番目のサンプルの候補 $\theta_{cand}$ をある $p({\theta}{\mid}{\theta}') = p({\theta}'{\mid}{\theta})$ を満たす確率分布、例えば $p(\theta\mid\theta') = \mathrm{Normal}(\theta\mid\theta', \sigma^2)$ を使って $\theta_{cand} \sim p(\theta\mid\theta_{k-1})$ と生成し、確率 $\min(1, \frac{f(\theta')}{f(\theta_{k-1})})$ で採択、不採択なら $\theta_k = \theta_{k-1}$ とする方法。このマルコフ連鎖がエルゴード的(ergodicity)であれば定常分布を持ち、更に釣り合い条件が成り立てばその定常分布と近似したい分布(確率密度関数)が一致する。このとき

1ステップごとの幅を広くとると採択率が低くなってしまう
1ステップごとの幅を狭くすると自己相関が長くなってしまうので、ステップ幅をうまく設定する必要がある。

Metropolis-Hastings法

Metropolis-Hastings法はMetropolis法の $p({\theta}{\mid}{\theta}') = p({\theta}'{\mid}{\theta})$ という条件を満たさない $p$ を用いる代わりに採択率をそれに応じて調整して詳細釣り合い条件を成り立たせる方法。これによってうまく $p$ を選べば元の位置から遠い、かつ採択率の高い $\theta'$ をサンプルすることができる。ただし、どのような $p$ を選べばよいかについては自明ではない。

Hamiltonian Monte-Carlo

Hamiltonian Monte-Carlo(HMC)はMetropolis-Hastings法の一種で、よい $p$ の構成法を与えてくれる。
以前HMCを理解しようとした際はこの辺りから理解が追い付かず

なぜ運動量を導入するの?
Hamiltonianって何?
Leap-Frog法って何?

みたいな疑問を全て咀嚼することなく"現在のサンプル $\theta_k$ を物体の位置として、近似する目標の関数から計算される位置エネルギーを考え、それに加えて正規分布から発生させた乱数を運動量とすることでその物体の動きが決まる。それを時間を進めてある程度動かした先の位置を次のサンプル候補 $\theta'$ とすると採択率の良い、離れたサンプルが取れるらしい"という程度の理解をしていた。とりあえずこれでそこまで遠い理解ではないものの、今回は1つ1つ確認したのでその手順を書いておく。

理解の手順

Hamiltonianを用いて採択率を決めると詳細釣り合い条件が成り立つということをひとまず受け入れる。
Hamiltonianとは何か導出してみる
HMCのアルゴリズムを見てみる
Leap-Frog法について理解する
Hamiltonianを用いて採択率を決めると詳細釣り合い条件が成り立つということを確認するという順番で追っていった。

Hamiltonianとは何か導出してみる

後々のため、位置、速度、加速度としてそれぞれ $\theta, \dot{\theta}, \ddot{\theta}$ , 運動量として $p=m\dot\theta$ , 運動エネルギーを $K$ , 位置エネルギーを $V$ とする。 Hamiltonianについて調べようとするとまず
"HamiltonianとはLagrangianをLegendre変換したものである"
といった表記が見つかる。全て知らないので一つ一つ見ていく。

Lagrangian

Lagrangian $L$ は $L = K - V$ で定義される。まずなんでこう定義されたのかについて調べる。ニュートンの運動方程式は
$F = m\ddot{\theta} = \frac{d}{dt}\left(m\dot\theta\right) = \frac{dp}{dt}$ と書け、また位置エネルギー $V = -\int{F}d\theta$ から $F = -\frac{dV}{d\theta}$ , 運動エネルギー $K = \frac{1}{2}m\dot{\theta}^2 \Leftrightarrow \frac{dK}{d\dot\theta} = m\dot\theta = p$ と書ける。

これをまとめて(以後、偏微分は本来各次元ごとに、という意味で $\theta_i$ とすべきなのだろうけど省略)
$\frac{d}{dt}\left(\frac{\partial K}{\partial\dot\theta}\right) - \left(-\frac{\partial V}{\partial\theta}\right) = 0$

ここで先ほどのように $L = K - V$ と定義すれば、 $K$ は $\dot\theta$ の関数で、 $V$ は $\theta$ の関数なので

$\frac{\partial L}{\partial\dot\theta} = \frac{\partial K}{\partial \dot\theta}$ , $\frac{\partial L}{\partial\theta} = -\frac{\partial V}{\partial \theta}$

であり、 $\frac{d}{dt}\left(\frac{\partial K}{\partial\dot\theta}\right) - \left(-\frac{\partial V}{\partial\theta}\right) = 0 \Leftrightarrow \frac{d}{dt}\left(\frac{\partial L}{\partial\dot\theta}\right) - \left(\frac{\partial L}{\partial\theta}\right) = 0$

と書き直すことができる。これは、直交座標→極座標のような座標変換に対して不変なのでうれしいらしい。(ニュートンの運動方程式は不変ではない)

Legendre変換

ある関数 $f(x,y,z)$ の全微分 $df = \frac{\partial f}{\partial x}dx + \frac{\partial f}{\partial y}dy + \frac{\partial f}{\partial z}dz$ に対し、 $a = \frac{\partial f}{\partial x}, b = \frac{\partial f}{\partial y}, c = \frac{\partial f}{\partial z}$ , $g = ax - f$ とおくと、 $g$ は $g(a,y,z)$ と書けて $dg = xda + adx - df = xda - bdy - cdz$ となる。( $df$ を代入した)
この変数変換をLegendre変換と呼ぶ。

Hamiltonian

ここまで来てやっとHamiltonianが出てくる。Lagrangian $L = K - V$ の全微分は
$dL = \frac{\partial L}{\partial \dot\theta} d\dot\theta + \frac{\partial L}{\partial \theta}d\theta$ とかけて、 $\frac{\partial L}{\partial \dot\theta} = \frac{\partial K}{\partial \dot\theta} = m\dot\theta = p$ だから

Hamiltonian $H = p\dot\theta - L$ と定義してやれば(ここも $p\dot\theta$ は本来次元の総和( $\sum_i^D{p_i\dot\theta_i}$ )とすべき)

$\frac{\partial H}{\partial p} = \dot{\theta}$ , $\frac{\partial H}{\partial \theta} = -\dot{p}$

ここで後者は $\frac{\partial H}{\partial \theta} = - \frac{\partial L}{\partial\theta} = \frac{\partial V}{\partial\theta} = -F = -m\ddot{\theta} = -\frac{d}{dt}\left(m\dot\theta\right) = -\dot p$ と変形した。これは $L(\theta, \dot\theta, t) \rightarrow H(\theta, p, t)$ というLegendre変換と解釈できる。

これらを用いると

$\frac{\partial \theta}{\partial t} = \dot{\theta} = \frac{\partial H}{\partial p}$ , $\frac{\partial p}{\partial t} = \dot{p} = -\frac{\partial H}{\partial \theta}$

であって、 $H$ の時間変化量が

$\frac{\partial H}{\partial \theta}d\theta + \frac{\partial H}{\partial p}dp = \frac{\partial H}{\partial \theta}\frac{\partial \theta}{\partial t}dt + \frac{\partial H}{\partial p}\frac{\partial p}{\partial t}dt = \frac{\partial H}{\partial \theta}\frac{\partial H}{\partial p}dt - \frac{\partial H}{\partial p}\frac{\partial H}{\partial\theta}dt = 0$

であることが確認できる。

HMCのアルゴリズムを見てみる

HMCの大まかなアルゴリズムは

1. 初期位置θを決める
2. 決められたステップ数繰り返したら終了
3. 運動量pを標準正規分布からサンプルする
4. 現在の位置(θ, p)でのHamiltonianを計算する
5. Leap-Frog法を使って物体の移動をシミュレーションする(この間、Hamiltonianは大体保存される。時間を離散化している分完全には保存されない。)
6. シミュレーション終了位置の(θ’, p')でのHamiltonianを計算する
7. 初期位置と終了位置でのHamiltonianを使って採択するかしないかを決め、2に戻る

次に見なければいけないのはLeap-Frog法に関して。

Leap-Frog法

これ自体はそこまで難しくなく、Euler法 $f(t+\Delta t) = f(t) + \frac{df}{dt}(t)\Delta t$ のように次のタイムステップでの状態を計算していく方法。
ただし詳細釣り合い条件を簡単にするためにはそれぞれのステップは可逆でないといけない*1。そのため、Leap-Frog法を使う必要がある。 Leap-Frog法は2つのパラメータを交互に時間を進めていく方法で、 $\theta(t), p(t)$ を $t$ ステップ後の状態として

$\theta(t+2) = \theta(t) + \frac{d\theta}{dt}(t+1)\Delta t = \theta(t) + p(t+1)\Delta t$
$p(t+2) = p(t) + \frac{dp}{dt}(t+1)\Delta t = p(t) - \frac{\partial V}{\partial\theta}(t+1)\Delta t$

ここで前者は $m=1$ とすることで $p = m\dot\theta = \dot\theta = \frac{\partial \theta}{\partial t}$ を使った。これは普通のDPで、初期状態を

$\theta(1) = \theta(0) + \frac{1}{2}p(0)\Delta t$

Tステップ後の終了状態を

$\theta(T) = \theta(T-1) + \frac{1}{2}p(T)\Delta t$

と計算してやることでTステップ先の $(\theta, p)$ が求められる。

f:id:isa_rentacs:20200706230751p:plain — pとθの更新

$\frac{\partial V}{\partial\theta}$ を計算するのに位置エネルギーが何になるか知らないといけないが、これは目的の関数を $f(\theta)$ として

$f(\theta) = \frac{exp(-S(\theta))}{Z}, Z$ は正規化定数としたときの $S(\theta)$ を使えばよいらしい。ここの理論的背景が解っていないものの、確率分布に対してはnegative log-likelihoodと同じなので計算自体はできる。

HMCの詳細釣り合い条件

ここまでくればあとは詳細釣り合い条件を確認すればよい。Leap-Frog法は可逆であるのでLeap-Frogの前後を $(\theta, p) \rightarrow (\theta', p')$ とすればこれは標準正規分布から $p$ がサンプルされる確率と等しく、これに採択率 $\min(1, \exp(H - H'))$ を掛ければ遷移確率が求まる。また、逆に $(\theta', -p') \rightarrow (\theta, -p)$ を考えると同様に遷移確率は標準正規分布から $-p'$ がサンプルされる確率と採択率 $\min(1, \exp(H' - H))$ の積で、これが $H \ge H', H \lt H'$ のどちらの場合も等しくなることを確認すればよい。