EMアルゴリズムで今更GMMを解く

趣旨

前回紹介したEMアルゴリズムを用いて超有名なグラフ混合分布(GMM)を解く。煩雑な計算については極力述べず、統計的な観点から重要な点だけおさえる。

usapyoi.hatenablog.com

なんで今更?

うるさいなぁ理解を深めるため。

では解説する。

表記法

1. 入力とハイパーパラメタ

  • D: 入力次元数
  • N: 入力データの数
  • K: 重ね合わせるガウス分布の数　未知だが推定不可なので決め打ちする(ハイパーパラメタ)
  • \boldsymbol{X} = (\boldsymbol{x}_1, \boldsymbol{x}_2, \cdots, \boldsymbol{x}_N) \in \mathbb{R}^{D \times N}: 総入力

2. モデルパラメタ

  • \boldsymbol{\pi} = (\pi_1, \pi_2, \cdots, \pi_K) \in \mathbb{R}^{K}: モデルパラメタその1　各ガウス分布の混合係数
  • \boldsymbol{\mu} = (\boldsymbol{\mu}_1, \boldsymbol{\mu}_2, \cdots, \boldsymbol{\mu}_K) \in \mathbb{R}^{D \times K}: モデルパラメタその2　各ガウス分布の平均
  • \boldsymbol{\Sigma} = (\boldsymbol{\Sigma}_1, \boldsymbol{\Sigma}_2, \cdots, \boldsymbol{\Sigma}_K) \in \mathbb{R}^{D \times D \times K}: モデルパラメタその3　各ガウス分布の分散共分散行列
  • \boldsymbol{\theta}: 上の3つのモデルパラメタを総称したもの。分けて書く必要性が特にない時はこっちを使う。

3. 潜在変数

  • \boldsymbol{Z} = (\boldsymbol{z}_1, \boldsymbol{z}_2, \cdots, \boldsymbol{z}_N) \in {\lbrace 0,1\rbrace}^{K \times N}: 潜在変数　対応する \boldsymbol{x}_nがどのガウス分布に所属するかを示す(後述)

問題設定

以下のようなガウス分布の線形重ね合わせを考える。

p(\boldsymbol{x}_n | \boldsymbol{\pi}, \boldsymbol{\mu}, \boldsymbol{\Sigma}) = \sum_{k=1}^K \pi_k \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_k, \boldsymbol{\Sigma}_k)

ただし、 \sum_{k=1}^K \pi_k = 1。これを最大にするようなモデルパラメタ \boldsymbol{\pi}, \boldsymbol{\mu}, \boldsymbol{\Sigma}を求めたい。

解法

Eステップ

負担率 r_{nk}^{(i)}を計算する。

r_{nk}^{(i)} = \frac{\pi_k^{(1)} \mathcal{N}(\boldsymbol{x}_n | \boldsymbol{\mu}_k^{(i)}, \boldsymbol{\Sigma}_k^{(i)})}{\sum_{\kappa=1}^K \pi_\kappa^{(i)} \mathcal{N}(\boldsymbol{x}_n | \boldsymbol{\mu}_\kappa^{(i)}, \boldsymbol{\Sigma}_\kappa^{(i)})}

期待値の式は省略。

Mステップ

さっき省略した期待値を最大にするようにパラメタを更新する。ここで N_kを以下のように定義する。

N_k^{(i)} = \sum_{n=1}^{N} r_{nk}^{(i)}

これを用いて、

\begin{align} \pi_k^{(i+1)} &= \frac{N_k^{(i)}}{N} \\\ \boldsymbol{\mu}_k^{(i+1)} &= \frac{\sum_{n=1}^N r_{nk}^{(i)} \boldsymbol{x}_n}{N_k^{(i)}} \\\ \boldsymbol{\Sigma}_k^{(i+1)} &= \frac{\sum_{n=1}^N r_{nk}^{(i)} (\boldsymbol{x}_n - \boldsymbol{\mu}_k^{(i)}){(\boldsymbol{x}_n - \boldsymbol{\mu}_k^{(i)})}^T }{N_k^{(i)}} \\\ \end{align}

なんでそうなるの?

EMアルゴリズムについては前の記事を参照のこと。

usapyoi.hatenablog.com

Eステップ

EMアルゴリズムによれば、Eステップでは期待値 Q^{(i)}(\boldsymbol{\theta})の計算を行う。

\begin{align} Q^{(i)}(\boldsymbol{\theta}) & = \mathbb{E}_{\boldsymbol{Z} \sim p(\boldsymbol{Z} | \boldsymbol{X}, \boldsymbol{\theta}^{(i)})} \left\lbrack \log p(\boldsymbol{X}, \boldsymbol{Z} | \boldsymbol{\theta}) \right\rbrack \\\ \end{align}

さて、潜在変数 \boldsymbol{Z}を導入し、まず p(\boldsymbol{X}, \boldsymbol{Z} | \boldsymbol{\theta})、次に p(\boldsymbol{Z} | \boldsymbol{X}, \boldsymbol{\theta}^{(i)})を求めよう。

1. 潜在変数の導入

潜在変数 \boldsymbol{z}_nを \boldsymbol{x}_nがどのガウス分布に所属するかを示すOne-hotベクトルとして定義する。つまり、

z_{nk} = \left\lbrace \begin{array}{cc} 1 & \text{if } \boldsymbol{x}_n \sim \mathcal{N} (\boldsymbol{\mu}_k, \boldsymbol{\Sigma}_k) \\ 0 & \text{otherwise} \end{array} \right.

そして、 \boldsymbol{z}_nを各 nについて並べたものを \boldsymbol{Z}とする。

2. p(\boldsymbol{X}, \boldsymbol{Z} | \boldsymbol{\theta})の計算

\sum_{k=1}^K \pi_k = 1より、 \boldsymbol{x}_nが k番目のガウス分布に所属する確率は \pi_kと書ける。よって、

\begin{align} p(\boldsymbol{z}_n | \boldsymbol{\pi}) &= \prod_{k=1}^K \pi_k^{z_{nk}} \\\ p(\boldsymbol{Z} | \boldsymbol{\pi}) &= \prod_{n=1}^N \prod_{k=1}^K \pi_k^{z_{nk}} \\\ \end{align}

一方、 \boldsymbol{z}_nを用いることで、以下のような関係性も成り立つ。

\begin{align} p(\boldsymbol{x}_n | \boldsymbol{z}_n, \boldsymbol{\mu}, \boldsymbol{\Sigma}) &= \prod_{k=1}^K {\mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_k, \boldsymbol{\Sigma}_k)}^{z_{nk}} \\\ p(\boldsymbol{X} | \boldsymbol{Z}, \boldsymbol{\mu}, \boldsymbol{\Sigma}) &= \prod_{n=1}^N \prod_{k=1}^K {\mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_k, \boldsymbol{\Sigma}_k)}^{z_{nk}} \\\ \end{align}

以上より、 p(\boldsymbol{X}, \boldsymbol{Z} | \boldsymbol{\theta})について

\begin{align} p(\boldsymbol{X}, \boldsymbol{Z} | \boldsymbol{\theta}) &= p(\boldsymbol{Z} | \boldsymbol{\pi}) p(\boldsymbol{X} | \boldsymbol{Z}, \boldsymbol{\mu}, \boldsymbol{\Sigma}) \\\ &= \prod_{n=1}^N \prod_{k=1}^K {\left\lbrace \pi_k \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_k, \boldsymbol{\Sigma}_k) \right\rbrace}^{z_{nk}} \end{align}

と書ける。

3. p(\boldsymbol{Z} | \boldsymbol{X}, \boldsymbol{\theta}^{(i)})の計算

ベイズの定理より、

\begin{align} p(\boldsymbol{Z} | \boldsymbol{X}, \boldsymbol{\theta}^{(i)}) &= \frac{p(\boldsymbol{X}, \boldsymbol{Z} | \boldsymbol{\theta}^{(i)})}{p(\boldsymbol{X} | \boldsymbol{\theta}^{(i)})} \\\ &= \frac{\prod_{n=1}^N \prod_{k=1}^K {\left\lbrace \pi_k^{(i)} \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_k^{(i)}, \boldsymbol{\Sigma}_k^{(i)}) \right\rbrace}^{z_{nk}}}{\prod_{n=1}^N \sum_{k=1}^K \pi_k^{(i)} \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_k^{(i)}, \boldsymbol{\Sigma}_k^{(i)})} \\\ \end{align}

と書ける。ややこしくなってきたが大丈夫。

4. 期待値 Q^{(i)}(\boldsymbol{\theta})の計算

以上の計算を踏まえ、いよいよ Q^{(i)}(\boldsymbol{\theta})、すなわち期待値の計算に移ろう。

\begin{align} Q^{(i)}(\boldsymbol{\theta}) & = \mathbb{E}_{\boldsymbol{Z} \sim p(\boldsymbol{Z} | \boldsymbol{X}, \boldsymbol{\theta}^{(i)})} \left\lbrack \log p(\boldsymbol{X}, \boldsymbol{Z} | \boldsymbol{\theta}) \right\rbrack \\\ & = \mathbb{E}_{\boldsymbol{Z} \sim p(\boldsymbol{Z} | \boldsymbol{X}, \boldsymbol{\theta}^{(i)})} \left\lbrack \log \prod_{n=1}^N \prod_{k=1}^K {\left\lbrace \pi_k \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_k, \boldsymbol{\Sigma}_k) \right\rbrace}^{z_{nk}} \right\rbrack \\\ &= \mathbb{E}_{\boldsymbol{Z} \sim p(\boldsymbol{Z} | \boldsymbol{X}, \boldsymbol{\theta}^{(i)})} \left\lbrack \sum_{n=1}^N \sum_{k=1}^K z_{nk} {\left\lbrace \log \pi_k + \log \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_k, \boldsymbol{\Sigma}_k) \right\rbrace} \right\rbrack \\\ &= \sum_{n=1}^N \sum_{k=1}^K \mathbb{E}_{\boldsymbol{Z} \sim p(\boldsymbol{Z} | \boldsymbol{X}, \boldsymbol{\theta}^{(i)})} \left\lbrack z_{nk} \right\rbrack \left\lbrace \log \pi_k + \log \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_k, \boldsymbol{\Sigma}_k) \right\rbrace \\\ \end{align}

ここで負担率 r_{nk}^{(i)}について

\begin{align} r_{nk}^{(i)} &= \mathbb{E}_{\boldsymbol{Z} \sim p(\boldsymbol{Z} | \boldsymbol{X}, \boldsymbol{\theta}^{(i)})} \left\lbrack z_{nk} \right\rbrack \end{align}

と定義すれば、

\begin{align} Q^{(i)}(\boldsymbol{\theta}) &= \sum_{n=1}^N \sum_{k=1}^K \mathbb{E}_{\boldsymbol{Z} \sim p(\boldsymbol{Z} | \boldsymbol{X}, \boldsymbol{\theta}^{(i)})} \left\lbrack z_{nk} \right\rbrack \left\lbrace \log \pi_k + \log \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_k, \boldsymbol{\Sigma}_k) \right\rbrace \\\ &= \sum_{n=1}^N \sum_{k=1}^K r_{nk}^{(i)} \left\lbrace \log \pi_k + \log \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_k, \boldsymbol{\Sigma}_k) \right\rbrace \\\ \end{align}

と多少は簡素に書ける。

4. r_{nk}^{(i)}って具体的にはいくらなの?

\begin{align} r_{nk}^{(i)} &= \mathbb{E}_{\boldsymbol{Z} \sim p(\boldsymbol{Z} | \boldsymbol{X}, \boldsymbol{\theta}^{(i)})} \left\lbrack z_{nk} \right\rbrack \\\ &= \sum_{\boldsymbol{Z}} z_{nk} \hspace{1mm} p(\boldsymbol{Z} | \boldsymbol{X}, \boldsymbol{\theta}^{(i)}) \\\ \end{align}

さて、前で示したような悍ましい式を p(\boldsymbol{Z} | \boldsymbol{X}, \boldsymbol{\theta}^{(i)})に入れても良いのだが、その前にある重要な性質を確認しておく: \boldsymbol{z}_nは統計的に独立している。つまり今回期待値を取りたいのは z_{nk}なので、 \nu \neq nにおける \boldsymbol{z}_\nuなどは考慮しなくても良いのである。よって先ほどの期待値は

\begin{align} r_{nk}^{(i)} &= \sum_{\boldsymbol{z}_n} z_{nk} \hspace{1mm} p(\boldsymbol{z}_n | \boldsymbol{X}, \boldsymbol{\theta}^{(i)}) \\\ &= \sum_{\boldsymbol{z}_n} z_{nk} \hspace{1mm} \frac{\prod_{\kappa=1}^K {\left\lbrace \pi_\kappa^{(i)} \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_\kappa^{(i)}, \boldsymbol{\Sigma}_\kappa^{(i)}) \right\rbrace}^{z_{n\kappa}}}{\sum_{\kappa=1}^K \pi_k^{(i)} \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_\kappa^{(i)}, \boldsymbol{\Sigma}_\kappa^{(i)})} \\\ \end{align}

と簡単化される(ここで、分母分子のインデックスを kと区別するため \kappaとした)。いやまだ簡単ではないだろと思うかもしれない。ここでさらにもう一つトリックを入れる: \boldsymbol{z}_nはOne-hotベクトルである。つまり z_{nk}の取りうる値は \lbrace 0, 1\rbraceしかない。 z_{nk} = 0の時は \sumの中身も 0になるのだから、 z_{nk}=1のときだけ考えれば良い。この時 k' \neq kを満たす k'について z_{nk'}=0となることにも注意しよう。

\begin{align} r_{nk}^{(i)} &= \sum_{\boldsymbol{z}_n} z_{nk} \hspace{1mm} \frac{\prod_{\kappa=1}^K {\left\lbrace \pi_\kappa^{(i)} \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_\kappa^{(i)}, \boldsymbol{\Sigma}_\kappa^{(i)}) \right\rbrace}^{z_{n\kappa}}}{\sum_{\kappa=1}^K \pi_\kappa^{(i)} \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_\kappa^{(i)}, \boldsymbol{\Sigma}_\kappa^{(i)})} \\\ &= {\left. \frac{\prod_{\kappa=1}^K {\left\lbrace \pi_\kappa^{(i)} \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_\kappa^{(i)}, \boldsymbol{\Sigma}_\kappa^{(i)}) \right\rbrace}^{z_{n\kappa}}}{\sum_{\kappa=1}^K \pi_\kappa^{(i)} \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_\kappa^{(i)}, \boldsymbol{\Sigma}_\kappa^{(i)})} \right|}_{z_{nk} = 1, z_{nk'}=0} \\\ &= \frac{ \pi_k^{(i)} \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_k^{(i)}, \boldsymbol{\Sigma}_k^{(i)})}{\sum_{\kappa=1}^K \pi_\kappa^{(i)} \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_\kappa^{(i)}, \boldsymbol{\Sigma}_\kappa^{(i)})} \\\ \end{align}

が成立する。なおこの負担率が何を意味するのかは触れない。ヨソ行け。

Mステップ

Mステップでは、Eステップで定義した期待値 Q^{(i)}(\boldsymbol{\theta}):

\begin{align} Q^{(i)}(\boldsymbol{\theta}) &= \sum_{n=1}^N \sum_{k=1}^K r_{nk}^{(i)} \left\lbrace \log \pi_k + \log \mathcal{N} (\boldsymbol{x}_n | \boldsymbol{\mu}_k, \boldsymbol{\Sigma}_k) \right\rbrace \\\ \end{align}

を、 \boldsymbol{\theta} (つまり、 \boldsymbol{\pi}、 \boldsymbol{\mu}、 \boldsymbol{\Sigma})について最大化する。偏微分の計算は面白くないので めんどくさいので もう疲れたので今ここで述べたい要旨からずれるので省略。計算的にはしんどいけど発想的にはそこまでしんどくないので多分大丈夫。なんなら他サイトあるし。しかし1点、拘束条件: \sum_k \pi_k = 1にだけ注意。これの最適化はLagrangeの未定乗数法を用いる。

まとめ

EMアルゴリズムをGMMに適用するという情報学徒なら1度はやることを改めて勉強しなおした。Pythonの実行コードもあるので(そんなものいくらでもありそうだが)参考にどうぞ。

github.com

おまけ

K*gglerの皆さんは理論なんてどうでもいいだろうね、だってこれでいいし。

import numpy as np
from sklearn.mixture import GaussianMixture

model = GaussianMixture(n_components=2)
model.fit(X)