はじめに
確率モデルから始めて、最近ようやくVAE系統までの流れを把握できたのでまとめておきます。
確率モデリング
確率モデリングは観測データ $X$ と 観測していない潜在変数 $Z$ が主な登場人物です。潜在変数 $Z$ というのはモデルの内部に潜んでいるのだが、データとして取り出すことはできない実体のある何かである可能性もありますし、実体として結びついていなくても、何かモデルの構造に入り込んでいるであろうものを抽象的に表しているだけの可能性もあります。これは設計者が決めることです。
また、確率モデルのパラメータ等も全て潜在変数というものだとしておきます。
例 ガウス分布から生起するデータ
観測データが
$$ x_i \sim {\mathcal N}(x \mid \mu, \sigma) $$
と生起すると考えましょう。$X = \{x_1, \cdots, x_N\}$ で $Z = \{\mu, \sigma \}$ としておけば観測データと潜在変数という形で分けることができます。要するに観測データは平均的には潜在的な値$\mu$になるはずなのだけども、はたまた別の潜在変数 $\sigma$ がこれを揺らがせているというふうに見るわけです。
これが、例えば $x$ は温度計の指す値で、$\mu$ は実際の空気の温度、$\sigma$ は温度計を取り巻く環境によるバラツキである、というような感じで物理的実体と照らし合わせてもいいですし、そうでなくても構いません（逆に物理的実体がわからない潜在変数を仮定したモデルを立てて、解析した結果から潜在変数について何かを言うのも一つの仕事です）。
回帰モデル
回帰モデルの場合は、観測データが $X = \{x_1, \cdots, x_N\}$ と $Y = \{y_1, \cdots, y_N\} $ として手元にあるとします。回帰モデルでは $y \simeq f(x)$ のような関係を想定しているので、例えば下記の用にモデル化したりします。
$$ y_i \sim {\mathcal N} (y_i | ax_i + b, \sigma) $$
すなわち、$y_i$ は平均的には $ax_i + b$ という $x_i$ に依存した形で表されるのですが、$\sigma$ くらいバラついてしまうということをモデル化しています。こちらは潜在変数を $Z = \{a, b, \sigma\}$ としていることになります。これも例えば、$x_i$ が起電力で$y_i$ が電流、 $a$ が抵抗の逆数で $bi$ が暗電流として、$\sigma$ が電流計測時の誤差という物理的実体と紐付いたモデルでもいいですし、$a, b$ が何かわからない潜在変数であると思っても構いません。
一般に上記ではよく知っている正規分布を使って観測データと潜在変数に関して説明しましたが、今後は潜在変数 $Z$ で観測変数を $X$ あるいは $Y$ などとして $p(X \mid Z)$ や $p(Y \mid X, Z)$ などを見ていくことにします。
あとはこれらの尤度関数を最大化するように潜在変数 $Z$ の値を点推定してやれば良いというのが確率モデリングの推定です。
ベイズモデリング
ベイズモデリングでは確率モデルの潜在変数 $Z$ に対して、事前分布 $p(Z)$ なるものを置きます。よく知られているベイズの定理
$$ p(Z \mid X) = \frac{p(X\mid Z)p(Z)}{p(X)} $$
によって、手元にある観測データ $X$ に依存した $Z$ の確率分布 $p(Z\mid X)$、すなわち事後分布を求めようというのがベイズ推論になります。 通常の確率モデリングでは $Z$ は尤度を最大化するような1つの値として求めていたところ、今回は分布として求めていることに大きな違いがあります。
MAP推定
よく知られているMAP推定は、ベイズ推論で事後分布を求めるのが困難であるときに（あるいはデータ $X$ のサイズが十分に大きく $p(Z \mid X) $ が局在しているであると思えるときに）、事後分布の最大値を示す $Z$ を点推定する手法です。
このとき事後分布の右辺、分母は不要で、
$$ p(Z \mid X) \propto p(X\mid Z)p(Z) $$
であるので、実質、尤度と事前分布の掛け算を最大化すれば良いということになります。もしも $p(Z)$ を無情報事前分布として構えるのであれば、MAP推定は最尤推定と同じ結果になります。
変分推論
真面目にベイズ推論で事後分布を求める場合、
$$ p(Z \mid X) = \frac{p(X\mid Z)p(Z)}{p(X)} $$
の右辺が都合よく計算できる必要があります。しかし、なかなかそれは上手く行きません。そこで上記の事後分布の形状を正確に計算することを諦め、適当な形状の確率分布 $q(Z)$ を仮定し、$p(Z \mid X)$ と一番似た形になるものを事後分布としてしまおうという方法を取ります。
ところで $p(Z \mid X) $ の形がわからないという話なのに、$q(Z)$ の形をどうやって決めれば良いのでしょうか。ここが変分推論の肝で、形はわからないが、どれくらい掛け離れていると言えるかを計算することはできるのです。それがKLダイバージェンス $KL[q:p]$ になります。（エビデンス→ELBO→KLダイバージェンスの流れが本流ですが割愛）
あとは $KL[q:p]$ が小さくなるように $q$ の形状を少しずつ変えていく方法があれば良いということになります。片っ端から $q$ を適当な形状にして $KL[q:p]$ の値を記録していくのも手かもしれませんが、ココは頭を少しだけ使って $q(Z; \eta)$ とパラメトライズしてしまいましょう。そして $q$ の形状を変える操作を $\eta$ の値を変える操作に肩代わりさせます。この値の変え方は $KL[q(Z; \eta):p(Z\mid X)]$ を $\eta$ で微分する勾配法に基づいて行うことにします。
これが近年の変分推論の主流となっています。
ここで再度強調するのは、元々観測変数 $X$ が手元にあり、諸々の観測できていない変数を潜在変数 $Z$ と置いており、観測できている $X$から$Z$ の分布を推論したいというのが事後分布の推論です。そして最後に登場した $\eta$ は事後分布の推論をお手伝いするパラメータであり、分布を求めようとしている潜在変数と異なり決定論的値です。
流れとしては
- $p(X \mid Z)$ と $p(Z)$ によってモデリング。
- $p(Z \mid X)$ を求めたい。
- $q(Z ; \eta)$ で肩代わりする。
- $KL[q(Z ; \eta) : p(Z \mid X)]$ を最小化する $\eta_{opt}$ を求める。
- $q(Z ; \eta_{opt})$ を求めたかった事後分布の近似解とする
というのが変分推論になります。
VAE
私はこの流れに乗って、VAEについて説明するのが一番分かりやすいと感じています。観測データ $X$ (大抵は画像であるケースが多いです) と、 潜在変数 $Z$ (画像のスタイルであるとか、文字そのものであったりとか) を想定しつつ
- $p(X \mid Z)$ と $p(Z)$ によってモデリング。
- $p(Z \mid X)$ を求めたい。
- 変分パラメータ（決定論的パラメータ）としてニューラルネットを導入 $$q(Z ; {\rm NN}(X, \eta)$$ で事後分布を肩代わりする。
- $KL[q(Z ; {\rm NN}(X, \eta)) : p(Z \mid X)]$ を最小化する $\eta_{opt}$ を求める。
- $q(Z ; {\rm NN}(X, \eta_{opt}))$ を求めたかった事後分布の近似解とする
というのがVAE系統の基本的な形です。決定論的な値である変分パラメータ $\eta$ を単体で使うのではなく、これをニューラルネットワークの形式で与えるという点が大きく異なります。ニューラルネットワークのパラメータが、元々の確率モデルの変分パラメータの役割をしているのですから、別にニューラルネットワーク自体がベイズ化されているわけではない点に注意してください（ニューラルネットワークのベイズ的発展ではなく、変分ベイズ推論へのニューラルネットワークの応用という具合でしょう）。
ちなみに、もう元々の生成モデル $p(X\mid Z)$ の方にも決定論的パラメータ噛ませちゃっても良いんじゃないの？という考えを導入すれば
- $p(X \mid {\rm}NN(Z, \phi))$ と $p(Z)$ によってモデリング。
- $p(Z \mid X)$ を求めたい。
- $q(Z ; {\rm NN}(X, \eta))$ で肩代わりする。
- 確率モデルの方にもニューラルネットを入れたことを陽に書き下すため $$ KL[q(Z ; {\rm NN}(X, \eta)) : \frac{p(X \mid {\rm}NN(Z, \phi))p(Z)}{p(X)}] $$ を最小化する $\eta_{opt}$ と $\phi_{opt}$ を求める、と表現し直す。単に$p(Z\mid X)$ をベイズの定理右辺で書き直しただけである。
- $q(Z ; {\rm NN}(X, \eta_{opt}))$ を求めたかった事後分布の近似解とする
と変更でき、諸々省略してしまえばVAEは
- 生成モデルであるデコーダー$p(X \mid {\rm}NN(Z, \phi))$とエンコーダー $q(Z ; {\rm NN}(X, \eta))$ を準備
- 生成モデルを使うときは、正規乱数生成器 $p(Z)$ で $Z$ を発生させる仕組みにする
- $KL[q(Z ; {\rm NN}(X, \eta)) : \frac{p(X \mid {\rm}NN(Z, \phi))p(Z)}{p(X)}]$ を最小化する $\eta_{opt}, \phi_{opt}$ を求める。
と書けてしまいます。いきなりこの状態で話をされても、思いつきで正規乱数に閉じ込めただけに感じてしまうかもしれませんが、変分推論にパラメータを無理矢理NNでねじ込んだというのが深層です。ちなみに、普通に変分推論しようという場合には $Z = \{z_1, z_2\}$ などと潜在変数が2つあるときに、これらを
$$ q(Z) = q(z_1)q(z_2) $$
と個々の潜在変数は独立（データの条件付き）であるという仮定をいれて解きやすくするのが平均場近似ですが、VAEの場合は変分モデル
$$ q(Z ; {\rm NN}(X, \eta)) $$
において $NN$の出力層の数を潜在変数の個数と合わせてしまえば、独立性を（形式上）は仮定せずに推論できるようになります。独立性を仮定しないため、潜在変数同士の相関を捨てること無く推論できるというのが一つの強みです。ただし、通常VAEでは変分モデル $q$ は具体的に
$$ q(Z ; {\rm NN}(X, \eta)) = {\mathcal N} (Z \mid \mu(X, \eta), {\rm diag} \sigma(X, \eta)^\2) $$
と対角共分散行列を持つ正規分布と決め込んで使うので、実を言うと、各成分の相関が無いという仮定を結局入れてしまっています。このように対角成分だけを考えることで、潜在変数の次元が$Z \in RD$ としたときに、$NN(X, \eta) \in R^{2D}$ を使って、
$$ \begin{align} \mu &= NN(X, \eta) [1 : D] \\ \sigma &= NN(X, \eta) [D+1 : 2D] \end{align} $$
と$Z$ の正規分布が従う平均と分散を決めさせられます（単に各成分の平均と分散を出力させればいいということ、ここに $[1 : D]$ というのは$1$ から $D$ 番目の成分という意味です）。この正規分布という仮定と、無相関であることの必然性は特にありません。
Normalizing Flow
潜在変数が正規分布である必然性はないことに着目し、正規分布に従う $Z$ に対して適当な全単写 $Z' = f(Z)$ を与えて、適当な分布 $q(Z')$ に従う $Z'$ を潜在変数の表現として使う方法もあります。ここで $f$ が全単写であればヤコビアンを通して、元の $q(Z')$ の分布と元の正規分布を行き来することができますので、尤度計算も正しく行えます。
一先ずメモのためにザッと書きました。
