TFPで階層モデルを書くときの便利なクラス tfd.JointDistributionCoroutine

はじめに

ただの自分用のメモです。

環境

google colabで使える最新のTFとTFPを下記でインストール。いずれも開発中のもの（stableではない）。

!pip install tf-nightly-gpu-2.0-preview
!pip install tfp-nightly

ついでにインポートまで。今回出てこないものも含まれていますが、colabで実行したセルのコピーなのであしからず…。

import tensorflow as tf
import tensorflow_probability as tfp
import numpy as np
import daft as daft
import matplotlib.pyplot as plt
from matplotlib import rc
import seaborn as sns

tfd = tfp.distributions
Root = tfd.JointDistributionCoroutine.Root

階層モデル

tfd.JointDistributionCoroutine がつい最近実装されました。Kerasライクに書くAPIばかりが整備され、どうも常に「順序」を気にせねばならず、柔軟性に欠いていたイメージでしたが、PyroやStanみたく自分で変数同士の関係を自由に書けるようなクラスとなっています。

例：モデル

pgm = daft.PGM([7, 6], node_unit=2)
pgm.add_node(daft.Node("e1", r"$e_1$",
                       1, 4))
pgm.add_node(daft.Node("e2", r"$e_2$",
                       1, 3))
pgm.add_node(daft.Node("g", r"$g$",
                       4, 4))
pgm.add_node(daft.Node("n", r"$n$",
                       2.5, 1.5))
pgm.add_node(daft.Node("m", r"$m$",
                       5, 1.5))
pgm.add_edge("e1", "g", True)
pgm.add_edge("e2", "g", True)
pgm.add_edge("g", "m", True)
pgm.add_edge("n", "m", True)
pgm.add_plate(daft.Plate([4, 0.5,
                          2, 2], "$i = 1:N$"))


pgm.render()
plt.show()

で記述されるモデルをtfd.JointDistributionCoroutineで書いてみます。

コード

ジェネレータを使うという仕様。Rootは階層モデルにおいて発生源となる確率変数を記述することになります（したがって Root で包んだ分布が事前分布となる）。

def model():
    e = yield Root(tfd.Independent(tfd.Exponential(rate=[2., 5.]), 1))
    g = yield tfd.Gamma(concentration=e[..., 0], rate=e[..., 1])
    n = yield Root(tfd.Normal(loc=0, scale=2.))
    m = yield tfd.Sample(tfd.Normal(loc=n, scale=g), 10)

tfd.Sample も比較的最近出てきた確率分布のラッパーで、確率分布とサンプルの数を与えて、独立に指定したサンプル数だけサンプリングを行ってくれます。いわばプレート表現を実現してくれるものとなっているわけですね。

joint = tfd.JointDistributionCoroutine(model)
x = joint.sample()
print(x)

#(<tf.Tensor: id=2326, shape=(2,), dtype=float32, numpy=array([0.78496206, 0.03862366], dtype=float32)>,
# <tf.Tensor: id=2354, shape=(), dtype=float32, numpy=34.197464>,
# <tf.Tensor: id=2378, shape=(), dtype=float32, numpy=1.5453969>,
# <tf.Tensor: id=2419, shape=(10,), dtype=float32, numpy=
# array([-23.206112, -20.230238,  14.247202,  -9.024473, -17.014738,
#         33.674152,  32.141914,  46.501278, -20.5491  ,  60.402077],
#       dtype=float32)>)

となり、$x$ にはモデルの中に含まれる全ての確率変数のサンプル値が格納されています。

joint.log_prob(x)

#<tf.Tensor: id=3449, shape=(), dtype=float32, numpy=-89.51276>

でxの対数同時確率を計算できます。