TFP: 確率モデルの書き方 with tfd.JointDistiburionSequential

はじめに

下記程度のグラフィカルモデルが読める前提です。

www.hellocybernetics.tech

TFPのモデリングはEdward2を利用しなければまともに書けない状態でありましたが、TensorFlow Probabilityの最もコアとなるモジュールにtfp.distributions.JointDistributionSequential が追加されたのでメモを残しておきます。

色々な分布

前提として（階層）ベイズモデルを扱います。なので（普通の意味での）パラメータに対して事前分布とハイパーパラメータが敷かれていることに注意してください。

正規分布を使った基本

観測データ $Y = (y _ i) _ {i=1} ^N$ とパラメータ $\mu, \sigma$ の同時分布のモデルは下記。

$$ \begin{align} p(\mu) &= {\rm Normal}(\mu \mid 0, 10000) \\ p(\sigma) &= {\rm HalfCauchy}(\sigma \mid 0, 5) \\ p(Y\mid \mu, \sigma) &= \prod_{i=1}^N {\rm Normal}(\mu, \sigma) \end{align} $$

として、$\mu, \sigma$ が独立であるという仮定を入れ

$$ p(Y, \mu, \sigma) = p(Y\mid \mu, \sigma)p(\mu)p(\sigma) $$

とします。グラフィカルモデルは下記。

TensorFlow Probabilityでは下記のようにコードを書きます。

def model(N):
    return tfd.JointDistributionSequential([
        tfd.Normal(0., 1e5), # mu
        tfd.HalfCauchy(0., 5.), # sigma
        lambda sigma, mu: tfd.Sample(
            tfd.Normal(loc=mu, scale=sigma),
            sample_shape=N
        )
    ])

何点か注意。

• tfd.Sample は概ねプレートの役割。ただし単に $N$ 個のサンプルではなく $(N, M)$型のように一般的なshapeで指定できます。
• tfd.JointDistibutionSequential にはtfd.Distribution を継承したクラスのインスタンスをリストで与えます。
• リストの $i$ 番目のインスタンスは $ i-1 $ 番目より小さなインデックスのインスタンスを認識できます。
• 認識する順序は $ i- 1, i-2, i-3, ...$ と降順になります。

条件付けられていない変数なら、何も考えずにリストに入れておけばよいのですが、条件付けられた変数は、リストに入れる順番を注意しなければなりません。特に「降順にインスタンスを認識する」というのが曲者で、上記の例でも分かるように、tfd.Sample(tfd.Normal) に lambda を使って条件付きの変数を与えるのですが、その際降順であるためにすぐ上の $sigma$ を第一引数に取っていること注意が必要です。

ちなみに自分のコードがモデルを意図したとおりに設計できているかをtfd.JointDistributionSequential._resolve_graph() によって簡易ですが確認することが出来ます。

model_ = model(100)
model_._resolve_graph()

# (('mu', ()), ('sigma', ()), ('x', ('sigma', 'mu')))

変数の接続が複雑な場合

次に下記のような同時分布を考えます。

$$ p(X, a, b, c, d) = p(a)p(b)p(c \mid a)p(d \mid a, b, c) \prod _ i p(x _ i \mid d, b) $$

具体的な確率分布は適当におくとして、グラフィカルモデルは下記です。

この場合は下記のようにダミーの引数を与えてコードを読んだときに、条件付きの変数が明確に分かるようにする方法が良いようです。

def model(N):
    return tfd.JointDistributionSequential([
        tfd.HalfCauchy(0., 5.), # a
        tfd.HalfCauchy(0., 5), # b
        lambda _, a: tfd.HalfCauchy(0., a), # c
        lambda c, b, a: tfd.StudentT(a, b, c), # d
        lambda d, _, b, __: tfd.Sample(
            tfd.Normal(loc=b, scale=d),
            sample_shape=N
        )
    ])

ただし、この場合は _resolve_graph()メソッドが

'''
(('a', ()),
 ('_', ()),
 ('c', ('_', 'a')),
 ('d', ('c', 'b', 'a')),
 ('x', ('d', '_', 'b', '__')))
'''

を返します。どうやらlambda で付けた引数の名前を変数名に持ってきているようです（そして生成される最後のデータに x を付けている模様）。するとどうしても本来 b と名付けられてほしい変数名が気持ち悪いです（といっても所詮、このメソッドのローカルの話なのでどうでも良いのだが…）。

その場合は、リストの順序を入れ替えることで「最初に条件付きとして引用されるときにダミー引数を使わなくて良いような順番」に変えてしまう方法があります。

def model(N):
    return tfd.JointDistributionSequential([
        tfd.HalfCauchy(0., 5.), # a
        lambda a: tfd.HalfCauchy(0., a), # c
        tfd.HalfCauchy(0., 5), # b
        lambda b, c, a: tfd.StudentT(a, b, c), # d
        lambda d, b, _, __: tfd.Sample(
            tfd.Normal(loc=b, scale=d),
            sample_shape=N
        )
    ])


'''
(('a', ()),
 ('c', ('a',)),
 ('b', ()),
 ('d', ('b', 'c', 'a')),
 ('x', ('d', 'b', '_', '__')))
'''

変数の順序は若干トリッキーですが、グラフの中で変数名がしっかりつけられています。また、ちゃんと変数名が付く順序で値を与えられているということ自体が、「有向グラフの順番に沿って確率変数の生成を書けている」ということにもなり、ある意味「読みやすいコード」になるような気がしないでもないです。

実はこのへん tfd.JointDistributionNamed という新しいクラスで解決される見込みですが（ただいまレビュー中）、ひとまず…。

sample

あとは普通の Distribution クラスのように扱えます。

a, c ,b, d, x = model_.sample()

とすれば、リストで与えたDistribution インスタンスと同じ順序でサンプルを返してくれます。

log_prob

当然 log_prob メソッドが使えます。 ここで log_prob メソッドの引数は JointDistributionSequential のインスタンス生成時に渡したリストの順番で、確率変数の実現値をリストにして与えれば良いです。最も簡単な例は

a, c ,b, d, x = model_.sample()
model_.log_prob([a, c, b, d, x])

例えばMAP推定ならば、データ $x$ を固定しつつ、$a, c, b, d$ を少しずつ勾配法で変えながらmodel_.log_probを評価していけば良いことになります。