らんだむな記憶

blogというものを体験してみようか!的なー

VAE (Variational Autoencoder) 論文を読んでみる

VAE (Variational Autoencoder) - らんだむな記憶を少し読んでみようかな。読んで見ようかなと書きつつも分からないところを省いたただの翻訳という説・・・。

たぶん以下で $p_\theta (x)$ みたいなのは例えば Bernoulli 分布の密度“函数” $p_\theta (x) = \theta^x (1-\theta)^{1-x}$ だとか、正規分布の密度函数 $p_{\mu,\sigma} (x) = \frac{1}{\sqrt{2\pi \sigma^2}} \exp \left\{ - \frac{(x - \mu)^2}{2 \sigma^2} \right\}$ のようなものをぼんやり念頭に置いておけば良いような気がする。

・・・

$x$ という連続あるいは離散的な変数がある。この変数について $N$ 個の独立同時分布なデータセット $X = \{x^{(i)}\}_{i=1}^N$ を考える。データはある確率過程によって、確認されないような連続確率変数 $z$ を伴いつつ生成される。この確率過程は 2 つのステップからなる:

(1) $z^{(i)}$ はある事前分布 $p_{\theta^*}(z)$ から生成される
(2) $x^{(i)}$ はある条件付き確率分布 $p_{\theta^*}(x|z)$ から生成される

仮定

事前分布 $p_{\theta^*}(z)$ と尤度 $p_{\theta^*}(x|z)$ は分布のパラメータ族 $p_\theta (z)$ と $p_\theta (x|z)$ から生じ、それらの PDFs (Probability Density Function; 確率密度函数) は $\theta$ と $z$ に関してほとんど至るところ微分可能であるとする。
残念ながら確率過程の大部分は我々は観測できない: 潜在変数 $z^{(i)}$ のみならず真のパラメータ $\theta^*$ も分からない。

重要なことだが、周辺確率や事後確率について単純化するような仮定はおかない。逆に、以下のような事情のもとでさえも効果的に機能するような一般的なアルゴリズムに興味がある:

  1. 重要なツールが使えないケース: 周辺尤度の積分 $\int p_\theta(z) p_\theta (x|z) dz$ が扱えない (したがって周辺尤度の評価や微分ができない)、真の事後密度 $p_\theta (z|x) = p_\theta (x|z) p_\theta(z) / p_\theta(x)$ が扱えない (したがって EM アルゴリズムが適用できない)、そして妥当な平均場近似・変分ベイズ法のアルゴリズムで必要な積分が扱えないケース。
  2. 巨大なデータセット

興味のある問題

  1. パラメータ $\theta$ に対して ML (最大尤度) あるいは MAP (最大事後確率) 評価を効率的に近似する。
    → パラメータ近似によって隠れた確率過程を模倣したり真のデータの模倣品を人工的に生成できる。
  2. パラメータ $\theta$ の選択に対して観測される値 $x$ が与えられた時に潜在変数 $z$ の事後予測を効率的に近似する。
    → 符号化やデータ表現のタスクで便利。
  3. $x$ の周辺予測を効率的に近似する。
    → $x$ に対する事前確率が必要な条件下で推測系のタスクが可能になる。ノイズ除去, 傷の除去や超解像を含むコンピュータビジョンに応用できる。

上記の問題をとくために、認識モデル $q_\phi (z|x)$ を導入する: これは直接は扱えない真の事後分布 $p_\theta (z|x)$ の近似である。生成モデルのパラメータ $\theta$ と一緒に認識モデルのパラメータ $\phi$ を学習する方法を導入する。

符号理論の観点からは、観測できない変数 $z$ は潜在表現あるいは符号としての解釈ができる。このため、同論文では認識モデル $q_\phi(z|x)$ を確率的エンコーダと見る。その心としてはデータポイント $x$ が与えられた時、それを生成する元になるであろうところの符号 $z$ の取り得る値全域にわたる分布 (e.g. ガウス分布) がつくられるからである。同様の気分で $p_\theta (x|z)$ を確率的デコーダと見る。その心は符号 $z$ が与えられた時に対応する可能性のある $x$ 全体にわたる分布がつくられるからである。

変分下限

周辺尤度は個別のデータポイントの周辺尤度上にわたる和 $\log p_\theta (x^{(1)}, \cdots, x^{(N)}) = \sum_{i=1}^N \log p_\theta (x^{(i)})$ から成る。この右辺の個々は
\begin{align}
\log p_\theta (x^{(i)}) = D_{KL} (q_\phi (z|x^{(i)}) \| p_\theta (z|x^{(i)})) + \mathcal{L}(\theta,\phi; x^{(i)}) \quad\quad\quad\quad\quad\quad (1)
\end{align}

と書き直すことができる。
右辺第一項は真の事後分布からの近似のカルバック・ライブラー・ダイバージェンスである。この KL ダイバージェンスは非負であるので、右辺第二項 $\mathcal{L}(\theta,\phi; x^{(i)})$ はデータポイント $i$ の周辺尤度上の (変分) 下限と呼ばれる。上式を書き直すと以下のようになる:
\begin{align}
\log p_\theta (x^{(i)}) \geq \mathcal{L}(\theta,\phi; x^{(i)}) = \mathbb{E}_{q_\phi (z|x^{(i)})} [ - \log q_\phi (z|x^{(i)}) + \log p_\theta (x^{(i)}, z) ] \quad (2)
\end{align}

これはJensenの不等式 - らんだむな記憶により
\begin{align}
\log p_\theta (x) = \log \int p_\theta (x,z) dz &= \log \int q_\phi(z|x) \frac{p_\theta (x,z)}{q_\phi(z|x)} dz \\
&= \log \mathbb{E}_{q_\phi (z|x)} \left[\frac{p_\theta (x,z)}{q_\phi(z|x)} \right] \geq \mathbb{E}_{q_\phi (z|x)} \left[\log \frac{p_\theta (x,z)}{q_\phi(z|x)} \right]
\end{align}

であることから従う。また
\begin{align}
\mathcal{L}(\theta,\phi; x^{(i)}) = - D_{KL} (q_\phi (z|x^{(i)}) \| p_\theta (z)) + \mathbb{E}_{q_\phi (z|x^{(i)})} \left[ \log p_\theta (x^{(i)}|z) \right] \quad\quad (3)
\end{align}

と書き直すことができる。これは $\mathcal{L}(\theta,\phi; x)$ の定義から直接計算で以下のようになることから従う。
\begin{align}
\log p_\theta (x) - \mathcal{L}(\theta, \phi; x) &= \mathbb{E}_{q_\phi (z|x)} \left[ - \log p_\theta (z|x) + \log q_\phi (z|x) \right] \\
&= D_{KL}(q_\phi (z|x) \| p_\theta (z)) + \mathbb{E}_{q_\phi (z|x)} \left[ - \log p_\theta (z|x) + \log p_\theta (z) \right] \\
&= D_{KL}(q_\phi (z|x) \| p_\theta (z)) + \mathbb{E}_{q_\phi (z|x)} \left[ - \frac{\log p_\theta (x|z)}{p_\theta(x)} \right] \\
& = D_{KL}(q_\phi (z|x) \| p_\theta (z)) - \mathbb{E}_{q_\phi (z|x)} [ \log p_\theta (x|z) ] + p_\theta(x)
\end{align}

微分することで変分下限 $\mathcal{L}(\theta,\phi; x^{(i)})$ を変分パラメータ $\phi$ と生成パラメータ $\theta$ について最適化したい。しかし $\phi$ に関する下限の勾配は少々問題となる。

SGVB 評価器と AEVB アルゴリズム

下限の実践的な評価器とパラメータに関する導函数を導入する。事後分布の近似を $q_\phi (z|x)$ の形式であると想定するが、 $x$ という条件を課さない $q_\phi (z)$ にも適用できることに注意する。
あとで導入するゆるやかな条件下では事後分布の適当な近似 $q_\phi (z|x)$ に対して、(補助の) ノイズ変数 $\varepsilon$ を用いた可微分変換 $g_\phi (\varepsilon, x)$ を使って確率変数 $\tilde{z} \sim q_\phi (z|x)$ と再パラメータ化できる。つまり:
\begin{align}
\tilde{z} \sim q_\phi (z|x) \quad \text{with} \quad \varepsilon \sim p(\varepsilon) \quad\quad\quad\quad\quad\quad (4)
\end{align}

である。

・・・と続くのだが、最終的には応用として (3) 式について $\mathbb{E}_{q_\phi (z|x^{(i)})} \left[ \log p_\theta (x^{(i)}|z) \right]$ における復号項 $\log p_\theta (x^{(i)}|z)$ の部分を Appendix C のように Bernoulli 分布なり Gauss 分布とするようだ。前者の場合、
\begin{align}
\log p (x|z) = \sum_{i=1}^D x_i \log y_i + (1-x_i) \log (1-y_i)
\end{align}

となる。$\varepsilon \sim \mathcal{N}(0,1)$ として $z^l = \mu + \sigma \odot \varepsilon$ として、変分限界
\begin{align}
&\mathcal{L}(\theta,\phi; x^{(i)}) = - D_{KL} (q_\phi (z|x^{(i)}) \| p_\theta (z)) + \mathbb{E}_{q_\phi (z|x^{(i)})} \left[ \log p_\theta (x^{(i)}|z) \right] \\
&\approx \frac{1}{2} \sum_{j=1}^J (1 + \log \sigma^2 - \mu^2 - \sigma^2) + \frac{1}{L} \sum_{l=1}^L \left\{ x^{(i)} \log \mathrm{dec}(z^l) + (1 - x^{(i)}) \log (1 - \mathrm{dec}(z^l)) \right\}
\end{align}

を最大化するというのが VAE の求める挙動ということになるようだ🤔 ここでモンテカルロ近似というのを使うようだが、その辺はよく分からない。