はじめに
ここで言う「実用」というのは、理論的な話をごちゃごちゃせず、変分ベイズ推論をチャチャッと回すためのモデルの書き方をPyroを使って説明しますという内容です。「おい！俺は株価を予想できる変分モデルが書きたいのに、全然実用的な内容じゃないぞ！」とか言われても、「個々のドメインで実用的なモデルはこれ！」とかいう解説は皆無ですのでご容赦ください。
尚、最尤推定だとかMAP推定だとかベイズ推論は知っている前提、あるいは言葉を自分で調べられる前提で書くので、記事内で一つ一つ解説はしません。
今回、記事を書こうと思ったモチベーションは、PyroがTF-Pに比べ使える機能を著しく制限しておりながらも、実はやりたいことを、概ね同じコードスタイルで書くことができるということに気づいたためです（というか、そのようにモジュールが整理されていると言える）。
データの分布形状が既知な場合の推論
問題設定
まず、変分推論が何をしているかを見るために、予め出来レース的な問題設定で話を進めてみます。今回扱いのは下記のような問題設定です。
コインを25回投げます。表が過半数である13回以上出た場合には、景品が貰えるという単純なゲームがあるとしましょう。 こうなってくると興味があるのは、コインを1回投げたときの表が出る確率です。これが単純に半分であるとしたならば、このゲームで勝てる確率は半々であろうと考えられます。一方、コインが仮に歪んでおり、表が出る確率が 0.5 であると言えない場合には注意が必要です。
さて、このゲームに使われているコインの表が出る確率を、300人のゲーム参加者の結果から推論してみましょう。
1人1人の参加者が25回コインを投げて、何回表が出たかを入念にメモしていけば、 「0〜25」までの数字が300個集まることになります。これはつまり、確率 $\theta$ で表が出るコインを25回投げたときに表が出た回数 $x$ を生成する二項分布
$$ {\rm Binomial} (x \mid 25, \theta) = {}_{25} C_x \theta^{x}(1-\theta)^{(25-x)} $$
を考えればいいということです。300個のデータを元に$ \theta $ が推論できれば良いという指針です。実はこのコインは歪んでおり、表が出る確率は 0.45 になっていたとしましょう（もちろんこのことは知らない設定で推論を実施する）。
N = 300 TRIAL = 25 THETA = 0.45 def get_coin(): return np.random.binomial(TRIAL, THETA, N) data = get_coin()
すると観測データは下記のようなヒストグラムになっています。概ね12回表が出るような人たちが最も多いように見えますが、このデータから コインを投げたときに表が出る確率を推論していこうというわけです。
ベイズ推論のためのモデリング
まず既に25回コインを投げたときに、コインがX回出るという話を二項分布を使ってモデル化しました。300人分のデータを集めているので、尤度関数
$$ \prod_{i=1}^{300}{\rm Binomial} (x_i \mid 25, \theta) $$
と表すことができ、 $x_i$ の方は既に手元にある確率変数の実現値であるので、上記はただの $\theta$ の関数になっています。ご存知の通りこの尤度関数を $\theta$ に関して最大化してやった場合の解を表が出る確率の推定値とするのが、最尤推定になります。 ベイズ推論では、これから推定するパラメータ $\theta$ の方に事前分布 $p(\theta)$ を予め考慮することでベイズモデリングを実施します。
私達は $\theta$ が $[0, 1]$ の区間に収まる連続値であることを知っているので、これを考慮した事前分布を構えてやりましょう。今回はベータ分布を使います。ベータ分布がどういうものであるかは細かい数式はさておいて、パラメータ $a, b$ を用いて $[0, 1]$ 上に確率分布を構成してくれるものです。例えば $a = 5, b = 5$ としたベータ分布からサンプルを10万個程取ってみると下記となります。
plt.hist(np.random.beta(5, 5, 100000), bins=100)
このベータ分布を用いて $\theta$ に事前分布を与えることにします。もしも上記のように $a=5, b=5$ と設定した場合は、推論をする前から、まあ概ねコインは表裏半々であろうと思っているわけです。しかし、半々からずれる確率もしっかり考慮してあることに注意してください。
一方で $a=1, b=1$ と設定した場合には下記のようなサンプルを発生させます。
plt.hist(np.random.beta(1., 1., 100000), bins=100)
この場合は、事前分布として真っ平らな分布を仮定しているわけですから、このコインは表が出る確率がどれくらい偏っていそうか検討もつかないという思いを抱いていることになります（とはいっても、半々であろうという思いと、いや全部表が出るようになっているんだ！という思いが同等に強いなんて、観測データからもありえないですわな…笑）
このような分布を無情報事前分布と言いますが、推論に自分の推測を一切乗せたくないのであれば無情報事前分布を用いることをおすすめします。（Stanなどでは、特別に事前分布を準備しない限り自動的に無情報事前分布が用いられる）
今回は下記の事前分布使うことにしましょう。
$$ p(\theta) = {\rm Beta} (\theta \mid a=1, b=1) $$
こうして、尤度と事前分布を準備したので、事後分布を一先ず書き下すことができます。これでベイズ推論を開始することができます。
$$ p(\theta \mid D) = \frac{\prod_{i=1}^{300}{\rm Binomial} (x_i \mid 25, \theta) p(\theta)}{p(D)} $$
の左辺を求めることができれば推論完了です。
共役事前分布を用いた解析的推論
さて、実は尤度関数に対して数学的に都合の良い事前分布を選択してやると、事後分布は解析的に求まります。 要するに、計算力の高い人間ならば、紙の上で推論が完了してしまうのです！
今回はそういうこともできるよ、ということだけ伝えて、変分推論はそれができないようなケースで使う手法であることを強調しておきます。
変分推論
さて、変分推論は事後分布を解析的に求められないときに、事後分布に適当な仮定を置いて最適化問題に変更してしまうような手法です。平均場近似のようにパラメータに条件付き独立などの制約を与えて上手に分解する方法もありますが、一先ず頭を使わなくても誰でもできる方法を紹介します。
それは事後分布、それ自体の形状を既に知っている確率分布で置いてしまうという方法です。今回でいうと、
$$ p(\theta \mid D) = \frac{\prod_{i=1}^{300}{\rm Binomial} (x_i \mid 25, \theta) p(\theta)}{p(D)} $$
という分布がどんな形状をしているのかは（まあ実は分かるのだが）、なかなか想像がつきません。よく分からんので、$p(\theta \mid D)$の形状がどんなものかを適当に想像して正規分布みたいな形になってんじゃないの？ といい加減に予想してみます。実際、$\theta$ は表が出る確率であるので、その事後分布は例えば $0.5$ を中心に減衰したような分布であったと思っても何もおかしくないですね。
つまり
$$ q(\theta) = \rm N (\theta \mid \mu, \sigma) $$
で事後分布を近似してしまおうと考えるのです。 $\mu, \sigma$ を上手に調整して $p(\theta \mid D)$ と帳尻の合う形状を探します。帳尻が合うというのは、$KL(q:p)$ が最も小さいという意味です。
Pyro で変分推論
Pyroを使うと上記の推論を簡単に記述できます。まず私達は尤度関数と事前分布を下記のように設定したのでした。
尤度関数 $$ {\rm Binomial} (x \mid 25, \theta) = {}_{25} C_x \theta^{x}(1-\theta)^{(25-x)} $$
事前分布 $$ p(\theta) = {\rm Beta} (\theta \mid a=1, b=1) $$
このような分布からデータを生成する流れは下記のようになります。 まず$\theta$がただ1つ、事前分布から生成されます。
$$ \theta \sim {\rm Beta} (\theta \mid a=1, b=1) $$
その後、そのパラメータ $\theta$ を使って、観測データが独立に300個手に入ります。
$$ x_i \sim {\rm Binomial} (x_i \mid 25, \theta) $$
観測データは300個あるのですが、各々のデータに別に $\theta$ を使うわけではないことに注意してください。最初に生成した $\theta$ と同一の$\theta$を用いて、独立に300個のデータが手に入ったという生成過程を考えます。
これをPyroで記述すると下記になります。
def model(data): # hyperparameter alpha0 = torch.tensor(1.0).type_as(data) beta0 = torch.tensor(1.0).type_as(data) # パラメータを生成 theta = pyro.sample("latent_fairness", dist.Beta(alpha0, beta0)) # 同一のパラメータthetaを用いて観測データを生成。 with pyro.plate("sample", len(data)): pyro.sample("obs", dist.Binomial(TRIAL, theta), obs=data)
私達はこのようなモデルを考え事後分布を下記のように一先ず書き下したのでした。
$$ p(\theta \mid D) = \frac{\prod_{i=1}^{300}{\rm Binomial} (x_i \mid 25, \theta) p(\theta)}{p(D)} $$
これに対して、変分推論モデルを正規分布でいい加減に作ったことをPyroで書くには下記のようになります。
def guide(data): # 最適化すべき変分パラメータ mu_q = pyro.param("mu_q", torch.tensor(0.5), constraint=constraints.positive).type_as(data) sigma_q = pyro.param("sigma_q", torch.tensor(0.1), constraint=constraints.positive).type_as(data) # 事後分布の肩代わりをしている正規分布から p が生成 pred_p = pyro.sample("latent_fairness", dist.Normal(mu_q, sigma_q)) return pred_p
ここで data を引数に取っているのは、model と guide の引数を合わせるというPyroの慣習のようなものでベイズ推論には本質的には関係ありません。私達は事後分布を勝手に近似した変分推論のモデルを素直に書いておけば良いのです（なぜPyroがこのような慣習を持つのかというと、特に深層生成モデルでは事後分布の近似モデル自体をデータに依存させるケースが多いためである。例えばVAEの変分モデルはデータに条件付けられているため、guide の中で引数データを使う）。
変分推論はKLダイバージェンスを最小化すると決まっているので、尤度と事前分布、変分モデルを書き終えた時点でやることは決まっています。Pyroではこの決まりきった作業はモジュールとして提供されているため、それをありがたく使えばいいです。オプティマイザーとしてはAdamを利用し、SVI クラスに作った model と guide とoptimizer を渡してやるだけでセッティングが完了します。学習は SVI.step()で1回進めることができます。
# setup the optimizer pyro.clear_param_store() adam_params = {"lr": 0.0005, "betas": (0.90, 0.999)} optimizer = Adam(adam_params) # setup the inference algorithm svi = SVI(model, guide, optimizer, loss=Trace_ELBO()) data = torch.tensor(data, dtype=torch.float) for step in range(2000): svi.step(data) if step % 100 == 0: print('.', end='')
学習完了後に、自分が勝手に作った変分モデル guide からサンプルを取り出せば、事後分布の近似が得られます。
sns.distplot([guide(data).detach().numpy() for _ in range(1000)])
概ね $p=0.45$ 付近に平均を持つ形状となりました。私達は本来仮定した尤度関数と事前分布の形状を無視して、適当に仮定した正規分布の中で、本来の形状にKLダイバージェンスの意味で最も近い形状を獲得したということを忘れてはなりません。
振り返り
さて、変分モデルを正規分布でなんとなく仮定してみましたが、これは良かったのでしょうか。 実は、計算の途中で標準偏差の方が大きくなっていくと変なことが起こります。元々推論すべきだった事後分布
$$ p(\theta \mid D) = \frac{\prod_{i=1}^{300}{\rm Binomial} (x_i \mid 25, \theta) p(\theta)}{p(D)} $$
は $[0, 1]$ にしか値を取りません。にも関わらず、正規分布で $\theta$ の分布を仮定すると、本来仮定したモデルと矛盾した生成が行われる可能性があるのです（そういうわけで、実は変分モデル guide を書くときには、平均の初期値を 0.5 で標準偏差の初期値を0.1 として変なことが起こりづらくしている。仮にわざと平均を 10 で初期化したりすればエラーが返ってくる。今回は分かりやすい問題で見ているのでいいが、複雑な問題やモデルになったときに、意外とこういう部分で躓いたりする）。
もしも、本当に形状が正規分布で近似できるのであれば、数学的にちょっとくらいおかしかろうが結果を取り出せて満足という場合もあります。しかし、今回の近似は実際はかなり危ういのです。
確率変数が $[0, 1]$ にしか値を取らないと分かっているのであれば、そのような確率分布を変分モデルとして仮定するほうが安全です。
定石としては、変分モデルの方にもベータ分布を仮定してしまいます。
def guide(data): alpha_q = pyro.param("alpha_q", torch.tensor(10.0), constraint=constraints.positive).type_as(data) beta_q = pyro.param("beta_q", torch.tensor(10.0), constraint=constraints.positive).type_as(data) pred_p = pyro.sample("latent_fairness", dist.Beta(alpha_q, beta_q)) return pred_p
そうすると、この変分モデルで得られた推論は、どう転んでも $[0, 1]$ にしか値を生成しません（実は今回のモデルの設定上、解析的にもこちらの形式が正しい）。
Pyroはとりあえずは動いてくれますが、あまりにいい加減な設定で動かしていると、変な挙動やエラーを返してくるので注意してください。 特に、チュートリアルや例題の多くは正規分布で変分モデルを作っているケースが多々あります。変分モデルに正規分布を作った時点で、その形状の中で最もまともなのを選ぶしかできなくなることに注意しましょう。
これに対してMCMCは変分モデルのように形状を仮定せず、自分が尤度関数と事前分布を指定することで作った形状を解析的に求められない事後分布をなんとか上手にサンプリングで作っていこうという試みになります。
追記：変分ベイズ推論を応用した最尤推定、MAP推定
さて、Pyroで提供されているのはこの変分モデルの最適化フレームワークです。 もしもあなたが、自分は十分にデータを持っているので、最尤推定やMAP推定で済ませたいと思った場合にはどうすればいいでしょうか。当然、頑張ってPyTorchでその最適化問題を記述するというのもありですが、せっかく確率分布や確率変数を楽に取り回せるフレームワークがあるのですから、これを利用しない手立てはありません。
MAP推定
MAP推定は事後分布 $p(\theta \mid D) $ で最も高い値を取る $\theta$ を推定値とする推定手法です。 例えば事後分布を求めてしまってから ${\rm argmax}_{\theta}p(\theta \mid D)$ とするのは一つの手でしょう。しかし変分推論では一般的に一つの $\theta$ に対して例えば正規分布を仮定すれば、推定すべきパラメータは平均と分散の2つになるなど、パラメータは増加します。MAP推定では $\theta$ だけを直接求めたいのに、わざわざ変分推論してから ${\rm argmax}$ を計算するなんて二度手間です。
そのためPyroにはモデルに対してMAP推定を実施するための変分モデルの書き方が提案されています。
それは guide を以下のように書き換えることです。
from pyro.contrib.autoguide import AutoDelta guide_map = AutoDelta(model)
AutoDeltaクラスは、モデルを引数に推定すべきパラメータ$\theta$ を自動で検出し、変分モデル$q(\theta)$ としてDelta関数を選んでくれます。$\delta (0) = \infty$ で $\delta({\rm not} 0) = 0$ となるものをデルタ関数と呼びます。$\delta(\theta)$ ならば $\theta=0$ でのみ値を持つので、 $q(\theta) = \delta(\theta - \eta)$ という変分モデルを仮定すると、要するに $\theta = \eta$ のときのみ値を取る確率分布であり、$\eta$を変分パラメータとして変分推論すれば、$\theta = \eta$ がMAP推定値になるという仕組み（だと思います）です。
AutoDelta クラスは辞書でパラメータの名前を key とそのテンソルを valueとして持っており、__call__ で辞書を返す仕組みになっています。
from pyro.contrib.autoguide import AutoDelta guide_map = AutoDelta(model) # setup the optimizer pyro.clear_param_store() adam_params = {"lr": 0.001, "betas": (0.90, 0.999)} optimizer = Adam(adam_params) # setup the inference algorithm svi = SVI(model, guide_map, optimizer, loss=Trace_ELBO()) data = torch.tensor(data, dtype=torch.float) for step in range(2000): svi.step(data) if step % 100 == 0: print('.', end='') guide_map() # -> {'latent_fairness': tensor(0.4540, grad_fn=<ExpandBackward>)}
最尤推定
こうなると、model内の事前分布に対して無情報事前分布置いて、上記の方法でMAP推定を実施すれば最尤推定が可能であるということになります。こうして、Pyroのフレームワークを通じてベイズ推論も点推定もできるようになりました。
