tf.contrib.bayesflow.hmcをちょっとだけ触った

はじめに

タイトルどおりbayesflowの中にあるhmcをちょっとだけ触りました。edwardがtensorflowに正式に吸収されるという話（以下）以降、（PyTorchを使っていたこともあり）全く触れられていませんでした。

PyTorchの方にもUberからPyroが出ていますが、やはりスタートの速さと開発者の数でEdwardの方が充実している印象でした（ただexampleやtutorialはPyroの方が充実してる）。 とりあえずtensorflowに吸収されたbayesflowに触れておこうということでちょっと触ったメモ用に残しておきます。

HMC(ハミルトニアンモンテカルロ法)

ハミルトニアンモンテカルロ法はマルコフ連鎖モンテカルロ法の一種です。 利用する条件としては尤度と事前分布が揃えばよく、正規化されている必要はありません（つまり全区間での積分が１でなくていい）。

これは何を言っているかというと、ガウス分布を例に上げれば

$$ N(\mu, \sigma) = \frac{1}{\sqrt{2\pi \sigma^2}} \exp \left( - \frac{(x-\mu)}{2\sigma^2} \right) $$

が本来の形ですが、必要なのは

$$ \exp \left( - \frac{(x-\mu)}{2\sigma^2} \right) $$

の部分だけになります。これをカーネルなどと呼びますが、畳み込みのフィルタをカーネルと読んだりするので ごちゃごちゃしてきますが、しっかり覚えましょう（TensorFlowだと、全結合にしても何にしてもニューラルネットのパラメータは全部kernelとbiasと表記しているんですね）。

要するに形状だけ分かればよくて、大きさはどうでも良いということですな。 HMCに関しては以下のスライドが端的にまとまっていて分かりやすいです。 簡単に述べれば、力学の運動方程式で位置をパラメータ、速度（運動量）をパラメータ移動（探索）に使い、高さ（ポテンシャル）を事後確率と置いて 物理法則に従わせてパラメータをサンプリングします。運動量をゴチャゴチャ変えてやればいろいろな範囲を探せるという雰囲気です。

基礎からのベイズ統計学第５章 from hiro5585

www.slideshare.net

ハミルトニアンという言葉はとりあえずエネルギーの総和くらいに思っておいても十分だと思います （もちろんそんなものをハミルトニアンと名付け特別扱いするのは物理学的には意味がある）。

まあ、ただサンプリング法って別にベイズ法それ自体ではなくて、あくまで確率分布の近似手法なので 今回はbayesflowというライブラリに慣れるため、普通にサンプリングをしてみます。

tf.contrib.bayesflow.hmcを使ってみた

２次元ガウス分布

まず平均&\mu&で標準偏差が&\sigma&のガウス分布の対数尤度を計算する関数を書きます。 今回は２次元ガウス分布を試すのですが、x軸もy軸も同じ値の平均で、当方的なガウス分布ってことにします。 というわけでプレースホルダはスカラー値です（一般のガウス分布をやりたければ、その次元のプレースホルダと式を書けばいい）。

mu = tf.placeholder(dtype=tf.float32)
scale = tf.placeholder(dtype=tf.float32)

def log_joint_gauss(x):
    return tf.reduce_sum(-0.5 * tf.square(x - mu) / scale ** 2)

これがサンプリングによって近似したい目標の分布としましょう。 HMCは上手くこの分布をサンプリングできるのでしょうか？という簡単な問題です。

chain, acceptance_probs = hmc.chain(n_iterations=10000,
                                    step_size=0.5,
                                    n_leapfrog_steps=2,
                                    initial_x=tf.zeros([2]),
                                    target_log_prob_fn=log_joint_gauss,
                                    event_dims=[0])

によって、サンプリングを行う計算グラフが準備出来ます（TensorFlowなのでこの段階では計算していません！）。

n_iterationsがサンプリングの回数で、step_sizeとn_leapfrog_stepsは微分方程式の数値計算手法であるリープフロッグ法のステップを指定する部分です。 step_sizeは大きいほど思い切ったパラメータの探索をしますが、大きすぎると棄却が指数的に増加するようです。 ちょっとそれぞれが何を表しているのか正確に把握していないので、分かり次第追記します。

initial_xは探索開始位置の初期値です。今回は$(0,0)$から開始します。

'target_log_prob_fn'は今回作成した対数尤度関数をpythonの関数として与えます。 event_dimsはリストで整数を与えることができて、独立だと'考えてはならない'Tensorのaxisを指定する形になります。 今回は２次元のデータが２次元正規分布に従っていると考えているので、一般には各次元は独立ではありません （今回のように当方的だと無相関かつ独立ですけども…。今思えばbayesflowの使い方を見るには例が悪かった…。）。

仮に各々の成分が１次元正規分布から独立に生起していると考える場合にはevent_dimsは指定しないという使い方をするのだと思います。

更に応用では、同時に複数のベクトル型のパラメータの事後分布を近似するときとかに使うのかなと思います （ARDとかだったら行列を構築する各基底ベクトルは独立に生起すると仮定する）。

以下ではサンプリングの最初の方を切り捨てます。burn-inとか言われるみたいです。

warmed_up = chain[2000:]
mean_est = tf.reduce_mean(warmed_up, 0)
var_est = tf.reduce_mean(tf.square(warmed_up), 0) - tf.square(mean_est)

mean_estとvar_estはburn in後のサンプルの平均と分散を計算するグラフを定義しています（ここではまだ計算していない！）。

以下で具体的に計算を実行します。sess.runの登場です。第一引数に計算グラフで定義した際の取り出したい変数をリストで複数指定できます。 ここらへんはTensorFlowの復習でした。結果は以下の通り

samples, mean, var = sess.run([warmed_up, mean_est, var_est], {mu: 2., scale: 2.})
print('mean:{}, var:{}'.format(mean, var))
sns.jointplot(samples[:,0], samples[:,1])

## mean:[1.9122131, 1.9952391], var:[3.8561068, 3.8079531]

真の分布の平均は$(2,2)$で分散は$4I$でしたからまずまず。

まあ今回やっているのは単に$x$のサンプリングなので、ベイズに使うならば本当は確率モデルのパラメータや潜在変数の事後分布を推論したいところです。

２次元ラプラス分布

式の書き換えが容易だったのでラプラス分布も試しておきました。

loc = tf.placeholder(dtype=tf.float32)
scale_lap = tf.placeholder(dtype=tf.float32)

def log_joint_laplace(x):
    return tf.reduce_sum(-tf.abs(x - loc) / scale_lap)

chain2, acceptance_probs2 = hmc.chain(n_iterations=10000,
                                    step_size=0.5,
                                    n_leapfrog_steps=2,
                                    initial_x=tf.zeros([2]),
                                    target_log_prob_fn=log_joint_laplace,
                                    event_dims=[0])
# Discard first half of chain as warmup/burn-in
warmed_up2 = chain2[2000:]
mean_est2 = tf.reduce_mean(warmed_up2, 0)
var_est2 = tf.reduce_mean(tf.square(warmed_up2), 0) - tf.square(mean_est2)

samples2 = sess.run(warmed_up2, {loc: -1, scale_lap: 1.})
sns.jointplot(samples2[:,0], samples2[:,1])


laplace = tf.distributions.Laplace(loc=-1., scale=1.)
laplace_sample = laplace.sample([8000,2])

sample_lap = sess.run(laplace_sample, {loc:0, scale:1})
sns.jointplot(sample_lap[:,0], sample_lap[:,1])

あれ、もっと尖ってくれていいかもしれません。が、まあちゃんと動いていることは確認できました。 この感じですと推論したい事後分布の計算式はtensorflowでしっかり書いて、Python関数でラッピングするという作業が必要そうです。

多分実際はいろいろな分布くっつけて事後分布を作って、そいつを投げれば勝手に推論してくれるAPIがあるのだろうけども、 どこにあるのか分かりませんでした…（Edwardはどこに行ったんだ…）。