KLダイバージェンス（Kullback-Leibler divergence）

KLダイバージェンス（Kullback-Leibler divergence）とは、２つの確率分布の差を数値化したもので、機械学習の分野では損失関数の項目として使用されます。特に生成モデルで重要な考え方になってきます。

KLダイバージェンスの式は、確率分布$p(x)$、$q(x)$があった場合に以下のようになります。

\[
\begin{eqnarray}
D_{KL}[p(x)][q(x)] &=& \int_{-\infty}^{\infty}p(x) \log \frac{p(x)}{q(x)}dx \\
&=& \int_{-\infty}^{\infty}p(x) \log p(x)dx – \int_{-\infty}^{\infty}p(x) \log q(x)dx
\end{eqnarray}
\]

KLダイバージェンスには、$D_{KL}[p(x)][q(x)] \geq 0$という性質があり、$D_{KL}[p(x)][q(x)] = 0$となるのは、$p(x)$と$q(x)$の分布が完全に一致する場合になります。

この性質から、機械学習では損失関数にKLダイバージェンスの項を追加して最小化することで、分布間を近づけるように学習させるということがよく行われます。

本記事では、このKLダイバージェンスに関して正規分布間の計算をするとどういった式が導出されるかを具体的に計算してみようと思います。生成モデルである変分自己符号化器(Variational Auto Encoder）では、まさに正規分布のKLダイバージェンスが使用されているため理解しておくと役に立ちます。

正規分布間のKLダイバージェンスの導出

１次元正規分布の場合

1次元正規分布の場合で正規分布間のKLダイバージェンスを導出してみようと思います。分布$p(x)$、$q(x)$を以下であるとします。

\[
\begin{eqnarray}
p(x) &=& \mathcal{N}(x|\mu_{1}, \sigma^{2}_{1}) \\
q(x) &=& \mathcal{N}(x|\mu_{2}, \sigma^{2}_{2})
\end{eqnarray}
\]

なお、1次元の正規分布の式は以下のようになります。

\[
\mathcal{N}(x|\mu, \sigma^{2}) = \frac{1}{\sqrt{2\pi\sigma^{2}}}\exp\left(-\frac{1}{2\sigma^{2}}(x – \mu)^{2}\right)
\]

では、KLダイバージェンスの式の中の$\log\frac{p(x)}{q(x)}$について、まず計算してみましょう。

\[
\begin{eqnarray}
\log\frac{p(x)}{q(x)} &=& \log p(x) – \log q(x) \\
&=& \log\left(\frac{1}{\sqrt{2\pi\sigma^{2}_{1}}}\exp\left(-\frac{1}{2\sigma^{2}_{1}}(x – \mu_{1})^{2}\right)\right)
– \log\left(\frac{1}{\sqrt{2\pi\sigma^{2}_{2}}}\exp\left(-\frac{1}{2\sigma^{2}_{2}}(x – \mu_{2})^{2}\right)\right)\\
&=& \log\frac{\sigma_{2}}{\sigma_{1}} + \frac{1}{2\sigma^{2}_{2}}(x – \mu_{2})^{2} – \frac{1}{2\sigma^{2}_{1}}(x – \mu_{1})^{2}
\end{eqnarray}
\]

KLダイバージェンスは、上記の式の$p(x)$に関する期待値になっているのが分かるかと思います。$p(x)$の期待値を$E_{p}[\cdot]$と書くと、以下のように書くことができます。

\[
\begin{eqnarray}
\int_{-\infty}^{\infty}p(x) \log \frac{p(x)}{q(x)}dx &=& E_{p}\left[\log\frac{p(x)}{q(x)}\right] \\
&=& E_{p}\left[ \log\frac{\sigma_{2}}{\sigma_{1}} \right] + E_{p}\left[\frac{1}{2\sigma^{2}_{2}}(x – \mu_{2})^{2} \right] – E_{p}\left[ \frac{1}{2\sigma^{2}_{1}}(x – \mu_{1})^{2} \right]
\end{eqnarray}
\]

まず第1項については、$x$に依存しない項目であるのでそのままになります。

\[
E_{p}\left[ \log\frac{\sigma_{2}}{\sigma_{1}} \right] = \log\frac{\sigma_{2}}{\sigma_{1}}
\]

第2項は、数式を分解して計算してみましょう。以下の式では$E_{p}[x] = \mu_{1}$、$E_{p}[x^{2}] = \sigma^{2}_{1} + \mu^{2}_{1}$を使用しています。後者は分散の定義から導かれます。

\[
\begin{eqnarray}
E_{p}\left[\frac{1}{2\sigma_{2}^{2}}(x – \mu_{2})^{2} \right] &=& \frac{1}{2\sigma_{2}^{2}} \left( E_{p}[x^{2}] -2\mu_{2}E_{p}[x] + \mu_{2}^{2} \right)\\
&=& \frac{1}{2\sigma_{2}^{2}} ( \sigma_{1}^{2} + \mu_{1}^{2} -2\mu_{1}\mu_{2} + \mu_{2}^{2})\\
&=& \frac{1}{2\sigma^{2}_{2}} \left( \sigma_{1}^{2} + (\mu_{1} – \mu_{2})^{2} \right)
\end{eqnarray}
\]

第3項も同じように計算することができます。

\[
\begin{eqnarray}
E_{p}\left[\frac{1}{2\sigma_{1}^{2}}(x – \mu_{1})^{2} \right] &=& \frac{1}{2\sigma_{1}^{2}} \left( E_{p}[x^{2}] -2\mu_{1}E_{p}[x] + \mu_{1}^{2} \right)\\
&=& \frac{1}{2\sigma_{1}^{2}} ( \sigma_{1}^{2} + \mu_{1}^{2} -2\mu_{1}^{2} + \mu_{1}^{2})\\
&=& \frac{1}{2}
\end{eqnarray}
\]

これで各項目の計算ができたので、最後に整理します。

\[
\begin{eqnarray}
\int_{-\infty}^{\infty}p(x) \log \frac{p(x)}{q(x)}dx &=& E_{p}\left[\log\frac{p(x)}{q(x)}\right] \\
&=& \log\frac{\sigma_{2}}{\sigma_{1}} + \frac{1}{2\sigma^{2}_{2}} \left( \sigma_{1}^{2} + (\mu_{1} – \mu_{2})^{2} \right) – \frac{1}{2}\\
&=& \frac{1}{2}\left( \log\frac{\sigma_{2}^{2}}{\sigma_{1}^{2}} + \frac{\sigma_{1}^{2}}{\sigma_{2}^{2}} + \frac{(\mu_{1} – \mu_{2})^{2}}{\sigma_{2}^{2}} – 1 \right)
\end{eqnarray}
\]

以上が、1次元正規分布間のKLダイバージェンスの式変形でした。

多次元正規分布の場合

次は、多次元の正規分布の場合で正規分布間のKLダイバージェンスを導出してみようと思います。分布$p(x)$、$q(x)$を以下であるとします。

\[
\begin{eqnarray}
p(\boldsymbol{x}) &=& \mathcal{N}(\boldsymbol{x}|\boldsymbol{\mu}_{1}, \boldsymbol{\Sigma}_{1}) \\
q(\boldsymbol{x}) &=& \mathcal{N}(\boldsymbol{x}|\boldsymbol{\mu}_{2}, \boldsymbol{\Sigma}_{2})
\end{eqnarray}
\]

なお、D次元の正規分布の式は以下のようになります。

\[
\mathcal{N}(\boldsymbol{x}|\boldsymbol{\mu}, \boldsymbol{\Sigma}) = \frac{1}{\sqrt{(2\pi)^{D}|\boldsymbol{\Sigma}|}}\exp\left(-\frac{1}{2}(\boldsymbol{x} – \boldsymbol{\mu})^{\top}\boldsymbol{\Sigma}^{-1} (\boldsymbol{x} – \boldsymbol{\mu}) \right)
\]

では、KLダイバージェンスの式の中の$\log\frac{p(x)}{q(x)}$について、まず計算してみましょう。

\[
\begin{eqnarray}
\log\frac{p(x)}{q(x)} &=& \log p(x) – \log q(x) \\
&=&
\log\left( \frac{1}{\sqrt{(2\pi)^{D}|\boldsymbol{\Sigma}_{1}|}}\exp\left(-\frac{1}{2}(\boldsymbol{x} – \boldsymbol{\mu_{1}})^{\top}\boldsymbol{\Sigma}_{1}^{-1} (\boldsymbol{x} – \boldsymbol{\mu_{1}})\right) \right) \\
&-& \log\left( \frac{1}{\sqrt{(2\pi)^{D}|\boldsymbol{\Sigma}_{2}|}}\exp\left(-\frac{1}{2}(\boldsymbol{x} – \boldsymbol{\mu_{2}})^{\top}\boldsymbol{\Sigma}_{2}^{-1} (\boldsymbol{x} – \boldsymbol{\mu_{2}})\right) \right) \\
&=& \frac{1}{2}\log\frac{|\boldsymbol{\Sigma}_{2}|}{|\boldsymbol{\Sigma}_{1}|}
+ \frac{1}{2}(\boldsymbol{x} – \boldsymbol{\mu_{2}})^{\top}\boldsymbol{\Sigma}_{2}^{-1} (\boldsymbol{x} – \boldsymbol{\mu_{2}})
-\frac{1}{2}(\boldsymbol{x} – \boldsymbol{\mu_{1}})^{\top}\boldsymbol{\Sigma}_{1}^{-1} (\boldsymbol{x} – \boldsymbol{\mu_{1}})
\end{eqnarray}
\]

1次元の時と同じように、$p(x)$の期待値を$E_{p}[\cdot]$と書くとKLダイバージェンスは、以下のように書くことができます。

\[
\begin{eqnarray}
\int_{-\infty}^{\infty}p(x) \log \frac{p(x)}{q(x)}dx = E_{p}\left[\log\frac{p(x)}{q(x)}\right] &=& E_{p}\left[ \frac{1}{2}\log\frac{|\boldsymbol{\Sigma}_{2}|}{|\boldsymbol{\Sigma}_{1}|} \right] \\
&+& E_{p}\left[ \frac{1}{2}(\boldsymbol{x} – \boldsymbol{\mu_{2}})^{\top}\boldsymbol{\Sigma}_{2}^{-1} (\boldsymbol{x} – \boldsymbol{\mu_{2}}) \right] \\
&-& E_{p}\left[ \frac{1}{2}(\boldsymbol{x} – \boldsymbol{\mu_{1}})^{\top}\boldsymbol{\Sigma}_{1}^{-1} (\boldsymbol{x} – \boldsymbol{\mu_{1}}) \right]
\end{eqnarray}
\]

まず第1項については、$x$に依存しない項目であるのでそのままになります。

\[
E_{p}\left[ \frac{1}{2}\log\frac{|\boldsymbol{\Sigma}_{2}|}{|\boldsymbol{\Sigma}_{1}|} \right] = \frac{1}{2}\log\frac{|\boldsymbol{\Sigma}_{2}|}{|\boldsymbol{\Sigma}_{1}|}
\]

第2項と第3項については、少し数学の特徴を利用して変形していきます。具体的には、以下の2次形式とトレースの関係やトレースの性質を使用します。2次形式というのは$\boldsymbol{x}^{\top}\boldsymbol{A}\boldsymbol{x}$のように表されるもので、正規分布の$\exp$の中に現れるような形式です。また、トレースというのは行列の対角成分の和をとる計算で、$\mathrm{tr}(\cdot)$のように表します。

上記の特徴を利用して第2項の式を変形してみましょう。なお、$E_{p}[\boldsymbol{x}] = \boldsymbol{\mu}_{1}$、共分散の定義より$E_{p}[\boldsymbol{x}\boldsymbol{x}^{\top}] = \boldsymbol{\Sigma}_{1} + \boldsymbol{\mu}_{1}\boldsymbol{\mu}_{1}^{\top}$を使用します。

\[
\begin{eqnarray}
E_{p}\left[ \frac{1}{2}(\boldsymbol{x} – \boldsymbol{\mu_{2}})^{\top}\boldsymbol{\Sigma}_{2}^{-1} (\boldsymbol{x} – \boldsymbol{\mu_{2}}) \right]
&=&
\frac{1}{2} E_{p}\left[\mathrm{tr}\left(\boldsymbol{\Sigma}_{2}^{-1}(\boldsymbol{x} – \boldsymbol{\mu}_{2})(\boldsymbol{x} – \boldsymbol{\mu}_{2})^{\top} \right) \right]\\
&=&
\frac{1}{2} \mathrm{tr}\left(\boldsymbol{\Sigma}_{2}^{-1} E_{p}\left[ (\boldsymbol{x} – \boldsymbol{\mu}_{2})(\boldsymbol{x} – \boldsymbol{\mu}_{2})^{\top} \right] \right)\\
&=&
\frac{1}{2} \mathrm{tr} \left(\boldsymbol{\Sigma}_{2}^{-1} E_{p}\left[ \boldsymbol{x}\boldsymbol{x}^{\top} -\boldsymbol{x}\boldsymbol{\mu}_{2}^{\top} -\boldsymbol{\mu}_{2}\boldsymbol{x}^{\top} + \boldsymbol{\mu}_{2}\boldsymbol{\mu}_{2}^{\top} \right] \right) \\
&=&
\frac{1}{2} \mathrm{tr} \left( \boldsymbol{\Sigma}_{2}^{-1} \left( E_{p} \left[ \boldsymbol{x}\boldsymbol{x}^{\top}\right] – 2 E_{p} \left[ \boldsymbol{x}\boldsymbol{\mu}_{2}^{\top} \right] + E_{p} \left [\boldsymbol{\mu}_{2}\boldsymbol{\mu}_{2}^{\top} \right] \right) \right) \\
&=&
\frac{1}{2} \mathrm{tr} \left( \boldsymbol{\Sigma}_{2}^{-1} \left( \boldsymbol{\Sigma}_{1} + \boldsymbol{\mu}_{1}\boldsymbol{\mu}_{1}^{\top} – 2\boldsymbol{\mu}_{1} \boldsymbol{\mu}_{2}^{\top} + \boldsymbol{\mu}_{2}\boldsymbol{\mu}_{2}^{\top} \right) \right) \\
&=&
\frac{1}{2} \mathrm{tr}\left( \boldsymbol{\Sigma}_{2}^{-1}\boldsymbol{\Sigma}_{1} \right) + \frac{1}{2} \mathrm{tr} \left( \boldsymbol{\Sigma}_{2}^{-1}(\boldsymbol{\mu}_{1} – \boldsymbol{\mu}_{2})(\boldsymbol{\mu}_{1} – \boldsymbol{\mu}_{2})^{\top} \right) \\
&=&
\frac{1}{2} \mathrm{tr}\left( \boldsymbol{\Sigma}_{2}^{-1}\boldsymbol{\Sigma}_{1} \right) + \frac{1}{2} (\boldsymbol{\mu}_{1} – \boldsymbol{\mu}_{2})^{\top}\boldsymbol{\Sigma}_{2}^{-1}(\boldsymbol{\mu}_{1} – \boldsymbol{\mu}_{2})
\end{eqnarray}
\]

同様に第3項を変形します。途中で$E_{p}\left[ (\boldsymbol{x} – \boldsymbol{\mu}_{1})(\boldsymbol{x} – \boldsymbol{\mu}_{1})^{\top}\right]$が出てきますが共分散の定義そのものなので$\boldsymbol{\Sigma}_{1}$になります。

\[
\begin{eqnarray}
E_{p}\left[ \frac{1}{2}(\boldsymbol{x} – \boldsymbol{\mu_{1}})^{\top}\boldsymbol{\Sigma}_{1}^{-1} (\boldsymbol{x} – \boldsymbol{\mu_{1}}) \right]
&=&
\frac{1}{2} E_{p}\left[\mathrm{tr}\left(\boldsymbol{\Sigma}_{1}^{-1}(\boldsymbol{x} – \boldsymbol{\mu}_{1})(\boldsymbol{x} – \boldsymbol{\mu}_{1})^{\top} \right) \right]\\
&=&
\frac{1}{2} \mathrm{tr}\left(\boldsymbol{\Sigma}_{1}^{-1} E_{p}\left[ (\boldsymbol{x} – \boldsymbol{\mu}_{1})(\boldsymbol{x} – \boldsymbol{\mu}_{1})^{\top} \right] \right)\\
&=&
\frac{1}{2}\mathrm{tr}\left( \boldsymbol{\Sigma}_{1}^{-1} \boldsymbol{\Sigma}_{1} \right) \\
&=&
\frac{1}{2}\mathrm{tr}\left( \boldsymbol{I}\right) \\
&=&
\frac{D}{2}
\end{eqnarray}
\]

これで各項目の計算ができたので、最後に整理します。

\[
\begin{eqnarray}
\int_{-\infty}^{\infty}p(x) \log \frac{p(x)}{q(x)}dx &=& E_{p}\left[\log\frac{p(x)}{q(x)}\right] \\
&=&
\frac{1}{2}\log\frac{|\boldsymbol{\Sigma}_{2}|}{|\boldsymbol{\Sigma}_{1}|}
+\frac{1}{2} \mathrm{tr}\left( \boldsymbol{\Sigma}_{2}^{-1}\boldsymbol{\Sigma}_{1} \right) + \frac{1}{2} (\boldsymbol{\mu}_{1} – \boldsymbol{\mu}_{2})^{\top}\boldsymbol{\Sigma}_{2}^{-1}(\boldsymbol{\mu}_{1} – \boldsymbol{\mu}_{2})
– \frac{D}{2}\\
&=&
\frac{1}{2}\left( \log\frac{|\boldsymbol{\Sigma}_{2}|}{|\boldsymbol{\Sigma}_{1}|} + \mathrm{tr}\left( \boldsymbol{\Sigma}_{2}^{-1}\boldsymbol{\Sigma}_{1} \right) + (\boldsymbol{\mu}_{1} – \boldsymbol{\mu}_{2})^{\top}\boldsymbol{\Sigma}_{2}^{-1}(\boldsymbol{\mu}_{1} – \boldsymbol{\mu}_{2}) – D \right)
\end{eqnarray}
\]

以上が、多次元であるD次元正規分布間のKLダイバージェンスの式変形でした。

まとめ

KLダイバージェンスに関して正規分布間の計算をするとどういった式が導出されるかを具体的に計算してみました。1次元正規分布間の場合と多次元であるD次元正規分布間の場合のそれぞれでKLダイバージェンスの式変形をしてみています。

具体的に計算したKLダイバージェンスの式は以下のようになります。

【1次元正規分布間のKLダイバージェンス】

$p(x) = \mathcal{N}(x|\mu_{1}, \sigma^{2}_{1})$、$q(x) = \mathcal{N}(x|\mu_{2}, \sigma^{2}_{2})$とするとき

\[
\begin{eqnarray}
D_{KL}[p(x)][q(x)] &=&
\int_{-\infty}^{\infty}p(x) \log \frac{p(x)}{q(x)}dx \\
&=&
\frac{1}{2}\left( \log\frac{\sigma_{2}^{2}}{\sigma_{1}^{2}} + \frac{\sigma_{1}^{2}}{\sigma_{2}^{2}} + \frac{(\mu_{1} – \mu_{2})^{2}}{\sigma_{2}^{2}} – 1 \right)
\end{eqnarray}
\]

【D次元正規分布間のKLダイバージェンス】

$p(\boldsymbol{x}) = \mathcal{N}(\boldsymbol{x}|\boldsymbol{\mu}_{1}, \boldsymbol{\Sigma}_{1})$、$q(\boldsymbol{x}) = \mathcal{N}(\boldsymbol{x}|\boldsymbol{\mu}_{2}, \boldsymbol{\Sigma}_{2})$とするとき

\[
\begin{eqnarray}
D_{KL}[p(x)][q(x)] &=&
\int_{-\infty}^{\infty}p(x) \log \frac{p(x)}{q(x)}dx \\
&=&
\frac{1}{2}\left( \log\frac{|\boldsymbol{\Sigma}_{2}|}{|\boldsymbol{\Sigma}_{1}|} + \mathrm{tr}\left( \boldsymbol{\Sigma}_{2}^{-1}\boldsymbol{\Sigma}_{1} \right) + (\boldsymbol{\mu}_{1} – \boldsymbol{\mu}_{2})^{\top}\boldsymbol{\Sigma}_{2}^{-1}(\boldsymbol{\mu}_{1} – \boldsymbol{\mu}_{2}) – D \right)
\end{eqnarray}
\]

生成モデルである変分自己符号化器(Variational Auto Encoder）等では、正規分布のKLダイバージェンスが使用されているため理解しておくと役に立ちます。

あわせて読みたい

【Python Tech】プログラミングガイド