はじめに
概要
現在開発が急ピッチで進んできている（ように私には見える）、TensorFlow Probabilityですが、 PyroやStanなどの先発組に比べて明らかに遅れを取っているように見えます。
このことに関して「ネット上に良いサンプルコードが見当たらない」ということと「ドキュメントを読んでもどのAPIを使えば良いか分からない」ということが大きな原因かなと思います。
特にtfp.distributionsとtfp.edward2の差異は明確ではなく、どっちをどのように使うのか…というのはかなりわかりづらいところです。
結論を述べると、tfp.edward2はtfp.distributionsの薄いラッパーです。tfp.distributionsを真面目に使ってガチャガチャ書かなければならないところを、省力化してくれる役割をになってくれる感覚です。従って、やろうとしていることがtfp.edward2で実行可能ならば迷わず使えばよく、若干低階層に降りなければできなさそうだな？と感じるときにtfp.distributionsを使う形になるかなと思います。
コードの前提
Jupyter などを使う際に下記のコードを実行していることを前提とします。
import numpy as np import tensorflow as tf import tensorflow_probability as tfp import matplotlib.pyplot as plt import seaborn as sns %matplotlib inline tfd = tfp.distributions tfe = tf.contrib.eager ed = tfp.edward2 tf.enable_eager_execution()
モジュールについてわかりづらいときは適宜、フルパスで記載します。
Edward2
Edward2 is a probabilistic programming language in TensorFlow and Python. It extends the TensorFlow ecosystem so that one can declare models as probabilistic programs and manipulate a model's computation for flexible training, latent variable inference, and predictions.
Edward2はもともとEdwardという3rd Party Libraryとして開発されていたものが、TensorFlowにマージされたものです。TFにマージされたときにAPIが変更されているようで、もともとのEdwardユーザーもついてきていないんじゃないか…？という印象を受けます。
しかし、ちゃんとEdward2を見てみると、利便性を損なわずTFとちゃんと融合していることが分かりました。
肝は tfp.edward2.RandomVariableクラス
実は最も重要なクラスが tfp.edward2.RandomVariableクラスです。edward2 が有している確率分布のクラス（例えば tfp.edward2.Normal）等も全てこのクラスを継承しています。細かいことはさておくとして、edward2が持つ確率分布クラスと、tfp.distributionsが持つ確率分布クラスは何が違うのでしょうか。ココが気になるところですね。
百聞は一見にしかずということで、実際にコードを見てみましょう。下記は一次元標準ガウス分布 $\mathcal N (0, 1)$ をインスタンス化しています。
print(tfp.distributions.Normal(loc=0., scale=1.)) print(tfp.edward2.Normal(loc=0., scale=1.)) # ->tfp.distributions.Normal("Normal/", batch_shape=(), event_shape=(), dtype=float32) # ->RandomVariable("-0.27558547", shape=(), dtype=float32, device=/job:localhost/replica:0/task:0/device:CPU:0)
tfp.distributions.Normalの方は、確率分布そのものを表すクラスとなっています。一方で、RandomVariableクラスの方は、何やら具体的な数値が入っているではありませんか（-0.27558547の部分）。ふむふむ、全く異なるものであることがわかりますね。
しかし、実はRandomVariableクラスのインスタンスはプロパティにtf.distributionsクラスのインスタンスを保持しており、RandomVariable.distributionでアクセスできます。今回の場合は、tfp.edward2.Normalクラスはtfp.distibutions.Normalクラスのインスタンスを保持しているということです。実際に見てみましょう。
print(tfp.edward2.Normal(loc=0, scale=1).distribution) print(tfp.distributions.Normal(loc=0, scale=1)) # ->tfp.distributions.Normal("Normal/", batch_shape=(), event_shape=(), dtype=float32) # ->tfp.distributions.Normal("Normal/", batch_shape=(), event_shape=(), dtype=float32)
はじめにでも述べた通り、Edward2の確率分布クラスは、TensorFlow Probabilityの確率分布の薄いラッパーだったのです。
逆にtfp.edward2.RandomVariableのコンストラクタにtfp.distributionsのインスタンスを食わせてやれば、RandomVariableクラスの確率分布を作ることもできます。
## 下記のコードは同じ処理を意味します print(tfp.edward2.Normal(loc=0, scale=1)) print(tfp.edward2.RandomVariable(tfp.distributions.Normal(loc=0, scale=1)))
ともかくEdward2では確率分布よりも確率変数を主体にしているようですね。たしかに結局私達が実際に手にするデータは全て確率変数（の実現値）の方であり、確率分布それ自体を見ているわけではありません。コードを書く時も確率分布を背景に持つ確率変数に関して計算処理を書いていくほうが、感覚的に慣れやすいような気もします。
ベイズロジスティック回帰
ロジスティック回帰を書く
下記のモデルをコードとして書いてみましょう。
$$ \begin{align} p & = \sigma({\bf w^T x} + b) \\ y & \sim {\rm Bern}(p) \end{align} $$
データ$\bf x$を1つ1つ扱うのではなく、横ベクトルとして格納した行列$\bf X$を使った数式を用いてコードを書くのが通例なので、
$$ \begin{align} {\bf p} & = \sigma ({\bf Xw + 1}*b) \\ {\bf y} & \sim {\rm Bern}({\bf p}) \end{align} $$
を実際にはコードとして書きます。${\bf 1}* b$はスカラー$b$をベクトル $\bf Xw$ に加算するときにブロードキャストしますという意味です。
# モデルを表す関数を作る def logistic_regression(X): # w.shape == (X.shape[1]) w = ed.Normal(loc=tf.zeros(X.shape[1]), scale=1., name="coeffs") # b.shape == (1,) b = ed.Normal(loc=0., scale=1., name="intercept") # y.shape == (X.shape[0]) y = ed.Bernoulli( logits=tf.tensordot(X, w, [[1], [0]]) + b, name="outcomes") return y # 実際にデータを流す num_features = 10 X = tf.random_normal([100, num_features]) y = logistic_regression(features) # -> y.shape == (100) のRandomVariable型
当然このモデルは学習しておらず、適当にガウス分布から生起させた $w$ と $b$ を使って値を計算しているにすぎません。
事後分布を書く
次に事後分布を書きます。$D$は手持ちのデータの集合として
$$ p({\bf w}, b \mid D) $$
が、私達の知りたいパラメータ $w, b$ に関する事後分布です。
def logistic_regression_posterior(num_features): posterior_w = ed.MultivariateNormalTriL( loc=tf.get_variable("w_loc", [num_features]), scale_tril=tfp.trainable_distributions.tril_with_diag_softplus_and_shift( tf.get_variable("w_scale", [num_features*(num_features+1) / 2])), name="w_posterior") posterior_b = ed.Normal( loc=tf.get_variable("b_loc", []), scale=tfp.trainable_distributions.softplus_and_shift( tf.get_variable("b_scale", [])), name="b_posterior") return posterior_w, posterior_b
posterior_w = ed.MultivariateNormalTriLは多次元ガウス分布を表すのですが、分散共分散行列は必ず対称行列になるという性質から、下三角行列だけわかれば情報としては十分なので、scaleとして下三角行列を渡しますというモデルです。もちろん普通に多次元ガウス分布を用いても良いでしょうけど、対称行列をフルで表現するのは単にメモリの無駄遣いです。
scale_tril=tfp.trainable_dstributions.tril_with_diag_softplus_and_shiftは下三角行列の成分をベクトルとして取り扱ってくれるクラスのようです。分散共分散行列の下三角行列だけを取り出した場合、データの次元 $d$とすれば、$d(d+1)/2$の成分を持ちます。
posterior_bの方は特に難しい話はなく、scale = tfp.trainable_distributions.softplus_and_shiftは値を正に保つために利用されているようです。
print(posterior_w) print(posterior_b)
とすれば、当然、デタラメな確率変数の実現値が得られるだけです。
事前分布のセッティング
さて各パラメータに事前分布を割り当てましょう。Pyroならnn.Moduleで普通にモデルを書いてから、事前分布を置きたいnn.Moduleクラス内のパラメータ名をキー、対応させたい事前分布をバリューとした辞書をpyroの特殊な関数に与えることでセッティングを行います。
一方で、この機能に相当する関数がTF-pには無さそうなので、下記のようなコードで対応するようです。
def set_prior_to_posterior_mean(f, *args, **kwargs): """Forms posterior predictions, setting each prior to its posterior mean.""" name = kwargs.get("name") if name == "w": return posterior_w.distribution.mean() elif name == "b": return posterior_b.distribution.mean() return f(*args, **kwargs) with ed.interception(set_prior_to_posterior_mean): predictions = logistic_regression(X)
ちょっと勉強中です…。
対数尤度
とりあえず適当にデータを作っておきます。この項目で劇的にedward2の恩恵を受けることになります（多分事前分布のセッティングもedward2の恩恵であるのだが理解不足で申し訳ないです）。
log_joint = ed.make_log_joint_fn(model)とすると、modelに含まれるRandomVariable型の対数確率の和を自動で計算してくれる関数log_jointが返ってきます。
X = tf.random_normal([100, 55]) y = tf.random_uniform([100], minval=0, maxval=2, dtype=tf.int32) log_joint = ed.make_log_joint_fn(logistic_regression)
としましょう。log_jointは関数であることに注意してください。どんな関数かと言うとlogistic_regressionの元々の引数と、更に関数内で使われるRandomVariable型の実現値を引数に取る関数です。
def logistic_regression(X):
w = ed.Normal(loc=tf.zeros(X.shape[1]), scale=1., name="coeffs") b = ed.Normal(loc=0., scale=1., name="intercept") y = ed.Bernoulli( logits=tf.tensordot(X, w, [[1], [0]]) + b, name="outcomes") return y
であったのでした。
続いて、MCMCに渡してやる前に下記のようにlog_joint関数を部分適用して別の関数を用意する必要があります。
TF-pのMCMCには、事後分布を求めたいパラメータであるwとbだけを流し込むようなAPIが提供されているため、
部分適用でデータの方は予め受け取った状態の関数を作っておくのです。
def target_log_prob_fn(w, b): """Target log-probability as a function of states.""" return log_joint(X, w=w, b=b, y=y)
MCMCの実行
準備完了です。
MCMCの遷移核としてはtfp.mcmc.HamiltonianMonteCarloを使うことにします。多分後々NUTSが出てくると思われるでしょう。
あとはtfp.mcmc.sample_chainに遷移核をわたし、サンプリングの回数やバーンインの回数を渡してやれば事後分布（からのサンプリング）が得られます。
hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=target_log_prob_fn,
step_size=0.1,
num_leapfrog_steps=5)
states, kernel_results = tfp.mcmc.sample_chain(
num_results=1000,
current_state=[tf.random_normal([55]), tf.random_normal([])],
kernel=hmc_kernel,
num_burnin_steps=500)
ここでcurrent_stateは、MCMCで遷移を始める初期状態です。なぜinitial_stateではなくcurrent_stateという名前なのかは、戻ってきたstateを引数にもう一度途中からサンプリング回したりできるだろ！という意味合いだと思われます。
事後分布を見る
statesはリストになっており、states[0]が $w$ のサンプル1000個が格納され、states[1]が $b$ のサンプル1000個が格納されています。よって下記のようにヒストグラムを見ることで、MCMCによって近似された事後分布を確認することができます。
plt.hist(states[1], bins=20) print('mean: ', states[1].numpy().mean()) print('std: ', states[1].numpy().std()) print('argmax: ', states[1].numpy().argmax()) # ->mean: 0.18845236 # ->std: 0.20227498
最後に
現状、立ち位置としてはtfp.distributionsとtorch.distributionsが、tfp.edward2とpyro.distributionsがそれぞれ対応しているように感じます（というかモジュールの構成としてTF側がPyTorch側を後追いする形になっている？）。また、コードの書き方としてもPyroの model と guide という書き方に似たやり方を見つけようとedward2が模索している感じにも見えます。
どうしてもgraphモードでもeagerモードでもシームレスなコードを実現しようとすると、標準的な書き方が若干回りくどくなってしまうのは仕方ないのでしょう。TF2.0からeagerがデフォルトになりますが、graphモードの立ち位置にも注目です。「学習モデルが決まって最適化に集中したい場合以外には顔を出さない」くらいならば、よりPythonicなAPIを提供することも出来るでしょうけども、あくまでeagerをデフォルトにしつつ、どちらも対等に使えるようにするという雰囲気な気がしますね。
eagerで書いていくにしても、高階関数の部分適用やカリー化などの概念を知っていたほうがTensorFlowは取り回しやすいでしょう。Pythonネイティブには受けが悪そうです...。
MCMCを利用する場合の書き方がTF-p単体ではまだまだ乱雑で、edward2がこの辺りに良いAPIをもたらそうとしているのは間違いないです。更にPyMC4もedward2の肩に乗りそうな雰囲気ですので、これからもっと高レベルAPIが整備されていくのは間違いないでしょう。MCMCはディープと組み合わせた場合には、どうしても計算量的に苦しいのですが計算基盤をガッチリ握っているTF陣営が、アプリケーションとしてどのような物を出してくるのかが楽しみですね。
