はじめに
タイトルの通り基本だけ書きます。 Pyroは解説できるほど触っていないので、大したことは書けませんが、何も知識が無いよりはとっつきやすくなるであろうことを書いておきます。
ベイズ推論
モデリング
ベイズモデリングとは観測データ $X$ を、パラメータだとか潜在変数だとか呼ばれる「未観測の確率変数」を使ってモデリングする試みです。例えば $X$ が正規分布に従うと思っているのならばパラメータ $\mu, \sigma$ を用いて、$ N(X | \mu, \sigma) $という確率モデルを考えてみたりします。
一般に $p(X | Z)$ とか書いておくことが多いです。 $Z$ というのはパラメータやら潜在変数やら、未観測の確率変数で$X$ の方が観測した手元のデータになります。ベイズモデリングでは未観測の確率変数 $Z$ に対して事前分布 $p(Z)$ を考え、確率モデルと事前分布のタプルで $(p(X|Z), p(Z))$をベイズモデルとして扱います。
事後分布
ベイズモデルを決めると、ベイズの定理から下記の事後分布を書き下すことが出来ます。
$$ p(Z | X) = \frac{p(X|Z)p(Z)}{p(X)} $$
分子は自身が仮定したベイズモデルによって構成されています（それが容易に計算できるかはさておいて）。分母は観測したデータ $X$ によってのみ構成されているので、すでに確定値です（確定値ではあるが計算できるかはこれも保留）。
このように未観測の確率変数 $Z$ に関する分布を観測済のデータ $X$ の条件付き分布（事後分布）として求めることをベイズ推論と言います。
予測分布
さて、 $Z$ とは具体的には仮定した正規分布の平均パラメータかもしれませんし、カテゴリ分布の確率を表すパラメータかもしれませんし、はたまた、物理的実体と照らし合わせた変数かもしれません。事後分布を求めることで、手元に在るデータから自分が仮定した未観測の確率変数に関する分布を得ることができるわけですが、私達が興味を示すのは、今後手に入る $X$ に関する予測であるケースが多いです。
あくまで $Z$ という未観測の確率変数は、 $X$ をモデリングするための仮定として置いたものであり、その潜在変数そのものよりも、置いた仮定によってどれだけ $X$ を上手く表せるかが重要になってくるケースが多いということです（というかそういう汎化性能の評価をすべきだと思う）。
ベイズで使われる予測分布は下記の式で表されます（すでに持っているデータ $X$ を使って $x_{new}$ の分布を表す）。
$$ p(x_{new} | X) = \int_Z p(x_{new}|Z)p(Z|X)dZ $$
これはすなわち、もともと仮定したベイズモデル $(p(X|Z), p(Z))$の確率モデル $p(X | Z)$ の $X$ の部分をこれから予測したい $x{new}$ に置き換え、 p(x{new} | Z) という確率モデルを事後分布 $p(Z|X)$ で重み付けして分布を構築するということです。
実際に使われる予測分布
多くのケースでは
$$ Z_i \sim p(Z|X) $$
とサンプリングを行って、
$$ p(x_{new} | Z_i) $$
という確率モデルを得る。ということをできるだけ大きな数 $N$ 回繰り返すことで、
$$ \frac{1}{N} \sum_{i=1}^N p(x_{new} | Z_i) $$
と $N$ 個の確率モデルの和で予測分布を近似して使います。 $Z_i \sim p(Z|X)$ は当然事後分布 $p(Z|X)$ により選ばれやすい $Z$ が頻出するはずですので、それに応じて選ばれやすい $Z_i$で構築される確率モデル $p(x_{new} | Z_i)$ の割合が多くなるので、元々の予測分布を近似できるという仕組みです（PPLでは数式として解くのが困難な積分はこの方法によって近似的に処理されることが頻繁にあるので、この部分を納得しているだけでだいぶ感覚が違うはずです）。
Pyroの基本
結局のPyroを始めとしたPPLでやろうとしていることは、
- データ $X$ に対して確率モデル $p(X|Z)$ を作る
- $Z$ に対して事前分布 $p(Z)$ を割り当てる
- 事後分布 $p(Z|X)$ を求める
- 予測分布 $p(x_{new} | X )$ を得る
という操作を手軽に実行する仕組みを与えたいというものです（もちろんもっとプリミティブに、確率変数というものを上手に扱う仕組み、例えば微分可能にする仕組みなどを入れることも重要な仕事のひとつである）。
それぞれの操作がどのようなAPIになっているのかを把握できれば、それだけでPPLの見通しがグッと良くなります。今回はPyroのAPIを見ていきましょう。
Pyroの確率変数の取扱
pyroでは確率変数は一貫してpyro.sample('name', distribution_fn)によって表すと決められています。では確率変数が観測データであるのか潜在変数であるのかをどのように見分けるのかというと、pyro.sample('name', distribution, obs=observed_data) として観測データを渡しておいたものだけが観測データとして扱われます。
単に distribution.sample() メソッドを利用した場合、それはPPLとしての確率変数という扱いはされず、言わば普通に numpy.random.randnなどを利用したのと同じになってしまいます（リッチな擬似乱数モジュールに過ぎないということ）。
pyro.sample()という関数は非常に重要なので必ず抑えましょう。
Pyroのハイパーパラメータの取扱
ハイパーパラメータのように固定値として使う物は普通にPythonの数値として、あるいはtorchにTensorとして与えれば問題ありません。個々に関しては普通のプログラミング同様なので困惑するところはないでしょう（困惑するとしたらベイズモデリング自体に困惑している可能性があります。ベイズモデリングでも決め打ちになるハイパーパラメータはあります。いくつか決め打ちを試してみて、その中で一番良さげなのを選ぶという手法もありますが）。
Pyroでの変分パラメータの取扱
Pyroで変分パラメータを扱う場合は、pyro.param("name", tensor)を用います。こいつで宣言された変数tensorは変分パラメータとしてPyroが認識するようになります。変分推論で事後分布を近似する際には、パラメトリックに置いた事後分布のパラメータを上記の関数で宣言することを忘れてはいけません。
変分ベイズ推論のコード：確率モデル
基本的なpyroのパーツとしては、本当に pyro.sample と pyro.param が分かれば一先ず…という感じです。あとは pyro.distributions を使うのですが、これは torch.distributions をラップしてる pyro用の確率分布モジュールで、特に特有の難しさはありません。
まず、ベイズ確率モデルは確率モデル $p(x|\theta)$ と事前分布 $p(\theta)$の組でできるのでした。ここで、事前分布 $p(\theta)$ は何らかの確率分布を仮定するのですがから、その確率分布を表すためのパラメータ $\alpha$ （俗にいうハイパーパラメータ）を決め打ちで準備する必要が在ることに注意してください。
すなわち、正確には
$$ p(x \mid \theta) $$
という確率モデルと
$$ p(\theta \mid \alpha) $$
という事前分布（$\alpha $ は決め打ち）を仮定することから始めます。
ここでは、データ $x$ に対する確率モデルとしてベルヌーイ分布 $p(x \mid f)$ と、事前分布としてベータ分布 $p(\f \mid \alpha, \beta)$ を選ぶこととしましょう。ハイパーパラメータは事前分布のパラメータ数に応じて数が変わることと、これに関しては単なる決め打ちの値になることに注意して下記のコードを読みましょう。
def model(data): alpha0 = torch.tensor(10.0) beta0 = torch.tensor(10.0) # sample f from the beta prior f = pyro.sample("latent_fairness", dist.Beta(alpha0, beta0)) # loop over the observed data for i in range(len(data)): # observe datapoint i using the bernoulli # likelihood Bernoulli(f) pyro.sample("obs_{}".format(i), dist.Bernoulli(f), obs=data[i])
上記のコードを数式で表すならば
$$ \alpha = 10 $$ $$ \beta = 10 $$ $$ f \sim \rm beta(f \mid \alpha=10, \beta = 10) $$ として、$f$ を使って下記のサンプリングを実施。 $$ x_i \sim {\rm Bern}(x\mid f) $$
ということになります。これによって
$$ {\rm Beta}(f\mid \alpha, \beta) \prod_{i}^{N} {\rm Bern}(x \mid f) $$
のような計算が実施されたことになるのです（実際には内部的にはこの対数を取って、対数同時確率を計算しているであろう）。ちなみに観測データである data は $model$関数の引数として渡して、内部の対応する確率変数を pyro.sample関数で生起させるときにobs 引数で紐づけておくことになります。
この辺りがpyroを扱う際の肝になってきます。
ちなみに今回はわかりやすく forループを用いていますが、 with pyro.plate が処理の並列化まで自動で行ってくれるため高速化に使えます（Stanでいうところのベクトル化）。
変分モデル
上記の確率モデル（正規化はされていないが）
$$ {\rm Beta}(f\mid \alpha, \beta) \prod_{i}^{N} {\rm Bern}(x \mid f) $$
に対して、事後分布 $p(f \mid x)$ が我々の知りたいものになります。今回、実はベルヌーイ分布に対してベータ分布を事前分布として選択しているため、解析的に事後分布を求めることができますが、今回は練習のため変分推論を行ってみます。
変分推論のポイントは、事後分布が解析的に解けなかったり、計算が面倒であったり、サンプリングすら計算量的によろしくなさそうな場合に、事後分布を適当なパラメトリックな関数（確率分布の形状として考えられそうなものを選ぶ）として置いてしまうというところです。
$$ q(f) = {\rm Beta} (\alpha', \beta') $$
というパラメトリックな関数として置いてみましょう。もちろんベータ関数で置かなくても良いです。事後分布は正規分布っぽいんじゃないか！？と思うならそれでも良いですし、実は減衰する正弦波みたいな形なんじゃないか！？と思うならそうしてください。そのような形状の者の中で、ELBOが最大化されるような形状が選ばれます。
今回は、解析的に解くとベータ関数で表されることを知っているのでインチキをして事後分布をベータ関数としておいてしまいます。
def guide(data): # register the two variational parameters with Pyro. 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) # sample latent_fairness from the distribution Beta(alpha_q, beta_q) pyro.sample("latent_fairness", dist.Beta(alpha_q, beta_q))
ここで、いま置いたベータ関数（で表されるベータ分布）によって $f \sim \rm Beta$ と出てくると決めたのであれば、それで表現される形の中でそれらしい事後分布が選ばれるだけで、それ以外の物を見つけることはできなくなります。今回は、事後分布の形状の仮定としてベータ分布を置いたことは（正解を知っているので）非常に筋の良い選択ということになりますが、通常はこの部分は手探り近いことを認識しておきましょう。
また、alpha_q と beta_q は変分パラメータなのでpyro.param で変数を作ることになります。引数は名前と初期値と、制約です（今回は正の値という制約が入っている。これは分散などをパラメトライズする場合にもしばしば使うので覚えておくこと）。
pyroの注意点としては、model が確率モデルと事前分布を記述し、 guide が変分モデルを記述すると明確に分けつつも、guide の引数をmodel と合わせる風習が在るようです。当然今回の場合 $q(f) $という変分モデルはデータを引数に取っていない単なるパラメトリックな関数としてなので、内部でデータは使われていないことに注意しましょう（ところがどっこい、VAEなどの場合、ニューラルネットワークのパラメータ自体が全て変分パラメータで、 decoderを確率モデル $p(x|z)$ 正規乱数を事前分布 $p(z)$ encoderを変分モデル $p(z | x)$ として推論を行うので、guide側もdata を内部で使う。 この辺はVAE自体を理解していない場合に混乱をするもととなるので注意が必要。そもそもVAEはパラメータを単にKL損失関数とした普通の勾配学習と思った方が実装は楽である。）。
学習コード
学習は驚くほど簡単です。
# set up the optimizer adam_params = {"lr": 0.0005, "betas": (0.90, 0.999)} optimizer = pyro.optim.Adam(adam_params) # setup the inference algorithm svi = pyro.infer.SVI(model, guide, optimizer, loss=pyro.infer.Trace_ELBO()) n_steps = 5000 # do gradient steps for step in range(n_steps): svi.step(data)
と記述するだけ！びっくりだね！！ SVIなるクラスが確率モデル model と変分モデル guide を渡せば学習をやってくれるというスグレモノになっています。 損失関数はELBOが使えますが、今後他のものも追加されていくのかどうかは謎…。
とりあえずモデルの設計者側が実施するのは基本的に、確率モデルを書いて、変分モデルを書いて、あとは変分推論用のクラスにそいつらを渡してやるという作業だけです。簡単ですね。もちろん確率モデルを階層ベイズモデルにしたければそうすればいいでしょう。
変分推論のカスタマイズ
変分推論のクラスを使うと学習がまるっきり隠蔽されます。少し低レベル側も見ましょう。
# define optimizer and loss function optimizer = torch.optim.Adam(my_parameters, {"lr": 0.001, "betas": (0.90, 0.999)}) loss_fn = pyro.infer.Trace_ELBO.differentiable_loss # compute loss loss = loss_fn(model, guide) loss.backward() # take a step and zero the parameter gradients optimizer.step() optimizer.zero_grad()
というコードがピンと来る方は、まさしくpyroがpytorchをラッピングしたものであるとすぐに分かるはずです。 変分推論は変分パラメータを普通の学習パラメータと思った場合に、勾配学習で最適化を行っているに過ぎません。ただし、その時の損失関数が確率モデル（＋事前分布）と勝手に決めた計算しやすい変分モデルとのKLダイバージェンスないしELBOで規定されているという代物になっています。
もちろん学習後は、変分モデル（事後分布の近似関数）を使って、ベイズ予測分布を構築したりすることができる点で、単なる確率モデルのパラメータの点推定とは異なるのですが、事後分布をパラメトライズして最適化を行うという形式は、計算テクニック上同じになってきます（ここらへんで混乱する人は、まずベイズ推論の基本を勉強する必要があるような気がします。かくいう自分もベイズに触れ始めたときは大層混乱した）。
更に損失関数として使っている pyro.infer.TraceELBO.differentiable_loss もよくよく見てみると
# note that simple_elbo takes a model, a guide, and their respective arguments as inputs def simple_elbo(model, guide, *args, **kwargs): # run the guide and trace its execution guide_trace = poutine.trace(guide).get_trace(*args, **kwargs) # run the model and replay it against the samples from the guide model_trace = poutine.trace( poutine.replay(model, trace=guide_trace)).get_trace(*args, **kwargs) # construct the elbo loss function return -1*(model_trace.log_prob_sum() - guide_trace.log_prob_sum()) svi = SVI(model, guide, optim, loss=simple_elbo)
こんな感じで（当然と言えば当然だが）数式レベルで確率モデルと変分モデルの対数確率の和（の負）を計算してることが丸見えになります。計算に使われている poutineモジュールやその他諸々のメソッドについては公式Docに任せますが、兎にも角にもpyro内でインスタンスを管理する上で色々上手に隠蔽しているということです。
pyroについて
仮に数式レベルで書くことに抵抗がそれほど無いのであれば、TensorFlow Probabilityをおすすめします。 個人的には pyro はTensorFlowで言うところのKeras的小回りの効かなさを感じます。やりたいことがシンプルであればpyroの枠組みで済みますが、損失関数レベルでごちゃごちゃ触るなら、隠蔽されている中身を紐解くのが帰って面倒です。
