14  潜在変数と推論 ── 見えない原因を捉える

第7章で、われわれは壁にぶつかったまま先へ進んだ。ボルツマンマシンの学習則は美しかった。正相から負相を引く。データの相関にモデルの相関を近づける。だが負相の計算に、すべての状態にわたる和が要る。ユニットが100個あれば 2^{100} 通り。宇宙の年齢をかけても終わらない。分配関数 Z という壁である。

第8章では、この壁を迂回した。エネルギー関数も分配関数も捨て、入力から出力への写像を教師あり学習で作り、連鎖律で勾配を計算する—バックプロパゲーションである。うまくいった。だが、あの問題の立て方には入っていないものがあった。ボルツマンマシンが持っていた「世界の確率的なモデルを内に持つ」という性格である。入力データの分布そのものを表さないので、新しい標本を生み出すこともない。断っておくと、これはバックプロパゲーションの限界ではない(第8章第1節)。本章の VAE も、同じ勾配法で学習する。

本章はその壁に、正面から挑む。ただし壁を破るのではなく、別の登り方をする。

鍵となる発想はこうだ—真の事後分布が計算できないなら、計算できる分布で近似し、近似の良さを測る量を最大化すればよい。これを変分推論と呼ぶ。そしてこの発想が、ヘルムホルツマシンを経て VAE に至り、最終的には第16章の自由エネルギー原理そのものになる。

ヒント本章のガイド

第3節の導出が本章の心臓部である。たった数行で、対数尤度が「下界」と「KL ダイバージェンス」の和に分解される。この一本の式が、本章のすべてであり、そして第16章のすべてでもある。符号を反転させると変分自由エネルギーになる。ここだけは飛ばさないでほしい。

第1〜2節は準備、第4〜5節が歴史的な系譜(EM とヘルムホルツマシン)、第6〜8節が現代的な実装(VAE)である。

急ぐなら第3節と第6節だけでもよい。ただし第5節の Wake-Sleep は、第7章の正相・負相と対応しているので、そこの気持ちよさは味わってほしい。


1. 潜在変数モデル

見えるものと見えないもの

手元にデータ x がある。顔写真だとしよう。この写真がどう生成されたかを、われわれはモデル化したい。

素朴には p(x) を直接モデル化すればよさそうだが、これは難しい。顔写真の空間は膨大で、その中で「顔らしい」画像が占める領域はごく一部の複雑な形をしている。それを直接書き下すのは無理がある。

発想を変える。顔写真の背後には、目に見えない要因があるはずだ。人物の同一性、表情、顔の向き、照明の方向。これらを表す潜在変数を z と呼ぼう。以下の連続値モデルでは、写真の画素値を並べた観測ベクトルを x \in \mathbb{R}^D、見えない要因を表す潜在ベクトルを z \in \mathbb{R}^d とする(D が観測の成分数、d が潜在変数の成分数である)。z が決まれば、そこから写真 x を生成するのは(比較的)易しい。

p_\theta(x) = \int p_\theta(x \mid z)\, p(z)\, dz

これが潜在変数モデルである。p(z) は潜在変数の事前分布(たとえば標準正規分布)、p_\theta(x\mid z) は生成過程(デコーダ)、\theta はパラメータである。

構造をよく見てほしい。まず単純な分布 p(z) から潜在変数を引き、次に p_\theta(x\mid z) に従って観測を作る。p_\theta(x\mid z) は「z ごとに一枚の画像を返す写像」ではなく、「その z からどんな画像がどれだけ生じやすいか」を表す条件付き分布である。上の積分は、それを p(z) で重み付けして足し合わせている。つまり単純な分布 p(z) を、p_\theta(x\mid z) を通して押し出すことで、複雑な分布 p_\theta(x) を作っているのである。第15章で扱うフローも拡散モデルも、根っこは同じ発想である。

二つの向き

潜在変数モデルには、二つの方向がある。

生成(generation)—z から x を作る。p_\theta(x\mid z) で記述される。トップダウンの向きである。

認識(recognition)または推論(inference)—x から z を推定する。p_\theta(z \mid x) で記述される。ボトムアップの向きである。

そして脳の話をするなら、この二つが視覚皮質の二つの経路に対応する、という見立てが自然に出てくる。網膜から高次野へ向かうフィードフォワード経路が認識、高次野から低次野へ戻るフィードバック経路が生成—という対応である。第16章の予測符号化で、この見立てを具体的な回路として扱う。

なぜ難しいのか

さて、\theta を学習したい。最尤法でいくなら \log p_\theta(x) を最大化すればよい。

これは観測が一つのときの書き方である。観測が N 個あるなら x_1, \dots, x_N(各 x_n \in \mathbb{R}^{D})と書き、同じモデルから独立に得られたと仮定して、共通の \theta について \sum_{n=1}^{N}\log p_\theta(x_n) を最大化する。第3節で導く下界も、観測ごとに求めて足し合わせればよい。いっぽう「見えない原因についての事後分布」は観測ごとに違う。第4節までは、その近似を観測ごとに個別に求める。第5節で、観測を受け取って近似分布を返す一つのネットワークを全観測で共有する形に変わる。以下ではまず、観測が一つの場合で計算の難しさを見よう。

\log p_\theta(x) = \log \int p_\theta(x\mid z) p(z)\, dz

この積分が計算できない。z が高次元なら、数値積分は絶望的である。ベイズの定理を使って事後分布を書くこともできる。

p_\theta(z \mid x) = \frac{p_\theta(x\mid z)\, p(z)}{p_\theta(x)}

だが分母に、まさに計算できない p_\theta(x) が出てくる。事後分布も計算できない。

第7章の分配関数 Z と、同じ構造の困難である。あのときも「すべての状態にわたる和」が壁になった。今度は「すべての潜在変数にわたる積分」が壁になっている。

ここがポイント

潜在変数モデルの困難は一点に集約される—見えない潜在変数について足し合わせた観測の尤度、すなわち周辺尤度 p_\theta(x) = \int p_\theta(x\mid z)p(z)dz が計算できない。だから事後分布 p_\theta(z\mid x) も計算できない。第7章の分配関数と同じ形の壁である。


2. 必要な数学:変分法の最小限

計算できない事後分布を、扱いやすい分布で近似したい。その準備として「変分」の発想だけを説明する。変分法の技術はほとんど使わない。

汎関数と変分

普通の関数は、数を受け取って数を返す。f(x) = x^2 のように。

汎関数(functional)は、関数を受け取って数を返す。たとえば

H[q] = -\int q(z) \log q(z)\, dz

はエントロピーであり、確率密度関数 q を受け取って数を返す。[\;] で書くのは、引数が関数であることを示す慣習である。

普通の最適化では、f(x) を最小にする x を探す。変分法では、H[q] を最小(最大)にする関数 q を探す。これが「変分」の意味である。

本章での使い方

われわれがやろうとしているのは、こうである。真の事後分布 p_\theta(z\mid x) は計算できない。そこで扱いやすい分布の族(たとえばガウス分布の族)を用意し、その中でいちばん近いものを探す。「分布の族の中で最適なものを探す」—これが変分推論であり、関数を探しているから「変分」なのである。

技術的な変分法(オイラー・ラグランジュ方程式など)は使わない。実際には q をパラメータで表して(q_\phi と書く)、\phi について普通に微分するだけである。\phi は近似分布の形を決めるパラメータの集まり(ガウス分布なら平均や分散を決める)、\theta は生成分布 p_\theta(x\mid z) を決めるパラメータの集まりである。\phi を変えると推定の仕方が変わり、\theta を変えるとデータの作られ方が変わる。だから身構える必要はない。

イェンセンの不等式

一つだけ道具を思い出しておく。第5章で使ったイェンセンの不等式である。

凹関数とは、グラフ上のどの二点を直線で結んでも、その二点のあいだでグラフが直線より下に来ない関数のことである。上にふくらんだ形をしている、と言ってもよい。この形のおかげで、「先に平均してから関数に入れた値」は「関数に入れてから平均した値」以上になる。以下の X は実数値をとる確率変数で、\mathbb{E}[X] はその値を確率で重み付けした平均である(観測ベクトル x とは別の記号である)。各期待値は有限に定まるとし、対数を使う場合は X > 0 とする。f が凹関数のとき、

f\big(\mathbb{E}[X]\big) \ge \mathbb{E}\big[f(X)\big]

\log は凹関数だから、

\log \mathbb{E}[X] \ge \mathbb{E}[\log X]

「対数の期待値」を「期待値の対数」で上から抑えられる—逆に言えば、\log を期待値の中に入れると値が下がる。次節でこれを使う。


3. 変分下界(ELBO)の導出

導出そのものは数行で終わる。だが一行ずつ、何をしているかを確かめながら進みたい。

何をしたいか

計算できない \log p_\theta(x) を、計算できる何かで下から抑えたい。下界が手に入れば、それを代理の目的関数として最大化できる。届かない山の高さの代わりに、確実に届く足場の高さを上げていく—そういう戦略である。

ただし、下界を上げれば元の量も必ず上がる、というわけではない。 下界と本体の差(あとで見るとおり KL ダイバージェンスである)も同時に動くからだ。両者の関係は次項で厳密に書き下す。そこがこの節の要点になる。

そのために、観測 x を一つ固定し、その観測のもとで潜在変数を推定するための確率分布 q(z) を導入する(事前分布 p(z) とは別物である。x ごとに選ぶ分布だが、いまは x を固定しているので条件を省いて q(z) と書く)。どんな分布でもよいというのが要点だ(後で「扱いやすい族」に制限するが、いまは任意である)。

導出その一 ── イェンセンを使う

以下では、q(z) > 0 となる z の範囲で p_\theta(x,z) > 0 とする—q が値を持つところでは p_\theta も値を持つ、ということである(そうでないと二行目の対数が -\infty になってしまう)。なお p_\theta(x,z) = p_\theta(x\mid z)p(z) は観測と潜在変数の同時密度で、逆向きの条件、つまり p_\theta(x,z) > 0 のところで q(z) > 0 とするのは、q を掛けて割る操作で積分の一部を落とさないためである。向きに注意してほしい。q(z) > 0 なのに p_\theta(x,z) = 0 となる領域を許すと、つまり q の台のほうを広く取ると、その領域で比の対数が定まらず、下界は -\infty\mathrm{KL}+\infty になりうる。あとで出てくる等式が意味を持たなくなる。密度の比と、その対数の期待値が定まることも仮定しておく。

\begin{aligned} \log p_\theta(x) &= \log \int p_\theta(x, z)\, dz &&\textsf{(同時分布で書いた)}\\ &= \log \int q(z)\, \frac{p_\theta(x,z)}{q(z)}\, dz &&\textsf{(} q(z) \textsf{ を掛けて割った)}\\ &= \log\, \mathbb{E}_{q(z)}\!\left[ \frac{p_\theta(x,z)}{q(z)} \right] &&\textsf{(積分を } q \textsf{ による期待値と読む)}\\ &\ge \mathbb{E}_{q(z)}\!\left[ \log \frac{p_\theta(x,z)}{q(z)} \right] &&\textsf{(イェンセンの不等式。} \log \textsf{ を中に入れた)} \end{aligned}

最後の量を \mathcal{L}(q,\theta) と書き、変分下界(variational lower bound)あるいは ELBO(evidence lower bound)と呼ぶ。

\mathcal{L}(q, \theta; x) = \mathbb{E}_{q(z)}\!\left[ \log \frac{p_\theta(x,z)}{q(z)} \right]

第二行の操作に注目してほしい。q(z) を掛けて割るという、それ自体は何もしていない変形である。だがこれによって積分が期待値の形になり、イェンセンが使えるようになった。「何もしない操作」が道を開く—数学ではよくあることだ。

導出その二 ── 差を測る

下界であることは分かった。だがどれだけ緩い下界なのかが知りたい。等号はいつ成り立つのか。

小さな離散の例で確かめよう。観測は x = 0, 1 の二通り、潜在変数は z = 1, 2, 3 の三通りとし、積分は三つの値についての和に読み替える。事前確率は p(z) = 1/3 ずつ、p_\theta(x=1\mid z)z = 1, 2, 3 の順に 0.3, 0.6, 0.3 とする。すると x = 1 が出る確率は (0.3 + 0.6 + 0.3)/3 = 0.4 で、その対数は約 -0.9163 である(以下、対数は自然対数)。

近似分布 q にも三つの値へ 1/3 ずつ割り当ててみる。下界の中の比は p_\theta(x=1,z)/q(z) = p_\theta(x=1\mid z) になるので、下界は (\log 0.3 + \log 0.6 + \log 0.3)/3 \approx -0.9729 である。-0.9163 より小さい。差は約 0.0566—これが「緩み」である。

では q を真の事後分布に合わせたらどうか。ベイズの定理から、p(z\mid x=1)(0.1, 0.2, 0.1)/0.4 = (0.25, 0.5, 0.25) である。この q を入れると、比 p_\theta(x=1,z)/q(z) はどの z でも 0.4 にそろう。だから下界は \log 0.4 そのもの、つまり緩みがゼロになる。比が z によらず一定になったとき等号が成り立つ—次にこれを式で確かめる。

そこで、別の道筋で同じ量に到達してみる。今度は不等式を使わない。

\begin{aligned} \log p_\theta(x) &= \int q(z) \log p_\theta(x)\, dz &&\textsf{(} z \textsf{ によらない)}\\ &= \int q(z) \log \frac{p_\theta(x,z)}{p_\theta(z\mid x)}\, dz &&\textsf{(ベイズの定理)}\\ &= \int q(z) \log \left[\frac{p_\theta(x,z)}{q(z)} \cdot \frac{q(z)}{p_\theta(z\mid x)}\right]\, dz &&\textsf{(対数の中で掛けて割った)}\\ &= \underbrace{\int q(z) \log \frac{p_\theta(x,z)}{q(z)}\, dz}_{= \,\mathcal{L}(q,\theta)} \;+\; \underbrace{\int q(z) \log \frac{q(z)}{p_\theta(z\mid x)}\, dz}_{= \,\mathrm{KL}\big(q(z) \,\|\, p_\theta(z\mid x)\big)} &&\textsf{(積を和に)} \end{aligned}

きれいな形が出た。

\log p_\theta(x) = \mathcal{L}(q,\theta) + \mathrm{KL}\big(q(z) \,\|\, p_\theta(z \mid x)\big)

この関係を図にすると 図 1 になる。

図 1: 対数尤度が下界と KL ダイバージェンスの和に分かれる関係(模式)。天井の高さ log p(x) は近似分布 q の選び方によらず一定である。だから下界を押し上げることと、KL を小さくして事後分布への近似を良くすることは、同じ一つの作業になる。

この一行が語ること

しばらく眺めてほしい。この式は三つのことを同時に言っている。第一に、なぜ \mathcal{L} が下界なのか。KL ダイバージェンスは常に 0 以上である。だから \log p_\theta(x) \ge \mathcal{L}。イェンセンを使わなくても下界性が出た。

第二に、緩みの正体。下界と真の値の差は、ちょうど「近似分布 q と真の事後分布 p_\theta(z\mid x) の KL ダイバージェンス」である。q が真の事後分布に一致すれば KL はゼロになり、等号が成り立つ。

第三に—これがいちばん大事なのだが—左辺は q に依存しない。\log p_\theta(x) は、q をどう選ぼうと変わらない。\theta だけで決まる量である。すると右辺の二項は、q を動かしたとき足して一定でなければならない。つまり

\underbrace{\mathcal{L}\ \text{を上げる}}_{\textsf{下界を押し上げる}} \quad\Longleftrightarrow\quad \underbrace{\mathrm{KL}\ \text{が下がる}}_{\textsf{近似が良くなる}}

ここで一つ、条件を確かめておきたい。左辺 \log p_\theta(x)q によらないのは \theta を止めているからである。だから「下界を上げること=事後分布の近似を良くすること」と言えるのは、\theta を固定して q だけを動かすあいだである。\theta も一緒に動かせば左辺も動くので、この言い換えはそのままでは使えない。

もう一つ。q を計算しやすい族(たとえばガウス)に制限すると、その族が真の事後分布を含んでいないかぎり、KL はゼロまで下がらない。下界と尤度のあいだには、その分の隙間が残り続ける。それでも、二つの目的—「尤度を上げたい」と「事後分布を近似したい」—が一つの量 \mathcal{L} の最大化に束ねられたことが、変分推論の核心である。

ここがポイント

\log p_\theta(x) = \mathcal{L}(q,\theta) + \mathrm{KL}\big(q \,\|\, p_\theta(z\mid x)\big) 左辺は(\theta を固定すれば)q によらない。だから \mathcal{L} を上げることと、事後分布への近似を良くすることは、同じ一つの作業である。そして q が真の事後分布に一致できたとき、\mathcal{L}\log p_\theta(x) に等しくなる—許した族の中に真の事後分布がなければ、そこには届かない。

ELBO の二通りの書き方

実装のために、\mathcal{L} を二つの項に分けておく。目指すのは、推定した潜在変数から観測 x を説明する対数尤度の平均を大きくし、同時に、近似分布 q事前分布 p(z) の KL を小さくすることである(引く相手は事前分布であって、下界の緩みを表していた事後分布ではない)。p_\theta(x,z) = p_\theta(x\mid z)p(z) を代入すると、

\begin{aligned} \mathcal{L}(q,\theta) &= \mathbb{E}_{q}\left[\log \frac{p_\theta(x\mid z)\, p(z)}{q(z)}\right] \\ &= \mathbb{E}_{q}\big[\log p_\theta(x\mid z)\big] + \mathbb{E}_q\left[\log \frac{p(z)}{q(z)}\right] &&\textsf{(} \log \textsf{ の積を分けた)}\\ &= \underbrace{\mathbb{E}_{q}\big[\log p_\theta(x\mid z)\big]}_{\textsf{再構成項}} - \underbrace{\mathrm{KL}\big(q(z)\,\|\,p(z)\big)}_{\textsf{正則化項}} \end{aligned}

この形が実装で使われる。二つの項には、はっきりした役割がある。

再構成項は「q から z をサンプルして x を再構成したとき、どれだけうまく再構成できるか」。大きいほどよい。

正則化項は「q が事前分布 p(z) からどれだけ離れているか」。離れすぎるとペナルティになる。

つまり ELBO の最大化とは、「よく再構成できて、かつ事前分布から離れすぎない q を探す」という作業である。オートエンコーダの目的関数に、正則化がついた形—第6節でこれがそのまま VAE になる。


4. EM アルゴリズム

変分下界 \mathcal{L}(q,\theta) には、事後分布を近似する q と、生成モデルのパラメータ \theta がある。片方を止めてもう片方を良くする。この二段階を交互に繰り返すのが、EM アルゴリズムである。

E ステップでは \theta を固定する。前節の分解では、q を動かしても対数尤度は動かない。だから ELBO を最大にするには、緩みを表す KL をゼロにすればよい。

q(z)=p_\theta(z\mid x)

M ステップでは、いま求めた q を固定して \theta を選び直す。\mathcal{L}=\mathbb{E}_q[\log p_\theta(x,z)]+H[q] の第二項は動かないので、

\theta\leftarrow\arg\max_\theta\;\mathbb{E}_{q(z)}[\log p_\theta(x,z)]

となる。\arg\max は最大値ではなく、それを実現するパラメータを選ぶ操作で、矢印はその値に更新することを示す。見えない z の候補を q で重みづけし、xz が両方見えたときの対数尤度を平均して、大きくするのである。

なぜ対数尤度が下がらないのだろうか。E ステップの直後は、下界が対数尤度に接する。M ステップで下界を上げれば、その上にある対数尤度も下がらない。対数尤度に有限の上限があれば値は収束するが、パラメータが最尤解に落ち着く保証とは別である。混合ガウスモデルのように、上限が有限でない場合もある。

ただし、E ステップは事後分布を計算できてこその一手だ。深層の非線形モデルでは、それが難しい。そこで q を扱いやすい分布の族に絞り、その中で KL を小さくする。これが第2節で予告した変分推論(変分ベイズとも呼ばれる)である。EM は、q の候補を全分布に広げ、事後分布そのものを選べる場合にあたる。


5. ヘルムホルツマシンと Wake-Sleep

推論の手間を、学習で前払いする

網膜に映るのは二次元の像なのに、われわれは三次元の世界を見る。ヘルムホルツが知覚を無意識的推論と呼んだのは、見えたものから原因を推し量る、この働きに注目したからである。世界の状態が z、感覚が x なら、知覚とは p(z\mid x) を求めることになる。

Dayan、Hinton、Neal、Zemel (1995) のヘルムホルツマシンは、その推論をするネットワークを別に用意した。

  • 生成モデル p_\theta(x\mid z) は、原因 z から感覚 x を作る
  • 認識モデル q_\phi(z\mid x) は、感覚 x から原因 z の分布を推定する

データ点ごとに q を毎回探す代わりに、認識モデルのパラメータ \phi を学んでおく。新しい x も一回通せば推論できる。この工夫を償却推論(amortized inference)と呼ぶ。推論の手間を学習で「前払い」するわけだ。

見たものと、夢で見たもの

どう学ぶのだろうか。Hinton、Dayan、Frey、Neal (1995) の Wake-Sleep は、二つの相を交互に回す。

Wake 相では、実データ x から認識モデルで z\sim q_\phi(z\mid x) を取り、\log p_\theta(x\mid z) を上げるように生成モデルの \theta を更新する。「見たものを、自分でも作れるようになる」段階だ。

Sleep 相では、z\sim p(z)x\sim p_\theta(x\mid z) と生成モデルから夢を見る。原因 z が分かっているので、\log q_\phi(z\mid x) を上げるように認識モデルの \phi を更新できる。「自分が作ったものを、正しく認識できるようになる」段階である。

第7章第5節の正相もデータ由来、負相もモデル由来の標本を使っていた。対応するのは、この標本の出どころである。ボルツマンマシンでは一つの勾配の二項を引き算したが、Wake-Sleep は別々のパラメータを別々の目的で更新する。睡眠中の自発活動との類比は魅力的だが、この対応だけで脳の学習則が分かったことにはならない。

目的が一つにそろわない

Sleep 相の標本の出どころに注意しよう。同じ x を生んだ z は、モデルの事後分布 p_\theta(z\mid x) に従う。KL は左側の分布で平均するので、Sleep 相が下げるのは

\mathrm{KL}\big(p_\theta(z\mid x)\,\|\,q_\phi(z\mid x)\big)

を、モデルが生む x について平均した量である。ELBO の緩みは逆向きの \mathrm{KL}(q_\phi\|p_\theta(z\mid x)) だった。しかも欲しいのは実データについての改善で、モデルの夢についての改善とは限らない。だから Wake-Sleep は、単一の ELBO を最適化する手続きになっていない。次節の VAE は、この目的を一つにそろえる。

生成と認識を対にする発想は、第16章の自由エネルギー原理にもつながる。ただし、式の共通性と生物の実際の働きは別に確かめる必要がある。


6. 変分オートエンコーダ

Kingma と Welling (2014) の変分オートエンコーダ(VAE)は、ヘルムホルツマシンの構成を受け継ぎつつ、ELBO という単一の目的関数を、勾配降下法で直接最適化する。

ここからは機械学習の慣例に従い、観測 x から近似事後分布 q_\phi(z\mid x) を決める認識ネットワークをエンコーダ(符号化器)、潜在変数 z から観測の分布 p_\theta(x\mid z) を決める生成ネットワークをデコーダ(復号器)と呼ぶ。観測から潜在表現へ進むのがエンコード、その逆がデコードである。第6章第1節で注意したように、神経科学での「刺激から脳活動へ」という基準とは向きの取り方が違う。この食い違いは第15章第10節で決着する。

何が問題だったか

やりたいのは、ELBO を \theta\phi の両方について最大化することである。

\mathcal{L}(\theta,\phi; x) = \mathbb{E}_{q_\phi(z\mid x)}\big[\log p_\theta(x\mid z)\big] - \mathrm{KL}\big(q_\phi(z\mid x)\,\|\,p(z)\big)

\theta については問題ない。期待値の中身だけが \theta に依存するので、サンプルして微分すればよい。

\phi が問題である。\phi は期待値を取る分布のほうに入っている。まず再構成項に注目し、x\theta を固定して f(z) = \log p_\theta(x\mid z) と略記しよう。f は潜在ベクトルを受け取り、観測 x の説明の良さを一つの数で返す関数である。

\nabla_\phi\, \mathbb{E}_{q_\phi(z\mid x)}\big[\, f(z) \,\big] = \;?

この書き方のままでは、標本を通して \phi へ勾配を流せない。乱数を引くところで、逆伝播の経路が途切れてしまう(勾配そのものが取れないわけではない。別の道は本節の「手を動かす」で見る)。

再パラメータ化トリック

Kingma と Welling の解決策は、驚くほど単純である。乱数を引く場所を、パラメータに依存しないところへ移す。

潜在変数が一つの場合で見よう。平均を \mu = 2、標準偏差を \sigma = 3 とし、標準正規分布(平均 0・分散 1)から引いた乱数が \epsilon = 0.5 だったとする。z = \mu + \sigma\epsilon と作れば、この回の値は 2 + 3 \times 0.5 = 3.5 である。ここで \epsilon を固定したまま \mu\sigma を動かせば、\partial z/\partial\mu = 1\partial z/\partial\sigma = \epsilon = 0.5 と、普通に微分できる。固定するのは一回の微分計算に使う乱数の値であって、乱数を繰り返し引けば z はちゃんと平均 \mu・標準偏差 \sigma の正規分布に従う。これを潜在変数の各成分について行う。

q_\phi(z\mid x) をガウス分布 \mathcal{N}\big(\mu_\phi(x), \mathrm{diag}(\sigma_\phi(x)^2)\big) とする。潜在変数の成分数を d とすると、ネットワークは平均 \mu_\phi(x) \in \mathbb{R}^d と、各成分が正の標準偏差 \sigma_\phi(x) \in \mathbb{R}^d を出力する(\sigma_\phi(x)^2 は成分ごとの二乗、\mathrm{diag} はそれを対角に並べた d \times d の共分散行列を作る操作である)。これは、観測 x を固定したときに潜在変数の各成分を独立として近似する、という選択であり、こうして分布を成分ごとの積に制限する方法を平均場近似という。真の事後分布でも各成分が独立だと主張しているわけではない。このとき、各成分が独立な標準正規ノイズ \epsilon \in \mathbb{R}^d を引けば、z を次のように書ける(\odot は成分ごとの積である)。

z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon, \qquad \epsilon \sim \mathcal{N}(0, I)

確かめよう。\epsilon が標準正規に従うなら、\sigma \odot \epsilon は平均 0・標準偏差 \sigma の正規分布に従う。それに \mu を足せば平均が \mu になる。たしかに z \sim \mathcal{N}\big(\mu_\phi, \mathrm{diag}(\sigma_\phi^2)\big) である。分布としてはまったく同じものだ。だが計算グラフの形が違う。

  • 元の書き方 — \phi → 分布 → (サンプリング)→ z。サンプリングで勾配が途切れる
  • 新しい書き方 — \epsilon\phi に依存しない外部からの乱数。z\mu_\phi, \sigma_\phi決定的な関数である。だから \partial z/\partial \phi が計算できる

これが再パラメータ化トリック(reparameterization trick)である。乱数を「外から注入されるノイズ」として扱うことで、微分可能な経路を確保した。期待値も書き換わる。

\mathbb{E}_{q_\phi(z\mid x)}\big[f(z)\big] = \mathbb{E}_{\epsilon \sim \mathcal{N}(0,I)}\Big[ f\big(\mu_\phi(x) + \sigma_\phi(x)\odot\epsilon\big) \Big]

右辺は \phi によらない分布での期待値になった。積分する相手が \phi と無関係になったので、あとは(被積分関数が \phi で微分でき、その微分が期待値を取れる程度に行儀がよければ)微分と期待値を交換できる。実際にはサンプルを一つ引いて、普通にバックプロパゲーションするだけである。この手が使えるための条件を、まとめておこう。

  • 潜在変数を「\phi に依存しない基本ノイズ \epsilon」と「\phi の微分可能な変換」に分けて書けること
  • その変換と、中身の関数 f が微分可能であること

だから離散的な潜在変数には、この形をそのまま使えない(\epsilon をどう変換しても、出口が飛び飛びになる)。そのときは、勾配を別の形で見積もる方法—本節の「手を動かす」で扱うスコア関数推定量や、離散変数を連続に緩める手法—を使う。「離散だと勾配法が使えない」のではなく、この経路が使えないのである。

ここがポイント

再パラメータ化トリックの要点は、乱数の発生源をパラメータから切り離すことにある。z = \mu_\phi + \sigma_\phi \odot \epsilon と書けば、\epsilon は外部からのノイズ、z\phi の決定的な関数になる。だから勾配が通る。

KL 項は解析的に計算できる

ELBO で差し引く \mathrm{KL}(q_\phi(z\mid x)\|p(z)) は、両方がガウス分布なら閉じた形で書ける。事前分布が標準正規分布なら、この KL を小さくすることは、近似分布の各成分の平均を 0 に、標準偏差を 1 に近づけるように働く。観測 x を固定して \mu_\phi(x)\sigma_\phi(x)\mu\sigma と略記し、j を潜在変数の成分を指す添字とする。p(z) = \mathcal{N}(0,I)q_\phi = \mathcal{N}(\mu, \mathrm{diag}(\sigma^2)) のとき、

\mathrm{KL} = \frac{1}{2}\sum_{j=1}^{d}\Big( \mu_j^2 + \sigma_j^2 - \log \sigma_j^2 - 1 \Big)

サンプリングが要らないので、こちらは分散なく計算できる。

各項の意味を読んでおこう。\mu_j^2 は「原点から離れるな」、\sigma_j^2 - \log\sigma_j^2 - 1\sigma_j = 1 で最小になる項で「分散を1に保て」。つまり KL 項は、潜在表現を標準正規分布の形に整えようとする力である。

この力には、行きすぎることがある。q_\phi(z\mid x) が入力によらず p(z) そのものになってしまえば、KL 項はゼロで最小になる。潜在変数は入力の情報を何も運ばなくなり、復号側はそれを無視して—たとえば、すでに生成した成分から次の成分を順に予測する強力な自己回帰モデルなら、自分の力だけで—入力を再構成しようとする。こうなると符号化側は潜在を放り出すのが得になる。後方崩壊(posterior collapse)と呼ばれる現象である。

再構成項があり、しかも的が入力に固定されているのに、途中の表現が使われなくなる。「的が潰れないこと」と「表現が使われること」は別である—この区別は、第17章第7節で予測符号化を検討するときに効いてくる。

全体像

まとめると、VAE の一回の更新はこうなる。

  1. データ x を取る
  2. エンコーダに通して \mu_\phi(x), \sigma_\phi(x) を得る
  3. \epsilon \sim \mathcal{N}(0,I) を引き、z = \mu_\phi + \sigma_\phi\odot\epsilon とする
  4. デコーダに通して p_\theta(x\mid z) を得る
  5. 再構成項 \log p_\theta(x\mid z) と KL 項を計算し、ELBO を作る
  6. -\mathcal{L} を損失としてバックプロパゲーション

ヘルムホルツマシンの構成(生成+認識の二本立て)に、ELBO という単一の目的関数と、それを微分可能にする仕掛けを与えたもの—それが VAE である。Wake-Sleep が二つの相で別々の目的を追っていたのに対し、VAE は一つの量を上げ続ける。

ヒント手を動かす

再パラメータ化トリックの効き目を、数値で確かめよう。比較するスコア関数推定量は、関数値 f(z) に対数密度のパラメータ微分 \partial_\phi\log q_\phi(z) を掛け、その標本平均で期待値の微分を見積もる方法である。第15章で登場する画像の値によるスコアとは、微分する相手が違う。コードは コード/14_reparam.py にある。

q_\phi(z) = \mathcal{N}(\phi,1)(一次元、平均だけがパラメータ)とし、f(z)=z^2 の期待値の \phi 微分を推定する。真の値は \mathbb{E}[z^2] = \phi^2+1 だから、微分は 2\phi である。\phi = 1、1回あたり100サンプル、それを1000回繰り返した。

PHI, M, REP = 1.0, 100, 1000  # φ、サンプル数、繰り返し

score, reparam = [], []
for _ in range(REP):
    eps = rng.standard_normal(M)
    z = PHI + eps
    # 1. スコア関数推定量: ∂_φ log q = (z - φ)
    score.append(np.mean(z**2 * (z - PHI)))
    # 2. 再パラメータ化: ∂_φ f(φ + ε) = 2(φ + ε)
    reparam.append(np.mean(2 * z))

score = np.array(score); reparam = np.array(reparam)
print(score.std(), reparam.std())

二つの行を見比べてほしい。上は f(z)=z^2 の値を使い、下は f の傾き 2z を使っている。差はそれだけである。

平均 標準偏差
スコア関数推定量 1.992 0.552
再パラメータ化 1.998 0.209

どちらも不偏である(平均が真の値 2 に近い)。だがばらつきが 2.6 倍違う。

なぜ差が出るのか。スコア関数推定量は f(z) の値そのものを重みに使い、再パラメータ化は f の傾きを使う。前者は f が大きい場所のサンプルに引きずられるが、後者は f の局所的な情報しか使わない。勾配を求めたいのだから傾きを使うほうが素直である—そう言われれば当たり前だが、この差が VAE を実用にした。


7. 生成モデルと認識モデルの関係

VAE では、エンコーダとデコーダを独立なネットワークとして持った。だが素朴には、こう思いたくなる—認識は生成の逆なのだから、片方から他方が決まるのではないか。

この直感が当たるかどうかを決めているのは、実は「線形か非線形か」ではない。単射(別の z が同じ x を生まないか)と、観測ノイズの有無である。ノイズがなく写像が全単射なら、非線形でも逆写像が存在する—第15章のフローがまさにそれだ。逆に線形でも、単射でなかったりノイズが乗っていれば、事後分布は逆写像ではない。

以下では、解析的に解ける線形ガウスの場合と、一般の非線形の場合を並べて確かめる。

線形の場合

潜在変数を重み付きで足し合わせ、そこに観測ノイズを加えてデータを作る—そういう生成過程が線形ガウスだとする。W \in \mathbb{R}^{D \times d} が生成の重みで、行が観測の成分、列が潜在変数の成分に対応する(各列は、一つの潜在成分が観測全体にどう効くかを表す)。\varepsilon \in \mathbb{R}^Dz と独立な観測ノイズで、第6節で潜在変数を作るために使った \epsilon \in \mathbb{R}^d とは別物である。\sigma > 0 は全観測成分に共通の標準偏差、下の二つの I はそれぞれ d \times dD \times D の単位行列である。

z \sim \mathcal{N}(0, I), \qquad x = Wz + \varepsilon, \quad \varepsilon\sim\mathcal{N}(0,\sigma^2 I)

これは確率的主成分分析(probabilistic PCA)と呼ばれるモデルである。このとき事後分布は解析的に求まり、やはりガウス分布になる。

p(z\mid x) = \mathcal{N}\big(M^{-1}W^\top x,\; \sigma^2 M^{-1}\big), \qquad M = W^\top W + \sigma^2 I

ここで M = W^\top W + \sigma^2 IId \times d の単位行列なので、Md \times d である。認識の重みが M^{-1}W^\top、つまり生成の重み W から決まっている。長方形の行列 W には、普通の逆行列は定義できない。代わりに使う擬似逆行列 W^+ \in \mathbb{R}^{d \times D}(行が潜在変数の成分、列が観測の成分)とは、観測 x に対して \|x - Wz\|^2 をいちばん小さくする z を返す行列である(最小にする z が複数あるときは、そのうち長さ \|z\| が最小のものを返す)。また列フルランクとは、W のどの列も、残りの列を重み付きで足して作ることはできない、という条件である。

\sigma \to 0 の極限では M^{-1}W^\top \to (W^\top W)^{-1}W^\top = W^+(擬似逆行列)である。まさに「逆」だ—ただしこの形が書けるのは W が列フルランクのとき、つまり潜在変数の各方向が観測に別々に効いているときである。

単射でないと壊れる

生成過程が x = g_\theta(z) + \varepsilon という非線形写像だとどうなるか。ここで g_\theta: \mathbb{R}^{d} \to \mathbb{R}^{D} は、潜在ベクトルから観測の平均ベクトルを作る関数で、\theta がその形を決める(線形の場合の Wzg_\theta(z) に置き換えたのである)。観測ノイズは直前と同じく z と独立な \varepsilon \sim \mathcal{N}(0, \sigma^2 I) とするので、p_\theta(x\mid z) = \mathcal{N}(g_\theta(z), \sigma^2 I) であり、g_\theta(z) はこの分布の平均にあたる。この仮定のもとで事後分布は

p(z\mid x) \propto p(z)\, \exp\!\left(-\frac{\|x - g_\theta(z)\|^2}{2\sigma^2}\right)

この分布は、一般にガウス分布ではない。多峰性を持つこともある。g_\theta が単射でなければ、複数の z が同じ x を生む—そのため事後分布は複数の峰を持ちうる。峰が分かれるかは、事前分布と観測ノイズにもよる。ここでも効いているのは非線形性そのものではなく、単射でないこと、そしてノイズがあることである。そして「認識モデルの重みが生成モデルの重みの逆行列」といった単純な関係は、まったく成り立たない。非線形写像の逆は、一般に閉じた形で書けないからである。だから VAE では、エンコーダを独立なネットワークとして用意し、学習で近似させる。「逆」は与えられるものではなく、学習するものになった。

設計で対称性を入れる

とはいえ、両者を完全に独立にするのは冗長でもある。実務では対称性を明示的に設計することがある。

  • 重み結合(tied weights)— デコーダの重みをエンコーダの重みの転置に取る。オートエンコーダで古くから使われる
  • アーキテクチャの鏡像化 — エンコーダとデコーダを層構成が対称になるように設計する

そして脳の側でも同じ問いが立つ。フィードフォワード結合とフィードバック結合は、互いの転置になっているのか。

解剖学的には、皮質の領野間結合はおおむね相反的である(A から B へ投射があれば、B から A へも投射がある)。だがシナプスの重みが数学的な意味で転置になっている保証はどこにもない。第8章で扱った重み輸送問題—バックプロパゲーションが生物学的に妥当でない理由の一つ—と、同じ論点がここにも現れる。


8. 離散潜在変数

最後に、第7章との接続をもう一つ確認しておく。

ボルツマンマシンと VAE の対比

両者とも潜在変数モデルである。だが潜在変数の性格が違う。

ボルツマンマシン VAE
潜在変数 離散\{0,1\} の二値) 連続(実ベクトル)
推論 MCMC(重い) 償却推論(一回のフォワード)
学習 正相 − 負相 ELBO の勾配
困難の所在 分配関数 Z 事後分布の近似

再パラメータ化トリックが使えるのは、潜在変数が連続だからである。z = \mu + \sigma\odot\epsilon という書き換えは、z が連続値でなければ意味をなさない。離散変数に対しては、\epsilon を少し動かしても z が動かない(あるいは飛ぶ)ので、この経路では勾配が流せない。これが、VAE が連続潜在変数を使う理由である。逆に言えば、離散的な潜在表現がほしいときには工夫が要る。

VQ-VAE

離散潜在変数を扱う代表的な方法が VQ-VAE(vector quantized VAE。van den Oord ら 2017)である。

エンコーダの出力を、あらかじめ用意した有限個のベクトル(コードブック)のうちいちばん近いものに置き換える(量子化する)。これで潜在表現が離散的になる。エンコーダの出力を z_e \in \mathbb{R}^d、候補を e_k \in \mathbb{R}^dK 個の候補の集合を \mathcal{C} = \{e_1, \dots, e_K\} と書く。次式の \arg\min は、距離を最小にする候補ベクトルそのものを選ぶ操作で、選ばれたものを z_q と書く。z_q 自体は実ベクトルだが、選択肢が有限個なので潜在表現は離散なのである。

z_q = \arg\min_{e_k \in \mathcal{C}} \| z_e - e_k \|

だが \arg\min は微分できない。ここでの工夫が straight-through 推定量(Bengio ら 2013)である—順伝播では量子化した値を使い、逆伝播では量子化がなかったことにして勾配をそのまま通す。数学的には正当化しづらい近似だが、実用上はよく機能する。

離散表現には利点がある。表現が有限個の「記号」になるので、その上で自己回帰モデルを学習できる。画像を離散トークンの列として扱い、言語モデルと同じ手法で生成する—近年の画像生成の一つの潮流は、この道筋にある。

次章へ

本章では、分配関数の壁を変分法で迂回し、VAE という実装まで辿り着いた。単純な分布 p(z) を写像で押し出して複雑な分布を作るという発想が、その根底にあった。

次章では、この発想を別の方向へ展開する。尤度を諦める道(GAN)、写像を可逆に制限する道(正規化フロー)、写像を少しずつにする道(拡散モデル)—三つの応答が並ぶ。そのうち正規化フローのところで、第5章第2節の変数変換の公式—|\det J| が密度を歪めるという話—が、いよいよ主役になる。

確認問題

  1. [確認]潜在変数モデル p_\theta(x) = \int p_\theta(x\mid z)p(z)\,dz において、最尤推定が困難になる理由を述べよ。またそれが第7章の分配関数の問題とどう似ているかを説明せよ。(第1節・第7章第6節)

  2. [導出]\log p_\theta(x) = \mathcal{L}(q,\theta) + \mathrm{KL}(q\|p_\theta(z\mid x)) を導出せよ。(第3節)

  3. [考える]\theta を固定すると、上の等式の左辺は q に依存しない。ここから何が結論できるか。二つ述べよ。(第3節)

  4. [導出]EM アルゴリズムの E ステップが「事後分布そのものを取る」ことになる理由を、第3節の等式から説明せよ。(第4節)

  5. [確認]Wake 相と Sleep 相が、第7章の正相・負相とそれぞれどう対応するかを述べよ。(第5節・第7章第5節)

  6. [考える]Wake-Sleep アルゴリズムの理論的な難点を述べよ。(第5節)

  7. [導出]p(z) = \mathcal{N}(0,I)q = \mathcal{N}(\mu, \mathrm{diag}(\sigma^2)) のときの KL 項の各項が、\mu\sigma に何を要求しているかを読み取れ。(第6節)

  8. [考える]「認識モデルは生成モデルの逆である」という主張が成り立つかどうかを決めているのは、線形か非線形かではない。何が決めているのかを述べよ。(第7節)


参考文献

  • Dayan, P., Hinton, G. E., Neal, R. M., & Zemel, R. S. (1995). The Helmholtz machine. Neural Computation, 7(5), 889–904. https://doi.org/10.1162/neco.1995.7.5.889 — 本章第5節の原典。第16章の系譜の出発点
  • Hinton, G., Dayan, P., Frey, B., & Neal, R. (1995). The wake-sleep algorithm for unsupervised neural networks. Science, 268(5214), 1158–1161. https://doi.org/10.1126/science.7761831
  • Neal, R. M., & Hinton, G. E. (1998). A view of the EM algorithm that justifies incremental, sparse, and other variants. In Learning in Graphical Models (pp. 355–368). Springer Netherlands. https://doi.org/10.1007/978-94-011-5014-9_12 — EM を自由エネルギーで書き直した仕事。第16章第2節の系譜で効いてくる[4節]
  • Kingma, D. P., & Welling, M. (2014). Auto-encoding variational bayes. arXiv:1312.6114. https://doi.org/10.48550/arxiv.1312.6114 — VAE の原典[6節]
  • Bishop, C. M. (2006). Pattern Recognition and Machine Learning. Springer. — EM と変分推論の教科書的な扱い
  • van den Oord, A., Vinyals, O., & Kavukcuoglu, K. (2017). Neural discrete representation learning. In Advances in Neural Information Processing Systems (NeurIPS). — VQ-VAE[8節]。straight-through 推定量は Bengio, Y., Léonard, N. & Courville, A. (2013)