1. VAE(変分オートエンコーダ)の損失関数
ざっくり言うと次の通り。
$$\sum_i\left( D_{KL}[f(N(\mu_i,\sigma_i))||f(N(0,1))]\right)+\alpha\times MSE(y_{true}, y_{pred})$$
但し$\alpha$はハイパーパラメータ。(KerasでVAEを実装する記事をいくつか拝読したが、こちら忘れがち。)
また、ここでの$f(分布hoge)$は、分布hogeの確率密度関数。
1-1. 意味
$D_{KL}[f({分布}_1)||f({分布}_2)]$は分布1と分布2の「近さ」を表現する指標。
可換律が成立しないため、厳密には距離の一種ではない。
$MSE$は、データ間の近さを表現する指標。
このことが分かれば、VAEの損失関数の意味は簡単だ。
$$(潜在空間における点の分布が、原点周りの多変量標準正規分布にどれだけ近いか) \\+ ハイパーパラメータ\times(欲しい出力と、実際のVAEの出力の近さ)$$
である。
こうしてみると、ハイパーパラメータをうまく調整することで、
「潜在変数を多変量標準正規分布させつつ、ちゃんとまともな出力を得ることのできるオートエンコーダ」が学習されることが容易に想像できるだろう。
1-2. なぜ正規分布にこだわるん？
多変量標準正規分布させる理由は、「潜在空間内の点の密度をある程度整えたい」という願いが存在することだ。
それと、分布の連続性を保証したいことも理由にあげられる。
というのも、潜在空間内の点の分布を機械任せにした標準のオートエンコーダでは、「点が密集する領域」と「点がほとんど現れない領域」が両極端に現れるからだ。都会と田舎みたいな感じだ。
これは、「潜在空間上の座標をランダムに決めてを選んでデコーダに入れれば、オートエンコーダのデコーダって生成モデルとして利用できるんじゃね？」というアイデアを試すときに問題となる。
このアイデアを試すために潜在空間からランダムに座標を指定したとき、そこが田舎であった場合、
それはデコーダにとっては「意味不明な入力」となってしまう。したがってデコーダは、説明不能なわけのわからないものを出力することになってしまう。
また、例え都会の座標を指定したとしても、オートエンコーダは潜在空間の分布が連続であることを保証しないため、やはり説明不能な出力を得る場合がある。
そこで、点の分布の仕方に「正解」を与えることで、分布が連続であること、および都会と田舎の格差が生じないことを学習させているのである。
2. KLダイバージェンス
2-1. 理論
定義式がWikipediaに乗っている。
(離散・連続の一致する)2種類の単変量分布$p(x),q(x)$にの間に定義される。
離散確率分布$p(x), q(x)$の間のKLダイバージェンス$D_{KL}[p(x)||q(x)]$は次の通り。
$$D_{KL}[p(x)||q(x)] = \sum_i{\left(p(x_i)\ln{\frac{p(x_i)}{q(x_i)}}\right)}$$
VAEでは$D_{KL}[f(N(\mu,\sigma))||f(N(0,1))]$を考えることになる。
$$f(N(\mu,\sigma)) = \frac{1}{\sqrt{2\pi\sigma^2}}\exp\left(-\frac{(x-\mu)^2}{2\sigma^2}\right)$$
であるから、
$$
\begin{array}{l}
D_{KL}[f(N(\mu,\sigma))||f(N(0,1))] \\=
\sum_i{\left(
\frac{1}{\sqrt{2\pi\sigma^2}}\exp\left(-\frac{(x_i-\mu)^2}{2\sigma^2}\right)
\ln{\frac{
\frac{1}{\sqrt{2\pi\sigma^2}}\exp\left(-\frac{(x_i-\mu)^2}{2\sigma^2}\right)
}{
\frac{1}{\sqrt{2\pi}}\exp\left(-\frac{x_i^2}{2}\right)
}}\right)}
\\=
\sum_i{\left(
\frac{1}{\sqrt{2\pi\sigma^2}}\exp\left(-\frac{(x_i-\mu)^2}{2\sigma^2}\right)
\ln{
\left(
\frac{1}{\sigma}
\exp\left(-\frac{(x_i-\mu)^2}{2\sigma^2}+\frac{x_i^2}{2}\right)
\right)
}\right)}
\\=
\sum_i{\left(
\frac{1}{\sqrt{2\pi\sigma^2}}\exp\left(-\frac{(x_i-\mu)^2}{2\sigma^2}\right)
\ln{
\left(
\exp\left(-\frac{(x_i-\mu)^2}{2\sigma^2}+\frac{x_i^2}{2}-\ln\sigma\right)
\right)
}\right)}
\\=
\sum_i{\left(
\frac{1}{\sqrt{2\pi\sigma^2}}\exp\left(-\frac{(x_i-\mu)^2}{2\sigma^2}\right)
\left(-\frac{(x_i-\mu)^2}{2\sigma^2}+\frac{x_i^2}{2}-\ln\sigma\right)\right)}
\\=
-\frac{1}{2}\sum_i{\left(
\frac{1}{\sqrt{2\pi\sigma^2}}\exp\left(-\frac{(x_i-\mu)^2}{2\sigma^2}\right)
\left(\frac{(x_i-\mu)^2}{\sigma^2}-x_i^2+2\ln\sigma\right)\right)}
\end{array}
$$
ここで、分散の定義式を思い出す。サンプルサイズが$n$の場合
$$\sigma^2=\frac{1}{n}\sum_{i=1}^n(x_i-\mu)^2$$
だ。故に
$$
\begin{array}{l}
\exp\left(-\frac{(x_i-\mu)^2}{2\sigma^2}\right)
\\=
\exp\left(-\frac{n}{2}\times\frac{(x_i-\mu)^2}{\sum_{j=1}^n(x_j-\mu)^2}\right)
\end{array}$$
だ。$n$が十分大きいとき、$\sum_{j=1}^n(x_j-\mu)^2$は$(x_i-\mu)^2$に対して非常に大きくなる。
したがって、$\exp{}$の中身は0に近づく。
故に
$$\exp\left(-\frac{(x_i-\mu)^2}{2\sigma^2}\right)\simeq1$$
といえる。
さて、KLダイバージェンスの式変形に戻ろう。
$$
\begin{array}{l}
D_{KL}[f(N(\mu,\sigma))||f(N(0,1))]
\\=
-\frac{1}{2}\sum_i{\left(
\frac{1}{\sqrt{2\pi\sigma^2}}\exp\left(-\frac{(x_i-\mu)^2}{2\sigma^2}\right)
\left(\frac{(x_i-\mu)^2}{\sigma^2}-x_i^2+2\ln\sigma\right)\right)}
\\\simeq
-\frac{1}{2}\sum_i{\left(
\frac{1}{\sqrt{2\pi\sigma^2}}
\left(\frac{(x_i-\mu)^2}{\sigma^2}-x_i^2+2\ln\sigma\right)\right)}
\\=
-\frac{1}{2}\frac{1}{\sqrt{2\pi\sigma^2}}\sum_i{
\left(\frac{(x_i-\mu)^2}{\frac{1}{n}\sum_{j=1}^n(x_j-\mu)^2}-x_i^2+2\ln\sigma\right)}
\\=
-\frac{1}{2}\frac{1}{\sqrt{2\pi\sigma^2}}\sum_i{
\left(n\times\frac{(x_i-\mu)^2}{\sum_{j=1}^n(x_j-\mu)^2}-x_i^2+2\ln\sigma\right)}
\\=
-\frac{1}{2}\frac{1}{\sqrt{2\pi\sigma^2}}\left(
n\times\frac{\sum_i(x_i-\mu)^2}{\sum_j(x_j-\mu)^2}
-\sum_ix_i^2+\sum_i\ln(\sigma^2)
\right)
\\=
-\frac{1}{2}\frac{1}{\sqrt{2\pi\sigma^2}}\left(
n
-n\times\bar{x^2}+n\ln(\sigma^2)
\right)
\\=
-\frac{1}{2}\frac{n}{\sqrt{2\pi\sigma^2}}\left(
1
-\bar{x^2}+\ln(\sigma^2)
\right)
\end{array}
$$
ここで、2乗平均$\bar{x^2}$は平均2乗と分散の和$\mu^2+\sigma^2$で表現できるから、
$$
D_{KL}[f(N(\mu,\sigma))||f(N(0,1))]\simeq
-\frac{1}{2}\frac{n}{\sqrt{2\pi\sigma^2}}\left(
1
-\mu^2-\sigma^2+\ln(\sigma^2)
\right)$$
となる。
2-2. 実際
書籍『生成Deep Learning』David Foster著 O'REILLY社 では次のように紹介されている。
$$D_{KL}[N(\mu,\sigma)||N(0,1)] = -\frac{1}{2}\sum(1+\log(\sigma^2)-\mu^2-\sigma^2)$$
ここに出てくる$\sum$は、各次元にわたって合計を計算するという意味と考えられる。
比較すると、係数$\frac{n}{\sqrt{2\pi\sigma^2}}$の消失が確認される。
$\frac{n}{\sqrt{2\pi}}$は定数であるため省略されうるものとしても、$\frac{1}{\sigma}$の消失は説明する方法が見当たらない。
ゆえに、書籍で説明されるKL情報量というのは、実際は**「本物のKLダイバージェンスを$\sigma$で割った近似値」**と理解されるべきであろう。
おそらく、KLダイバージェンスを$\sigma$で割っているのは、分散が大きくなることによる損失の増大を緩めることで、点の分布の広さをある程度確保する狙いがあるのではないかと愚考する。