Variational AutoEncoder (VAE) の気持ちを理解したい

Variational AutoEncoder (VAE) の気持ちを完全に理解すべく書いてみます
間違いや誤解を招く表現などがあったらご指摘いただければさいわいです

TL;DR
記法
はじめに
Variational Autoencoder (VAE)
実験
おわりに

TL;DR

データをボトルネックに通して，非線形な「本質」を抽出する．
「本質」を確率的に破壊してロバストにする．

記法

めんどいので，スカラーもベクトルも確率変数も全部イタリック体で書きます
$\mathcal{N}$ はガウス分布を表します

はじめに

人工知能 (Artificial Intelligence; AI)

人間の知能を機械で模倣しようという取り組み．
色々定義があるが，一般的に知能は「学習」「推論」の2つを行う能力とされる．
過去の入力から環境について特徴を捉え（学習），そこから未来の状態について予測を行う（推論）．

機械学習 (Machine Learning; ML)

人工知能の一ジャンルで，学習データの統計的な性質を使って学習・推論を行う手法の総称．
Support Vector Machine (SVM) や，後述する深層学習などが有名．

深層学習 (Deep Learning; DL)

ニューラルネットワーク (Neural Network; NN) を使った機械学習手法の総称．
人間が特徴量を用意しなくても，大量のデータから自動で特徴量を抽出してくれる．
多層パーセプトロン (Multilayer Perceptron; MLP) をはじめとして，畳み込みニューラルネットワーク (Convolutional Neural Network; CNN) が画像分類コンペで高精度をたたき出した*1ので有名になった．
入力から出力への有向非巡回グラフになるように計算グラフを構築することで，グラフの逆辺をたどってチェーンルールに従い自動微分をして得た勾配をもとに，何かしらの損失関数に対して最適化を行う（誤差逆伝播法）．
基本，大量のデータで大量に計算をして高い精度を出す．計算自体は行列積や成分ごとの関数などの比較的単純なものが多いので，GPUを使うと速い．

多様体仮説

データ $x$ は多くの場合，その内容よりも次元数がいくらか大きい空間 (ambient space) で表現されているため，より小さい次元の空間 (latent space) でも表現できるという仮説*2．

例

640x480のRGB画像はラスタ表現で $640 \times 480 \times 3=921600$ 次元のベクトルとして表されているが，近傍の画素が相関しまくっているのでもっと圧縮できる．自然画像だけなら，砂嵐ノイズのようなランダムな点を考えなくていいので，次元数はもっと小さくできる（はず）．

Autoencoder (AE)

入力 $x$ を変換する1層の隠れ層 $h$ から，元の入力 $x$ を再構成する $x'$ を出力する*3．
$x \to h \to x'$ のように変換するので， $h$ の次元数を $x$ より小さくする（ボトルネック型）と次元削減ができる．
多様体仮説を採用すると，画像などのデータ $x$ に対してもっと小さい次元の隠れ層 $h$ で表現できる（はず）．
エンコーダ $E$ とデコーダ $D$ を使うと，

$\displaystyle h=E(x) \\ x'=D(h)$

と表せる．損失関数 $\mathcal{L}$ は，

$\displaystyle \mathcal{L} = d(x, x')$

となる．ただし， $d(x,x')$ は入力 $x$ と再構成 $x'$ の誤差．

Denoising Autoencoder (DAE)

入力 $x$ に少しノイズを加えて破壊して，復元・再構成することでロバストな推論を行うAE*4．
データが埋まっている多様体から少し離れた点を多様体上に戻してやる問題を解いているため，AEより明示的に多様体学習をしているといえる．

潜在変数モデル (Latent Variable Model; LVM)

観測によって直接得られる観測変数 (observed variables) から，直接得ることのできない潜在変数 (latent variables) を推論する．
うまく高次元データを表現できる低次元の潜在変数を取ってくれば，データに顕著に表れる特徴＝「本質」？が抽出できるので，その「本質」のもつ構造を多様体学習で求めたい．

Variational Autoencoder (VAE)

AEに変分ベイズ法を適用する*5．DAEと異なり，確率的モデルとしてノイズで潜在変数（隠れ層）を破壊する．
観測変数 $x$ に対する潜在変数 $z$ を仮定し，その関係を表す事後分布 $p(z|x)$ や $p(x|z)$ を求めたい．

導出

まず，尤度 $p_\theta(x)$ の最大化をする．モデルパラメータを $\theta$ と書くことにする
たいてい指数分布族なので，上に凸で単調増加するlogをかけた対数尤度 $\log p_\theta(x)$ を最大化することにする．

$\displaystyle \begin{align*} \log p_\theta(x) &= \log \int p_\theta(x,z) dz \end{align*}$

$x$ の分布は複雑な多様体を仮定しているので， $p_\theta(x)$ や $p_\theta(x,z)$ は陽に求めることができない．
そのため，別の分布 $q_\phi(z|x)$ ( $\phi$ はそのパラメータ) を使って近似することにする．

$\displaystyle \begin{align*} \log p_\theta(x) &= \log \int q_\phi(z|x) \cfrac{p_\theta(x,z)}{q_\phi(z|x)} dz \\ &= \log \int q_\phi(z|x) \cfrac{p_\theta(x|z) p_\theta(z)}{q_\phi(z|x)} dz \end{align*}$

ここで，イェンセンの不等式から，logを外に出すと対数尤度の下限 (Evidence Lower Bound; ELBO) が求まる：

$\displaystyle \begin{align*} \log p_\theta(x) &\ge \int q_\phi(z|x) \log \cfrac{p_\theta(x|z) p_\theta(z)}{q_\phi(z|x)} dz \\ &= \int q_\phi(z|x) \left[ \log p_\theta(x|z) - \log \cfrac{q_\phi(z|x)}{p_\theta(z)} \right] dz \\ &= \int q_\phi(z|x) \log p_\theta(x|z) dz - \int q_\phi(z|x) \log \cfrac{q_\phi(z|x)}{p_\theta(z)} dz \\ &= \mathbb{E}_{z \sim q_\phi (z|x)} \left[ \log p_\theta(x|z) \right] - D_\mathrm{KL} \left[ q_\phi(z|x) \| p_\theta(z) \right] \\ &= \mathrm{ELBO}(x; \theta, \phi) \end{align*}$

ただし， $D_\mathrm{KL}$ はKLダイバージェンスとする．
これを最大化することにより，対数尤度を最大化することができる（詳細は割愛）．

事後分布の具体形

ELBOが算出できたところで，具体的な分布について考える．
近似分布 $q_\phi(z|x)$ は，入力 $x$ が条件づけられたときの潜在変数 $z$ の近似分布だった．これを各成分独立なガウス分布とする． $q_\phi(z|x)$ のパラメータ $\mu, \sigma$ を推定するNNを，エンコーダ $E(x; \phi)$ とする．
事後分布 $p_\theta(x|z)$ は，潜在変数 $z$ が条件づけられたときの $x$ の近似分布だった．これは，二値画像ならベルヌーイ分布とし，そうでない画像ならガウス分布とする．この分布のパラメータ $x'$ を推定するNNを，デコーダ $D(z; \theta)$ とする．
以上まとめると，

$\displaystyle \begin{align*} q_\phi(z|x) &= \mathcal{N}(z; \mu(x), \sigma(x)), & [\mu(x), \sigma(x)] = E(x; \theta) \\ p_\theta(x|z) &= \mathcal{N}(x; x'(z), 1)~\mathrm{or}~\mathrm{Bernoulli}(z; x'(z)), & x'(z) = D(z; \phi) \\ \end{align*}$

Reparameterization Trick

一方で，このままでは $x \xrightarrow{E} z \xrightarrow{D} x'$ のように計算グラフを構築することができない． $q_\phi(z|x)$ から $z$ をサンプリングしてくるところは自動微分できないので，誤差逆伝播法が使えなくなってしまう．そこで，"reparameterization trick"をする．先ほど， $q_\phi(z|x)$ を各成分独立なガウス分布とした．そのため，以下のように誤差逆伝播法を保ったサンプリングができる：

$\displaystyle \begin{align*} z = \mu(x) + \sigma(x) \odot \epsilon, \\ \text{where}~\epsilon \sim \mathcal{N}(0, I) \end{align*}$

ただし， $\odot$ はアダマール積を表す．ばらつきのところだけサンプリングして，中心と広がりを微分できるようにする．

ELBOの実装

また，最大化するELBOの具体的な実装について考える．先ほどのELBOの式は以下の通りだった：

$\displaystyle \begin{align*} \mathrm{ELBO}(x; \theta, \phi) &= \mathbb{E}_{z \sim q_\phi (z|x)} \left[ \log p_\theta(x|z) \right] - D_\mathrm{KL} \left[ q_\phi(z|x) \| p_\theta(z) \right] \end{align*}$

第1項について考える：

$\displaystyle \begin{align*} \mathbb{E}_{z \sim q_\phi (z|x)} \left[ \log p_\theta(x|z) \right] \end{align*}$

$p_\theta(x|z)$ を正規分布 $\mathcal{N}(x;x'(z),I)$ とすると，これは入力と再構成誤差の二乗誤差を $-1$ 倍したものに等しくなる．

$\displaystyle \begin{align*} \mathbb{E}_{z \sim q_\phi (z|x)} \left[ \log p_\theta(x|z) \right] &= \mathbb{E}_{z \sim q_\phi (z|x)} \left[ \log \mathcal{N}(x;x'(z),I) \right] \\ &= \mathbb{E}_{z \sim q_\phi (z|x)} \left[ -\cfrac{1}{2} (x - x'(z))^\top(x - x'(z)) \right] + C \\ &= \mathbb{E}_{z \sim q_\phi (z|x)} \left[ -\cfrac{1}{2} \| x - x'(z) \|^2 \right] + C \end{align*}$

ただし， $C$ は正規化定数である．
第2項について考える：

$\displaystyle \begin{align*} D_\mathrm{KL} \left[ q_\phi(z|x) \| p_\theta(z) \right] \end{align*}$

まず，事前分布 $p_\theta(z)$ について仮定が必要なので，標準ガウス分布 $p_\theta(z)=\mathcal{N}(z;0,I)$ をおく．
そうすると，事後分布 $q_\phi(z|x)$ もガウス分布だったので，KLダイバージェンスは解析的に求まる：

$\displaystyle \begin{align*} D_\mathrm{KL} \left[ q_\phi(z|x) \| p_\theta(z) \right] &= \sum_{i \in \lambda} \cfrac{1}{2} \Big[ \mu_i(x)^2 + \sigma_i(x)^2 - \log \sigma_i(x)^2 - 1 \Big] \end{align*}$

サンプルごとのELBOは分散が大きいので，ミニバッチ学習で平均をとって最大化する関数として使う．元論文*6ではこの関数を Stochastic Gradient Variational Bayes (SGVB) 推定量と呼んでいる．
最小化問題を解く機械学習の慣習に従って，以降はELBOを $-1$ 倍して考えると，以下のように実装できる：

$\displaystyle \begin{align*} \mathcal{L} &:= -\cfrac{1}{|\mathcal{B}|} \sum_{i \in \mathcal{B}} ELBO(x^{(i)}; \theta, \phi) \\ &= -\cfrac{1}{|\mathcal{B}|} \sum_{i \in \mathcal{B}} \left[ \mathbb{E}_{z \sim q_\phi (z|x^{(i)})} \left[ \log p_\theta(x^{(i)}|z) \right] - D_\mathrm{KL} \left[ q_\phi(z|x^{(i)}) \| p_\theta(z) \right] \right] \\ &= -\cfrac{1}{|\mathcal{B}|} \sum_{i \in \mathcal{B}} \left[ -\cfrac{1}{2} \| x^{(i)} - x'(z(x^{(i)})) \|^2 - \sum_{i \in \lambda} \cfrac{1}{2} \Big( \mu_i(x^{(i)})^2 + \sigma_i(x^{(i)})^2 - \log \sigma_i(x^{(i)})^2 - 1 \Big) \right] \\ &= \cfrac{1}{|\mathcal{B}|} \sum_{i \in \mathcal{B}} \left[ \cfrac{1}{2} \| x^{(i)} - x'(z(x^{(i)})) \|^2 + \sum_{i \in \lambda} \cfrac{1}{2} \Big( \mu_i(x^{(i)})^2 + \sigma_i(x^{(i)})^2 - \log \sigma_i(x^{(i)})^2 - 1 \Big) \right] \end{align*}$

ただし， $x^{(i)}$ はミニバッチ $\mathcal{B}$ 中の $i (i \in \mathcal{B})$ 番目のサンプルを表し，学習に影響を及ぼさない定数項は無視した．

ELBOの各項について

$\displaystyle \begin{align*} {-}\mathrm{ELBO}(x; \theta, \phi) &= \mathbb{E}_{z \sim q_\phi (z|x)} \left[ -\log p_\theta(x|z) \right] + D_\mathrm{KL} \left[ q_\phi(z|x) \| p_\theta(z) \right] \end{align*}$