TFPの変分ベイズの書き方2つ
tensorflow probability(TFP)では変分ベイズによる推定ができますが書き方に幅があります。
またBayesian DNNを書くのに使用できる関数があります。変分ロジスティック回帰のコードでそれを比較します。
素朴な方法(ed.make_log_joint_fn を使う)
https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Probabilistic_PCA.ipynb
などの例で使われている方法です。
統計モデル(ここではlogistic回帰)を
def logistic_regression(features,coefs): coeffs = ed.Normal(loc=0., scale=1., sample_shape=features.shape[1], name="coeffs") outcomes = ed.Bernoulli(logits=tf.tensordot(features, coeffs, [[1], [0]]), name="outcomes") return outcomes
と定義します。推定すべき変数(stanのparameter)を入力とし、出力はデータの確率分布関数(stanのdataに相当)です。
変分推論のために新たにモデルを定義します。
def variational_model(qw_mean, qw_stddv, qz_mean, qz_stddv): qfeature = ed.Normal(loc=qw_mean, scale=qw_stddv, name="qw") qcoeffs = ed.Normal(loc=qz_mean, scale=qz_stddv, name="qz") return qw, qz
変分推論における統計モデルと変分モデルの関係はVAEではこれらはEncoder,Decoderに相当します。
変分推論では近似モデルの因子分解が仮定されているため各確率変数は互いに関係を持たない形になります。
モデルに代入する確率変数(paremeter)は以下のように定義されます。
features = tf.random_normal([3, 2]) coeffs_value = tf.random_normal([2]) outcomes_value = tf.to_int32(tf.random_uniform([3]))
ここでdataに相当するoutcome_valueはtf.placeholderからデータ値を供給されるような形になっています。
各変数に対応した確率分布関数の対数の和を変分推論では最大化するのですが、Edward2ではモデルの複数のパラメータの合計から対数尤度を得るための関数としてed.make_log_joint_fn() が用意されています。
この関数はモデルとその入力となるデータ、確率変数を引数として取り、確率変数を引数とする関数を返します。
logistic regressionの例ではモデルと変分モデルに対して
log_joint = ed.make_log_joint_fn(logistic_regression) target = log_joint(features, coeffs=coeffs_value, outcomes=outcomes_value) log_joint_q = ed.make_log_joint_fn(variational_model) target_q=log_joint_q(features, coeffs=coeffs_value)
のようになります。
このモデルを最適化問題に落とし込む実装をします。ELBO(evidence lower bound)を
]
を(pはモデルの分布関数、qは近似モデルの)書くとその最大化はKL divergence(モデルがデータに合致する度合いと)の最小化と等価になり(())、
ELBO = energy + entropy
とも解釈できます。上で得られたモデルに対応した尤度をとELBOをparameter features, coeffsの関数として
energy = target(features, coeffs)
entropy = -target_q(features, coeffs)
elbo = energy + entropy
optimizer = tf.train.AdamOptimizer(learning_rate = 0.05)
train = optimizer.minimize(-elbo)
という風にかけます。Edward1の頃と比べるとtensorflowむき出しのコードです。このtrainをSession内で実行していきます。
Flipoutを使う
TFPのlogistic regression, VAEのexampleではDenseFlipoutという関数が使われています。
https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/logistic_regression.py
https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/vae.py
重みパラメータの事前分布が多変量正規分布の時のみ使えるものでReparameterization Trickを内包したものでDNNの1つの層として提供されています(
https://www.tensorflow.org/probability/api_docs/python/tfp/layers/DenseFlipout
https://www.tensorflow.org/probability/api_docs/python/tfp/layers/Convolution2DFlipout
など)。
[1803.04386] Flipout: Efficient Pseudo-Independent Weight Perturbations on Mini-Batches
linear regressionを1層のニューラルネットとして見ています(これはpyroの例と同じ構成です)。
layer = tfp.layers.DenseFlipout(
units=1, activation=None,
kernel_posterior_fn=tfp.layers.default_mean_field_normal_fn(),
bias_posterior_fn=tfp.layers.default_mean_field_normal_fn())
logits = layer(features)
labels_distribution = tfd.Bernoulli(logits=logits)
KL divergenceは各Flipoutの出力のlossであり(分布の差)、それにモデルの出力labels_distributionの負の対数尤度をを足したものがELBOになります。
neg_log_likelihood = -tf.reduce_mean(labels_distribution.log_prob(labels))
kl = sum(layer.losses) / FLAGS.num_examples
elbo_loss = neg_log_likelihood + kl
VAEの例を見ると分布が多変量正規分布など解析的に計算できる場合はtfd.kl_divergenceが使え、そうでない場合は近似事後分布からサンプルするので以下のようにサンプロを使った平均値になります。
neg_log_likelihood = -tf.reduce_mean(labels_distribution.log_prob(labels)) if params["analytic_kl"]: rate = tfd.kl_divergence(approx_posterior, latent_prior) else: rate = (approx_posterior.log_prob(approx_posterior_sample) - latent_prior.log_prob(approx_posterior_sample)) avg_rate = tf.reduce_mean(rate) elbo_loss = neg_log_likelihood + kl elbo = tf.reduce_mean(elbo_local) loss = -elbo
という風に書かれます。
