読者です 読者をやめる 読者になる 読者になる

めもめも

このブログに記載の内容は個人の見解であり、必ずしも所属組織の立場、戦略、意見を代表するものではありません。

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}

    =\sum_{\mathbf Z}q(\mathbf Z)\ln q(\mathbf Z) - \sum_{\mathbf Z}q(\mathbf Z)\ln p(\{\mathbf X_n\},\mathbf Z \mid \theta) ――― (4)

したがって、(4)を最小にする \theta を求めることは、

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

を最大にする \theta を求めることに他ならず、EMアルゴリズムにおけるMステップ(事後分布の下での対数尤度の期待値の最大化)と同じことになります。