Conditional Variational Auto Encoder

はじめに

　先の2回の投稿（こことここ）では、Variational Auto Encoder（VAE）をBayes推論の枠組みで解説した。今回は、Conditional Variational Auto Encoder（CAVE）をBayes推論の枠組みで説明する。

問題設定

　回帰問題を考え、$N$個のペア$(\vec{x}_n, \vec{y}_n)$が観測されているとする。$X=\{\vec{x}_1,\cdots,\vec{x}_N\}, Y=\{\vec{y}_1,\cdots,\vec{y}_N\}$と置いたとき、未観測データ$\vec{x}_{\alpha}$に対応する$\vec{y}_{\alpha}$を生成する確率分布$p(\vec{y}_{\alpha}|\vec{x}_{\alpha},X,Y)$を求めたい。潜在変数$\vec{z}$を導入し、$X$、$Y$、$\vec{z}$の同時確率分布$p(X,Y,\vec{z})$を考え、Bayesの定理を適用すると次式を得る。 $$\begin{equation} p(\vec{z}|X,Y) = \frac{p(Y|X,\vec{z})p(\vec{z})}{p(Y|X)} \label{eq9} \end{equation}$$ ただし、式変形の途中で$p(X|\vec{z})=p(X)$を用いた。事後確率$p(\vec{z}|X,Y)$が求まれば、次式により$\vec{y}_{\alpha}$を生成する確率分布を求めることができる。 $$\begin{equation} p(\vec{y}_{\alpha}|\vec{x}_{\alpha},X,Y)=\int d\vec{z}\;p(\vec{y}_{\alpha}|\vec{x}_{\alpha},\vec{z})p(\vec{z}|X,Y) \label{eq3} \end{equation}$$ 事後確率$p(\vec{z}|X,Y)$を求めることが目的である。

最適化すべき量

　$p(\vec{z}|X,Y)$を直接求めることはせず、パラメータ$\phi$を持つ関数$q_{\phi}(\vec{z}|X,Y)$を導入し、次のKullback Leibler divergenceを最小にすることを考える。 $$\begin{equation} D_{KL} \left[ q_{\phi}(\vec{z}|X,Y)||p(\vec{z}|X,Y) \right]=\int d\vec{z}\;q_{\phi}(\vec{z}|X,Y) \ln{ \frac{ q_{\phi}(\vec{z}|X,Y) } { p(\vec{z}|X,Y) } } \end{equation}$$ これを変形すると次式を得る。 $$\begin{equation} D_{KL} \left[ q_{\phi}(\vec{z}|X,Y)||p(\vec{z}|X,Y) \right] = D_{KL} \left[ q_{\phi}(\vec{z}|X,Y)||p(\vec{z}) \right]-E_{q_{\phi}(\vec{z}|X,Y)}\left[\ln{p(Y|X,\vec{z})}\right]+\ln{p(Y|X)} \label{eq1} \end{equation}$$ ただし、式変形の途中で式($\ref{eq9}$)を用いた。式($\ref{eq1}$)右辺にある$\ln{p(Y|X)}$は$\phi$に依存せず、観測値だけから決まる定数である。従って、次式が成り立つ。 $$\begin{equation} \min_{\phi} D_{KL} \left[ q_{\phi}(\vec{z}|X,Y)||p(\vec{z}|X,Y) \right] = \min_{\phi} {\left[ D_{KL} \left[ q_{\phi}(\vec{z}|X,Y)||p(\vec{z}) \right]-E_{q_{\phi}(\vec{z}|X,Y)}\left[\ln{p(Y|X,\vec{z})}\right] \right] } \label{eq2} \end{equation}$$ 式($\ref{eq2}$)の右辺第1項を小さく、第2項の期待値を大きくすれば良い。第1項は$q_{\phi}(\vec{z}|X,Y)$をできるだけ$p(\vec{z})$に近い形の分布にすることを要求し、この分布の下で対数尤度$\ln{p(Y|X,\vec{z})}$の期待値を大きくすることを第2項は要求する。第1項は正則化項に相当する。

KL divergenceの計算

　式($\ref{eq2}$)の右辺第1項を考える。いま次の仮定をおく。 $$\begin{eqnarray} q_{\phi}(\vec{z}|X,Y)&=&\mathcal{N}(\vec{z}|\vec{\mu}_{\phi}(X,Y),\Sigma_{\phi}(X,Y)) \\ p(\vec{z})&=&\mathcal{N}(\vec{z}|\vec{0},I_D) \end{eqnarray}$$ ここで、$\vec{z}$の次元を$D$とした。$I_D$は$D\times D$の単位行列である。どちらの分布も正規分布とし、$q_{\phi}(\vec{z}|X,Y)$の平均と共分散行列は$\phi,X,Y$から決まる量とする。これらは、入力$X,Y$、パラメータ$\phi$のニューラルネットワークを用いて計算される。一方、$p(\vec{z})$の方は平均0、分散1の標準正規分布である。このとき、$D_{KL} \left[ q_{\phi}(\vec{z}|X,Y)||p(\vec{z}) \right]$は解析的に計算することができる。 $$\begin{equation} D_{KL} \left[ q_{\phi}(\vec{z}|X,Y)||p(\vec{z}) \right]=\frac{1}{2}\left[ -\ln{|\Sigma_{\phi}(X,Y)|} -D +\mathrm{Tr}\left(\Sigma_{\phi}(X,Y)\right)+\vec{\mu}_{\phi}(X,Y)^T\vec{\mu}_{\phi}(X,Y) \right] \label{eq4} \end{equation}$$

ここまでの処理の流れ

　式($\ref{eq2}$)を計算する際の手順は以下のようになる。
分布$q_{\phi}(\vec{z}|X,Y)$は$X$と$Y$から$\vec{z}$を生成するEncoder、$p(Y|X,\vec{z})$は$X$と$\vec{z}$から$Y$を生成するDecoderとみなすことができる。青色で示した部分は最小化すべき量であり、赤字はサンプリングするステップである。青色の式の和を勾配降下法により最小にするが、その際、誤差逆伝播ができなければならない。$q_{\phi}(\vec{z}|X,Y)$はその$\phi$依存性のため誤差逆伝播時の微分鎖の中に組み込まれるが、サンプリングという処理の勾配を定義することができない。対数尤度の期待値の計算に工夫が必要である。

対数尤度の期待値の計算

　計算したい式は次式である。 $$\begin{equation} E_{q_{\phi}(\vec{z}|X,Y)}\left[\ln{p(Y|X,\vec{z})}\right]=\int d\vec{z}\;q_{\phi}(\vec{z}|X,Y)\ln{p(Y|X,\vec{z})} \end{equation}$$ この式に再パラメータ化トリック（re-parametrization trick）を適用する。すなわち $$\begin{equation} \vec{z}\sim\mathcal{N}(\vec{z}|\vec{\mu}_{\phi}(X,Y),\Sigma_{\phi}(X,Y)) \end{equation}$$ の代わりに $$\begin{eqnarray} \vec{\epsilon}&\sim&\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\\ \vec{z}&=& \vec{\mu}_{\phi}(X,Y)+\Sigma_{\phi}^{1/2}(X,Y)\vec{\epsilon} \label{eq7} \end{eqnarray}$$ を用いてサンプリングを行う。これを用いて期待値を書き直すと次式を得る。 $$\begin{equation} E_{q_{\phi}(\vec{z}|X,Y)}\left[\ln{p(Y|X,\vec{z})}\right]=\int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\ln{p(Y|X, \vec{z}=\vec{\mu}_{\phi}(X,Y)+\Sigma_{\phi}^{1/2}(X,Y)\vec{\epsilon})} \label{eq11} \end{equation}$$ 処理の流れは以下のように変更される。
上図であれば誤差逆伝播が可能となる。

未観測データの生成

　未観測データ$\vec{y}_{\alpha}$を生成する確率分布は次式で与えられた。 $$\begin{equation} p(\vec{y}_{\alpha}|\vec{x}_{\alpha},X,Y)=\int d\vec{z}\;p(\vec{y}_{\alpha}|\vec{x}_{\alpha},\vec{z})p(\vec{z}|X,Y) \end{equation}$$ 事後確率$p(\vec{z}|X,Y)$の近似解$q_{\phi}(\vec{z}|X,Y)$を用いると $$\begin{equation} p(\vec{y}_{\alpha}|\vec{x}_{\alpha},X,Y)\approx\int d\vec{z}q_{\phi}(\vec{z}|X,Y)p(\vec{y}_{\alpha}|\vec{x}_{\alpha},\vec{z}) \end{equation}$$ を得る。先と同様に再パラメータ化トリックを適用すると $$\begin{equation} p(\vec{y}_{\alpha}|\vec{x}_{\alpha},X,Y)\approx\int d\vec{\epsilon}\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)p(\vec{y}_{\alpha}|\vec{x}_{\alpha}, \vec{z}=\vec{\mu}_{\phi}(X,Y)+\Sigma_{\phi}^{1/2}(X,Y)\vec{\epsilon}) \label{eq10} \end{equation}$$ となる。

実装に向けた詳細な計算

　最初に$\vec{\mu}_{\phi}(X,Y)$と$\Sigma_{\phi}(X,Y)$を次のように置く。 $$\begin{eqnarray} \vec{\mu}_{\phi}(X,Y)&=&(\mu_{\phi,1}(X,Y),\cdots,\mu_{\phi,D}(X,Y))^T \label{eq5}\\ \Sigma_{\phi}(X,Y)&=&\mathrm{diag}(\sigma^2_{\phi,1}(X,Y),\cdots,\sigma^2_{\phi,D}(X,Y)) \label{eq6} \end{eqnarray}$$ このとき式($\ref{eq4}$)は次式となる。 $$\begin{equation} D_{KL} \left[ q_{\phi}(\vec{z}|X,Y)||p(\vec{z}) \right]= \frac{1}{2} \sum_{d=1}^{D}\left\{ -\ln{\sigma^2_{\phi,d}(X,Y)}-1+\sigma^2_{\phi,d}(X,Y)+\mu_{\phi,d}^2(X,Y) \right\} \label{eq8} \end{equation}$$ また、$\vec{z}$の成分は次式で与えられる。 $$\begin{equation} z_d=\mu_{\phi,d}(X,Y)+\sigma_{\phi,d}(X,Y)\epsilon_d \end{equation}$$ 観測値が独立同分布に従うと仮定すると、式($\ref{eq11}$)の中にある対数尤度は以下のように変形される。 $$\begin{equation} \ln{p(Y|X,\vec{z})}= \sum^{N}_{n=1}\ln{p(\vec{y}_n|\vec{x}_n,\vec{z})} \end{equation}$$ さらに計算を進めるには、具体的に$X,Y$として、何を与えるか決定しなければならない。 ここでは、$X$として0から9までのラベルを、$Y$としてMNISTの画像（２値画像）を与えることにする。$X$の各観測値$\vec{x}_n$は$9$次元のone-hotベクトルで表現される。各画素が独立同分布に従うと仮定すると、$\vec{y}_n$の次元数を$M$として $$\begin{equation} \ln{p(\vec{y}_n|\vec{x}_n,\vec{z})}=\sum_{m=1}^{M}\ln{p(y_{n,m}|\vec{x}_n,\vec{z})} \end{equation}$$ と書くことができる。いま考える画像は0と1から構成されるから、$p(y_{n,m}|\vec{x}_n,\vec{z})$として0と1を生成するBernoulli分布を仮定する。 $$\begin{eqnarray} p(y_{n,m}|\vec{x}_n,\vec{z})&=&\mathrm{Bern}\left(y_{n,m}|\eta_{\theta,m}\left(\vec{x}_n,\vec{z}\right)\right) \\ \mathrm{Bern}(x|\eta)&=&\eta^{x}(1-\eta)^{1-x} \end{eqnarray}$$ $\eta_{\theta,m}\left(\vec{x}_n,\vec{z}\right)$は、入力が$\vec{x}_n$と$\vec{z}$、パラメータとして$\theta$を持つニューラルネットワークで学習される。以上を踏まえて処理の流れを書き直すと下図となる。

次に式($\ref{eq10}$)を考える。これは、観測値$X,Y$とラベル$\vec{x}_{\alpha}$が与えられときの$\vec{y}_{\alpha}$の実現確率である。 $$\begin{equation} p(\vec{y}_{\alpha}|\vec{x}_{\alpha},X,Y)\approx \int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\;p(\vec{y}_{\alpha}|\vec{x}_{\alpha},\vec{z}) \end{equation}$$ ここで、$z_d=\mu_{\phi,d}(X)+\sigma_{\phi,d}(X)\epsilon_d$である。上式は以下のように書くことができる。 $$\begin{equation} \prod_{m=1}^M p(y_{\alpha,m}|\vec{x}_\alpha,X,Y)\approx \int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\;\prod_{m=1}^M p(y_{\alpha,m}|\vec{x}_\alpha,\vec{z}) \end{equation}$$ すなわち、要素$y_{\alpha,m}$ごとに次式が成り立つ。 $$\begin{equation} p(y_{\alpha,m}|\vec{x}_\alpha,X,Y)\approx \int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\;p(y_{\alpha,m}|\vec{x}_\alpha,\vec{z}) \end{equation}$$ $p(y_{\alpha,m}|\vec{x}_\alpha,\vec{z})$としてBernoulli分布を仮定したから $$\begin{equation} p(y_{\alpha,m}|\vec{x}_\alpha,X,Y)\approx \int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\;\mathrm{Bern}(y_{\alpha,m}|\eta_{\theta,m}(\vec{x}_\alpha,\vec{z})) \end{equation}$$ となる。確率分布$p(y_{\alpha,m}|\vec{x}_\alpha,X,Y)$の下での$y_{\alpha,m}$の期待値は $$\begin{eqnarray} <y_{\alpha,m}>&=&\sum_{y_{\alpha,m}=0,1} y_{\alpha,m}\;p(y_{\alpha,m}|\vec{x}_\alpha,X,Y) \\ &\approx& \int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\;\sum_{y_{\alpha,m}=0,1} y_{\alpha,m} \mathrm{Bern}(y_{\alpha,m}|\eta_{\theta,m}(\vec{x}_\alpha,\vec{z})) \\ &=& \int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\;\eta_{\theta,m}(\vec{x}_\alpha,\vec{z}) \end{eqnarray}$$ となる。$\vec{z}$は$\vec{\epsilon}$に依存する項であることに注意する。$\eta_{\theta,m}(\vec{x}_\alpha,\vec{z})$はDecoderの出力である。上式から、復号化した結果を得るには、$\eta_{\theta,m}(\vec{x}_\alpha,\vec{z})$を標準正規分布に従ってサンプリングすれば良いことが分かる。さらに、式($\ref{eq8}$)のKullback Leibler divergenceを十分小さくできれば、すなわち、$\sigma_{\phi,d}(X,Y)\rightarrow 1,\mu_{\phi,d}(X,Y)\rightarrow 0$とできれば、$\vec{z}=\vec{\epsilon}$とすることができるので、標準正規分布から生成した値$\vec{\epsilon}$と$\vec{x}_\alpha$からDecoderの出力を直接得ることができる。

まとめ

　今回は、CVAEをBayes推定の枠組みで説明した。前回のVAEの論法とほとんど同じである。VAEでは未観測データ$\vec{x}$が従う確率分布$p(\vec{x}|X)$を求める過程でVAEの構造を見出した。一方、CVAEでは未観測データ$\vec{x}$に対応する$\vec{y}$が従う確率分布$p(\vec{y}|\vec{x},X,Y)$を求める過程でCVAEの構造が現れることを見た。その構造は、VAEに少し手を加えれば実現できる程度のものである。ChainerのVAEのサンプルコードをベースにすればすぐに実装できそうである。

Variational Auto Encoder 〜その2〜

はじめに

　先のページで、Variational Auto Encoder（VAE）をBayes推論の枠組みで解説し、Chainerのサンプルコードをみた。今回は、サンプルコードを実際に動かし、その動作を確認する。

前回の補足

　前回の式(16)は、観測値$X$の下での未知変数$\vec{x}$の実現確率であった。 $$\begin{equation} p(\vec{x}|X)\approx \int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\;p(\vec{x}|\vec{z}) \end{equation}$$ ここで、$z_d=\mu_{\phi,d}(X)+\sigma_{\phi,d}(X)\epsilon_d$である。$\vec{x}$の各成分が独立同分布で生成されると仮定すると、以下のように書き換えることができる。 $$\begin{equation} \prod_{m=1}^M p(x_{m}|X)\approx \int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\;\prod_{m=1}^M p(x_m|\vec{z}) \end{equation}$$ すなわち、要素$x_m$ごとに次式が成り立つ。 $$\begin{equation} p(x_{m}|X)\approx \int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\;p(x_m|\vec{z}) \end{equation}$$ $p(x_m|\vec{z})$としてBernoulli分布を仮定したから $$\begin{equation} p(x_{m}|X)\approx \int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\;\mathrm{Bern}(x_m|\eta_{\theta,m}(\vec{z})) \end{equation}$$ となる。確率分布$p(x_{m}|X)$の下での$x_m$の期待値は $$\begin{eqnarray} <x_m>&=&\sum_{x_m=0,1} x_m\;p(x_{m}|X) \\ &\approx& \int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\;\sum_{x_m=0,1} x_m \mathrm{Bern}(x_m|\eta_{\theta,m}(\vec{z})) \\ &=& \int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\;\eta_{\theta,m}(\vec{z}) \end{eqnarray}$$ となる。$\vec{z}$は$\vec{\epsilon}$に依存する項であることに注意する。ここまでの議論で言えることは、復号化した結果を得るには、$\eta_{\theta,m}(\vec{z})$を標準正規分布に従ってサンプリングすれば良いということである。$\eta_{\theta,m}(\vec{z})$はDecoderの出力である。

Chainerの実装の確認

　前回は、net.pyを見たので、今回はtrain_vae.pyを見る。trainerを用いた実装部分は特に指摘することはないので、結果を描画している箇所を解説する。最初は訓練データに関わる描画部分である。 適当に画像を9枚選択し、これを関数__call__で符号化・復号化している。前回指摘したように、復号化の際、$\sigma_{\phi,d}(X)$は無視されている。計算のあと元画像と復号化画像を保存している。epochを100とした結果は以下の通りである。

訓練画像


復号化した訓練画像

次はテスト画像に関わる部分である。 ここでも適当に9枚の画像を選択し、元画像と復号化画像を保存している。epochを100とした結果は以下の通りである。

テスト画像


復号化したテスト画像

最後に、標準正規分布に従う値$\vec{z}$から復号化する部分である。 9個の乱数$\vec{z}$を作り、関数decodeを呼び出している。epochを100とした結果は以下の通りである。
訓練・テスト画像を符号化・復号化した結果と比べるとかなり精度の悪い結果である。関数decodeは$\eta_{\theta,m}(\vec{z})$を出力する。上で見たように本来の$\vec{z}$は、$z_d=\mu_{\phi,d}(X)+\sigma_{\phi,d}(X)\epsilon_d$として与えらるべきものである。精度が悪いのは、$\vec{z}$を標準正規分布に置き換えたことが原因である。ところで、勾配降下法で最小にすべき式の1つが次式であった（前回の式(19)）。 $$\begin{equation} D_{KL} \left[ q_{\phi}(\vec{z}|X)||p(\vec{z}) \right]= \frac{1}{2} \sum_{d=1}^{D}\left\{ -\ln{\sigma^2_{\phi,d}(X)}-1+\sigma^2_{\phi,d}(X)+\mu_{\phi,d}^2(X) \right\} \end{equation}$$ これを十分小さくできれば、すなわち、$\sigma_{\phi,d}\rightarrow 1, \mu_{\phi,d}\rightarrow 0$とできれば、標準正規分布による置き換えは意味があるものになる。残念ながらepochを1000としても大して精度は良くならない。Chainerのサンプル実装を変更する必要があるかもしれない。

まとめ

　今回は、$<x_m>$が$\eta_{\theta,m}(\vec{z})$を求めることに帰着すること示し、Chainerのサンプルコードの計算結果を考察した。 さらに、標準正規分布による$\vec{z}$から計算した画像の精度が悪い理由についても述べた。改善するには、epoch数を増やすだけではなくコードの見直し（多層化、初期化関数の変更）も必要であろう。次回はこの辺りのことについてまとめたい。

Variational Auto Encoder

はじめに

　Variational Auto Encoder（VAE）をBayes推論の枠組みで解説し、Chainerのサンプルコードを読解する。

問題設定

　観測値$X=\{\vec{x}_1,\cdots,\vec{x}_N\}$が与えられたとき、未知の値$\vec{x}_*$を生成する確率分布$p(\vec{x}_*|X)$を求めたい。潜在変数$\vec{z}$を導入し、$X$と$\vec{z}$の同時確率分布$p(X,\vec{z})$を考え、Bayesの定理を適用すると次式を得る。 $$\begin{equation} p(\vec{z}|X) = \frac{p(X|\vec{z})p(\vec{z})}{p(X)} \label{eq9} \end{equation}$$ 事後確率$p(\vec{z}|X)$が求まれば、次式により$\vec{x}_*$を生成する確率分布を求めることができる。 $$\begin{equation} p(\vec{x}_*|X)=\int d\vec{z}\;p(\vec{x}_*|\vec{z})p(\vec{z}|X) \label{eq3} \end{equation}$$ 事後確率$p(\vec{z}|X)$を求めることが目的である。

最適化すべき量

　$p(\vec{z}|X)$を直接求めることはせず、パラメータ$\phi$を持つ関数$q_{\phi}(\vec{z}|X)$を導入し、次のKullback Leibler divergenceを最小にすることを考える。 $$\begin{equation} D_{KL} \left[ q_{\phi}(\vec{z}|X)||p(\vec{z}|X) \right]=\int d\vec{z}\;q_{\phi}(\vec{z}|X) \ln{ \frac{ q_{\phi}(\vec{z}|X) } { p(\vec{z}|X) } } \end{equation}$$ これを変形すると次式を得る。 $$\begin{equation} D_{KL} \left[ q_{\phi}(\vec{z}|X)||p(\vec{z}|X) \right] = D_{KL} \left[ q_{\phi}(\vec{z}|X)||p(\vec{z}) \right]-E_{q_{\phi}(\vec{z}|X)}\left[\ln{p(X|\vec{z})}\right]+\ln{p(X)} \label{eq1} \end{equation}$$ ただし、式変形の途中で式($\ref{eq9}$)を用いた。式($\ref{eq1}$)右辺にある$\ln{p(X)}$は$\phi$に依存せず、観測値だけから決まる定数である。従って、次式が成り立つ。 $$\begin{equation} \min_{\phi} D_{KL} \left[ q_{\phi}(\vec{z}|X)||p(\vec{z}|X) \right] = \min_{\phi} {\left[ D_{KL} \left[ q_{\phi}(\vec{z}|X)||p(\vec{z}) \right]-E_{q_{\phi}(\vec{z}|X)}\left[\ln{p(X|\vec{z})}\right] \right] } \label{eq2} \end{equation}$$ 式($\ref{eq2}$)の右辺第1項を小さく、第2項の期待値を大きくすれば良い。第1項は$q_{\phi}(\vec{z}|X)$をできるだけ$p(\vec{z})$に近い形の分布にすることを要求し、この分布の下で対数尤度$\ln{p(X|\vec{z})}$の期待値を大きくすることを第2項は要求する。第1項は正則化項に相当する。

KL divergenceの計算

　式($\ref{eq2}$)の右辺第1項を考える。いま次の仮定をおく。 $$\begin{eqnarray} q_{\phi}(\vec{z}|X)&=&\mathcal{N}(\vec{z}|\vec{\mu}_{\phi}(X),\Sigma_{\phi}(X)) \\ p(\vec{z})&=&\mathcal{N}(\vec{z}|\vec{0},I_D) \end{eqnarray}$$ ここで、$\vec{z}$の次元を$D$とした。$I_D$は$D\times D$の単位行列である。どちらの分布も正規分布とし、$q_{\phi}(\vec{z}|X)$の平均と共分散行列は$\phi$と$X$から決まる量とする。これらは、入力$X$、パラメータ$\phi$のニューラルネットワークを用いて計算される。一方、$p(\vec{z})$の方は平均0、分散1の標準正規分布である。このとき、$D_{KL} \left[ q_{\phi}(\vec{z}|X)||p(\vec{z}) \right]$は解析的に計算することができる。 $$\begin{equation} D_{KL} \left[ q_{\phi}(\vec{z}|X)||p(\vec{z}) \right]=\frac{1}{2}\left[ -\ln{|\Sigma_{\phi}(X)|} -D +\mathrm{Tr}\left(\Sigma_{\phi}(X)\right)+\vec{\mu}_{\phi}(X)^T\vec{\mu}_{\phi}(X) \right] \label{eq4} \end{equation}$$

ここまでの処理の流れ

　式($\ref{eq2}$)の最適化を行う際に勾配降下法を用いる。処理の流れは以下のようになる（下図参照）。
分布$q_{\phi}(\vec{z}|X)$は$X$から$\vec{z}$を生成するEncoder、$p(X|\vec{z})$は$\vec{z}$から$X$を生成するDecoderとみなすことができる。青色で示した部分は最小化すべき量であり、赤字はサンプリングするステップである。勾配降下法を実現するには、誤差逆伝播ができなければならない。$q_{\phi}(\vec{z}|X)$はその$\phi$依存性のため誤差逆伝播時の微分鎖の中に組み込まれるが、サンプリグという処理の勾配を定義することができない。対数尤度の期待値の計算に工夫が必要である。

対数尤度の期待値の計算

　計算したい式は次式である。 $$\begin{equation} E_{q_{\phi}(\vec{z}|X)}\left[\ln{p(X|\vec{z})}\right]=\int d\vec{z}\;q_{\phi}(\vec{z}|X)\ln{p(X|\vec{z})} \end{equation}$$ この式に再パラメータ化トリック（re-parametrization trick）を適用する。すなわち $$\begin{equation} \vec{z}\sim\mathcal{N}(\vec{z}|\vec{\mu}_{\phi}(X),\Sigma_{\phi}(X)) \end{equation}$$ の代わりに $$\begin{eqnarray} \vec{\epsilon}&\sim&\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\\ \vec{z}&=& \vec{\mu}_{\phi}(X)+\Sigma_{\phi}^{1/2}(X)\vec{\epsilon} \label{eq7} \end{eqnarray}$$ を用いてサンプリングを行う。これを用いて期待値を書き直すと次式を得る。 $$\begin{equation} E_{q_{\phi}(\vec{z}|X)}\left[\ln{p(X|\vec{z})}\right]=\int d\vec{\epsilon}\;\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)\ln{p(X| \vec{z}=\vec{\mu}_{\phi}(X)+\Sigma_{\phi}^{1/2}(X)\vec{\epsilon})} \end{equation}$$ このときの処理の流れは以下のようになる（下図参照）。
上図であれば誤差逆伝播が可能となる。

未知変数の生成

　未知変数を生成する確率分布は次式で与えられた。 $$\begin{equation} p(\vec{x}_*|X)=\int d\vec{z}\;p(\vec{x}_*|\vec{z})p(\vec{z}|X) \end{equation}$$ 事後確率$p(\vec{z}|X)$の近似解$q_{\phi}(\vec{z}|X)$を用いると $$\begin{equation} p(\vec{x}_*|X)\approx\int d\vec{z}q_{\phi}(\vec{z}|X)p(\vec{x}_*|\vec{z}) \end{equation}$$ を得る。先と同様に再パラメータ化トリックを適用すると $$\begin{equation} p(\vec{x}_*|X)\approx\int d\vec{\epsilon}\mathcal{N}(\vec{\epsilon}|\vec{0},I_D)p(\vec{x}_*| \vec{z}=\vec{\mu}_{\phi}(X)+\Sigma_{\phi}^{1/2}(X)\vec{\epsilon}) \end{equation}$$ となる。

Chainer実装の確認

　ChainerのサンプルコードにVAEがある。これはMNISTデータセットにVAEを適用したものである。MNISTは2値画像であるから$\vec{x}_n$として0と1が784(=28$\times$28)個並んだベクトルを考えることになる。実際にコードを見て行く前に先に導出した式をもう少し詳細に計算しておく。

　最初に$\vec{\mu}_{\phi}(X)$と$\Sigma_{\phi}(X)$を次のように置く。 $$\begin{eqnarray} \vec{\mu}_{\phi}(X)&=&(\mu_{\phi,1}(X),\cdots,\mu_{\phi,D}(X))^T \label{eq5}\\ \Sigma_{\phi}(X)&=&\mathrm{diag}(\sigma^2_{\phi,1}(X),\cdots,\sigma^2_{\phi,D}(X)) \label{eq6} \end{eqnarray}$$ このとき式($\ref{eq4}$)は次式となる。 $$\begin{equation} D_{KL} \left[ q_{\phi}(\vec{z}|X)||p(\vec{z}) \right]= \frac{1}{2} \sum_{d=1}^{D}\left\{ -\ln{\sigma^2_{\phi,d}(X)}-1+\sigma^2_{\phi,d}(X)+\mu_{\phi,d}^2(X) \right\} \label{eq8} \end{equation}$$ また、$\vec{z}$の成分は次式で与えられる。 $$\begin{equation} z_d=\mu_{\phi,d}(X)+\sigma_{\phi,d}(X)\epsilon_d \end{equation}$$ 観測値が独立同分布に従うと仮定すると対数尤度は以下のように変形される。 $$\begin{equation} \ln{p(X|\vec{z})}= \sum^{N}_{n=1}\ln{p(\vec{x}_n|\vec{z})} \end{equation}$$ $\vec{x}_n$の次元数を$M$(=784)とすると $$\begin{equation} \ln{p(\vec{x}_n|\vec{z})}=\sum_{m=1}^{M}\ln{p(x_{n,m}|\vec{z})} \end{equation}$$ となる。いま考える画像は0と1から構成される。従って、$p(x_{n,m}|\vec{z})$として0と1を生成する確率分布であるBernoulli分布を仮定する。 $$\begin{eqnarray} p(x_{n,m}|\vec{z})&=&\mathrm{Bern}\left(x_{n,m}|\eta_{\theta,m}\left(\vec{z}\right)\right) \\ \mathrm{Bern}(x|\eta)&=&\eta^{x}(1-\eta)^{1-x} \end{eqnarray}$$ $\eta_{\theta,m}\left(\vec{z}\right)$は、入力$\vec{z}$、パラメータ$\theta$のネットワークで学習される量である。 以上を踏まえて処理の流れを書き直すと下図となる。

以上で準備が整ったので順にコードを見ていく。まず最初にネットワークを定義したクラスVAEを見る。コンストラクタは以下の通り。 Encoderとして$\vec{\mu}_{\phi}(X)$と$\Sigma_{\phi}(X)$を生成する層がそれぞれ1層ずつ定義されている。Decoderとして$\vec{\eta}_{\theta}(\vec{z})$を生成する2層が定義されている。次に関数encodeを見る。 $\mu_{\phi,d}$と$\ln{\sigma_{\phi,d}^2}$を生成する処理が記述されている。次は関数decodeである。 $\vec{z}$を受け取り$\vec{\eta}_{\theta}(\vec{z})$を返す処理が記述されている。次は__call__である。 Encoderで計算した平均を入力としてDecoderを呼び出している。分散を無視して復号化している。次はget_loss_funcである。

13行目：Encoderで平均と対数分散を計算する。

17行目：ここから始まるループは$\vec{z}$のサンプリングのためのものである。

18行目：正規分布からサンプリングする。関数F.gaussianの中で再パラメータ化トリックが実行されていることに注意する。

19行目：Bernoulli分布の対数にマイナスを付けたものが計算される。

23行目：式($\ref{eq8}$)が計算される。

まとめ

　今回は、VAEをBayes推論の枠組みで解説し、Chainerのサンプルコードを見た。ニューラルネットワークで計算されるのは確率分布のパラメータであることを明確に示した。言い換えると他の手法でパラメータを計算できるのであればそれでも構わないということである。計算の仮定で確率分布にいくつかの仮定（ガウス分布やBernoulli分布）をしていることに注意しなければならない。今回Bernoulli分布を使用したのはターゲットとした観測値が0と1から構成される2値画像であるためである。2値でない観測値を対象とするならそれに見合った確率分布を導入する必要がある。次回はChainerのコードを動かして得られる結果を考察したい。

参考文献

Tutorial on Variational Autoencoders

Pattern Recognition and Machine Learning

ベイズ推論による機械学習入門

Variational Autoencoder徹底解説