PyTorch : Pyro SVI (2) 条件付き独立性、サブサンプリング 及び Amortization

by

PyTorch : Pyro SVI (2) 条件付き独立性、サブサンプリング 及び Amortization (翻訳)

翻訳 : (株)クラスキャット セールスインフォメーション
更新日時 : 11/08/2018 (v0.2.1)
作成日時 : 10/13/2018 (v0.2.1)

* 本ページは、Pyro のドキュメント SVI Part II: Conditional Independence, Subsampling, and Amortization を
動作確認・翻訳した上で適宜、補足説明したものです：

* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

目標: SVI を巨大なデータセットにスケールする

$N$ 観測を持つモデルに対してモデルとガイドを実行して ELBO を構築することは log pdf の評価を伴い、その複雑さは $N$ で酷くスケールします。巨大なデータセットにスケールすることを望む場合これは問題です。幸い、モデル/ガイドが (利用できる) ある条件付き独立性構造を持つ条件下では ELBO 目的はサブサンプリングを自然にサポートします。例えば、latent が与えられたとき観測が条件的に独立である場合には ELBO の log 尤度項目は次で近似されます :

\[
\sum_{i=1}^N \log p({\bf x}_i | {\bf z}) \approx \frac{N}{M}
\sum_{i\in{\mathcal{I}_M}} \log p({\bf x}_i | {\bf z})
\]

ここで $\mathcal{I}_M$ は $M<N$ であるサイズ $M$ のインデックスのミニバッチです (議論についてはリファレンス [1,2] 参照)。Great, problem solved! しかしこれを Pyro でどのように行なうのでしょう ?

 

Pyro で条件付き独立性をマークする

ユーザがこの類のことを Pyro で行なうことを望むのであれば、彼または彼女は最初にモデルとガイドが Pyro が関連する条件付き独立性を活用できるような方法で書かれることを確実にする必要があります。これがどのように成されるのか見てみましょう。Pyro は条件付き独立性をマークするために 2 つの言語プリミティブを提供します : irange と iarange です。2 つのより単純なものから始めましょう。

 

irange

前のチュートリアルで使用した例に戻りましょう。便利のために、モデルの主ロジックをここで複製しておきます :

def model(data):
    # sample f from the beta prior
    f = pyro.sample("latent_fairness", dist.Beta(alpha0, beta0))
    # loop over the observed data using pyro.sample with the obs keyword argument
    for i in range(len(data)):
        # observe datapoint i using the bernoulli likelihood
        pyro.sample("obs_{}".format(i), dist.Bernoulli(f), obs=data[i])

このモデルについて潜在的確率変数 latent_fairness が与えられたとき観測は条件付き独立です。これを Pyro で明示的にマークするために基本的には Python 組み込み range を Pyro construct irange で置き換える必要があるだけです :

def model(data):
    # sample f from the beta prior
    f = pyro.sample("latent_fairness", dist.Beta(alpha0, beta0))
    # loop over the observed data [WE ONLY CHANGE THE NEXT LINE]
    for i in pyro.irange("data_loop", len(data)):
        # observe datapoint i using the bernoulli likelihood
        pyro.sample("obs_{}".format(i), dist.Bernoulli(f), obs=data[i])

pyro.irange は一つの主な違いとともに range に非常に類似しています: irange の各発動はユーザに一意な名前を提供することを要求します。2 番目の引数はちょうど range のように整数です。

So far so good. Pyro は今では潜在的確率変数が与えられたとき観測の条件付き独立性をを活用できます。しかしこれは実際にはどのように動作するのでしょう？基本的には pyro.irange はコンテキスト・マネージャを利用して実装されています。for ループの本体の総ての実行において新しい (条件付き) 独立性コンテキストに入り、それから for ループ本体の最後に抜け出します。これについてはっきりと明示しましょう :

  • 各観測された pyro.sample ステートメントは for ループ本体の異なる実行内で発生しますので、Pyro は各観測を独立であるとマークします。
  • latent_fairness が与えられたときこの独立性は正しく条件付き独立性です、何故ならば latent_fairness は data_loop のコンテキストの外側でサンプリングされるからです。

先に進む前に、irange を使用するときに避けるべき幾つかの落とし穴に言及しましょう。上のコード・スニペットの次の変形を考えます :

# WARNING do not do this!
my_reified_list = list(pyro.irange("data_loop", len(data)))
for i in my_reified_list:
    pyro.sample("obs_{}".format(i), dist.Bernoulli(f), obs=data[i])

これは望まれる挙動を達成しません、何故ならば単一の pyro.sample ステートメントが呼び出される前に list() は data_loop コンテキストに完全に入って抜けるからです。同様に、ユーザはコンテキスト・マネージャの境界に渡って可変な計算を漏らさないように気をつける必要があります、何故ならばこれは微妙なバグに繋がるからです。例えば、pyro.irange は時間モデル (= temporal model) のためには適切ではありません、そこではループの各 iteration は前の iteration に依存します; この場合代わりに range が使用されるべきです。

 

iarange

概念的には iarange はそれがベクトル化された演算であることを除いて irange と同じです (torch.arange が range に対するように)。そのようなものとしてそれは irange に伴って出現する明示的な for ループに比較して潜在的に膨大なスピードアップを可能にします。私達の実行例に対してこれがどのように見えるか見てみましょう。最初に tensor の形式にあるデータが必要です :

data = torch.zeros(10)
data[0:6] = torch.ones(6)  # 6 heads and 4 tails

それから次を持ちます :

with iarange('observe_data'):
    pyro.sample('obs', dist.Bernoulli(f), obs=data)

これを類似の irange コンストラクションとポイントごとに比較してみましょう: – ちょうど irange のように、iarange はユーザに一意な名前を指定することを要求します。- このコードスニペットは単一の (観測される) 確率変数 (つまり obs) だけを導入することに注意してください、何故ならば tensor 全体が一度に考慮されるからです。- このケースでは iterator の必要性はありませんので、irange コンテキストに伴う tensor の長さを指定する必要はありません。

irange で言及した落とし穴は iarange にもまた当てはまりますので注意してください。

 

サブサンプリング

Pyro で条件付き独立性をどのようにマークするかを今では知っています。これはそれ自体有用ですが (SVI Part III の 依存性追跡セクション 参照)、巨大なデータセット上で SVI を行えるようにサブサンプリングもまた行ないたいです。モデルとガイドの構造に依拠して、Pyro はサブサンプリングを行なう幾つかの方法をサポートします。これらを一つずつ通り抜けましょう。

 

irange と iarange で自動サブサンプリング

最初に最も単純なケースを見ましょう、そこでは irange と iarange への 1 つか 2 つの追加引数でサブサンプリングをただで (= for free) 得ます :

for i in pyro.irange("data_loop", len(data), subsample_size=5):
    pyro.sample("obs_{}".format(i), dist.Bernoulli(f), obs=data[i])

これだけのことです: 単に引数 subsample_size を使用します。model() を実行したときはいつでも今ではデータの 5 つのランダムに選択されたデータポイントのために対数尤度を評価するだけです; 更に、対数尤度は $\tfrac{10}{5} = 2$ の適切な因数で自動的にスケールされます。iarange についてはどうでしょう？呪文は全く類似しています :

with iarange('observe_data', size=10, subsample_size=5) as ind:
    pyro.sample('obs', dist.Bernoulli(f),
                obs=data.index_select(0, ind))

重要なこととして、iarange は今ではインデックスの tensor ind を返します、これはこの場合は長さ 5 になります。引数 subsample_size に加えて irange が tensor データのフルサイズを知るように引数 size も渡すことにも注意してください、その結果それは正しいスケーリング・ファクタを計算できます。ちょうど irange のためのように、ユーザは iarange により提供されるインデックスを使用して正しいデータポイントを選択する責任を負います。

最後に、データが GPU 上であるうならばユーザは引数 use_cuda=True を irange または iarange に渡さなければならないことに注意してください。

 

irange と iarange によるカスタム・サブサンプリング・ストラテジー

上の model() が実行されるたびに irange と iarange は新しいサブサンプル・インデックスをサンプリングします。サブサンプリングはステートレスですから、これは幾つかの問題に繋がる可能性があります: 基本的に十分に巨大なデータセットについて巨大な数の iteration の後でさえもデータポイントの一部が決して選択されないであろう無視できない確率があります。これを回避するためにユーザは irange と iarange の subsample 引数を利用してサブサンプリングを制御できます。詳細は ドキュメント を見てください。

 

ローカル・ランダム変数だけがあるときのサブサンプリング

次で与えられる同時確率密度を持つモデルを持つことを念頭におきます :

$$p({\bf x}, {\bf z}) = \prod_{i=1}^N p({\bf x}_i | {\bf z}_i) p({\bf z}_i)$$

この依存構造を持つモデルについてサブサンプリングで導入されるスケール・ファクターは ELBO の総ての項を同じ総量でスケールします。これは例えば、vanilla VAE の場合です。これは VAE に対してサブサンプリングを完全に制御してミニバッチを直接モデルとガイドに渡すことを何故ユーザに許容するかを説明しています; iarange は依然として使用されますが、subsample_size と subsample はそうではありません。これがどのように見えるかを詳細に見るためには、VAE チュートリアル を見てください。

 

グローバルとローカル・ランダム変数の両者があるときのサブサンプリング

上のコインフリップの例では irange と iarange はモデルに現れましたがガイドには現れません、何故ならばサブサンプリングされる唯一のものは観測だったからです。より複雑な例を見てみましょう、そこではサブサンプリングはモデルとガイドの両者に現れます。簡潔にするために、議論を幾分抽象的なものにして完全なモデルとガイドを書くことを回避しましょう。

次の同時分布で指定されるモデルを考えます :

$$p({\bf x}, {\bf z}, \beta) = p(\beta)
\prod_{i=1}^N p({\bf x}_i | {\bf z}_i) p({\bf z}_i | \beta)$$

N 観測 $\{ {\bf x}_i \}$ と $N$ ローカル潜在ランダム変数 $\{ {\bf z}_i \}$ があります。グローバル潜在ランダム変数 $\beta$ もあります。ガイドは次のように分解できます :

$$q({\bf z}, \beta) = q(\beta) \prod_{i=1}^N q({\bf z}_i | \beta, \lambda_i)$$

ここで $N$ ローカル変分パラメータ $\{\lambda_i \}$ を導入することは明白でしたが、一方で他の変分パラメータは暗黙的です。モデルとガイドの両者は条件付独立性を持ちます。特に、モデル側では、$\{ {\bf z}_i \}$ が与えられたとき観測 $\{ {\bf x}_i \}$ は独立です。更に、$\beta$ が与えられたとき潜在ランダム変数 $\{\bf {z}_i \}$ は独立です。ガイド側では、変分パラメータ $\{\lambda_i \}$ と $\beta$ 潜在ランダム変数 $\{\bf {z}_i \}$ は独立です。Pyro でこれらの条件付独立性をマークしてサブサンプリングを行なうためにはモデルとガイドの両者で irange か iarange のいずれかを利用する必要があります。irange を使用して基本的なロジックの概略を述べましょう (コードのより完全なピースは pyro.param ステートメント, etc. を含むでしょう)。最初に、モデル :

def model(data):
    beta = pyro.sample("beta", ...) # sample the global RV
    for i in pyro.irange("locals", len(data)):
        z_i = pyro.sample("z_{}".format(i), ...)
        # compute the parameter used to define the observation
        # likelihood using the local random variable
        theta_i = compute_something(z_i)
        pyro.sample("obs_{}".format(i), dist.MyDist(theta_i), obs=data[i])

コインフリップを実行する例と対称的に、ここでは irange コンテキストの内側と外側の両者で pyro.sample ステートメントを持つことに注意してください。次にガイドです :

def guide(data):
    beta = pyro.sample("beta", ...) # sample the global RV
    for i in pyro.irange("locals", len(data), subsample_size=5):
        # sample the local RVs
        pyro.sample("z_{}".format(i), ..., lambda_i)

インデックスはガイドでは一度だけサブサンプリングされることに極めて注意してください ; Pyro バックエンドは、モデルの実行中インデックスの同じセットが使用されることを確かなものにします。この理由で、subsample_size はガイドでのみ指定される必要があります。

 

Amortization

グローバルとローカル潜在ランダム変数とローカル変分パラメータを持つモデルを再度考えましょう :

$$p({\bf x}, {\bf z}, \beta) = p(\beta)
\prod_{i=1}^N p({\bf x}_i | {\bf z}_i) p({\bf z}_i | \beta) \qquad \qquad
q({\bf z}, \beta) = q(\beta) \prod_{i=1}^N q({\bf z}_i | \beta, \lambda_i)$$

スモールからミディアムサイズの $N$ に対してこのようなローカル変分パラメータを使用することは良いアプローチです。けれども、$N$ がラージであれば、それに渡り最適化しようとしている空間が $N$ で増大する事実は現実問題となり得ます。このデータセットのサイズによる嫌な増大を回避する一つの方法は amortization です。

これは次のように動作します。ローカル変分パラメータを導入する代わりに、single パラメトリック関数 $f(\cdot)$ を学習して次の形式を持つ変分分布で作業していきます :

$$q(\beta) \prod_{n=1}^N q({\bf z}_i | f({\bf x}_i))$$

関数 $f(\cdot)$ — これは基本的には与えられた観測を (そのデータポイントに適合された) 変分パラメータのセットにマップします — は事後を正確に補足するために十分にリッチである必要がありますが、今では変分パラメータの不愉快な数を導入しなければならないことなく巨大なデータセットを処理できます。このアプローチは他の恩恵もまたあります: 例えば、学習の間 $f(\cdot)$ は異なるデータポイント内の統計的パワーを共有することを効果的に可能にします。これは正確に VAE で使用されるアプローチであることに注意してください。

 

Tensor shape と pyro.iarange

このチュートリアルの pyro.iarange の使用方法は比較的単純なケースに制限されていました。例えば、どの iarange も他の iarange の内側にネストされませんでした。iarange をフル活用するためには、ユーザは Pyro tensor shape セマンティクスを使用するために注意しなければなりません。議論のためには tensor shape チュートリアル を見てください。

 

References

  1. Stochastic Variational Inference, Matthew D. Hoffman, David M. Blei, Chong Wang, John Paisley
  2. Auto-Encoding Variational Bayes, Diederik P Kingma, Max Welling
 

以上