Adversarial Variational Bayes 〜理論（その2）〜

はじめに

　前回の投稿では、論文Adversarial Variational Bayesの前半部分をまとめた。今回は後半部分をまとめる。

前半部分の結論

　従来のVAEの目的関数は $$\begin{equation} (\theta^{*},\phi^{*})=\arg\max_{\theta,\phi} {\rm E}_{p_{d}(x)}\left[ -{\rm KL} \left( q_{\phi}(z|x), p(z) \right) + {\rm E}_{q_{\phi}(z|x)} \left[\ln{p_{\theta}(x|z)}\right] \right] \label{eq1} \end{equation}$$ である。本論文の前半部分では式($\ref{eq1}$)から次の2つの目的関数を導出した。 $$\begin{eqnarray} T^{*}(x,z)&=&\arg{\max_{T}}\; {\rm E}_{p_{d}(x)} \left[ {\rm E}_{q_{\phi}(z|x)} \left[ \ln{\sigma\left(T(x, z)\right)} \right] + {\rm E}_{p(z)} \left[ \ln{\left(1-\sigma\left(T(x, z)\right)\right)} \right] \right] \label{eq2}\\ (\theta^{*},\phi^{*})&=&\arg\max_{\theta,\phi} {\rm E}_{p_{d}(x)}\;{\rm E}_{q_{\phi}(z|x)} \left[ \ln{p_{\theta}(x|z)}-T^{*}(x,z) \right] \label{eq3} \end{eqnarray}$$ ここで、$\sigma(\cdot)$はシグモイド関数である。実際に計算を行うと、式($\ref{eq2}$)から最適な値$T^*(x,z)$を求めることは難しいことが分かる。その原因は、式($\ref{eq1}$)の右辺にあるKullback Leibler divergence にある。この量は、事後分布$q_{\phi}(z|x)$の形状を事前分布$p(z)$の形状に近づけようとする。一般にこれら2つの分布は大きく異なる形状を持つため、$q_{\phi}(z|x)\approx p(z)$の実現は難しい。また、事後分布を事前分布に近づける操作は、事後分布から観測値$x$の依存性を失くすことに相当するため、Bayes推論の立場から考えても妥当ではない。本論文の後半部分ではこの操作を改善する。

後半部分の内容

　事後分布$q_{\phi}(z|x)$に良く似た形状を持ち、かつ、解析計算が可能な補助分布（auxiliary distribution）$r_{\alpha}(z|x)$を導入する。例えば、$r_{\alpha}(z|x)$として正規分布を考えることができる。この補助分布を用いて、式($\ref{eq1}$)の右辺を次のように書き換える。 $$\begin{equation} {\rm E}_{p_{d}(x)}\left[ -{\rm KL} \left( q_{\phi}(z|x), r_{\alpha}(z|x) \right) + {\rm E}_{q_{\phi}(z|x)} \left[ -\ln{r_{\alpha}(z|x)} +\ln{p_{\theta}(x,z)}\right] \right] \label{eq4} \end{equation}$$ $r_{\alpha}(z|x)$を$q_{\phi}(z|x)$を近似する関数と置いたので、式($\ref{eq4}$)に現れるKullback Leibler divergenceの値は、式($\ref{eq1}$)のそれよりもずっと容易に小さな値とすることができる。

　いま、補助分布として次の正規分布を仮定する。 $$\begin{equation} r_{\alpha}(z|x)=\mathcal{N}(z|\mu(x),\sigma_s^2(x)) \end{equation}$$ これは、変数変換 $$\begin{equation} \tilde{z}=\frac{z-\mu(x)}{\sigma_s(x)} \end{equation}$$ の下で $$\begin{equation} r_{\alpha}(z|x)=\frac{1}{\sigma_s}\mathcal{N}(\tilde{z}|0,1)\equiv \frac{1}{\sigma_s}r_0(\tilde{z}) \end{equation}$$ となるから $$\begin{eqnarray} {\rm KL} \left( q_{\phi}(z|x), r_{\alpha}(z|x) \right) &=& \int dz\;q_{\phi}(z|x)\ln{\frac{q_{\phi}(z|x)}{r_{\alpha}(z|x)}} \label{eq5} \\ &=& \int d\tilde{z}\;\sigma_s q_{\phi}(\sigma_s \tilde{z}+\mu|x) \ln{ \frac {\sigma_s q_{\phi}(\sigma_s\tilde{z}+\mu|x)} {r_0(\tilde{z})} }\\ &=& {\rm KL} \left( \tilde{q}_{\phi}(\tilde{z}|x), r_0(\tilde{z}) \right) \label{eq6} \end{eqnarray}$$ となる。ただし、$\tilde{q}_{\phi}(\tilde{z}|x)\equiv \sigma_s q_{\phi}(\sigma_s\tilde{z}+\mu|x)$とした。$\tilde{q}_{\phi}(\tilde{z}|x)$は $$\begin{equation} \int d\tilde{z}\tilde{q}_{\phi}(\tilde{z}|x)=1 \end{equation}$$ を満たす新たな分布である。さらに、次式が成り立つ。 $$\begin{eqnarray} {\rm E}_{q_{\phi}(z|x)} \left[ \ln{r_{\alpha}(z|x)} \right] &=& \int dz\;q_{\phi}(z|x)\ln{r_{\alpha}(z|x)} \\ &=& \int d\tilde{z}\;\tilde{q}_{\phi}(\tilde{z}|x)\ln{r_{0}(\tilde{z})}-\ln{\sigma_s} \\ &=& {\rm E}_{\tilde{q}_{\phi}(\tilde{z}|x)} \left[ \ln{r_{0}(\tilde{z})} \right]+{\rm const.} \end{eqnarray}$$ 以上から、式($\ref{eq4}$)は、定数項を無視して、次式に変換される。 $$\begin{equation} {\rm E}_{p_{d}(x)}\left[ -{\rm KL} \left( \tilde{q}_{\phi}(\tilde{z}|x), r_{0}(\tilde{z}) \right) + {\rm E}_{\tilde{q}_{\phi}(\tilde{z}|x)} \left[ -\ln{r_{0}(\tilde{z})} \right]+ {\rm E}_{q_{\phi}(z|x)} \left[ \ln{p_{\theta}(x,z)}\right] \right] \label{eq7} \end{equation}$$ ここで、$-{\rm KL} \left( \tilde{q}_{\phi}(\tilde{z}|x), r_{0}(\tilde{z}) \right)$に注目して、論文の前半部分で用いた議論を繰り返し、$T(x,\tilde{z})$を導入すると、式($\ref{eq7}$)は次の2式と等価となる。 $$\begin{eqnarray} T^{*}(x,\tilde{z})&=&\arg{\max_{T}} \;{\rm E}_{p_{d}(x)} \left[ {\rm E}_{\tilde{q}_{\phi}(\tilde{z}|x)} \left[ \ln{\sigma\left(T(x, \tilde{z})\right)} \right] + {\rm E}_{r_0(\tilde{z})} \left[ \ln{\left(1-\sigma\left(T(x, \tilde{z})\right)\right)} \right] \right] \label{eq8}\\ (\theta^{*},\phi^{*}) &=& \arg\max_{\theta,\phi} \;{\rm E}_{p_{d}(x)} \left[ {\rm E}_{q_{\phi}(z|x)} \left[ \ln{p_{\theta}(x,z)} \right] + {\rm E}_{\tilde{q}_{\phi}(\tilde{z}|x)} \left[ -T^*(x,\tilde{z})-\ln{r_0(\tilde{z})} \right] \right] \label{eq9} \end{eqnarray}$$ これらに、再パラメータ化トリックを適用する。すなわち、 $$\begin{equation} z\sim q_{\phi}(z|x) \end{equation}$$ を $$\begin{eqnarray} \epsilon&\sim&p(\epsilon)\\ z&=&z_{\phi}(x,\epsilon) \end{eqnarray}$$ に置き換える。このとき $$\begin{equation} \tilde{z}\sim \tilde{q}_{\phi}(\tilde{z}|x) \end{equation}$$ は $$\begin{eqnarray} \tilde{z}&=&\frac{z_{\phi}(x,\epsilon)-\mu}{\sigma_s}\equiv\tilde{z}_{\phi}(x,\epsilon) \end{eqnarray}$$ に置き換わる。また、分布$r_0(\tilde{z})$は標準正規分布であるから $$\begin{equation} \tilde{z} \sim r_0(\tilde{z}) \end{equation}$$ は $$\begin{eqnarray} \eta&\sim&\mathcal{N}(\eta|0,1)\equiv p(\eta) \\ \tilde{z}&=&\eta \end{eqnarray}$$ と書くことができる。以上から式($\ref{eq8}$),($\ref{eq9}$)は次式に変形される。 $$\begin{eqnarray} T^{*}(x,\tilde{z}) &=& \arg{\max_{T}} \;{ \rm E}_{p_{d}(x)} \Bigl[ {\rm E}_{p(\epsilon)} \left[ \ln{\sigma\left(T(x, \tilde{z}_{\phi}(x,\epsilon))\right)} \right] + {\rm E}_{p(\eta)} \left[ \ln{\left(1-\sigma\left(T(x, \eta)\right)\right)} \right] \Bigr] \label{eq10}\\ (\theta^{*},\phi^{*}) &=& \arg\max_{\theta,\phi} \;{\rm E}_{p_{d}(x)} \Bigl[ {\rm E}_{p(\epsilon)} \left[ \ln{p_{\theta}(x,z_{\phi}(x,\epsilon))} \right] + {\rm E}_{p(\epsilon)} \left[ -T^*(x,\tilde{z}_{\phi}(x,\epsilon))-\ln{r_0(\tilde{z}_{\phi}(x,\epsilon))} \right] \Bigr] \label{eq11} \end{eqnarray}$$ $r_0$は標準正規分布であるから、式($\ref{eq11}$)の最後の項は $$\begin{equation} -\ln{r_0(\tilde{z}_{\phi}(x,\epsilon))}=\frac{1}{2}\|\tilde{z}_{\phi}(x,\epsilon)\|^2+{\rm const.} \end{equation}$$ となる。最適化を行う際は定数項は無視すれば良い。また、$p_{\theta}(x,z)$の計算は $$\begin{equation} p_{\theta}(x,z)=p_{\theta}(x|z)p(z) \end{equation}$$ を利用すれば良い。ここまでの議論を反映したアルゴリズムが論文に掲載されている。