Pyro 0.3.0 : SVI (1) 確率的変分推論へのイントロダクション (翻訳)

翻訳 : (株)クラスキャットセールスインフォメーション
作成日時 : 12/16/2018 (v0.3.0)

* 本ページは、Pyro のドキュメント SVI Part I: An Introduction to Stochastic Variational Inference in Pyro を
翻訳した上で適宜、補足説明したものです：

http://pyro.ai/examples/svi_part_i.html

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

SVI (1) 確率的変分推論へのイントロダクション

Pyro は確率的変分推論を汎用目的推論アルゴリズムとしてサポートするために特別な注意が払われてデザインされています。Pyro で変分推論を行なうことをどのように取り扱うかを見てみましょう。

セットアップ

私達は既にモデルを Pyro で定義したものと仮定します (これがどのように成されるかの詳細はイントロ (1) 参照)。簡潔な確認として、モデルは確率関数 model(*args, **kwargs) として与えられて、これは一般的には引数を取ります。model() の異なるピースはマッピングを通してエンコードされます :

観測 ⟺ pyro.sample with obs 引数
潜在確率変数 ⟺ pyro.sample
パラメータ ⟺ pyro.param

さて幾つかの表記を定めましょう。モデルは観測 $x$ と潜在確率変数 $z$ そしてパラメータ $\theta$ を持ちます。それは次の形の同時確率密度を持ちます :
\[
p_{\theta}({\bf x}, {\bf z}) = p_{\theta}({\bf x}|{\bf z}) p_{\theta}({\bf z})
\]

$p_{\theta}({\bf x}, {\bf z})$ を作り上げる様々な確率分布 $p_i$ は次の特性を持つことを仮定します :

各々の $p_i$ からサンプリングできます。
pointwise な対数確率密度関数 (= log pdf) $p_i$ を計算できます。
$p_i$ はパラメータ $\theta$ に関して微分可能です。

モデル学習

このコンテキストで良いモデルを学習するための規準は対数エビデンスを最大化することです, i.e. 次で与えられる $\theta$ の値を見つけることを望みます :
\[
\theta_{\rm{max}} = \underset{\theta}{\operatorname{argmax}} \log p_{\theta}({\bf x})
\]

ここで対数エビデンス $\log p_{\theta}({\bf x})$ は次で与えられます :
\[
\log p_{\theta}(x) = \log \int\! d{\bf z}\; p_{\theta}({\bf x}, {\bf z})
\]

一般的なケースではこれは二重に難しい問題です。これは何故ならば (固定された $\theta$ に対してさえも) 潜在確率変数 $z$ に渡る積分がしばしば扱いにくいからです。逆に、$\theta$ の総ての値に対して対数エビデンスをどのように計算するかを知る場合でさえも、$\theta$ の関数として対数エビデンスを最大化することは一般に困難な非凸 (= non-convex) 最適化問題でしょう。

$\theta_{\rm{max}}$ を見つけることに加えて、潜在変数 $z$ に渡る事後分布を計算したいです :
\[
p_{\theta_{\rm{max}}}({\bf z} | {\bf x}) = \frac{p_{\theta_{\rm{max}}}({\bf x} , {\bf z})}{
\int \! d{\bf z}\; p_{\theta_{\rm{max}}}({\bf x} , {\bf z}) }
\]

この式の分母は (通常は扱いにくい) エビデンスであることに注意してください。変分推論は $\theta_{\rm{max}}$ を見つけて事後分布 $p_{\theta_{\rm{max}}}({\bf z} | {\bf x})$ への近似を計算するためのスキームを提供します。それがどのように動作するかを見てみましょう。

ガイド

基本的なアイデアはパラメータ化された分布 $q_{\phi}({\bf z})$ を導入することです、ここで $\phi$ は変分パラメータとして知られています。殆どの文献ではこの分布は変分分布 (= variational distribution) と呼ばれ、そして Pyro のコンテキストではそれはガイドと呼ばれます (1音節 instead of nine!)。ガイドは事後分布への近似としてサーブします。

モデルのように、ガイドは確率関数 guide() としてエンコードされ、これは pyro.sample と pyro.param ステートメントを含みます。それは観測されたデータを含みません、何故ならばガイドは正しく正規化された分布である必要があるからです。Pyro は model() と guide() が同じ call シグネチャを持つことを強制することに注意してください、i.e. 両者の callable は同じ引数を取るべきです

ガイドは事後分布 $p_{\theta_{\rm{max}}}({\bf z} | {\bf x})$ への近似ですから、ガイドはモデルの総ての潜在確率変数に渡る妥当な同時確率密度を提供する必要があります。Pyro で確率変数がプリミティブ・ステートメント pyro.sample() で指定されたとき最初の引数は確率変数の名前を示すことを思い出してください。これらの名前はモデルとガイドで確率変数を整列させるために使用されます。非常に明示するためには、モデルが確率変数 $z_1$ を含むならば :

def model():
    pyro.sample("z_1", ...)

def guide():
    pyro.sample("z_1", ...)

2 つのケースで使用される分布は異なるかもしれませんが、名前は 1-対-1 で並ばなければなりません。

ひとたびガイドを指定したのであれば (下に幾つかの明示的な例を与えます)、推論に進む準備ができました。学習は最適化問題としてセットアップされてそこでは訓練の各反復は $\theta-\phi$ 空間でステップを取り、それはガイドを正確な事後分布に近づけます。これを行なうためには適切な目的関数を定義する必要があります。

ELBO

単純な導出 (例えば参照 [1] を見てください) は私達が追い求めているものを生成します : エビデンス下限 (ELBO, evidence lower bound) です。ELBO は $\theta$ と $\phi$ の両者の関数で、ガイドからのサンプルに関する期待値として定義されます :
\[
{\rm ELBO} \equiv \mathbb{E}_{q_{\phi}({\bf z})} \left [
\log p_{\theta}({\bf x}, {\bf z}) – \log q_{\phi}({\bf z})
\right]
\]

仮定により期待値内の対数確率が計算できます。そしてガイドはサンプリングできるパラメトリック分布と仮定されていますので、この量のモンテカルロ推定を計算できます。重要なことは、ELBO は対数エビデンスへの下限であること, i.e. $\theta$ と $\phi$ の総ての選択に対して次を持ちます :
\[
\log p_{\theta}({\bf x}) \ge {\rm ELBO}
\]

そして ELBO を最大化するために (確率的) 勾配ステップを取れば、(期待値の) 対数エビデンスもまたより高く押し上げていきます。更に、ELBO と対数エビデンスの間の隔たりはガイドと事後分布間の KL ダイバージェンスにより与えられることも示すことができます :
\[
\log p_{\theta}({\bf x}) – {\rm ELBO} =
\rm{KL}\!\left( q_{\phi}({\bf z}) \lVert p_{\theta}({\bf z} | {\bf x}) \right)
\]

この KL ダイバージェンスは 2 つの分布間の「近さ」の特定の (非負な) 尺度です。従って、固定された $\theta$ に対して、$\phi$ 空間で ELBO を増大させるステップを取るとき、ガイドと事後分布の間の KL ダイバージェンスを減少させます, i.e. ガイドを事後分布に移動させます。一般的なケースでは、ガイドが移動する事後分布 $\log p_{\theta}({\bf z} | {\bf x})$ を追跡して、ガイドとモデルが追い掛けっこをするように、$\theta$ と $\phi$ 空間の両者で同時に勾配ステップを取ります。おそらく幾分驚くべきことに、移動するターゲットにも拘らず、この最適化問題は多くの異なる問題に対して (近似の適切なレベルに) 解くことができます。

従って高いレベルで変分推論は容易です : 行なうことが必要な総てはガイドを定義して ELBO の勾配を計算することです。実際には、一般的なモデルとガイドのペアのために勾配を計算することは幾つかの複雑さに繋がります (議論のためにはチュートリアル SVI パート III 参照)。このチュートリアルの目的のためには、その解かれた問題を熟慮して変分推論を行なうために Pyro が提供するサポートを見ましょう。

SVI クラス

Pyro では変分推論を行なうための機構は SVI クラスにカプセル化されています (現在のところ SVI は ELBO 目的関数のためのサポートだけを提供しますが、将来的には Pyro は代替の変分目的関数のサポートを提供します)。

ユーザは 3 つのものを提供する必要があります : モデル、ガイド、そして optimizer です。モデルとガイドは上で議論して optimizer は下で多少議論しますので、3 つの構成要素総てを手元に持つものと仮定しましょう。ELBO 目的関数を通して最適化を行なう SVI のインスタンスを構築するために、ユーザは次のように書きます :

import pyro
from pyro.infer import SVI, Trace_ELBO
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())

SVI オブジェクトは 2 つのメソッド step () と evaluate_loss() を提供します、これらは変分学習と評価のためのロジックをカプセル化します :

メソッド step() は単一の勾配ステップを取りそして損失 (i.e. minus the ELBO) の見積もりを返します。提供される場合、step() への引数は model() と guide() へパイプされます。
メソッド evaluate_loss() は勾配ステップを取ることなしに損失の見積もりを返します。ちょうど step() のためのように、提供されれば、evaluate_loss() への引数は model() と guide() にパイプされます。

損失が ELBO であるケースについては、両方のメソッドはまたオプション引数 num_particles を受け取ります、これは損失 (evaluate_loss の場合) を計算するためと損失 (step の場合) と勾配を計算するために使用されるサンプル数を示します。

Optimizer

Pyro では、以下の条件でモデルとガイドは任意の確率関数であることが許容されます :

ガイドは obs 引数を持つ pyro.sample ステートメントを含まない。
モデルとガイドは同じ call シグネチャを持ちます。

これは幾つかの挑戦を提示します、何故ならばそれは、e.g. 特定の潜在確率変数とパラメータが時々だけ出現するような、model() と guide() の異なる実行がまったく異なる挙動を持つかもしれないことを意味するからです。実際にパラメータは推論の過程において動的に作成されるかもしれません。換言すれば、$\theta$ と $\phi$ によりパラメータ化される、最適化を行おうとする空間は動的に増大して変わることができます。

この挙動をサポートするために、Pyro は各パラメータに対して学習中にそれが最初に出現したとき optimizer を動的に生成する必要があります。幸い、PyTorch は動的なケースのために容易に再目的化できる軽量最適化ライブラリ (torch.optim 参照) を持ちます。

これら総ては optim.PyroOptim により制御され、これは基本的には PyTorch optimizer まわりの薄いラッパーです。PyroOptim は 2 つの引数を取ります : PyTorch optimizer optim_constructor のためのコンストラクタと optimizer 引数 optim_args の仕様です。高いレベルでは、最適化の過程で、新しいパラメータが見られるときはいつでも、optim_args により与えられる引数とともに与えられたタイプの新しい optimizer をインスタンス化するために optim_constructor が使用されます。

殆どのユーザは多分 PyroOptim と直接には相互作用しません、そして代わりに optim/__init__.py で定義されるエイリアスと相互作用します。それがどのように進むかを見てみましょう。optimizer 引数を指定する 2 つの方法があります。より単純なケースでは、optim_args は固定された辞書でこれは総てのパラメータのために PyTorch optimizer をインスタンス化するために使用される引数を指定します :

from pyro.optim import Adam

adam_params = {"lr": 0.005, "betas": (0.95, 0.999)}
optimizer = Adam(adam_params)

引数を指定する 2 番目の方法はより洗練された制御レベルを可能にします。ここではユーザは新しく見られるパラメータのために optimizer の作成時に Pyro により起動される callable を指定しなければなりません。この callable は次のシグネチャを持たなければなりません :

module_name: (もしあれば、) パラメータを含むモジュールの Pyro 名
param_name: パラメータの Pyro 名

これは例えば、異なるパラメータのために学習率をカスタマイズする機能をユーザに与えます。この類の制御のレベルが有用である例については、baseline の議論を見てください。API を示す単純な例がここにあります :

from pyro.optim import Adam

def per_param_callable(module_name, param_name):
    if param_name == 'my_special_parameter':
        return {"lr": 0.010}
    else:
        return {"lr": 0.001}

optimizer = Adam(per_param_callable)

これは、Pyro パラメータ my_special_parameter に対しては 0.010 の学習率をそして他の総てのパラメータのためには 0.001 の学習率を使用することを Pyro に単に知らせます。

単純な例

単純な例で終わりにします。両面があるコインが与えられました。貴方はコインが公正であるか否か, i.e. それが表か裏に同じ頻度で落下するかを決定することを望みます。貴方は、次の 2 つの観測をもとにコインの尤もらしい公正さについて事前信念を持ちます :

それは US Mint により発行された標準的なクォーターである。
それは長年の使用により少し傷があります (= banged up)。

そのためコインが最初に生成されたときにそれが極めて公正であったことを想定する一方で、その公正さが完全な 1:1 比率から偏位することを考慮します。従ってコインが 11:10 の比率で表を裏よりも選択したと判明した場合、貴方は驚かないでしょう。反対にコインが 5:1 の比率で表を裏よりも選択したことが判明すれば、貴方は非常に驚くでしょう — それはそこまで傷ついてはいません。

これを確率モデルに変えるために表と裏を 1 と 0 にエンコードします。コインの公正さを実数 $f$ としてエンコードします、ここで $f$ は $f \in [0.0, 1.0]$ を満たして $f=0.50$ は完全に公正なコインに対応します。$f$ についての私達の事前信念はベータ分布、特に $\rm{Beta}(10,10)$ としてエンコードされ、これは区間 $[0.0, 1.0]$ 上の対称確率分布で $f=0.5$ で最大になります。

Figure 1: コインの公正さについての事前信念をエンコードする分布 Beta

幾分曖昧な事前 (信念) よりもより正確な、コインの公正さについての何かを学習するために、実験を行なって幾つかのデータを集める必要があります。コインを 10 回弾いて各フリップの結果を記録するとしましょう。実際には多分 10 試行以上を行なうことを望むでしょうが、しかしこれはチュートリアルですので。

リストデータでデータを集めたと仮定すると、対応するモデルは次で与えられます :

import pyro.distributions as dist

def model(data):
    # define the hyperparameters that control the beta prior
    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])

ここで単一の潜在確率変数 (‘latent_fairness’) を持ち、これは $\rm{Beta}(10, 10)$ に従って分布します。その確率変数に条件付けられて、bernoulli 尤度を使用してデータポイントの各々を観測します。各観測には Pyro で一意な名前が割り当てられることに注意してください。

次のタスクは対応するガイド i.e. 潜在確率変数 $f$ のための適切な変分分布を定義することです。ここでの唯一の実際的な要件は $q(f)$ が範囲 $[0.0, 1.0]$ に渡る確率分布であるべきことで、何故ならば $f$ はその範囲外では無意味だからです。単純な選択は、2 つの訓練可能なパラメータ $\alpha_q$ と $\beta_q$ によりパラメータ化されたもう一つの beta 分布を使用することです。実際に、この特定のケースではこれは「正しい」選択です、何故ならば bernoulli と beta 分布の共役は正確な事後分布が beta 分布であることを意味します。Pyro では次のように書きます :

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))

ここで注意すべき幾つかのことがあります :

確率変数の名前は model と guide の間で正確に並ぶように注意しています。
model(data) と guide(data) は同じ引数を取ります。
変分パラメータは torch.tensor です。requires_grad フラグは pyro.param により自動的に True に設定されます。
最適化の間 alpha_q と beta_q が非負であり続けることを確かなものにするために constraint=constraints.positive を使用します。

さて確率的変分推論の遂行へと進むことができます。

# set up the optimizer
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())

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

step() メソッドで私達はデータを渡し、それからこれは model と guide に渡されることに注意してください。

この時点で見逃している唯一のものは何某かのデータです。そこで幾つかデータを作成して上の総てのコードスニペットを完全なスクリプトに集めましょう :

from __future__ import print_function
import math
import os
import torch
import torch.distributions.constraints as constraints
import pyro
from pyro.optim import Adam
from pyro.infer import SVI, Trace_ELBO
import pyro.distributions as dist

# this is for running the notebook in our testing framework
smoke_test = ('CI' in os.environ)
n_steps = 2 if smoke_test else 2000

# enable validation (e.g. validate parameters of distributions)
assert pyro.__version__.startswith('0.3.0')
pyro.enable_validation(True)

# clear the param store in case we're in a REPL
pyro.clear_param_store()

# create some data with 6 observed heads and 4 observed tails
data = []
for _ in range(6):
    data.append(torch.tensor(1.0))
for _ in range(4):
    data.append(torch.tensor(0.0))

def model(data):
    # define the hyperparameters that control the beta prior
    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
        pyro.sample("obs_{}".format(i), dist.Bernoulli(f), obs=data[i])

def guide(data):
    # register the two variational parameters with Pyro
    # - both parameters will have initial value 15.0.
    # - because we invoke constraints.positive, the optimizer
    # will take gradients on the unconstrained parameters
    # (which are related to the constrained parameters by a log)
    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))

# setup the optimizer
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())

# do gradient steps
for step in range(n_steps):
    svi.step(data)
    if step % 100 == 0:
        print('.', end='')

# grab the learned variational parameters
alpha_q = pyro.param("alpha_q").item()
beta_q = pyro.param("beta_q").item()

# here we use some facts about the beta distribution
# compute the inferred mean of the coin's fairness
inferred_mean = alpha_q / (alpha_q + beta_q)
# compute inferred standard deviation
factor = beta_q / (alpha_q * (1.0 + alpha_q + beta_q))
inferred_std = inferred_mean * math.sqrt(factor)

print("\nbased on the data and our prior belief, the fairness " +
      "of the coin is %.3f +- %.3f" % (inferred_mean, inferred_std))

サンプル出力:

based on the data and our prior belief, the fairness of the coin is 0.532 +- 0.090

推定は正確な事後平均と比較されます、これはこの場合 16/30=0.53¯ で与えられます。コインの公正さの最後の推定は事前 (すなわち 0.5) により選択された公正さと生の実験上の頻度 (6/10=0.60).により提示された公正さの間にあることに注意してください。

References

Automated Variational Inference in Probabilistic Programming, David Wingate, Theo Weber
Black Box Variational Inference, Rajesh Ranganath, Sean Gerrish, David M. Blei
Auto-Encoding Variational Bayes, Diederik P Kingma, Max Welling

以上

月	火	水	木	金	土	日
					1	2
3	4	5	6	7	8	9
10	11	12	13	14	15	16
17	18	19	20	21	22	23
24	25	26	27	28	29	30
31