Jensenの不等式 - らんだむな記憶で触れたカルバック・ライブラー情報量について特定の確率分布の時に計算してみたい。PRML の 1.6. Relative entropy and mutual information にて relative entropy 或いは Kullback-Leibler divergence という名で知られるものとして
\begin{align}
\mathrm{KL}(p\|q) &= - \int p(X) \ln q(x) dx - \left(- \int p(x) \ln p(x) dx\right) \\
&= - \int p(x) \ln \left\{ \frac{q(x)}{p(x)} \right\} dx
\end{align}
Exercise 2.13 にて確率分布に 2 つのガウス分布を仮定する場合に計算せよというものがある。オリジナルの場合は多次元を想定しているようだが面倒なので 1 次元のガウス分布として扱う。
キーとなる定積分
\begin{align}
\int_{-\infty}^\infty e^{-x^2} dx = \sqrt{\pi}
\end{align}
これを変数変換することで
\begin{align}
\int_{-\infty}^\infty \exp \left( - \frac{x^2}{2} \right) dx = \sqrt{2\pi}
\end{align}
を得る。
さて、 $p(x) = N(x | \mu_1, \sigma_1^2)$, $q(x) = N(x | \mu_2, \sigma_2^2)$ と置こう。ここで
\begin{align}
N(x | \mu, \sigma^2) = \frac{1}{\sqrt{2\pi} \sigma} \int_{-\infty}^\infty \exp \left( - \frac{(x-\mu)^2}{2 \sigma^2} \right) dx
\end{align}
である。これを素直に $KL(p\|q)$ に代入すると
\begin{align}
KL(p\|q) &= \frac{1}{\sqrt{2\pi} \sigma_1} \int \exp \left( - \frac{(x-\mu_1)^2}{2 \sigma_1^2} \right) \left( \log \frac{\sigma_2}{\sigma_1} - \frac{(x-\mu_1)^2}{2 \sigma_1^2} + \frac{(x-\mu_2)^2}{2 \sigma_2^2} \right) dx \\
&= I + J + K
\end{align}
となるが、3 つに分けて計算する。まず
\begin{align}
I = \frac{1}{\sqrt{2\pi} \sigma_1} \int \exp \left( - \frac{(x-\mu_1)^2}{2 \sigma_1^2} \right) \log \frac{\sigma_2}{\sigma_1} dx = \log \frac{\sigma_2}{\sigma_1} = -\frac{1}{2} \log \frac{\sigma_1^2}{\sigma_2^2}
\end{align}
となる。これは $y = \frac{x-\mu_1}{\sigma_1}$ という変数変換のもと「キーとなる定積分」で処理をすれば良い。最後の項は合算時の見栄えの都合で少し面倒な形に式変形している。次に、同じ変数変換を用いて
\begin{align}
J &= - \frac{1}{\sqrt{2\pi} \sigma_1} \int \exp \left( - \frac{(x-\mu_1)^2}{2 \sigma_1^2} \right) \frac{(x-\mu_1)^2}{2 \sigma_1^2} dx \\
&= - \frac{1}{\sqrt{2\pi} \sigma_1} \frac{\sigma_1}{2} \int \exp \left( - \frac{y^2}{2} \right) y^2 dy \\
&= - \frac{1}{\sqrt{2\pi} \sigma_1} \frac{\sigma_1}{2} \int y \left( y \exp \left( - \frac{y^2}{2} \right) \right) dy \\
& = - \frac{1}{\sqrt{2\pi} \sigma_1} \frac{\sigma_1}{2} \int \exp \left( - \frac{y^2}{2} \right) dy = - \frac{1}{\sqrt{2\pi} \sigma_1} \frac{\sigma_1}{2} \sqrt{2\pi} = - \frac{1}{2}
\end{align}
が部分積分により求まる。最後に、またまた同じ変数変換を用いて $(x-\mu_2)^2 = \sigma_1^2 y^2 + 2(\mu_1-\mu_2) \sigma_1 y + (\mu_1 - \mu_2)^2$ となることに注意して展開すると、
\begin{align}
K &= \frac{1}{\sqrt{2\pi} \sigma_1} \int \exp \left( - \frac{(x-\mu_1)^2}{2 \sigma_1^2} \right) \frac{(x-\mu_2)^2}{2 \sigma_2^2} dx \\
&= \frac{1}{\sqrt{2\pi} \sigma_1} \frac{\sigma_1}{2 \sigma_2^2} \left( \sigma_1^2 \!\!\int y^2 e^{- \frac{y^2}{2}} dy + 2(\mu_1-\mu_2) \sigma_1 \!\!\int y e^{- \frac{y^2}{2}} dy + (\mu_1-\mu_2)^2 \!\!\int e^{- \frac{y^2}{2}} dy \right) \\
& = \frac{1}{\sqrt{2\pi} \sigma_1} \frac{\sigma_1}{2 \sigma_2^2} \left( \sigma_1^2 \sqrt{2\pi} + (\mu_1-\mu_2)^2 \sqrt{2\pi} \right) = \frac{\sigma_1^2}{2 \sigma_2^2} + \frac{(\mu_1-\mu_2)^2}{2 \sigma_2^2}
\end{align}
が部分積分により求まる。 $I$ と $J$ と $K$ を合算して
\begin{align}
KL(p\|q) = \frac{1}{2} \left( - \log \frac{\sigma_1^2}{\sigma_2^2} - 1 + \frac{\sigma_1^2}{\sigma_2^2} + \frac{(\mu_1-\mu_2)^2}{\sigma_2^2} \right)
\end{align}
が求まった。特に $q$ のほうが標準正規分布 $q(x) = N(x|0,1)$ の時には
\begin{align}
KL(p\|q) = \frac{1}{2} \left( - \log \sigma_1^2 - 1 + \sigma_1^2 + \mu_1^2 \right)
\end{align}
となる。
以上を踏まえると VAE (Variational Autoencoder) における損失関数は
- エンコーダ・デコーダを通して再構成した画像と入力画像の交差エントロピー
- エンコーダから出てくる出力を確率分布と見た場合に標準正規分布とのカルバック・ライブラー情報量
の合算で構成され、これの最小値を求める形で勾配降下法を適用しているようだ・・・。なんで標準正規分布との擬距離をとっているのかは知らないけど、[1312.6114] Auto-Encoding Variational Bayesの Appendix B でそんな感じだから？ということらしい？？
