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

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 08/05/2020 (1.4.0)

* 本ページは、Pyro の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

• SVI Part I: An Introduction to Stochastic Variational Inference in Pyro

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

 

 

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

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

 

セットアップ

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

1. 観測 ⟺ pyro.sample with obs 引数
2. 潜在確率変数 ⟺ pyro.sample
3. パラメータ ⟺ 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$ は以下の特性を持つことを仮定します :

1. 各々の $p_i$ からサンプリングできます。
2. pointwise な log pdf (対数確率密度関数) $p_i$ を計算できます。
3. $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}({\bf 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() で指定されたとき最初の引数は確率変数の名前を示すことを思い出してください。これらの名前はモデルとガイドで確率変数を整列 (= align) させるために使用されます。非常に明示的にするためには、モデルが確率変数 $z_1$ を含むならば :

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

then、ガイドは一致する sample ステートメントを持つ必要があります :

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() を提供します、これらは変分学習と評価のためのロジックをカプセル化します :

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

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

 

Optimizer

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

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

これは幾つかの挑戦を提示します、何故ならばそれは、(例えば、) 特定の潜在確率変数とパラメータが時々だけ出現するような、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 は次のシグネチャを持たなければなりません :

1. module_name: (もしあれば、) パラメータを含むモジュールの Pyro 名
2. 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 により発行された標準的なクォーター (25セント硬貨) である。
• それは長年の使用により少し傷があります (= 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):
    # beta 事前分布を制御するハイパーパラメータ
    alpha0 = torch.tensor(10.0)
    beta0 = torch.tensor(10.0)
    # beta 事前分布からのサンプル f
    f = pyro.sample("latent_fairness", dist.Beta(alpha0, beta0))
    # 観測データに渡るループ
    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):
    # Pyro で 2 つの変分パラメータを登録する
    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)
    # 分布 Beta(alpha_q, beta_q) から latent_fairness をサンプリングする
    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 を使用します。

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

# optimizer をセットアップする
adam_params = {"lr": 0.0005, "betas": (0.90, 0.999)}
optimizer = Adam(adam_params)

# 推論アルゴリズムをセットアップする
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())

n_steps = 5000
# 勾配ステップを行なう
for step in range(n_steps):
    svi.step(data)

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

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

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

# 検証を有効にする (e.g. 分布のパラメータを検証する)
assert pyro.__version__.startswith('1.4.0')
pyro.enable_validation(True)

# REPL にある場合に param ストアをクリアする
pyro.clear_param_store()

# 6 つの観測された表と 4 つの観測された裏で幾つかのデータを作成する
data = []
for _ in range(6):
    data.append(torch.tensor(1.0))
for _ in range(4):
    data.append(torch.tensor(0.0))

def model(data):
    # beta 事前分布を制御するハイパーパラメータを定義する
    alpha0 = torch.tensor(10.0)
    beta0 = torch.tensor(10.0)
    # beta 事前分布から f をサンプリングする
    f = pyro.sample("latent_fairness", dist.Beta(alpha0, beta0))
    # 観測データに渡りループする
    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):
    # 2 つの変分パラメータを Pyro で登録する
    # - 両者のパラメータは初期値 15.0 を持ちます。
    # - constraints.positive を起動しますので、optimizer は非制約のパラメータ上で勾配を取ります
    # (これは対数により制約されたパラメータに関連します)
    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)
    # 分布 Beta(alpha_q, beta_q) から latent_fairness をサンプリングする
    pyro.sample("latent_fairness", dist.Beta(alpha_q, beta_q))

# optimizer をセットアップする
adam_params = {"lr": 0.0005, "betas": (0.90, 0.999)}
optimizer = Adam(adam_params)

# 推論アルゴリズムをセットアップする
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())

# 勾配ステップを行なう
for step in range(n_steps):
    svi.step(data)
    if step % 100 == 0:
        print('.', end='')

# 学習された変分パラメータを手に入れる
alpha_q = pyro.param("alpha_q").item()
beta_q = pyro.param("beta_q").item()

# ここで beta 分布についての幾つかの事実を利用します。
# コインの公正さの推論された平均を計算する
inferred_mean = alpha_q / (alpha_q + beta_q)
# 推論された標準偏差を計算する
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.5\bar{3}\) で与えられます。コインの公正さの最終的な推定は事前分布 (すなわち \(0.50\)) により選択された公正さと生の実験上の頻度 (\(6/10 = 0.60\)) により提示された公正さの間にあることに注意してください。

 

References

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

以上