Variational Autoencoder を勉強してみる

# はじめに

Variational Autoencoder をお勉強して動かしてみたのでメモっておきます。
大昔、AutoEncoder だと特徴量として使えそうで面白そうだなー、と大昔思ったことがありますが、動かしてみたのは今回が初めてです。あと、勉強として下記を読みました。

https://arxiv.org/abs/1312.6114

実際に勉強してみてわかったのは、

- VAEの原理のコードそのものはシンプル。AEとの差分はきちんとある。
- 普通に次元圧縮するのと、どの程度違うの？、何故いいの？については、定性的な説明はできますが、いまいちアピールしにくいかもと思います。一応、確率の枠組みで制御を入れていますが、制御項を入れることは、別に他にいろいろありますし。

私の実装はこちら：
https://github.com/torupati/study_autoencoder


# ＶＡＥの定式化

入力、というか、観測ベクトルをx とします。これを、低次元の空間 z に射影します。その後、ｚからｘに復元します。
それぞれencoder, decoder と呼びます。関数として

$$
z = f(x)
$$

と

$$
x = g(z)
$$

と書くことにします。
最小化したい誤差の設計は、できるだけ情報を落とさずに潜在空間に射影したいので

$$
|| x - g(z) || ^2 = || x - g(f(x))||^2
$$

とするのが自然です。全ての訓練データについて、潜在空間の表現ｚからの復元誤差を最小化するのがauto-encoder です。

これは、通常のＮＮの教師あり問題の学習であり、通常の方法で学習を実行することができます。

### 観測$x$ が与えられたときの確率分布 p(z|x) を考えてみる

ここでｚを確率変数とし、確率分布$p(z)$ に従うと仮定します。これは観測について情報を持っていない状態で与えられる情報であり、事前分布です。我々にとって有用なのは、観測ｘが与えられたときのｚの分布、すなわち事後分布 $P(z|x)$ です。

この事後分布$P(z|x) $ が得られたら、復元誤差の計算に使用できます。復号化された値 $x'=g(z)$ と入力値について、潜在変数で積分して復元誤差の期待値を計算できます。

$$
E_{z\sim p(z|x)} [||x - g(z)||^2] = \int ||x - g(z)||^2 p(z|x) dz
$$

これをどうやって求めましょうか？結論を簡単に言いますと、ずばり

1. $p(z|x)$ を関数 $q_\phi(z)$ で近似する。ただし、この分布は$x$に依存する。
1. 潜在変数 $z$ での積分計算は $q_\phi(z)$ からのサンプリングに置き換える

ということをします。

事前分布 $P(z)$ があるので、尤度 $P(x|z)$ を設計してから事後確率を計算し、周辺化して確率分布を正規化して、、、というのは、計算が複雑でできません。なので、ここでは強引に$p(z|x)$ の形を決めてしまって計算してしまおう、というのが、この変分法の考え方になります。

### パラメータ推定

本来求めたいパラメータは、

- Encoder: 観測 $x$ を潜在変数に変換する関数のパラメータ
- Decoder: 潜在変数 $z$ を観測と同じ次元の値 $x'$に復元する関数のパラメータ

なのですが、ここでは潜在変数 $z$ を確率分布として扱うので変わってきます。特に、$p(z|x)$ を扱いやすい関数 $g(z|\phi)$ で扱おうとしています。なので、

- Encoder: 観測$x$ から関数 $g(z|\phi)$のパラメータ$\phi$への変換する関数を考える。その関数のパラメータ$\phi$ 
- Decoder: 潜在変数 $z$ を観測と同じ次元の値 $x'$に復元する関数のパラメータ（変わらず）

となります。

### 対数尤度の最大化

対数尤度の最大化の計算を考えます。実際には最大化ではなく、局所最適解なのですが、まぁ、細かいことは言わず。

$$
p(x) = \int p(x,z)dz
$$

に対して$q_\phi(z) \geq  0$ かつ $\int q_\phi (z) = 1$ を満たす $q(z)$ を用いると、対数尤度の下限となる関数を求めることができます。

$$
\log p_\theta(x) \geq \int q(z) \log \frac{p(x,z)}{q(z)}dz 
$$

実際に計算すれば、等号成立は確かめられます。おじさんは、凸関数の不等式として、重み和の対数、が、対数の重み和で下から抑えられる、なんて思って覚えています。等号成立条件は$q(z) = p(z|x)$ で、不等式の差分はKL divergence です。

今、$p(z|x)$ を $q_\phi(z|x)$で近似するので、

$$
\log p_\theta(x) = KL(q_\phi(z|x) || p_\theta(z|x)) + \int q_\phi(z|x) \log \frac{p_\theta(x,z)}{q_\phi(z|x)}dz
$$

この下限を

$$
\mathcal{L}(\theta,\phi;x) = \int q_\phi(z|x) \log \frac{p_\theta(x,z)}{q_\phi(z|x)}dz
$$

を最大化したいです。この下限は下記のように分解できます。

$$
\mathcal{L}(\theta,\phi;x) = - KL(q_\phi(z|x)||p_\theta(z)) + E_{z \sim q_\phi(z|x)} [\log p_\theta(x|z)]
$$

尤度の期待値が可適応で高くなりそうになると、ＫＬdivergence による項が事前分から離れることに対するペナルティを与える構造になっています。

### 下限の最大化の計算

今、２つのNNモデルが下限の計算に関係しています。

1. エンコーダに関するもの：観測 $x$ を入力として受け付けて、潜在変数の確率分布のパラメータ $\phi$ を出力するモデルです。これはミニバッチで複数のx がある場合、出力はそれぞれの $q(z|\mu(x),\Sigma(x))$になります

1. デコーダ：潜在変数 $z$ を入力として観測と同じ次元の値 $x'$を出力するNNモデルがあります。この出力が理想的には入力した$x$と等しくなってほしく、これを教師データとすることができます。つまり、$\log p_\theta(x|z)$を最大化したい、ということです。

があります。

そして結論からいうと、具体的な実装はNNのSGD(back propagation)です。どんなNNと誤差関数かというと、

- NN モデル (Variational Autoencoder)
  1. encoder に対応するＮＮモデルでパラメータ $\mu$, $\Sigma$ を求める。
  1. 求めたパラメータを用いて、$q(z|\mu,\Sigma)$ でサンプリングする。
  1. その後、z をもとにdecoder に入力して $x'$を計算し、出力とする

- 誤差関数 $\mathcal{L}(\theta,\phi;x)$
  1. $KL(q_\phi(z|x)||p_\theta(z))$ は計算式を導出できるので、直接計算する
  1. $\mathbb{E}_{z \sim q_\phi(z|x)} [\log p_\theta(x|z)]$は $q(z|\mu(x), \Sigma(x))$からサンプリングして計算する

## 感想

当初、観測を入れてパラメータが出てくるＮＮ、というのに驚いたのですが、よく考えるとそれほど普段扱っているものとギャップはありませんでした。今まで直接潜在変数 $z$ を出していたとすると、その$z$を分布の平均$\mu$ と見方を変えることにギャップはないと思います。そして分散がでてきてもいいかな、なんて思ったりします。実際は分散の対数ですが。

もと論文だと$\theta$、$\phi$を書いたり、サンプリングをreparameterization と言ってみたり、なかなか複雑ですが、コードをみると何を計算しているかすっきり分かってよいです。

あと、$\theta$や$\phi$が何を指すのかあいまいに書いてしまった。元論文だとNNの重みを指していて、それはそれでいいのだけれど。確率分布のパラメータを直接指定するものと誤解しやすいと思います。平均、分散を求めるところ、$\mu = \mu(x;\phi)$ みたいに書く方がいいのかな。自分は途中から$\mu$を登場させたが。

# まとめ

VAEの計算内容について調査内容をメモしました。実際に動かしたものも書きたかったのだが、時間切れ。orz

今後
- 正規分布使っているところをきちんと書く。もと論文はベルヌーイ分布も書いていたが。
- VAEの実行結果を書き足す（たくさん実装例が無料で掲載されていますが。）
- 記事を書く前はＰＣＡ、ＡＥなどについても手を出そうかと思ったが、ちょっと間に合わなかった。できれば潜在空間での意味について比較したいな。

おじさん疲れた。病み上がりなので寝る。(2023/08/16)