確率モデルのフレームワーク,pyroを勉強中です.MCMCなどの勉強とともに実装のフレームワークを探したところ,pyroにたどり着きました.

pymcをまず始めようと思いましたが,更新が古く自分の中で却下になりました.
次に,edwardを始めようと思いましたが,現在,edawardはedward2として,tensorflow-probabilityに取り込まれました.tenforflow-probabilityはまだ発展途上で,exampleも少なく,自分がtensorflowに慣れてないということもあり,やめました.

ということで,pytorchを勉強しているということもあり,pytorchバックエンドのpyroを始めることにしました.

pyro
パイロキネシスのpyroらしい

コイントスの例

ベイズ推定の最初の例といえばやはりこれ.コイントスが公平かを調べるという問題です.

コインが表が出る確率を という潜在変数で表します.0.5だと公平なコイントスと言えます.
コインが表か裏かを表現する確率を で表します.

の確率分布を求めるというのが今回の目標になります.そのためには,観測データ に対して,次の対数尤度を最大化することを目指します.

一般にこの積分を解析的に求めることは難しいです.(共役事前分布を設定するとできるということはご存知だと思いますが,今は忘れましょう)

さて, を計算するのは難しいのですが,その下限を考えてそれを最大化させることを考えます.この下限をELBOと言います.
をパラメーターとする分布として,ELBOを次のように定義します:

どんなでもが成立します.
これは,であることから確認できます.

このは変分ベイズ事後分布などと呼ばれています.(どのような名称が一般的かはわかってないです…)

このELBOについて, について勾配を計算し(期待値の積分はモンテカルロサンプリングで近似します)勾配法でELBOを最大化させます.
これについては,pyroの枠組みで公式チュートリアルがあります.(もうちょっとよく勉強してから自分の言葉で記事にしたいと思っています)

pyroの枠組み

さて問題設定について述べたところで,pyroにおける書き方について解説したいと思います.
pyroでは主に,モデルパート, を記述する部分と,
ガイドパート,を記述する部分に分かれます.

この記事は,チュートリアル・SVI-1がベースになっていますが,自分なりに補いつつ解説していきたいと思います.
また,pyroのバージョンは0.3.0を現時点では使っています.

pyro (モデルパート)

このうち, は事前分布と呼ばれ,今回の問題では次のように設定します.

モデルはpyroの枠組みで記述すると次のようになります.

def model(data):
    # set prior
    alpha0 = torch.tensor(5.0)
    beta0 = torch.tensor(5.0)
    prior = dist.Beta(alpha0, beta0)

    # sample z from the prior
    z = pyro.sample("z", prior)

    with pyro.plate("data_plate", len(data)):
        # observe data
        pyro.sample("obs", dist.Bernoulli(z), obs=data)

pyro.sample で分布からサンプルができます.この時に確率変数の名前を一意に設定しなければなりません.

pyro.plate は独立な確率変数をfor文のようにサンプルする時に用います.
上の例では,obsdataのサイズと同じサイズのテンソルになっていますが,
各成分obs[0], obs[1], ...zが与えられた上での条件付き独立になっています.

このように明示的に条件付き独立を宣言することは,変分推論の効率を良くします.
詳しくは,チュートリアル・SVI-2チュートリアル・SVI-3を参照してみてください.

pyro (ガイドパート)

について記述する部分はguideと呼ばれます.(これが一般的なのかはよくわかりません)
これをpyroの枠組みで記述すると次のようになります.

def posterior():
    alpha_q = pyro.param("alpha_q", torch.tensor(15.0),
                         constraint=constraints.positive)
    beta_q = pyro.param("beta_q", torch.tensor(15.0),
                        constraint=constraints.positive)

    return dist.Beta(alpha_q, beta_q)

def guide(data):
    # sample z from the posterior
    pyro.sample("z", posterior())

を定義(posterior() の箇所)して,サンプルします.
最適化の対象のパラメーターはpyro.paramで生成します.勾配計算などはpytorchのテンソル計算を援用しています.
パラメーターの初期化もこの時に行います.torch.tensor(15.0) の箇所が該当します.

チュートリアルにもありますが,いくつか注意点があります.

  • modelとguideの引数は使わないとしても,完全に一致させないといけません.
  • 潜在変数である”z”は,modelとguideで同じ名前”z”でサンプルしなければなりません.
  • (今回は) alpha_q,beta_q は正の値でなければならないので,constarintで制約することができます.

pyro(推論パート)

ここまでくると,パラメーターの最適化は簡単で,pyroの枠組みで実行できます.

adam_params = {"lr": 0.001, "betas": (0.90, 0.999)}
optimizer = Adam(adam_params)

# setup the inference algorithm
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())

# do gradient steps
for step in range(n_steps):
    svi.step(data)

ほぼブラックボックスですが,ELBOを目的関数として勾配法が実行できます.

pyro(評価パート)

上で定義したposterior()は,パラメーター最適化した分布を返すので平均や分散などを計算できます.

inferred_mean = posterior().mean
inferred_std = math.sqrt(posterior().variance)

print("mean: %.3f, std: %.3f"% (inferred_mean, inferred_std) )

分布からサンプリングして,ヒストグラムを作成することも可能です.

plt.hist([posterior().sample() for _ in range(1000)], bins=20, range=(0.05,1), density=True)

実験結果

最後に実験結果を載せておきます.まあこれは入門書でいつもみるやつですね.
サンプル数が増えるに従って,の分布が真の値に峰をもつような分布に近づいていきます.
表が出る真の確率を0.6しました.

データ数10 (mean: 0.552, std: 0.093)
pyro_coin0
データ数30 (mean: 0.577, std: 0.087)
pyro_coin1
データ数50 (mean: 0.586, std: 0.085)
pyro_coin2

まとめ

上の例は,最も簡単な例であり,共役事前分布をとることでELBOなどを考えることなく,解析解を求めることができてしまいます.
しかし,上のコードでは共役事前分布であることを全く用いずに数値的に推論ができているという点が重要です.
それぞれのパーツを簡潔に記述できてしまうというのもpyroの利点ですね.

コード全文は gist に載せました.上のコードだけではライブラリのインポートなどが不十分です.

次回は,ELBOによる最適化とベイズ線形回帰に取り組んでいこうと思っています.