EMアルゴリズムとは、その導出

はじめに

• EMアルゴリズム、混合正規分布モデル (GMM) と一緒に語られるせいで、説明が煩雑になりがち。だが、EMアルゴリズムの本質はGMMには無い。
• もっとEMアルゴリズムだけをクリアに説明して、その応用としてGMMを紹介する記事とかほしいので、この記事を書き始めました。

EMアルゴリズムの分かりづらさの原因

先に話しておくと、EMアルゴリズムの分かりづらさは次の3つにあると思います：

1. 解きたい問題の説明がまず難しい
2. 上界（下界）で最適化する議論が分かりづらい
3. 突然の分布$q$の登場で混乱

この記事では順番に説明していきます。そのためちょっと長くなっていますが、この3つが分かればだいたい大丈夫だと思います。

解きたい問題

EMアルゴリズムの文脈で解きたい問題は、ざっくり言うと「潜在変数があるモデルの最尤推定をしたい」になります。下で詳しく解説します。

状況説明

• 推定したいモデルがあり、そのモデルは次のことが分かっています：
  • 「原因$A$」から「中間$B$」が生成される
  • 「中間$B$」から「結果$C$」が生成されて値が観測される
• 「原因$A$」を司るパラメータを$\theta_1$,
  「中間$B$」の状態を表す変数を$z$, 「中間$B$」を司るパラメータを$\theta_2$¹,
  「結果$C$」で観測できる値を$x$とします。

  \underbrace{「原因A」}_{\theta_1}→\underbrace{「中間B」}_{z,\ \theta_2}→\underbrace{「結果C」}_{x}

• 観測できるのは「結果$C$」の観測値$x$だけです。「中間B」の状態$z$は観測できないので、潜在変数と呼ばれます。

仮定

• 「原因$A$」から「中間$B$」へのモデル（つまり、$\theta$から$z$がどうやって生成されるか）が分かっているものとします。この確率モデルを$p(z|\theta_1)$とします。
• 「中間$B$」から「結果$C$」へのモデル（つまり、$z, \theta_2$から$x$がどうやって生成されるか）が分かっているものとします。この確率モデルを$p(x|z, \theta_2)$とします。

  \underbrace{「原因A」}_{\theta_1}
  \hphantom{10}
  \underbrace{→}_{p(z|\theta_1)}
  \hphantom{10}
  \underbrace{「中間B」}_{z,\ \theta_2}
  \hphantom{10}
  \underbrace{→}_{p(x|z, \theta_2)}
  \hphantom{10}
  \underbrace{「結果C」}_{x}

問題

観測のデータセット$D = (x_1, x_2, \ldots, x_n)$が与えられたとき、パラメータ$\theta = (\theta_1, \theta_2)$を最尤推定せよ。すなわち、

\def\argmax{\mathop{\rm arg~max}\limits}
\def\argmin{\mathop{\rm arg~min}\limits}

\theta^* = \argmax_{\theta}\ p(D|\theta)

となる$\theta^*$を求めよ。このとき、$p(z|\theta_1),\ p(x|z,\theta_2)$を利用してよい。
ただし、各$x_i$に対応する潜在変数$z_i$は観測できないので与えられていない。

例: GMMの問題
GMMのモデルと変数の対応に関して説明します。
紹介したい本質の部分ではないのですが、イメージがつかない人向けに後ろに書いておきます。

解き方の考え方

この問題の難しい点

潜在変数の組を$S=(z_1,z_2, \ldots, z_n)$と略して表記します。
（本や他の記事などだと、$D = (x_1, x_2, \ldots, x_n)$を別表記で$x_{1:n}$と書いていたり、単に$X$などと書いていることが多いです。同様に$S=(z_1,z_2, \ldots, z_n)$を別表記で$z_{1:n}$と書いていたり、単に$Z$などと書いていることが多いです。表記は違いますが、同じものと思ってください）

この最尤推定をしようとすると

\begin{align}
  \theta^* &= \argmax_\theta\ p(D|\theta)\\
           &= \argmin_\theta\ -\log p(D|\theta)\\
           &= \argmin_\theta\ -\log \int p(D | S, \theta_2)p(S|\theta_1) dS
\end{align}

のような計算ができます。
$-\log$を付けたのは確率分布の最尤推定をする際に簡単になるケースが多いからです。
最後の式に出てくる$p(D | S, \theta_2),\ p(S | \theta_1)$はデータの独立性が担保されていれば、
$p(D | S, \theta_2) = \prod_i p(x_i|z_i,\theta_2),\ p(S | \theta_1) = \prod_i p(z_i|\theta_1)$でそれぞれ計算できます。

問題なのは$\log\int$の形が出ていることです。
$\log\int$は解析的に計算できることはほぼ無く、手計算で$\theta^*$を計算することはまず無理です。
(一方で、$\int\log$の形は計算しやすいことが多いです）

じゃあ何とか数値計算に持っていこうと思うわけですが、どうやって$\argmin$になる$\theta$を数値的に取り出す手法を考える必要があります。そこで次に紹介する”キモ”を基に出てくる手法がEMアルゴリズムです。

解き方のキモ

上で述べた問題に対して数値計算での最適化方法を考えます。
ここでのキモとなる考え方は、目的関数の上界を作り、その上界の中で最小化することを繰り返すことです。²

この操作を可視化したものが次のgifです：

この操作を繰り返して$\theta_t$を構成すると「目的関数の値が単調に減少していくこと」が一般的に証明されています。³⁴

このやり方の問題は2点あります：
 問題1. 「上界はどうやって作る？」
 問題2. 「上界の最小値はどうやって計算する？」
このアルゴリズムで中心の問題になってくるのは上界の作り方です。上界をテキトーに導出しても、その上界の最小値の計算は困難なことがほとんどです。導出した上界が複雑な式だったりするとあまり意味がありません。

そこでいい感じの上界を作るときに結果として導き出されたのが、次に紹介するELBOです。

上界（& ELBO）の導出

上で述べた通り、最小値の計算が簡単な上界を出したいです。
$\int\log$が出るように調整していくと良さそうです（そのため、一旦上の式変形は忘れて作っていきます）。

作った結論だけ述べますが⁵、$\int\log$が出るように目的関数の$-\log p(D|\theta)$を式変形してみると

\begin{align}
 - \log p(D|\theta) 
 &=  -\int q(S) \log\frac{p(D,S|\theta)}{q(S)} dS\ \underbrace{-\ \mathrm{KL}(q(S)||p(S|D,\theta))}_{\leq 0}\\
 &\leq  -\int q(S) \log\frac{p(D,S|\theta)}{q(S)} dS
\end{align}

という上界が取得できます。ここで$q$は任意の分布です。

突然出てきた任意の分布$q$に面を食らってしまう方がほとんどでしょう。
大丈夫です。すぐにこの正体は説明します。

ELBOとは

用語の説明をします。$\log p(D|\theta)$はエビデンスと呼ばれます。
紹介した上界に関して符号を反転すると、$\log p(D|\theta) \geq \int q(S) \log\frac{p(D,S|\theta)}{q(S)} dS$ですから、右辺はエビデンスの下界になっています。
そこで右辺の$\int q(S) \log\frac{p(D,S|\theta)}{q(S)} dS$はELBO（Evidence Lower Bound) という名前が付けられています。

今後書くのが面倒なので$\mathrm{ELBO}(q(S), p(D,S|\theta)) := \int q(S) \log\frac{p(D,S|\theta)}{q(S)}dS$ とします。
これを使えば上の式は次の形で書け、すっきり（？）します。

\begin{align}
 - \log p(D|\theta) 
 &=  -\mathrm{ELBO}(q(S), p(D,S|\theta))\ \underbrace{-\ \mathrm{KL}(q(S)||p(S|D,\theta))}_{\leq 0}\\
 &\leq  -\mathrm{ELBO}(q(S), p(D,S|\theta))
\end{align}

重要なのはELBOが$\int\log$の形をしていることです。つまり計算が容易な形であり、$\theta$に対する最適化も容易です。

任意の分布qはどうするのか

結局$q$は何者なのでしょうか。
考え方としては「$q$は上界$-\mathrm{ELBO}(q(S), p(D,S|\theta))$のパラメータで、上界がどれだけ目的関数に近いかを決めるもの」と理解できそうです。
そうなると、上界が最も目的関数に近いように$q$を設計したくなります。

上の等式の各項の関係をよく見てみると、$-\mathrm{KL}(q(S)||p(S|D,\theta))$の項が「目的関数と上界のギャップ」になっています。
つまり、KLの項が0になるように$q$をとれば良さそうです。

すなわち、更新$t$回目で$\theta_t$が与えられているとき、上界を最も目的関数に近くするためには

q(S) := p(S|D, \theta_t)

と設定してあげれば良さそうです。

EMアルゴリズム

上の流れをまとめると、「目的関数の上界を作り、その上界の中で最小化する」は下のようなステップで計算できそうです。

愚直に書き下したタイプ

1. 上界を計算
   a. $q(S) = p(S|D, \theta_t)$を計算
   b. 上界$-\mathrm{ELBO}(q(S), p(D,S|\theta))$を計算
2. $\theta_{t+1} = \argmin {}_\theta\ -\mathrm{ELBO}(q(S), p(D,S|\theta))$で更新

いい感じにまとめたタイプ

1. $\mathrm{ELBO}(p(S|D, \theta_t), p(D,S|\theta))$を計算
   • $p(D,S|\theta)$は$p(D,S|\theta) = p(D | S, \theta_2)p(S|\theta_1) = \prod_i p(x_i|z_i,\theta_2) \prod_i p(z_i|\theta_1)$で計算可
   • $p(S|D, \theta)$は$p(S|D, \theta) = \frac{p(D,S|\theta)}{\int p(D, S|\theta)dS}$で計算可
2. $\theta_{t+1} = \argmax {}_\theta\ \mathrm{ELBO}(p(S|D, \theta_t), p(D,S|\theta))$で更新

上のアルゴリズムをEM (Expectation-Maximization) アルゴリズムと呼びます。
第1ステップはELBOを計算するため、期待値計算を含んでいます。そのためEステップ (Expectation step) と呼ばれます。第2ステップは最大化を行うため、Mステップ (Maximization step) と呼ばれます。

例: GMMに対するEMアルゴリズム
紹介したい本質の部分ではないのですが、イメージがつかない人向けに後ろに書いておきます。

おわりに

• ぶっちゃけ、まとめてようやく理解した
• MMアルゴリズムの枠組みで導出したかったので一部冗長だったかも

おまけ：GMMでの理解

GMMで解きたい問題

GMMのモデルと変数の対応は次のようになります:

• 「原因$A$」は「複数のガウス分布のどれからデータ点を生成するかを決める、確率値のパラメータ」です。このパラメータが$\theta_1 = \lbrace\pi_k\rbrace_{k=1}^K$です。
• 「原因$A$」から「中間$B$」での生成では、生成するガウス分布としてどれが選ばれたかを示すone-hotベクトル$\boldsymbol{z} = (z_1, \ldots, z_K)$を潜在変数として生成します。頑張って計算すると、$p(\boldsymbol{z}|\theta_1) = \prod_k \pi_k^{z_k}$となります。
• 「中間$B$」は潜在変数$\boldsymbol{z}$と「どのガウス分布からデータ点を生成するかが決まっているとき、どんなガウス分布を使うかを表すパラメータ」です。このパラメータは$\theta_2=\lbrace (\mu_k, \Sigma_k) \rbrace_{k=1}^K$です。
• 「中間$B$」から「結果$C$」の生成では、実際に選択されたガウス分布からデータ点$x$を生成します。頑張って計算すると、$p(x|\boldsymbol{z},\theta_2) = \prod_k \mathcal{N}(x|\mu_k, \Sigma_k)^{z_k}$となります。
• 「結果$C$」は「以上のプロセスでサンプリングされた点」です。すなわち混合ガウス分布からサンプリングされた点になります。
• GMMで解くべき問題は、「混合ガウス分布からサンプリングされた点が複数与えられたとき、その混合ガウス分布を決めるパラメータ$\theta = \{(π_k, \mu_k, \Sigma_k)\}_{k=1}^K$は何だったか？最尤推定せよ」です。

※ 詳細は@kenmatsu4さんの記事などを参考にしてください。

GMMに対するEMアルゴリズム

あとで書く

1. ここの$\theta_2$の説明はかなり苦しい内容なのでわかりにくいです。すぐ下の仮定まで読んだら$\theta_2$が何者かわかると思います。 ↩

2. これは上界最小化アルゴリズム（MMアルゴリズム）という一般的な枠組みで用いられる考え方です。つまりEMアルゴリズムはMMアルゴリズムを特定の問題（つまり「解きたい問題」の章で述べた問題）に適用することで導出できます。 ↩

3. 例えば「金森敬文, 鈴木大慈, 竹内一郎, 佐藤一誠. 機械学習のための連続最適化（機械学習プロフェッショナルシリーズ）, 講談社, 2016」のpp.223-224を参照のこと。
   MMアルゴリズムで単調減少性を担保するには一定の条件が必要ですが、このページではEMアルゴリズムの導出の紹介だけにとどめておきます。一定の条件とは、構成した下界$u_{\theta_t}(\theta)$が$\theta=\theta_t$で$f(\theta)$と一致することです。すなわち、$f(\theta_t) = u_{\theta_t}(\theta_t)$を確認する必要があります。難しくないので計算してみるといいと思います。つまり、本ページの最後に紹介しているアルゴリズムの$\mathrm{ELBO}$の$\theta$に$\theta_t$を代入してみてください。 ↩

4. 単調減少性しか担保していないため、EMアルゴリズムは局所最小解で停止する可能性があります。すなわち、必ずしも最適解にはたどり着かないです。 ↩

5. 証明は簡単です。等式の右辺を変形すると左辺と一致するのは簡単に確認できると思います。問題は「何でこの式が出てきたか」なのですが、まずは「最適化しやすい上界を先人が頑張って作った」くらいの理解でいいかと思います。実はELBOは$\log\int$の形式を$\int\log$の形式にするために、イェンセンの不等式を用いて導出が簡単にできます。ここでは流れを重視して割愛しました。 ↩