あの論文を検証してみた! - シリーズ第4回 - Neural Processes 論文(解説編)

こんにちは!ブレインズコンサルティングの大下です。

今回は、「あの論文を検証してみた!」のシリーズ第4回、Neural Processes の論文について解説します。 Neural Processes の論文を選んだモチベーションは、もともと確率とニューラルネットワークの合わせ技のモデルに興味があったところ、 Deepmind 社が、ICMLでガウス過程の深層学習版として提案した当論文が目に入って来たというところです。

検証環境

まずは、動作確認に使った検証環境を明記しておきます。

  • Ubuntu 18.04.1 LTS (Bionic Beaver)
    • CPU: Intel(R) Core(TM) i7-6700K CPU @ 4.00GHz / 4 cores / 8 processors
    • Memory: 48GiB
    • HDD: 約193GiB (Available なサイズ)
  • Python 3.6.5
  • pycharm-community-2018.3.3
  • pytorch 1.0.0
  • Colaboratory(GPU)

本検証におけるアウトプット

  • 当記事
  • メモ用紙(極秘(笑))

省略語

Abbreviated Full words
NN Neural Network
GP Gaussian Process
NP Neural Process
VAE Variational Auto Encoder
ELBO Evidence Lower BOund (変分下界)

概要

Neural Processes(NPs)は、ニューラルネットワーク(NN)とガウス過程(GP)のいいとこどりをしたニューラル潜在変数モデルの1種です。 ガウス過程(GP)のように、NPも関数(カーネル)上の分布を定義づけ、新しい観測データに対して、高速に適応する能力を持っているようです。(メタラーニングとしての側面を持つようです。)

NPは、よくある教師あり学習のように、学習・予測をすることもできますが、 複数の小データセットから、予測分布を推定することもできます。(以下は、Kaspar Martensより抜粋)

f:id:bci-oshita:20190225174339p:plain
Two Scenarios

グラフィカルモデルと、計算グラフは、以下のように定義づけられています。(論文より抜粋)

f:id:bci-oshita:20190225172640p:plain
Neural Processes Model

ここで、 C は、Context(文脈) を意味し、 T は、Target(目的)を意味しています。  C は、学習時にも予測時にも X_C, Y_Cのいずれも使用可能(既に取得済、観測済)で、  T については、 X_T は、学習時にも予測時にも使用可能(既に取得済、観測済)ですが、  Y_Tは、学習時にしか使えない(取得できていない)ものとして考えます。

VAEの拡張

NPのELBO は、VAEをベースに拡張して構築されています。 ここでは、VAEからどのように、NPのELBOを導いているかを解明していきます。 (論文では、ロジックが飛んでいるため、その補足をします。)

VAEのELBO

基本的なVAEのELBO( L)は、以下のように表現できます。

\begin{align} \log p(Y) \geqq L = E_{q(z|Y)}[\log p(Y|z)] - KL[q(z|Y) || p(z)] \end{align}

ここで、 q は、任意の確率分布(意味としては、p(z)の近似分布)でよいです。

この等式は、グラフィカルモデル(確率変数の接続構造)に依存しない点に注意します。

導出方法は、外部サイトの記事を参照ください。

条件付きVAEのELBO

条件付き分布 p(Y|X) を使ったVAEでは、以下のように、 \log p(Y|X) の変分下界(ELBO)を表現できます。

\begin{align} \log p(Y|X) \geqq L = E_{q(z|A)}[\log p(Y|z, X)] - KL[q(z|A) || p(z|X)] \end{align}

ここで、 q は、任意の確率分布(意味としては、p(z|X)の近似分布)でよく、  A は、任意の確率変数(ベクトル or リスト)でよい点に注意します。

また、基本的なVAEと同様に、この等式は、グラフィカルモデル(確率変数の接続構造)に依存しない点に注意します。

NPのELBOの導出

条件付き分布を使ったVAEのELBO の、 Y, X, A を以下のように定義します。 (はてブの数式の制約を回避するために、遠回りな定義になっています。)

\begin{align} Y &:= Y_T \\ \end{align}

\begin{align} X &:= X_G, Y_C = X_T, C \\ \end{align}

\begin{align} A &:= X_G, Y_G = C, T \\ \end{align}

ただし、

\begin{align} X_C &:= x_{1:m} \\ \end{align}

\begin{align} Y_C &:= y_{1:m} \\ \end{align}

\begin{align} X_T &:= x_{m+1:n} \\ \end{align}

\begin{align} Y_T &:= y_{m+1:n} \\ \end{align}

\begin{align} X_G &:= x_{1:n} = X_C, X_T \\ \end{align}

\begin{align} Y_G &:= y_{1:n} = Y_C, Y_T \\ \end{align}

\begin{align} C &:= X_C, Y_C \\ T &:= X_T, Y_T \\ \end{align}

ELBOの式に代入すると、

\begin{align} \log p(Y|X) \geqq L = E_{q(z|A)}[\log p(Y|z, X)] - KL[q(z|A) || p(z|X)] \end{align}

\begin{align} \log p(Y_T|X_T, C) \geqq L = E_{q(z|C, T)}[\log p(Y_T|z, X_T, C)] - KL[q(z|C, T) || p(z|X_T, C)] \end{align}

グラフィカルモデルを踏まえたELBOの検討

以降は、NPのグラフィカルモデルをシンプルにした下記のモデルを使って考えます。(本質的には、論文のグラフィカルモデルと同等である点に注意します。)

f:id:bci-oshita:20190226150047p:plain
graphical model

このグラフィカルモデルによる同時分布 p(C, z, Y_t, X_t)の分解式は、以下の通りです。

\begin{align} p(C, z, Y_t, X_t) = p(C)p(z|C)p(Y_t | z, X_t)p(X_t) \\ \end{align}

グラフィカルモデルと、導出したELBOの式(右辺)を比較参照すると、以下の点が疑問になります。

  • ELBO( L)の第1項の p(Y_T|z, X_T, C) について、 Y_T の分布に、 z の観測(条件)の下で C が影響するのか?
    •  C, z, Y_t は、head-to-tail モデルであるので、 z の観測の下では、 C, Y_tは独立。つまり、 p(C, Y_t | z) = p(C|z)p(Y_t|z)
    •  p(C, Y_t | z) = p(C|z)p(Y_t|z) は、同時分布 p(C, z, Y_t)の分解式とベイズの定理から導出できます。
  • ELBO( L)の第2項のKLの第2引数の p(z|X_T, C) について、 z の分布に、 Y_Tの観測(条件)なく、 X_T が影響するのか?
    •  z, Y_t, X_t は、head-to-head モデルであるので、 Y_t が未観測の(条件にない)場合、 z, X_t は、独立である点に注意します。つまり、 p(z, X_t) = p(z)p(X_t)
    •  p(z, X_t) = p(z)p(X_t) は、同時分布 p(z, Y_t, X_t)の分解式を Y_t に対して周辺化することで導出できます。

そこで、この2点の疑問について解消していきます。

ELBOの第1項

上記、グラフィカルモデルを用いて、 L の第1項を簡略化できるか検討していきます。

グラフィカルモデルの同時分布(分解式)から、 z, X_t の観測の下で、 Y_t, C が独立であることが導出できます。具体的な計算は、以下の通りです。

\begin{align} p(Y_t, C | z, X_t) &= \frac{p(Y_t, C, z, X_t)}{p(z, X_t)} \\ &= \frac{p(C)p(z|C)p(Y_t | z, X_t)p(X_t)}{p(z, X_t)} \\ &= \frac{p(z)p(C|z)p(Y_t | z, X_t)p(X_t)}{p(z, X_t)} \\ &= p(C|z)p(Y_t | z, X_t) \\ &= p(Y_t | z, X_t)p(C|z) \\ \end{align}

両辺を、 Y_t に対して周辺化すると( Y_t の変数で、積分すると) \begin{align} \int p(Y_t, C | z, X_t)dY_t &= p(C|z) \int p(Y_t | z, X_t)dY_t \\ p(C | z, X_t) &= p(C|z) \\ \end{align}

この式を、1つ前の式に代入すると

\begin{align} p(Y_t, C | z, X_t) &= p(Y_t | z, X_t)p(C|z) \\ p(Y_t, C | z, X_t) &= p(Y_t | z, X_t)p(C|z, X_t) \\ \end{align}

以上から、 z, X_t の観測の下で、 Y_t, C が独立になります。 この導出方法は、 X_t, Y_t(1変数)を X_T, Y_T(ベクトル変数、リスト、同時確率)に置き換えても成立するので、 X_T, Y_T についても同様に、 z, X_T の観測の下で、 Y_T, C は、独立になります。

この条件付き独立性から、 p(Y_T|z, X_T, C) = p(Y_T|z, X_T) が導けます。(つまり、 C を省略できます。)

実際、計算すると、

\begin{align} p(Y_t | z, X_t)p(C | z, X_t) = p(Y_t, C | z, X_t) = p(Y_t | C, z, X_t)p(C | z, X_t) \\ \end{align}

両辺を、共通する p(C | z, X_t) で割って、以下の等式を得ます。(厳密には、 p(C | z, X_t) > 0 を保証するか、非ゼロになるように定義する必要があることに注意。)

\begin{align} p(Y_t | z, X_t) &= p(Y_t | z, C, X_t)) \\ \end{align}

また、 X_t := X_T Y_t := Y_T と確率変数をベクトル(リスト、同時確率)とみなしても同じ議論が成立するので、

\begin{align} p(Y_T|z, X_T, C) = p(Y_T | z, X_T) \\ \end{align}

つまり、 z の観測の下での  Y_T の確率分布においては、( Y_Tとは、 zにおける条件付き独立である) C を条件から省略することができることを示しています。

ELBOの第2項

グラフィカルモデルを用いて、 L の第2項を簡略化できるか検討していきます。

第2項・ KL の第2引数  p(z | X_T, C) = p(z | C, X_T) について、周辺化することで、 p(z|C, X_T) を計算します。

\begin{align} p(z|C, X_T) &= \int p(z, y_T | C, X_T) dy_T \\ &= \int p(z|C)p(y_T|z, X_T)) dy_T \\ &= p(z|C) \int p(y_T|z, X_T)) dy_T \\ &=p(z|C) \\ \end{align}

ここで、1行目から2行目の式変形は、グラフィカルモデルによる同時分布の分解式から導出できる点に注意します。

以上より、 Y_T 未観測の下での  z の確率分布においては、( Y_T未観測の下で zとは独立である) X_T を条件から省略することができることを示しています。

ELBO の簡略化

以上2つの結果から、第1項の p(Y_T|z, X_T, C)では、条件から C を省略でき、第2項の p(z|X_T, C)では、条件から X_T を省略できるとわかりました。

これらを、 L の式に代入し、以下の式を得ます。

\begin{align} \log p(Y_T|X_T, C) \geqq L = E_{q(z|C, T)}[\log p(Y_T|z, X_T, C)] - KL[q(z|C, T) || p(z|X_T, C)] \end{align}

\begin{align} \log p(Y_T|X_T, C) \geqq L = E_{q(z|C, T)}[\log p(Y_T|z, X_T)] - KL[q(z|C, T) || p(z|C)] \end{align}

この得られた不等式(or ELBO)が、今回の論文の 式(8) に一致することに注意します。

あとは、論文に記載されている通り、一般に、条件付き事前分布  p(z|C) の計算は困難であるため、近似分布  q(z|C) に置き換えたELBOを用いて、実装します。 (細かく言うと、 q_n(z|C) p(z|C) に弱収束することを前提にするイメージですが、ここは、未確認&未証明です。)

先の不等式を、近似分布 q(z|C)に置き換えて、以下の式を得ます。

\begin{align} \log p(Y_T|X_T, C) \geqq L = E_{q(z|C, T)}[\log p(Y_T|z, X_T)] - KL[q(z|C, T) || q(z|C)] \end{align}

この式が、論文の式(9) です。

学習に必要な分布

変分下界(ELBO)の式を構築できたので、具体的に学習に必要な分布の定義を行います。

学習に必要な分布は、以下の通りです。

  •  p(Y_T|z, X_T) = \prod_{t=m+1}^n p(Y_t | z, X_t) : 各  Y_t の真の条件付き分布(仮定)
  •  q(z|C, T) : NNによる近似分布
  •  q(z|C) : NNによる近似分布

分布の定義

以下のように、分布を定義します。

  •  p(z) := N(0, I) (初期分布)
  •  p(Y_t | z, X_t) := N(Y_t | m, s^{2})
  •  q(z | \emptyset) := q(z) := p(z)
  •  q(z | C) := N(z|m_C, {s_C}^2)
  •  q(z | C, T) := N(z|m_G, {s_G}^2)

ただし、各記号は以下のように定義します。

  •  G := C, T = \{x_i, y_i\}_{i=1}^{n}
  •  m := \hat{y}_{t} := g(X_t, z)
  •  s := Const or Learnable (Deepmind社のコードでは、NNで学習させています。)
  •  m_C :=u(a(h(X_C, Y_C))) の出力1
  •  \log s_{C} : u(a(h(X_C, Y_C))) の出力2
  •  m_G : u(a(h(X_G, Y_G))) の出力1
  •  \log s_{G} : u(a(h(X_G, Y_G))) の出力2
  •  g : NNで実装
  •  h : NNで実装
  •  u : NNで実装
  •  a := mean(平均する関数)

上記定義に合わせた、計算グラフは、以下のようになります。(論文の計算グラフをもう少し、詳細化しました。)

f:id:bci-oshita:20190304101818p:plain
computational diagram

補足

先の記述では、 p(Y_T|z, X_T) = \prod_{t=m+1}^n p(Y_t | z, X_t) が成り立つことを前提として定義していたので、ここで、導出しておきます。

以下のような、サブグラフを使って(に注目して)、導出します。

f:id:bci-oshita:20190227155718p:plain
sub graphical model of NP 2

このグラフィカルモデルの同時分布の分解式を計算すると、以下のようになります。

\begin{align} p(z, Y_T, X_T) &= \prod_{t \in T} p(Y_t|z, X_t)p(z)p(X_t) \\ p(Y_T | z, X_T)p(z, X_T) &= \prod p(Y_t|z, X_t)p(z)\prod p(X_t) \\ &= \prod p(Y_t|z, X_t)p(z, X_T) \\ \end{align}

上記式を、共通項 p(z, X_T)で、割って、無事、求める等式を得ます。

\begin{align} p(Y_T | z, X_T) &= \prod_{t \in T} p(Y_t|z, X_t) \\ \end{align}

どこが確率過程なの??

上記のようなグラフィカルモデル、計算グラフで、NNを使って確率モデルを実現できそうなことは、うすうす わかってきたのですが、 一体、どの部分を指して確率過程なのか?というところが疑問になります。

実は、当論文の「2. Model - 2.1. Neural processes as stochastic processes」に、記載されているのですが、個人的には結構わかりにくかったです。おそらく、わかる人には、すぐわかる書き方っぽい。。(わかりにくいけど、内容がちゃんとしているのは、やはりすごいです。)

ポイントを絞って、ひも解いてみたいと思います。(詳細は、論文を追ってみてください。)

まず、2つの確率過程によって生成される確率変数の集合 X=\{X_t\}, Y=\{Y_t\}と、その間の写像 F:X \rightarrow Y を考えます。 ただし、写像 F は、 F(X_t) = Y_t を満たすような写像に限るとします。 この時、 \{F(X_t)\}_t は、確率過程になります。( \{F(X_t)\}=Yで、 Yが確率過程(と同一視可能)より)

この F(X_t) = Y_t の確率過程をモデル化したのが、このNeural Processes になります。

では、この確率過程とNeural Processes のモデルがどうつながるかを、論文の2.1 節のポイントを絞って概説します。

確率過程と結合分布

ここで、 F(X_t) = Y_t の確率分布をうまく定義することが目的です。 このうまく定義する確率分布の前提となるモデルが、これまで議論してきたグラフィカルモデルに繋がっていることを示します。

まずは、 F(X_t) = Y_t に対して、 \{Y_t\}_{t=1}^T のように、 任意に選んだ有限個の確率変数に対する結合分布 \rho_Tを以下のように記述します。 ( \{Y_t\} 上では、無限次元の結合分布を考えますが、まずは任意に選んだ有限個の確率変数に対する有限次元上の結合分布を考えます。)

\begin{align} \rho_T(y_1, \cdots, y_T) &:= \rho_T(Y_1=y_T, \cdots, Y_T = y_T) \\ &= \rho_T( F(X_1) = y_1, \cdots, F(X_T) = y_T) \\ \end{align}

ここで、 \rho_T は、 F, X に依存した有限結合分布です。

Kolmogorov's Extension Theorem より、この \rho=\{\rho_T\}_{T} が、exchangeability と consistency を持つならば、 以下を満たす確率分布 p pは、確率過程 \{Y_t\} = \{F(X_t)\} 上の結合分布)が一意に存在することが言えます。 (exchangeability と consistency については、論文等をご参照ください。)

\begin{align} \rho_T(y_1, \cdots, y_T) = p(Y_1=y_1, \cdots, Y_T=y_T, \prod_{t=T+1}^{\infty}\mathbb{R}) \end{align}

つまり、 \{Y_t\}に対する結合分布 pに対して、  \{Y_t\} 上の興味がある任意の部分列 \{Y_1, \cdots, Y_T\} 以外を周辺化した結合分布が、  \{Y_1, \cdots, Y_T\} 上の結合分布 \rho_Tに一致することを意味していいます。

結合分布とグラフィカルモデル

ここで、具体的に、 \rho_Tをexchangeability と consistency を持つように定義することを考えます。

まず、consistency が成り立つように、 \rho_T, f を、以下のように定義します。( p は、 \{Y_t\} 上の任意の確率測度としておきます。)

  •  \rho_T(y_1, \cdots, y_T) = \rho_T(y) := p(y | x) = \int p(y, f | x)df
  •  y := (y_1, \cdots, y_T)
  •  x := (x_1, \cdots, x_T)
  •  f:= (y_{T+1}, \cdots)

ここで、 \{\rho_T\}_{T} が、 \{y_t\} に対して exchangeability を持つために、以下のようなグラフィカルモデルを考えます。 ( x_t y_t が、 tで対応づくようにします。)

f:id:bci-oshita:20190301193036p:plain
graphical model 2

このモデルにおける各確率変数間の同時分布の分解式は、

\begin{align} p(y, f, x) = p(f)\prod_{t}p(y_t|f, x_t)p(x_t) \end{align}

 xの条件付き分布を計算すると、( p(x) = \prod_{t}p(x_t) を使います)

\begin{align} p(y, f | x) = p(f)\prod_{t}p(y_t|f, x_t) \end{align}

先の式に代入することで、このモデルでは、 \rho_T を以下のように記述できます。

\begin{align} \rho_T(y_1, \cdots, y_T) &= \rho_T(y) \\ &= \int p(y, f | x)df \\ &= \int p(f)\prod_{t}p(y_t|f, x_t) df \\ \end{align}

このモデルでは、 \{y_t\} は、 f, \{x_t\} の条件の下で、独立です。 (つまり、 p(y|f, x) = \prod_{t}p(y_t|f, x)

結局、このように定義した、 \rho_T は、exchangeability と consistency を持ちます。

結果、Kolmogorov's Extension Theorem より、上記のように定義した結合分布の集合(Collection) \{\rho_T\}_{T} に対して、 唯一の p が存在するため、 \rho_T の右辺がこの p に一致します。 (これは、Finetti's theorem (の条件付きバージョン)に対応するようです。)

具体的なモデル(実装)は、GPをベースに考えているため、 p(y_t|f, x_t) := N(y_t|g(x_t, f), s^{2}) とします。

当節のグラフィカルモデルの f を、 z とすると、NPのグラフィカルモデルの右側が得られます。 ちなみに、残りのNPのグラフィカルモデルの左側は、 f(=z) の生成過程をモデル化したものになります。 (結果的には、近似分布  q(z|C), q(z|C, T) を計算するためのモデルになります。)

関数のサンプリングとは?

いくつかの解説サイトで、関数のサンプリングというような言い回しがありますが、それは、  z \sim N(0, I) または、 z \sim q(z|C) z \sim q(z|C, T) による ランダムな z に対して、 g(x_t, z) を算出することを意味しているようです。

感想

この論文は、数学的な理論が背景にあり、読んでいて何か芸術性(美しさ)を感じました。 そして、グラフィカルモデルは、楽しかったです。

まとめ

  • NP のELBO は、VAEを条件付きに拡張し、グラフィカルモデルの構造から分布を簡略化することで得られる
  • NPにおける確率過程とは、確率過程 \{Y_t\} のことを指している
    • ただし、確率過程 \{Y_t\}は、確率過程 \{X_t\} から、なんらかの変換(未知の変換)をした列とみなすことが前提
    • つまり、 F(X_t)=Y_tが前提
  • NPは、数学的な理論づけが、どこか美しい
  • NPの実装・検証では、GPをベースにモデル化している
    • 理論上は、GP以外の確率過程にも適用可能
  • NPは、よくある教師あり学習の枠組みと、複数の小データセットからの推定(Few-Shot Learning)の枠組みも持つ
    • Few-Shot の枠組みとしては、ベイズ更新を使う様子(これについては、またいつか・・・)

参考リンク