ハミルトニアン・モンテカルロ(HMC)法のざっくりとした解説とPythonによる実装

ベイズモデリングが流行っている中で多くのRユーザーはStanを使って解析をしているんではないかと思います。そして、Stanはハミルトニアン・モンテカルロ(HMC)法と呼ばれる方法で事後分布からのサンプルを得ています。色々と解説記事はありますが、超ざっくりとHMCの原理をメモとして残しておくことにします。

ここで、基本的に私はHMCを伊庭先生のハミルトン4.pdf - Google ドライブこの資料で勉強したので、伊庭先生の資料を読めば私のこの記事は必要ないことをあらかじめ断っておきます((´^ω^)
なんでこれでいいの？という疑問は伊庭先生の資料で解決することでしょう((´^ω^)

記号の定義

まず、以下のように記号を定義します：

X$X$: データ
θ=(θ1,…,θd)′$\theta =({\theta }_1,\dots ,{\theta }_d{)}^{{}^{\prime }}$: d$d$次元パラメータベクトル。′${}^{\prime }$は転置を表します。
r(θ|X)$r(\theta |X)$: 事後分布(後にp$p$が別に出てくるので、ここではr$r$としました)

上記の事後分布からHMCでサンプリングをすることを考えます。

ハミルトニアンとハミルトン方程式

まっっっったく詳しくないのですが、ハミルトニアンとは運動エネルギーV(p)$V(p)$とポテンシャルエネルギー(位置エネルギー)U(θ)$U(\theta )$の和で定義される以下の物理量のことです。ここでp=(p1,…,pd)′$p=(p_1,\dots ,p_d{)}^{{}^{\prime }}$はθ$\theta$と同じd$d$次元のベクトルとします：

H(θ,p)=V(p)+U(θ)$H(\theta ,p)=V(p)+U(\theta )$

また、ハミルトン方程式と呼ばれるものが以下のように表現されます：

dθj(t)dtdpj(t)dt==∂V(p)∂pj−∂U(θ)∂θj$\begin{array}{rcl}\frac{d{\theta }_j(t)}{dt}& =& \frac{\mathrm{\partial }V(p)}{\mathrm{\partial }{p}_j}\\ \frac{d{p}_j(t)}{dt}& =& -\frac{\mathrm{\partial }U(\theta )}{\mathrm{\partial }{\theta }_j}\end{array}$

物理の言葉を借りると、p$p$は運動量でθ$\theta$はある物体の位置を表しています。上記のハミルトン方程式はある物質に運動する力p$p$を与えてやることで、その物質の位置θ$\theta$が変化するという現象を表したモデルになっています。このアナロジーは以下の説明でHMCを理解する上でとっても役立ちます。

さて、このとき以下のようにp$p$とθ$\theta$の同時確率密度関数q(p,θ)$q(p,\theta )$をハミルトニアンH(θ,p)$H(\theta ,p)$だけの関数だけの関数、つまり

q(p,θ)=f(H(p,θ))$q(p,\theta )=f(H(p,\theta ))$

という形で定義したとき、上記のハミルトン方程式に従ってp$p$の値とθ$\theta$の値を変化させれば、それは確率密度関数q(p,θ)=f(H(p,θ))$q(p,\theta )=f(H(p,\theta ))$からのサンプリングをしたこととみなせるという素晴らしい性質を持っています。ここがHMCのポイントだと思います。

ハミルトニアン方程式の解き方

一般に、リープ・フロッグ法と呼ばれる方法で解くようです。ここで次のように記号を導入します：

p(t)$p(t)$：t$t$番目のステップにおけるp$p$の値
θ(t)$\theta (t)$：t$t$番目のステップにおけるθ$\theta$の値

また、次のようにパラメータを導入します：

T$T$：状態変化を行う回数
ϵ$ϵ$：1ステップで状態変化をどの程度許すか

詳しい解説はグーグル先生にまかせて、以下に計算の流れを記します：

===========================================================

Step 0. t←0$t\leftarrow 0$、またp(0)$p(0)$とθ(0)$\theta (0)$に初期値を設定。

Step 1. 以下のようにθ$\theta$を更新：
θj(t+0.5)←θj(t)−ϵ2∂U(p)∂pj(t)${\theta }_j(t+0.5)\leftarrow {\theta }_j(t)-\frac{ϵ}{2}\frac{\mathrm{\partial }U(p)}{\mathrm{\partial }{p}_j(t)}$

Step 2. 以下のようにp$p$を更新：
pj(t+1)←pj(t)+ϵ2∂V(θ)∂θj(t)$p_j(t+1)\leftarrow p_j(t)+\frac{ϵ}{2}\frac{\mathrm{\partial }V(\theta )}{\mathrm{\partial }{\theta }_j(t)}$

Step 3. 以下のようにθ$\theta$を更新：
θj(t+1)←θj(t+0.5)−ϵ2∂U(p)∂pj(t+1)${\theta }_j(t+1)\leftarrow {\theta }_j(t+0.5)-\frac{ϵ}{2}\frac{\mathrm{\partial }U(p)}{\mathrm{\partial }{p}_j(t+1)}$

Step 4. t←t+1$t\leftarrow t+1$ としてStep. 1 - Step. 4をt=T$t=T$になるまで繰り返す。

===========================================================

ここで、t+0.5$t+0.5$ステップという普段見ない表記を導入していますが、中間のステップということを表現したいがために便宜的に導入したものです。

具体的なq(p,θ)$q(p,\theta )$の形

上記のハミルトニアンの右辺に出てくるV(p)$V(p)$とU(θ)$U(\theta )$を、具体的に以下のように指定してやります：
V(p)U(θ)==∑i=1Np2i2τ2−logr(θ|X)$\begin{array}{rcl}V(p)& =& \sum _{i=1}^N\frac{p_i^2}{2{\tau }^2}\\ U(\theta )& =& -\log r(\theta |X)\end{array}$

またq(p,θ)=f(H(p,θ))$q(p,\theta )=f(H(p,\theta ))$のf$f$について、指数関数の逆数
f(y)=exp(−y)$f(y)=\exp (-y)$
を指定してやります。

このときq(p,θ)$q(p,\theta )$は以下のように：

q(p,θ)====exp(−∑i=1Np2i2τ2+logr(θ|X))exp(−∑i=1Np2i2τ2)exp(logr(θ|X))(∏i=1Nexp(−p2i2τ2))r(θ|X)(∏i=1N12π−−√τexp(−p2i2τ2))r(θ|X)$\begin{array}{rcl}q(p,\theta )& =& \exp (-\sum _{i=1}^N\frac{p_i^2}{2{\tau }^2}+\log r(\theta |X))\\ & =& \exp (-\sum _{i=1}^N\frac{p_i^2}{2{\tau }^2})\exp (\log r(\theta |X))\\ & =& (\prod _{i=1}^N\exp (-\frac{p_i^2}{2{\tau }^2}))r(\theta |X)\\ & =& (\prod _{i=1}^N\frac{1}{\sqrt{2\pi }\tau }\exp (-\frac{p_i^2}{2{\tau }^2}))r(\theta |X)\end{array}$

上記の分布をよくみてください。平均0、分散τ2${\tau }^2$の正規分布と目的の事後分布の掛け算になっています。掛け算で表現されるということは、つまりp$p$とθ$\theta$は独立しているということです。

さて、q(p,θ)$q(p,\theta )$はハミルトニアン方程式の状態変化の定常分布でした。加えて、上記の式をみるとp$p$は平均0、分散τ2${\tau }^2$の正規分布に従います。つまり、θ$\theta$とは独立にp$p$を平均0、分散τ2${\tau }^2$の正規分布からサンプリングし、その値を用いてハミルトン方程式を解くことで新たなθ$\theta$の値を得ます。

ここで、ハミルトン方程式と解く際に、現在のθ$\theta$の値から状態を変化させるために以下のようにリープ・フロッグ法のStep 0. を変更します：

===========================================================

Step 0.
t←0$t\leftarrow 0$と設定
p(0)$p(0)$に平均0、分散τ2${\tau }^2$の正規分布から発生させた値を設定
θ(0)$\theta (0)$を現在のθ$\theta$の値を設定

Step 1. 以降は上で述べた方法と同じで、ハミルトン方程式を解いて新たなθ$\theta$を得ます。

===========================================================

以上の流れで得たθ$\theta$を目的の事後分布からのサンプルとみなしたいところですが、上記のハミルトン方程式をリープ・フロッグ法を使って得ということは、コンピュータで”近似的”に解くということで、厳密解ではありません。そこで、ハミルトン方程式を解いて得たθ$\theta$を、目的の分布からのサンプルの”候補”とみて、メトロポリス法のように以下のようなサンプルの受理・棄却のステップを設けてやります。ここでθ$\theta$とp$p$を現在のθ$\theta$の値、θ∗${\theta }^{\ast }$とp∗$p^{\ast }$をハミルトン方程式を解いてθ$\theta$とp$p$の状態変化を行った後の値とします：

===========================================================

確率 min(exp(H(p∗,θ∗))expH(p,θ),1)$min(\frac{\exp (H(p^{\ast },{\theta }^{\ast }))}{\exp H(p,\theta )},1)$ で状態変化後の値θ∗${\theta }^{\ast }$に移動

===========================================================

このように、ハミルトン方程式を決定論的に値を決定するパートと上記の受理・棄却パートのように確率的に値を決定するパートがあるので、HMCはハイブリッド・モンテカルロと呼ばれたりもしています。

HMCのパラメータ

今まで見てきた中で、分析者が指定しなければならないパラメータが以下のように3つあります：

τ$\tau$：運動量p$p$が従う正規分布の分散
T$T$：リープ・フロッグ法でハミルトン方程式を解くときの状態変化の回数
ϵ$ϵ$；リープ・フロッグ方でハミルトン方程式を解くときの1ステップあたりの状態変化の大きさ

τ$\tau$を大きくすれば0から離れた値のp$p$が出やすくなり(物理的にはθ$\theta$の運動が大きくなり)、それにともなってθ$\theta$の値は大きく変わりやすく(物理的にはθ$\theta$の位置が大きく移動しやすく)なります。T$T$の値を大きくすれば、状態変化をする回数が多くなるということなので、前の値のθ$\theta$から異なる値へより変化していきます。ϵ$ϵ$の値を変更すると、状態変化の具合が変わります。

Stanに実装されているNUTSという手法は、上記のパラメータT$T$を自動化するもので、HMCの敷居を大きく下げました。

Pythonによる実装

PythonでHMCを実装します。
今回は、簡単に平均100、分散22$2^2$の正規分布からサンプリングをしたいと思います。

pythonのコードは以下のようになりました：

%matplotlib inline
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# Log probability
# ここでは平均mu, 標準偏差sigmaの正規分布
# -1をかけると物理でいうポテンシャルエネルギー
def log_normal(x, mu, sigma):
    return -0.5*np.log(2*np.pi*sigma**2) - (x-mu)**2/(2*sigma**2)

# Log probabilityの導関数
def d_log_normal(x, mu, sigma):
    return -(x-mu)/sigma**2

# 運動エネルギー
def momentum(p, tau):
    return p**2/(2*tau**2)

# 運動エネルギーの導関数
def d_momentum(p, tau):
    return  p/tau**2

# ハミルトニアン
# 運動エネルギーとポテンシャルエネルギーの和
def Hamiltonian(x, p, tau):
    global mu, sigma
    
    return momentum(p, tau) + (-1.*log_normal(x, mu, sigma))

# リープ・フロッグ法を事項する関数
def proceed_leapflog(epsilon, x, p, tau):
    global mu, sigma
    
    x += -0.5*epsilon*(-1.*d_momentum(p, tau))
    p += epsilon*d_log_normal(x, mu, sigma)
    x += -epsilon*(-1.*d_momentum(p, tau))
    
    return x, p

# HMCを１ステップ実行するサブルーチン
def proceed_HMC_iteration(x, tau, epsilon, T):
    p = np.random.normal(0, tau, size=1)[0]
    p_new = p
    x_new= x
    for t in range(T):
        x_new, p_new = proceed_leapflog(epsilon, x_new, p_new, tau)
        
    alpha = np.exp(Hamiltonian(x, p, tau) - Hamiltonian(x_new, p_new, tau))
    u = np.random.uniform()
    if u < alpha:
        x_accepted = x_new
    else:
        x_accepted = x
        
    return x_accepted

# HMCを実行する関数
def proceed_HMC(tau, epsilon, T, ite, init):
    # Initialize
    x = [init]
    for i in range(ite):
        x.append(proceed_HMC_iteration(x[i], tau, epsilon, T))
    
    return x

以下のように目的の正規分布の平均を初期値と設定して実行してみます。

init = 100
theta = proceed_HMC(tau=2, epsilon=0.1, T=10, ite=2000, init=init)

plt.plot(theta)

ちょっと自己相関が高い感じになってるのでリープ・フロッグ法でハミルトニアン方程式を解く際のステップ数T$T$を30にして、もっとθ$\theta$の状態を変化を待ってみましょう。

init = 100
theta = proceed_HMC(tau=2, epsilon=0.1, T=30, ite=2000, init=init)

plt.plot(theta)

自己相関が小さくなった気がしますね。

ここでちょっと意地悪をして、初期値の値を目的の分布のテールがほとんどない0からスタートしてみます。

init = 0
theta = proceed_HMC(tau=2, epsilon=0.1, T=30, ite=5000, init=init)

plt.plot(theta)

すごいぜHMC。メトロポリス・ヘイスティングではこうはいかないと思います。

最後に

ベイズ推論をするにあたって、多くの場合で事後分布が解析的に解けないために事後分布に従う乱数（事後サンプル）を発生させ、それ使用して色々な分析をします。よく聞くマルコフ連鎖モンテカルロ(MCMC)法は、事後サンプルを得る方法群の名前です。ギブスサンプラーやメトロポリス・ヘイスティング法といった色々な種類のMCMC法が存在します。ベイズ推論において、Stanに実装されているHMCも事後サンプルを得るために使われる一つの方法です。要はどんな方法であれ、最終的に手元に"理論的に正しい"事後サンプルが手に入っていれば解析ができるので、それでよいのです。ただしそれぞれの方法には特徴があって、問題に合わせて使い分ける必要があるように思います。

世の中にはStanやBUGS、JAGSといったベイズ推論を実行することのできるとっても便利な言語があって、我々はそれらを使ってすぐにベイズ推論を行うことができるようになっています。モデリングの試行錯誤も楽ちんです。一から実装していては、モデルが変わるたびにコードを書き直さなければならないので試行錯誤が大変です。ただし、例えば最終的に構築されたモデルをに対してギブスサンプラーが構築できる時にはギブスサンプラーを使ったほうがサンプリング時間が圧倒的に短く、収束も早いことがあります。昔、以下の本の内容をStanで実装していて思い知らされました（この本の著者はマーケティング分野ではかなり有名な研究者です）。

Bayesian Statistics and Marketing (Wiley Series in Probability and Statistics)

特定のモデルについてですが、Rossiは線形代数や分布の特性を駆使してギブスサンプラーによって効率的にサンプリングできる関数をパッケージ化して{bayesm}というパッケージを作っています。Stanとは比べ物にならないスピードで収束します。
Stan便利だぜ、すげえぜ、と思いますが、他の方法の方が良いこともあるということは頭に入れて置かなければなりません。