画像生成AIの勉強中に、VAEが出てきたのでそれについてまとめます。
VAEの概要
端的にVAE（Variational AutoEncoder）とは
・エンコーダーとデコーダーから構成されている
・潜在変数に正規分布などの確率分布を仮定したモデリング
・画像生成などで用いられる
です。
エンコーダーとデコーダについて補足すると
・エンコーダーは入力データからより良い潜在変数zを見つけ出すためのモデル
・デコーダーは得られた潜在変数zからできる限り入力データの確率分布に近しい事後確率分布を推定するモデル
となっています。AutoEncoderとの違いは潜在変数に確率分布を仮定することにあります。
潜在変数とは
因子分析では「潜在変数とは直接観測できない観測であり、かつ観測できた変数から推測できる変数」という説明がなされます。しかし、VAEの文脈では以下の説明の方が潜在変数をよりわかりやすく理解できるのではないかと思います。
VAEの文脈において、潜在変数とは情報の本質的な部分を表現したものと考えられます。
高次元データはノイズが多くデータとして大して大事じゃない部分があります。それらのノイズを排除した（いわゆる次元圧縮）よりデータの本質の部分を低次元データで表現したものが潜在変数となります。
VAEの構造
以下のようになっています。
データがどのように処理されているのかについて説明すると
１、画像データX1をエンコーダー($\theta$)に入力し、正規分布のパラメータ（$\muと\sigma^2$）を出力する。
２、変数変換トリックを用い、潜在変数Zを出力する
３、Zをデコーダー($\phi$)に入力し、出力結果から学習する。デコーダは潜在変数が与えられた時の事後分布を出力するように学習している。
エンコーダーとデコーダー
ニューラルネットワークを用いることがほとんどだと思いますが、情報を圧縮・復元できるならどのようなモデルを用いても問題ないと思います。
変数変換トリックとは
変数変換トリックとはVAE全体で学習させるための工夫です。
エンコーダーに直接確率分布を出力させる場合、潜在変数Zを出力した確率分布からサンプリングすることになります。このとき、サンプリングした際に計算グラフが途切れ、誤差の逆伝播ができなくなりパラメータの学習ができなくなってしまいます。そこで$\epsilon ~ N(0,1)$から値をサンプリングし、その値に分散を掛けた上で、平均を足し、潜在変数Zを出力することで、計算グラフを切れず、エンコーダーのパラメータまで誤差逆伝播ができ、パラメータを学習できるようになります。
目的関数
VAEは入力されたデータの確率分布を学習することが目的です。したがって$p_\phi(x)$の対数尤度を最大化することが目的になります。
途中式などは先人たちの素晴らしい資料があるのでそちらの説明をご覧ください。
最終的に最大化したい関数はELBOと呼ばれ、ELBOは以下のようになります。
ELBO(x; \theta, \phi) = E_{p_{\theta}(z|x)}[log p_{\phi}(x|z)] - D_{KL}[q_{\theta}(z|x) || p(z)]
各項について説明します
再構成誤差
$E_{p_{\theta}(z|x)}[log p_{\phi}(x|z)]$の部分を再構成誤差と呼びます。
上記式はエンコーダー$\theta$を通して、得た潜在変数の確率分布がp(z|x)であった時、この確率分布からサンプリングされたzをデコーダー$p_{(\phi)}$に入力して得られたものの期待値です。これが最大になるようにパラメータ$\thetaと\phi$を学習します。
ここで注意点は$p_{\phi}(x|z)のx$が$\hat{x}(=\muと\sigmaで推定される確率分布)$ではない点です。$x$は実際に入力したデータのことであり、$\hat{x}$は$p_{\phi}(z)=\hat{x}$とエンコーダーが推定した値のことです。ここで、$xと\hat{x}$がズレると（つまり、エンコーダーの推定値と入力したデータの値が異なる）どうなるでしょうか？
$xと\hat{x}$が正規分布だと仮定します。この時$log p_{\phi}(x|z)$は以下のように書き換えられます。$p_{\phi}$を多変量正規分布の式で書き換えただけです。
E_{p_{\theta}(z|x)}[log p_{\phi}(x|z)] = E_{p_{\theta}(z|x)}[log \frac{1}{\sqrt{(2\pi)^D * I}}exp(-\frac{1}{2}(x - \hat{x})^TI^{-1}(x - \hat{x}))]
logにつてい展開すると
=E_{p_{\theta}(z|x)}[-\frac{1}{2}(x - \hat{x})^T(x - \hat{x}) + log\frac{1}{\sqrt{(2\pi)^D * I}}]
$p_{\theta}(z|x)$でモンテカルロ法によりサンプリングし、期待値の近似値を求めます。サンプル数をLとした時以下の式に展開できます。よく見ると２乗誤差になっていますね。
=-\frac{1}{2*L} \sum^L_l ((x - \hat{x})^T(x - \hat{x}) + const)
この時、$xと\hat{x}$が同値の時（全体が上に凸の関数であるため）、再構成誤差は最大になります。
KLダイバージェンス
$D_{KL}[q_{\theta}(z|x) || p(z)]$の部分をKLダイバージェンスです。
$q_{\theta}(z|x)$はエンコーダー$\theta$を利用し、入力データxを潜在変数zに圧縮した確率分布です。p(z)は私たちが任意に設定する確率分布で正規分布など（タスクに依存する）になります。
つまり、ここでは次元削減後の潜在変数zの確率分布が正規分布p(z)に近いほど、KLダイバージェンスの値は小さくなり、目的関数を最大化できるということになります。
余談ですが、ここでのKLダイバージェンスはReverseKLDを用いています。画像生成タスクではいい感じの画像となる座標が複数あると考えられます。これらの座標の中の1つでも学習できれば十分なのでReverseKLDを利用していると考察しています。
よくある勘違い
よくある勘違いとして、VAEはエンコーダーに入力する確率分布とデコーダーが出力する確率分布を近付けるように学習しているというものがあります。しかし、目的関数を見ると「デコーダー（＝関数）の尤度を最大にしている」部分＋「エンコード後の潜在変数の分布が設定した分布にどれだけ近づいているか」部分に分かれます。デコーダーは私の知る限り表現力が高いNNを用いますが、デコーダーをその他の次元を復元する方法で代用することも可能だと思われます。
潜在変数をなぜ正規分布にするのか
最後に
上記の疑問を含め色々、調べているうちに潜在変数をなぜ確率分布と仮定しているのかを解説している面白い記事に出会いました。
記事の詳細はこちら：
いくつかの疑問を要約します
Q1：潜在変数を連続値である確率分布にする理由は？
A1：未学習のデータに対しても出力が可能になるため。例えば離散値の潜在変数だった場合、座標(2, 2)の出力結果が画像の４だったとする。この時学習していない座標(2.1, 2.1)はどうするのか。潜在変数を連続値である確率分布で学習している場合、4に似たような画像が出力されやすくなる。
Q2：なぜ多変量正規分布を潜在変数の確率分布として仮定するのか？
A2：明確な答えはありませんでした。ただ、メリットとして学習がうまくいった場合、潜在変数の各軸は画像の独立した特徴を捉えたベクトルになるそうです。つまり、画像の複雑な構成要素を複数の独立した個別の要素で説明できるようになるということです。具体的には顔の写真を個別の要素に分解すると、「肌色、年齢、彩度」などになりますが、それらの特徴を潜在変数の各軸が表現しているということです。
参考文献
・【超初心者向け】VAEの分かりやすい説明とPyTorchの実装：
・「潜在変数」（z）を多変量正規分布に近づけるβ-VAEモデルはなぜ、データが持つ特徴のもつれを解く（disentangle）ことができるのか？：
・ゼロから作るDeep Learning ❺：
変分ベイズ（潜在変数をモデルに組み込んだ上で事後分布を推定する手法）の枠組みでVAEを説明したもの：
KLダイバージェンスについて：