本記事は「これなら分かる！はじめての数理統計学」シリーズに含まれます。
不適切な内容があれば，記事下のコメント欄またはお問い合わせフォームよりご連絡下さい。
目次
正規分布同士のKLダイバージェンス
確率変数$X=\{ X_{1}, \ldots, X_{D} \}$に対して定義される
\begin{align}
p(X) &\sim \N (\vmu_{p}, \Sigma_{p}) \\[0.7em]
q(X) &\sim \N (\vmu_{q}, \Sigma_{q})
\end{align}
p(X) &\sim \N (\vmu_{p}, \Sigma_{p}) \\[0.7em]
q(X) &\sim \N (\vmu_{q}, \Sigma_{q})
\end{align}
に対して，以下が成り立つ。
\begin{align}
&\left. \KL \left[ q(X) \right\| p(X) \right] \notag \\[0.7em]
&= \frac{1}{2} \left\{ \log \frac{|\Sigma_p|}{|\Sigma_q|}
+ \Tr \left[ \Sigma_{p}^{-1} \Sigma_{q} \right]
+ \left( \vmu_q-\vmu_p \right)^T \Sigma_{p}^{-1} \left( \vmu_q-\vmu_p \right)-D
\right\}
\end{align}
&\left. \KL \left[ q(X) \right\| p(X) \right] \notag \\[0.7em]
&= \frac{1}{2} \left\{ \log \frac{|\Sigma_p|}{|\Sigma_q|}
+ \Tr \left[ \Sigma_{p}^{-1} \Sigma_{q} \right]
+ \left( \vmu_q-\vmu_p \right)^T \Sigma_{p}^{-1} \left( \vmu_q-\vmu_p \right)-D
\right\}
\end{align}
ただし，$|\cdot|$は行列式を表す。
変分推論の計算でよく利用される計算です。というのも，多くの変分推論では真の分布に正規分布，近似分布にも正規分布を仮定することが多く，そうなると正規分布同士のKLダイバージェンスを解析的に計算する必要が出てくるのです。特にVAE（変分オートエンコーダ: variational autoencoder）の下限導出で大活躍する公式です。
証明
まずは，KLダイバージェンスの定義を変形します。
\begin{align}
\left. \KL \left[ q(X) \right| p(X) \right] &= E_{q(X)} \left[ \log \frac{q(X)}{p(X)} \right] \\[0.7em]
&= E_{q(X)} \left[ \log q(X)-\log p(X) \right] \\[0.7em]
&= E_{q(X)} \left[ \log q(X) \right]-E_{q(X)} \left[ \log p(X) \right] \label{eq_KL_1}
\end{align}
\left. \KL \left[ q(X) \right| p(X) \right] &= E_{q(X)} \left[ \log \frac{q(X)}{p(X)} \right] \\[0.7em]
&= E_{q(X)} \left[ \log q(X)-\log p(X) \right] \\[0.7em]
&= E_{q(X)} \left[ \log q(X) \right]-E_{q(X)} \left[ \log p(X) \right] \label{eq_KL_1}
\end{align}
多変量正規分布の確率密度関数より，$\log p(X)$と$\log q(X)$は以下のように表されます。
\begin{align}
\log p(X) &= \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{p}|-(\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) \right\} \\[0.7em]
\log q(X) &= \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{q}|-(\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q) \right\}
\end{align}
\log p(X) &= \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{p}|-(\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) \right\} \\[0.7em]
\log q(X) &= \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{q}|-(\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q) \right\}
\end{align}
定数$c$の期待値は$E[c]=c$となることから，$\log p(X)$と$\log q(X)$の期待値は以下のように表されます。
\begin{align}
E_{q(X)} \left[ \log p(X) \right] &= E_{q(X)} \left[ \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{p}|-(\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) \right\} \right] \\[0.7em]
&= \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{p}|-E_{q(X)} \left[ (\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) \right] \right\} \\[0.7em]
E_{q(X)} \left[ \log q(X) \right] &= E_{q(X)} \left[ \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{q}|-(\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q) \right\} \right] \\[0.7em]
&= \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{q}|-E_{q(X)} \left[ (\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q) \right] \right\}
\end{align}
E_{q(X)} \left[ \log p(X) \right] &= E_{q(X)} \left[ \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{p}|-(\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) \right\} \right] \\[0.7em]
&= \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{p}|-E_{q(X)} \left[ (\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) \right] \right\} \\[0.7em]
E_{q(X)} \left[ \log q(X) \right] &= E_{q(X)} \left[ \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{q}|-(\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q) \right\} \right] \\[0.7em]
&= \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{q}|-E_{q(X)} \left[ (\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q) \right] \right\}
\end{align}
これらを式($\ref{eq_KL_1}$)に代入します。
\begin{align}
\left. \KL \left[ q(X) \right| p(X) \right]
&= \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{q}|-E_{q(X)} \left[ (\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q) \right] \right\} \notag \\[0.7em]
&\quad\quad-\frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{p}|-(\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) \right\} \\[0.7em]
&= \frac{1}{2} \biggl\{ \log \frac{|\Sigma_{p}|}{|\Sigma_{q}|}-E_{q(X)} \left[ (\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q)\right] \notag \\[0.7em]
&\quad\quad+E_{q(X)} \left[(\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) \right] \biggr\}
\end{align}
\left. \KL \left[ q(X) \right| p(X) \right]
&= \frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{q}|-E_{q(X)} \left[ (\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q) \right] \right\} \notag \\[0.7em]
&\quad\quad-\frac{1}{2}\left\{-D\log 2\pi-\log |\Sigma_{p}|-(\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) \right\} \\[0.7em]
&= \frac{1}{2} \biggl\{ \log \frac{|\Sigma_{p}|}{|\Sigma_{q}|}-E_{q(X)} \left[ (\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q)\right] \notag \\[0.7em]
&\quad\quad+E_{q(X)} \left[(\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) \right] \biggr\}
\end{align}
ここからは，第二項目と第三項目の計算を分けて行っていきます。
第二項目の計算
まずは，第二項目を計算していきます。二次形式の行列部分が対称行列である場合，行列のトレースを利用することで計算を簡潔にすることができます。
二次形式とトレースの性質
$A$が対称行列であるとき，二次形式は以下のようにトレースを用いて記述することができる。
\begin{align}
\vx^T A \vx &= \Tr\left[ A \vx \vx^T \right]
\end{align}
\vx^T A \vx &= \Tr\left[ A \vx \vx^T \right]
\end{align}
$\Sigma_q$は多変量正規分布の分散共分散行列ですから対称行列です。したがって，第二項目に二次形式とトレースの性質を適用することができます。
\begin{align}
E_{q(X)} \left[ (\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q) \right]
&= \Tr \left[ \Sigma_{q}^{-1} E_{q(X)} \left[(\vx-\vmu_q)(\vx-\vmu_q)^T\right] \right] \\[0.7em]
&= \Tr \left[ \Sigma_{q}^{-1} \Sigma_{q} \right] \\[0.7em]
&= \Tr\left[ I_{D} \right] \\[0.7em]
&= D
\end{align}
E_{q(X)} \left[ (\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q) \right]
&= \Tr \left[ \Sigma_{q}^{-1} E_{q(X)} \left[(\vx-\vmu_q)(\vx-\vmu_q)^T\right] \right] \\[0.7em]
&= \Tr \left[ \Sigma_{q}^{-1} \Sigma_{q} \right] \\[0.7em]
&= \Tr\left[ I_{D} \right] \\[0.7em]
&= D
\end{align}
ただし，$I_{D}$は$D$次元単位行列を表し，以下の分散共分散行列の定義を利用しました。
\begin{align}
\Sigma_{q} &= E_{q(X)} \left[(\vx-\vmu_q)(\vx-\vmu_q)^T\right]
\end{align}
\Sigma_{q} &= E_{q(X)} \left[(\vx-\vmu_q)(\vx-\vmu_q)^T\right]
\end{align}
第三項目の計算
続いて，第三項目を計算していきます。$\Sigma_p$と$\Sigma_q$は多変量正規分布の分散共分散行列ですから，対称行列です。したがって，第二項目と同様に第三項目に二次形式とトレースの性質を適用することができます。簡単のため，まずは期待値計算の中身だけを考えます。
\begin{align}
(\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) &= \Tr \left[ \Sigma_{p}^{-1}(\vx-\vmu_p)(\vx-\vmu_p)^T\right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1}(\vx\vx^T-\vx\vmu_p^T-\vmu_p\vx^T+\vmu_p\vmu_p^T)\right] \label{eq_tr_p}
\end{align}
(\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) &= \Tr \left[ \Sigma_{p}^{-1}(\vx-\vmu_p)(\vx-\vmu_p)^T\right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1}(\vx\vx^T-\vx\vmu_p^T-\vmu_p\vx^T+\vmu_p\vmu_p^T)\right] \label{eq_tr_p}
\end{align}
続いて，式($\ref{eq_tr_p}$)に期待値を施すことを考えます。ここで，期待値とトレースが交換可能であることを利用します。
期待値とトレースの交換
$D$次元正方行列$A$を考える。
\begin{align}
A &=
\begin{pmatrix}
a_{11} & \cdots & \\
\vdots & \ddots & \\
& & a_{DD}
\end{pmatrix}
\end{align}
A &=
\begin{pmatrix}
a_{11} & \cdots & \\
\vdots & \ddots & \\
& & a_{DD}
\end{pmatrix}
\end{align}
期待値の線形性とトレースの定義により，期待値とトレースが交換可能であることが分かる。
\begin{align}
E\left[ \Tr\left[ A \right] \right] &= E\left[ \sum_{d=1}^{D}a_{dd} \right] \\[0.7em]
&= \sum_{d=1}^{D} E\left[ a_{dd} \right] \\[0.7em]
&= \Tr\left[
\begin{pmatrix}
E[a_{11}] & \cdots & \\
\vdots & \ddots & \\
& & E[a_{DD}]
\end{pmatrix}
\right] \\[0.7em]
&= \Tr\left[ E\left[ A \right] \right]
\end{align}
E\left[ \Tr\left[ A \right] \right] &= E\left[ \sum_{d=1}^{D}a_{dd} \right] \\[0.7em]
&= \sum_{d=1}^{D} E\left[ a_{dd} \right] \\[0.7em]
&= \Tr\left[
\begin{pmatrix}
E[a_{11}] & \cdots & \\
\vdots & \ddots & \\
& & E[a_{DD}]
\end{pmatrix}
\right] \\[0.7em]
&= \Tr\left[ E\left[ A \right] \right]
\end{align}
したがって，式($\ref{eq_tr_p}$)に期待値を施すと以下のように変形できます。
\begin{align}
&E_{q(X)} \left[ \Tr \left[ \Sigma_{p}^{-1}(\vx\vx^T-\vx\vmu_p^T-\vmu_p\vx^T+\vmu_p\vmu_p^T)\right] \right] \notag \\[0.7em]
&= \Tr \left[ E_{q(X)} \left[ \Sigma_{p}^{-1}(\vx\vx^T-\vx\vmu_p^T-\vmu_p\vx^T+\vmu_p\vmu_p^T)\right] \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} E_{q(X)} \left[\vx\vx^T-\vx\vmu_p^T-\vmu_p\vx^T+\vmu_p\vmu_p^T \right] \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \left\{ E_{q(X)} \left[\vx\vx^T\right]-E_{q(X)} \left[\vx\vmu_p^T\right]-E_{q(X)} \left[\vmu_p\vx^T\right]+E_{q(X)} \left[\vmu_p\vmu_p^T \right] \right\} \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \left\{ E_{q(X)} \left[\vx\vx^T\right]-E_{q(X)}\left[\vx\right]\vmu_p^T-\vmu_pE_{q(X)} \left[\vx^T\right]+\vmu_p\vmu_p^T \right\} \right]
\end{align}
&E_{q(X)} \left[ \Tr \left[ \Sigma_{p}^{-1}(\vx\vx^T-\vx\vmu_p^T-\vmu_p\vx^T+\vmu_p\vmu_p^T)\right] \right] \notag \\[0.7em]
&= \Tr \left[ E_{q(X)} \left[ \Sigma_{p}^{-1}(\vx\vx^T-\vx\vmu_p^T-\vmu_p\vx^T+\vmu_p\vmu_p^T)\right] \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} E_{q(X)} \left[\vx\vx^T-\vx\vmu_p^T-\vmu_p\vx^T+\vmu_p\vmu_p^T \right] \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \left\{ E_{q(X)} \left[\vx\vx^T\right]-E_{q(X)} \left[\vx\vmu_p^T\right]-E_{q(X)} \left[\vmu_p\vx^T\right]+E_{q(X)} \left[\vmu_p\vmu_p^T \right] \right\} \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \left\{ E_{q(X)} \left[\vx\vx^T\right]-E_{q(X)}\left[\vx\right]\vmu_p^T-\vmu_pE_{q(X)} \left[\vx^T\right]+\vmu_p\vmu_p^T \right\} \right]
\end{align}
ここで，$E[\vx\vx^T]$に対して以下の分散共分散の定義を利用します。
分散共分散行列の定義
分散の定義からも類推される通り，分散共分散行列は以下で定義される。
\begin{align}
\Sigma &= E\left[ (\vx-\vmu)(\vx-\vmu)^T \right] \\[0.7em]
&= E\left[ \vx\vx^T \right]-E[\vx]E[\vx^T]
\end{align}
\Sigma &= E\left[ (\vx-\vmu)(\vx-\vmu)^T \right] \\[0.7em]
&= E\left[ \vx\vx^T \right]-E[\vx]E[\vx^T]
\end{align}
実際に適用してみましょう。
\begin{align}
&\Tr \left[ \Sigma_{p}^{-1} \left\{ E_{q(X)} \left[\vx\vx^T\right]-E_{q(X)}\left[\vx\right]\vmu_p^T-\vmu_pE_{q(X)} \left[\vx^T\right]+\vmu_p\vmu_p^T \right\} \right] \notag \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \left\{ \left(\Sigma_q+\vmu_q\vmu_q^T\right)-\vmu_q\vmu_p^T-\vmu_p\vmu_q^T+\vmu_p\vmu_p^T \right\} \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \left\{ \Sigma_q+\vmu_p\vmu_p^T-\vmu_p\vmu_q^T-\vmu_q\vmu_p^T+\vmu_q\vmu_q^T \right\} \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \left\{ \Sigma_q+\left( \vmu_p-\vmu_q \right)\left( \vmu_p-\vmu_q \right)^T \right\} \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \Sigma_q \right] + \Tr\left[ \Sigma_{p}^{-1} \left(\vmu_p-\vmu_q \right)\left( \vmu_p-\vmu_q \right)^T \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \Sigma_q \right] + \left(\vmu_p-\vmu_q \right)^{T}\Sigma_{p}^{-1}\left( \vmu_p-\vmu_q \right)
\end{align}
&\Tr \left[ \Sigma_{p}^{-1} \left\{ E_{q(X)} \left[\vx\vx^T\right]-E_{q(X)}\left[\vx\right]\vmu_p^T-\vmu_pE_{q(X)} \left[\vx^T\right]+\vmu_p\vmu_p^T \right\} \right] \notag \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \left\{ \left(\Sigma_q+\vmu_q\vmu_q^T\right)-\vmu_q\vmu_p^T-\vmu_p\vmu_q^T+\vmu_p\vmu_p^T \right\} \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \left\{ \Sigma_q+\vmu_p\vmu_p^T-\vmu_p\vmu_q^T-\vmu_q\vmu_p^T+\vmu_q\vmu_q^T \right\} \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \left\{ \Sigma_q+\left( \vmu_p-\vmu_q \right)\left( \vmu_p-\vmu_q \right)^T \right\} \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \Sigma_q \right] + \Tr\left[ \Sigma_{p}^{-1} \left(\vmu_p-\vmu_q \right)\left( \vmu_p-\vmu_q \right)^T \right] \\[0.7em]
&= \Tr \left[ \Sigma_{p}^{-1} \Sigma_q \right] + \left(\vmu_p-\vmu_q \right)^{T}\Sigma_{p}^{-1}\left( \vmu_p-\vmu_q \right)
\end{align}
KLダイバージェンスの計算
以上から，KLダイバージェンスは以下のように計算されます。
\begin{align}
&\left. \KL \left[ q(X) \right| p(X) \right] \notag \\[0.7em]
&= \frac{1}{2} \biggl\{ \log \frac{|\Sigma_{p}|}{|\Sigma_{q}|}-E_{q(X)} \left[ (\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q)\right]\notag\\[0.7em]
&\quad\quad+E_{q(X)}\left[(\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) \right] \biggr\} \\[0.7em]
&= \frac{1}{2} \left[ \log \frac{|\Sigma_{p}|}{|\Sigma_{q}|}-D+\left\{\Tr \left[ \Sigma_{p}^{-1} \Sigma_q \right] + \left(\vmu_p-\vmu_q \right)\Sigma_{p}^{-1}\left( \vmu_p-\vmu_q \right)^T\right\} \right] \\[0.7em]
&= \frac{1}{2} \left\{ \log \frac{|\Sigma_p|}{|\Sigma_q|}
+ \Tr \left[ \Sigma_{p}^{-1} \Sigma_{q} \right]
+ \left( \vmu_q-\vmu_p \right)^T \Sigma_{p}^{-1} \left( \vmu_q-\vmu_p \right)-D
\right\}
\end{align}
&\left. \KL \left[ q(X) \right| p(X) \right] \notag \\[0.7em]
&= \frac{1}{2} \biggl\{ \log \frac{|\Sigma_{p}|}{|\Sigma_{q}|}-E_{q(X)} \left[ (\vx-\vmu_q)^T\Sigma_{q}^{-1}(\vx-\vmu_q)\right]\notag\\[0.7em]
&\quad\quad+E_{q(X)}\left[(\vx-\vmu_p)^T\Sigma_{p}^{-1}(\vx-\vmu_p) \right] \biggr\} \\[0.7em]
&= \frac{1}{2} \left[ \log \frac{|\Sigma_{p}|}{|\Sigma_{q}|}-D+\left\{\Tr \left[ \Sigma_{p}^{-1} \Sigma_q \right] + \left(\vmu_p-\vmu_q \right)\Sigma_{p}^{-1}\left( \vmu_p-\vmu_q \right)^T\right\} \right] \\[0.7em]
&= \frac{1}{2} \left\{ \log \frac{|\Sigma_p|}{|\Sigma_q|}
+ \Tr \left[ \Sigma_{p}^{-1} \Sigma_{q} \right]
+ \left( \vmu_q-\vmu_p \right)^T \Sigma_{p}^{-1} \left( \vmu_q-\vmu_p \right)-D
\right\}
\end{align}
参考文献
本稿の執筆にあたり参考にした文献は，以下でリストアップしております。
コメント
コメント一覧 （2件）
分かりやすい導出で大変助かりました。どうも有難うございます。二点ほど気付いたのですが、
【1】(３８) の２項目が２次形式となるためには、転置の位置が最初の項にかかるかと思いました。
【2】(9)-(20)において期待値をとる分布が$q(X)$ではなく$q(Z)$になっている箇所がありました。またこれは個人的な好みになってしまいますが、(1)(2)の$N$が$\mathcal{N}$、期待値が$\mathbb{E}$でも良いのかなと思いました。以上念の為シェアさせていただきました。重ね重ね御礼申し上げます。
T.O.様
ご指摘ありがとうございます。
【1】について
ご指摘の通りです。本文を修正致しました。
【2】について
$q(Z)$に関して、ご指摘の通りです。本文を修正致しました。ノーテーションに関して、当サイトで一貫して正規分布は$\N(\cdot)$を用い、期待値は$E[\cdot]$を用いる方針となっております。こちらは専門書によっても表記揺れがある箇所であり、違和感がある場合でもご理解いただけますと幸いです。