今更聞けないEMアルゴリズムの解説

今更ながらEMアルゴリズムとは何かについて、調査・勉強したのでまとめておく。Andrew氏のレクチャノートとパターン認識と機械学習下巻第9章を参考にした。

EMアルゴリズムの目的

観測変数$\boldsymbol{x}$と観測できない潜在変数$\boldsymbol{z}$を含む確率モデルの尤度関数$P( \boldsymbol{x} \mid \boldsymbol{\theta})$ を最大化するパラメータ$\boldsymbol{\theta}$ を見つけることである。

比較的複雑な観測変数の(周辺)分布$P( \boldsymbol{x})$を、より扱いやすい観測変数と潜在変数の同時分布$P( \boldsymbol{x}, \boldsymbol{z})$によって表すことができることがある（例えば混合正規分布などがそれに相当する）。このように、モデルを簡単に扱うために潜在変数$\boldsymbol{z}$を導入し、EMアルゴリズムの効果的な適用を図るケースがしばしばある。

EMアルゴリズムの解説

E-Stepの説明

E-Stepでは、$\boldsymbol{\theta}$固定の下、尤度関数の下界の分布を最大化する。以下で尤度関数を変形し、下界を求める手順を示す。

$$\begin{eqnarray} \displaystyle \sum_{ i = 1 }^{ N } \ln P( \boldsymbol{x}_i \mid \boldsymbol{\theta}) &=& \displaystyle \sum_{ i = 1 }^{ N } \ln \displaystyle \sum_{\boldsymbol{z}_i} P( \boldsymbol{x}_i, \boldsymbol{z}_i \mid \boldsymbol{\theta})\\ &=& \displaystyle \sum_{ i = 1 }^{ N } \ln \displaystyle \sum_{\boldsymbol{z}_i} Q(\boldsymbol{z}_i)\frac{P( \boldsymbol{x}_i, \boldsymbol{z}_i \mid \boldsymbol{\theta})}{Q(\boldsymbol{z}_i)}\\ &\geq&  \displaystyle \sum_{ i = 1 }^{ N }\displaystyle \sum_{\boldsymbol{z}_i} Q(\boldsymbol{z}_i) \ln \frac{P( \boldsymbol{x}_i, \boldsymbol{z}_i \mid \boldsymbol{\theta})}{Q(\boldsymbol{z}_i)} \end{eqnarray}$$

(1)式から(2)式への変形は、$\boldsymbol{z}_i$の任意の確率分布$Q(\boldsymbol{z}_i)$でかけて割っただけである。
(2)式から(3)式への変形は、Jensen's Inequalityを利用した。以下の(4)式を参照されたい。$\ln$が凹関数なので、期待値を関数に入れた値のほうが、関数に入れた値の期待値より大きい。今回の場合、$\frac{P( \boldsymbol{x}_i, \boldsymbol{z}_i \mid \boldsymbol{\theta})}{Q(\boldsymbol{z}_i)}$を値と捉え、$Q(\boldsymbol{z}_i)$によって期待値を算出することを考える。

$$\begin{eqnarray}  \ln E_{\boldsymbol{z}_i 〜 Q(\boldsymbol{z}_i) } \left[ \frac{P( \boldsymbol{x}_i, \boldsymbol{z}_i \mid \boldsymbol{\theta})}{Q(\boldsymbol{z}_i)}\right] &\geq& E_{\boldsymbol{z}_i 〜 Q(\boldsymbol{z}_i) } \ln \left[ \frac{P( \boldsymbol{x}_i, \boldsymbol{z}_i \mid \boldsymbol{\theta})}{Q(\boldsymbol{z}_i)}\right] \end{eqnarray}$$

(2)式から(3)式の変形は、$Q(\boldsymbol{z}_i)$がどのような確率分布であっても成立する。
尤度関数最大化の目的を考えると、$\boldsymbol{\theta}$固定の下、下界(3)式を最大化、つまり(3)の等号の成立を図るのが自然であろう。
Jensen's Inequalityの等号の成立条件は、凸関数が線形関数の場合か、狭義凸関数（線形ではない）の場合は、凸関数の中身が１点分布となる場合である。つまり以下を満たす。

$$\begin{eqnarray} \displaystyle \frac{P( \boldsymbol{x}_i, \boldsymbol{z}_i \mid \boldsymbol{\theta})}{Q(\boldsymbol{z}_i)} &=& const \end{eqnarray}$$

また、$\sum_{\boldsymbol{z}_i }Q(\boldsymbol{z}_i) =1$より、以下が成り立つ。

$$\begin{eqnarray} Q(\boldsymbol{z}_i)  &=& \displaystyle \frac{P( \boldsymbol{x}_i, \boldsymbol{z}_i \mid \boldsymbol{\theta})}{\sum_{\boldsymbol{z}_i }P( \boldsymbol{x}_i, \boldsymbol{z}_i \mid \boldsymbol{\theta})}\\ &=&  \displaystyle \frac{P( \boldsymbol{z}_i \mid \boldsymbol{x}_i, \boldsymbol{\theta})P( \boldsymbol{x}_i \mid \boldsymbol{\theta})}{P( \boldsymbol{x}_i \mid \boldsymbol{\theta})}\\ &=& P( \boldsymbol{z}_i \mid \boldsymbol{x}_i, \boldsymbol{\theta}) \end{eqnarray}$$

すべての$i$について(8)式を実施することが、E-Stepにて行うことである。

（補足）E-Stepのその他の説明

(3)式の下界の各々$i$について、以下のように式変形する。
$$\begin{eqnarray} \ln P( \boldsymbol{x}_i \mid \boldsymbol{\theta}) &\geq& \displaystyle \sum_{\boldsymbol{z}_i} Q(\boldsymbol{z}_i) \ln \frac{P( \boldsymbol{x}_i, \boldsymbol{z}_i \mid \boldsymbol{\theta})}{Q(\boldsymbol{z}_i)}\\ &=& \displaystyle \sum_{\boldsymbol{z}_i} Q(\boldsymbol{z}_i) \ln \frac{P( \boldsymbol{z}_i \mid \boldsymbol{x}_i, \boldsymbol{\theta})P( \boldsymbol{x}_i \mid \boldsymbol{\theta})}{Q(\boldsymbol{z}_i)}\\ &=&  \ln P( \boldsymbol{x}_i \mid \boldsymbol{\theta}) + \displaystyle \sum_{\boldsymbol{z}_i} Q(\boldsymbol{z}_i) \ln \frac{P( \boldsymbol{z}_i \mid \boldsymbol{x}_i, \boldsymbol{\theta})}{Q(\boldsymbol{z}_i)}\\ &=& \displaystyle \ln P( \boldsymbol{x}_i \mid \boldsymbol{\theta}) - KL(Q(\boldsymbol{z}_i)||P( \boldsymbol{z}_i \mid \boldsymbol{x}_i, \boldsymbol{\theta})) \end{eqnarray}$$

(10)式から(11)式への変形は、$\boldsymbol{z}_i$に依存しない$P( \boldsymbol{x}_i \mid \boldsymbol{\theta})$を前に出し、$Q(\boldsymbol{z}_i)$の$\boldsymbol{z}_i$に関する積分が1になることによる。

さて、下界である右辺の最大化を考えた場合、(12)式第2項$Kullback–Leibler$ divergenceの最小化が必要となる。$KL$ divergence $\geq 0$より、  $KL(Q(\boldsymbol{z}_i)||P( \boldsymbol{z}_i \mid \boldsymbol{x}_i, \boldsymbol{\theta})) = 0$となる$Q(\boldsymbol{z}_i)$を求めることと同義となる。よって、$Q(\boldsymbol{z}_i) = P( \boldsymbol{z}_i \mid \boldsymbol{x}_i, \boldsymbol{\theta})$を求めることとなり、これは(8)式と一致する。

また、$\boldsymbol{z}_i$が連続変数の場合のE-Stepについては、以下に変分法を用いた説明を加えたので、参照していただきたい。
今更聞けないEMアルゴリズムの解説〜潜在変数が連続変数の場合のEステップの説明〜

M-Stepの説明

今度は、$Q(\boldsymbol{z}_i)$を固定し、$\boldsymbol{\theta}$を動かし、尤度関数の下界分布の最大化を図る。

$$\begin{eqnarray} \hat{\theta}&=&\mathop{\arg\,\max}\limits_\boldsymbol{\theta} \displaystyle \sum_{\boldsymbol{z}_i} Q(\boldsymbol{z}_i) \ln \frac{P( \boldsymbol{x}_i, \boldsymbol{z}_i \mid \boldsymbol{\theta})}{Q(\boldsymbol{z}_i)}\\ &=& \mathop{\arg\,\max}\limits_\boldsymbol{\theta} \displaystyle \sum_{\boldsymbol{z}_i} Q(\boldsymbol{z}_i) \ln P( \boldsymbol{x}_i, \boldsymbol{z}_i \mid \boldsymbol{\theta})\\ &=& \mathop{\arg\,\max}\limits_\boldsymbol{\theta} \displaystyle \sum_{\boldsymbol{z}_i} P( \boldsymbol{z}_i \mid \boldsymbol{x}_i, \boldsymbol{\theta}_{old}) \ln P( \boldsymbol{x}_i, \boldsymbol{z}_i \mid \boldsymbol{\theta}) \end{eqnarray}$$

(13)式から(14)式の変形は$\boldsymbol{\theta}$によって影響のない部分を排除した。
(15)式はE-Stepで求めた、$Q(\boldsymbol{z}_i) = P( \boldsymbol{z}_i \mid \boldsymbol{x}_i, \boldsymbol{\theta}_{old})$を代入した。

このように、E-Stepにて潜在変数の事後確率を求め、その事後確率によってM-Stepにて期待値計算を行い、目的のパラメータを更新する。このプロセスを収束するまで繰り返す。イメージを下図に示す。

パターン認識と機械学習 下 (ベイズ理論による統計的予測)  第9章 図9.14

以上、EMアルゴリズムの全貌である。複雑な分布を簡単な分布の構成として扱うことで、簡単に尤度関数の最適化ができるようになった。EMアルゴリズム適用時のポイントは、何を潜在変数の分布として、何を観測変数として扱うのかというモデル化であろう。卑近な混合正規分布への適用例をしっかりと学び、その類推で応用していくのが簡単に思える。

EMアルゴリズムを拡張した変分推論については、また今度。

参考文献

CS229 Lecture notes Andrew Ng

パターン認識と機械学習 下 (ベイズ理論による統計的予測)  第9章, C.M. ビショップ (著), 元田 浩, 栗田 多喜夫, 樋口 知之, 松本 裕治, 村田 昇 (監訳)