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

今回は、「あの論文を検証してみた！」のシリーズ第４回、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）のいいとこどりをしたニューラル潜在変数モデルの１種です。ガウス過程（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（）の第１項のについて、の分布に、の観測（条件）の下でが影響するのか？
- $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（）の第２項のKLの第２引数のについて、の分布に、の観測（条件）なく、が影響するのか？
- $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$ に対して周辺化することで導出できます。

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

ELBOの第１項

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

グラフィカルモデルの同時分布（分解式）から、 $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}

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

\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の第２項

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

第２項・ $KL$ の第２引数 $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}

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

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

ELBO の簡略化

以上２つの結果から、第１項の $p(Y_T|z, X_T, C)$ では、条件から $C$ を省略でき、第２項の $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\}$ は、確率過程 $\{X_t\}$ から、なんらかの変換（未知の変換）をした列とみなすことが前提
- つまり、 $F(X_t)=Y_t$ が前提
NPは、数学的な理論づけが、どこか美しい
NPの実装・検証では、GPをベースにモデル化している
- 理論上は、GP以外の確率過程にも適用可能
NPは、よくある教師あり学習の枠組みと、複数の小データセットからの推定（Few-Shot Learning）の枠組みも持つ
- Few-Shot の枠組みとしては、ベイズ更新を使う様子（これについては、またいつか・・・）

ブレインズコンサルティング・テックブログ

あの論文を検証してみた！ - シリーズ第４回 - Neural Processes 論文（解説編）