EMアルゴリズムの幾何学的解釈

「EMアルゴリズムが幾何学的に説明できる」と聞いた驚きが情報幾何に興味を持ったきっかけということで、そこのところを綺麗に整理してみます。議論の元ネタは、「情報幾何学の新展開」の第12章です。

幾何を用いない（普通の）EMアルゴリズムの説明は、こちらを参照ください。

確率分布が構成する空間

観測可能な変数 $\mathbf X$ と観測できない変数（Latent variable） $\mathbf Z$ を持つ確率分布について、考えうるすべての分布 $p(\mathbf X,\mathbf Z)$ を集めた空間 $S$ を用意します。

この中で特に、パラメータ $\theta$ で特徴づけられたモデルの分布 $p(\mathbf X,\mathbf Z \mid \theta)$ を集めると、これは、空間 $S$ の部分空間 $M$ を構成します。これを「モデル空間」と呼びます。

一方、観測データ $\{{\mathbf X}_n\}_{n=1}^N$ が与えられた場合、この観測データが得られる確率が 1 になる（この観測データにオーバーフィッティングした）確率分布が構成できます。

　 $\hat q(\mathbf X,\mathbf Z) = \prod_{n=1}^N \delta(\mathbf X-\mathbf X_n)q(\mathbf Z)$

ここに、 $q(\mathbf Z)$ は、 $\sum_{\mathbf Z}q(\mathbf Z)=1$ を満たす任意の関数です。

このような $\hat q(\mathbf X,\mathbf Z)$ をすべて集めたものは、やはり、空間 $S$ の部分空間 $D$ となります。これを「データ空間」と呼びます。

最短距離を用いた推定方法

Latent variableを持たないモデルでは、観測データの確率を 1 にする分布は1つにきまるので、これは、 $S$ 上の1点 $q$ を表します。そこで、 $q$ からモデル空間 $M$ に何らかの意味で「正射影」して得られる $M$ 上の点を推定分布として採用することが可能になります。「正射影」の意味（つまり、空間 $S$ に導入する接続）のとり方によって、推定方法が変わることになります。

一方、今の場合は、Latent variableがあるために、観測データに対応する分布は1つに決まらず、部分空間 $D$ を構成しました。そこで、データ空間 $D$ とモデル空間 $M$ の「最短距離」を実現する2点を見つけ出して、空間 $M$ 側の点を推定分布として採用するという方法が考えられます。

ここで特に「距離」として、データ空間上の点からモデル空間上の点に対するKLダイバージェンスを採用します。

　 ${\rm KL}\left(\hat q(\mathbf X,\mathbf Z)\,||\,p(\mathbf X,\mathbf Z \mid \theta)\right) = {\rm KL}\left(\prod_{n=1}^N \delta(\mathbf X-\mathbf X_n)q(\mathbf Z)\,||\, p(\mathbf X,\mathbf Z \mid \theta)\right)$

上記を $\theta$ および $q(\mathbf Z)$ の関数（汎関数）と見なして、これを最小にする $\theta$ を推定パラメータとして採用します。

実は、これによって得られる $\theta$ は下記の確率を極大にします。すなわち、EMアルゴリズムと同じ最尤推定になっていることが証明されます。（証明は後ほど。）

　 $\ln p(\mathbf X \mid \mathbf \theta) = \ln \sum_{\mathbf Z}p(\mathbf X,\mathbf Z \mid \mathbf \theta)$

emアルゴリズム

「emアルゴリズム」とは、EMアルゴリズムの幾何学バージョンで次のような繰り返し操作になります。

はじめにモデル空間 $M$ の一点 $p_{\rm old}$ を任意に選択して、データ空間 $D$ 上で、ここへのKLダイバージェンスが最小になる点を探します。

　 $q = {\arg\min}_{q \in D} {\rm KL}(q\,||\,p_{\rm old})$ 　――― (1)

$p_{\rm old}$ から上記の $q$ への対応を「e-射影」と呼びます。

続いて、モデル空間 $M$ 上で、 $q$ からのKLダイバージェンスが最小になる点を探します。

　 $p_{\rm new} = {\arg\min}_{p \in M} {\rm KL}(q\,||\,p)$ 　――― (2)

$q$ から上記の $p_{\rm new}$ への対応を「m-射影」と呼びます。

$p_{\rm new}$ を新たに $p_{\rm old}$ として、上記の「e-射影」と「m-射影」の操作を繰り返します。それぞれの射影でKLダイバージェンスは単調に減少していくので、最終的にデータ空間 $D$ とモデル空間 $M$ の間のKLダイバージェンスを極小にする点の組に収束します。

EMアルゴリズムとの対応

(1)(2)の計算を実際に行ってみると、それぞれ、EMアルゴリズムの操作、すなわち、事後分布

　 $q(\mathbf Z) = p(\mathbf Z \mid \{\mathbf X_n\}, \mathbf \theta_{\rm old})$

の計算（Eステップ）と、事後分布の下での対数尤度の期待値

　 $Q(\mathbf \theta,\mathbf\theta_{\rm old}) = \sum_{\mathbf Z}q(\mathbf Z)\ln p(\{\mathbf X_n\}, \mathbf Z \mid \mathbf \theta)$

の最大化（Mステップ）に一致することが分かります。

従って、emアルゴリズムは、EMアルゴリズムと同じ結果（最尤推定）を導くことが証明されます。

それでは、(1)(2)を実際に計算してみます。まずは(1)です。

　 ${\rm KL}(q\,||\,p_{\rm old}) = \int \sum_{\mathbf Z}\prod_{n=1}^N \delta(\mathbf X-\mathbf X_n)q(\mathbf Z)\ln \frac{\prod_{n=1}^N \delta(\mathbf X-\mathbf X_n)q(\mathbf Z)}{p(\mathbf X,\mathbf Z \mid \theta_{\rm old})}\,d{\mathbf X}$

　　　　 $=\sum_{\mathbf Z}q(\mathbf Z)\left\{\ln q(\mathbf Z)-\ln p(\{\mathbf X_n\},\mathbf Z \mid \theta_{\rm old})\right\}$

束縛条件 $\sum_{\mathbf Z}q(\mathbf Z)=1$ を考慮して、ラグランジュの未定乗数項を加えて $q(\mathbf Z)$ の変分を取ります。

　 $\delta\left\{{\rm KL}(q\,||\,p_{\rm old})+\lambda(\sum_{\mathbf Z}q(\mathbf Z)-1)\right\}$

　　　　 $=\sum_{\mathbf Z}\left\{\ln q(\mathbf Z)-\ln p(\{\mathbf X_n\},\mathbf Z \mid \theta_{\rm old})+1+\lambda\right\}\delta q(\mathbf Z)$

上記が任意の $\delta q(\mathbf Z)$ に対して 0 になることから、

　 $q(\mathbf Z) = e^{-(\lambda+1)}p(\{\mathbf X_n\},\mathbf Z \mid \theta_{\rm old})$ 　――― (3)

さらに束縛条件より、

　 $1=\sum_{\mathbf Z}q(\mathbf Z)=e^{-(\lambda+1)}\sum_{\mathbf Z}p(\mathbf Z \mid \{\mathbf X_n\}, \theta_{\rm old})p(\{\mathbf X_n\} \mid \theta_{\rm old})$

　　　　 $=e^{-(\lambda+1)}p(\{\mathbf X_n\} \mid \theta_{\rm old})$

これから決まる $e^{-(\lambda+1)}$ を(3)に代入して、

　 $q(\mathbf Z) = \frac{p(\{\mathbf X_n\},\mathbf Z \mid \theta_{\rm old})}{p(\{\mathbf X_n\} \mid \theta_{\rm old})} = p(\mathbf Z \mid \{\mathbf X_n\}, \theta_{\rm old})$

これは、EMアルゴリズムにおけるEステップ（事後分布の計算）と同じ計算になります。

続いて、上記で決まる $q(\mathbf Z)$ を用いて (2) を計算します。

　 ${\rm KL}(q\,||\,p) = \int \sum_{\mathbf Z}\prod_{n=1}^N \delta(\mathbf X-\mathbf X_n)q(\mathbf Z)\ln \frac{\prod_{n=1}^N \delta(\mathbf X-\mathbf X_n)q(\mathbf Z)}{p(\mathbf X,\mathbf Z \mid \theta)}\,d{\mathbf X}$