xiangze's sparse blog

機械学習、ベイズ統計、コンピュータビジョンと関連する数学について

Exhaustive Hamilton Monte Carloの紹介

stan MCMC

Stan Advent Calendar12/20の記事です。
Stan version2.10.0で実装されたExhaustive HMC(xHMC)というアルゴリズムについて
論文Identifying the Optimal Integration Time in Hamiltonian Monte Carlo
の内容に基づいて説明します。

定式化

論文の1章ではHMCが定義されるような状況を一般化するための数学的定式化について説明しています。
サンプルされる値の集合をQとすると分布関数 π:Q->R(実数)は
\pi==e^{-V}q_1 \wedge \dots \wedge q_n
と書かれます。Vは物理的にはポテンシャルエネルギーに相当します。
運動量に相当する量は余接バンドル(cotangent bundle) T*QからQへの関数として
\xi==e^{-K}p_1 \wedge \dots \wedge p_n
というn-形式で書かれ、Kに相当する部分は運動エネルギーに相当します。
ここからハミルトニアンは形式的には
H=K+V=-\log \frac{d\pi_H}{d\Omega}
d\Omega=p_1 \wedge \dots \wedge p_n \wedge q_1 \wedge \dots \wedge q_n
と書かれます。
このξからサンプルした運動量pを q\in Qにくっつけるliftと呼ばれる操作で余接バンドルに写します。
l:Q->T*Q
q|->(q,p)
ハミルトン方程式の解軌道は
\phi_t^H:T^*Q \rightarrow T^*Qと書かれていて
これに射影w:T*Q->Qでもとに戻す操作を施してHMCの中の一回の積分操作gは
g=w \cdot \phi^H_t \cdot l
と書かれます。

f:id:xiangze:20161220040513p:plain
論文のfig.1 Q=S1(円周)の場合の余接バンドル(b)とその上の分布関数(c )、等エネルギー面の集合(d)の関係

停止基準

ハミルトンモンテカルロ法(HMC)ではハミルトン方程式に従った数値積分と確率的な遷移を交互に行うことになります。
指定した分布関数に従ったサンプリングをする為には2点z,z'間の遷移の詳細釣り合い条件
P(z)*K(z',z)=P(z')*K(z,z')

P(z)=\frac{d\pi_H}{d\Omega}(z)
K(z_L,z_0)=\frac{ \frac{d\pi_H}{d\Omega}(z_L) }{ \sum_{m=0}{L} \frac{d\pi_H}{d\Omega}(z_m ) }

を満たす必要があります。このためには相空間内のある点zから始まる軌道だけでなく、点zを通るような軌道も取得しておく必要があり、時間的に逆方向の積分も行わなければいけません。これはすでにStanで使われているNUTS(no-U-turn sampler)に含まれている機能です。

ハミルトン積分を長い期間行えばそれだけ正確なサンプリングができますが、サンプリングという点では非効率です。
ある場合には長い期間積分すれば相空間内の同じエネルギーの点の集合(level set)の任意の点に近い点に到達することは可能かもしれませんが(動的エルゴード性 1.1.2) 。あるいは周期軌道になっているような場合はまんべんなくサンプリングすることは出来ないことになりますが、ハミルトニアンで決まるlevel setと方程式の軌道がそのような形をしているかを事前に知ることは困難です*1。
そこで積分の期間(数値積分のステップ)をある基準に従って決めてそこで積分を打ち切って確率的遷移を行うのが効率的と考えられます。またその打ち切る期間の最適な値がありそうです(fig5, 3.EXPLICIT TERMINATION CRITERIA)。
ここで提唱されているxHMCの既存の手法(普通のNUTS)との違いはこの積分の停止基準を明示しているところにあります。
エネルギー以外の相空間内の点z=(p,q)に対するスカラー値関数u(z)の時間微分の平均
\kappa(T,z)=\frac{1}{T}\int_0^T dt \frac{du}{dt} \dot \phi_t^H(z)= \frac{u \dot \phi^H_T(z)-u(z)}{T}
は積分期間を長くしていくと0に近づいていく筈です。それが一定の値以下になったら積分を停止するのが良さそうです。
ここではビリアル(運動量と位置の積)
G=\sum_i q^ip_i
をスカラー値関数として使っています。ビリアルは常に有限の値をとり、その時間微分を長時間積分すると0になるため
(ビリアル定理)
exhaustionの値として用いこれが一定値δ以下となることを停止基準としています。
場所ごとに計量(運動量部分の共分散行列)が定義されるようなリーマン多様体上でもHMCは定義でき、この場合には積分の停止基準としてパラメータの軌道上の平均ρTに対して計量g(p,ρT)の逆行列が負になるという条件が使えるそうです*2。これは一定値δの様なパラメータを含まないため厳密性が増しているようですが、gの計算コストが大きく実用性は必ずしも高くないようです(3.3)。
4章の数値実験ではxHMCがうまくいく例(多変量正規分布で相関の大きい場合)とそうでない例を挙げています。

実装

stanのコードではxHMCはRiemannian HMCとともに実装されています。
https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/hmc/xhmc/base_xhmc.hpp#L58
ビリアルの時間微分dG_dtはbase_hamiltonian で定義されており、
https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/hmc/hamiltonians/base_hamiltonian.hpp#L41
それを継承するような形になっており、Stanで使うことの出来るいくつかの計量ごとに実装されています(https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/hmc/hamiltonians/dense_e_metric.hpp#L34)。これをbase_xhmc::transitionで実行しています。

追記

NUTSの実装にバグがあったらしく、ガウス分布の分散がわずかに小さく推定されてしまうということなのですが、HMCの部分ではなくサンプリングの部分に問題があったようです。
Stan 2.10 through Stan 2.13 produce biased samples - Statistical Modeling, Causal Inference, and Social Science

*1:分かっていたら数値的にサンプリングする必要がありません

*2:原理がまだよくわかりません