EMアルゴリズム入門

テックタヌキです。
皆さんは「EMアルゴリズム」という言葉を聞いたことがありますか？本記事では、私が日々のデータ分析業務を通じて学んだ、この強力な統計的手法についてご紹介します。

本記事では、EMアルゴリズムの基本概念から実際の応用例まで、できるだけわかりやすく解説していきます。もし記事の中で不明点や疑問があれば、ぜひコメントでお知らせください。

それでは、次から本編です！

この記事に書いてあること

• EMアルゴリズムの概要
• 混合正規モデルのパラメータ推定

EMアルゴリズムのモチベーション

EMアルゴリズム(expectation-maximization algorithm)は、確率モデルのパラメータに対する最尤推定を行う方法の一つです。

EMアルゴリズムの前に、指数分布のパラメータ推定の例を考えてみましょう。
$n$個の標本$X_1 \cdots, X_n$が独立に、指数分布によって得られたとします。
このとき、未知なパラメータ$\lambda$に対して、以下の尤度関数$L(\lambda)$を最大とするような$\hat{\lambda}$を$\lambda$として推定するのが最尤法でした。

$$L(\lambda) = \prod_{i=1}^n p(x_i|\lambda) = \lambda^n \exp (- \lambda \sum_{i=1}^n x_i)$$

実際に、$\hat{\lambda}$を求めてみましょう。このまま、$\lambda$に関して微分して0となるものを求めてもよいのですが、尤度関数の対数を取ったもの（対数尤度関数とよぶ）を考えることが多いです。対数尤度関数は、次のように書けます。

$$\log L(\lambda) = n \log \lambda - \lambda \sum_{i=1}^n x_i$$

これを$\lambda$に関して微分すると、^[1]

$$\frac{d L(\lambda)}{d \lambda} = \frac{n}{\lambda} - \sum_{i=1}^n x_i$$

です。これを0と置いて、$\lambda$に対して解くと、^[2]

$$\hat{\lambda} = \frac{\sum_{i=1}^n x_i}{n}$$

となります。つまり、最尤推定によって得られた、未知のパラメータ$\lambda$に対する推定値はデータの平均値になります。^[3]

ここまでの手順を簡単に整理すると、

1. 尤度関数を計算する。
2. (対数)尤度関数を微分して、最大となる点をもとめる。

となります。しかし尤度関数が複雑の場合、微分しても陽に解けないこともありそうです。このようなときに、威力を発揮するのが、EMアルゴリズムです。

混合正規分布モデル

EMアルゴリズムについて考えるために、本記事では混合正規分布モデルについて考えてみましょう。

構成されるクラス数をK、構成比率を$p=(p_1, \cdots,p_K)$、各クラスの正規分布のパラメータを$(\mu_1, \cdots, \mu_K, \sigma^2_1, \cdots, \sigma^2_K)$とすると、混合正規分布は以下のように表せます。

$$f_x = \sum_{k=1}^K p_k \frac{\exp{-(x-\mu_k)^2/2\sigma^2_k}}{\sqrt{2 \pi \sigma_k^2}}$$

ここで、$\sum_k p_k = 1$とします。

対数尤度を考えてみると、

$$\log L(\theta) = \log \prod_{n=1}^N \sum_{k=1}^K p_k \frac{\exp{-(x-\mu_k)^2/2\sigma^2_k}}{\sqrt{2 \pi \sigma_k^2}}= \sum_{n=1}^N \log \sum_{k=1}^K p_k \frac{\exp{-(x-\mu_k)^2/2\sigma^2_k}}{\sqrt{2 \pi \sigma_k^2}}$$

となります。^[4][5]
先ほどの手順に従うと、$\mu, \sigma^2, p$について最大化をする必要がありますが、微分しようにも、$\log$の中に$\sum$が入っていて解析的に解くのは難しそうです。

一般的なEMアルゴリズムの導入

一般的なEMアルゴリズムについて考えていきましょう。^[6]元々の目的は、

$$L(\theta) = \log p(X|\theta)$$

を最大化する$\theta$を見つけることでした。^[7]（推定したいパラメータをすべて、$\theta$と書いています。）
ここで、この$\theta$を推定するために、近似列$\theta_n$を構成することを考えてみましょう。
つまり、

$$L(\theta) - L(\theta_n) = \log p(X|\theta) - \log p(X| \theta_n) \geq 0$$

で、$n \rightarrow \infty$で$L(\theta_n) \rightarrow L(\theta)$となるようなものです。

潜在変数

ここで、潜在変数というものを考えてみましょう。
潜在変数とは、直接は観察されませんが、観測データを制御する何らかの確率変数$z$のことです。
これはXとの同時分布を考えることでモデル化します。^[8]

$$p(X|\theta) = \sum_z p(X,z| \theta) = \sum_z p(X|z, \theta)p(z|\theta)$$

この式を、上の式に代入すると、

$$L(\theta) - L(\theta_n) = \log \sum_z p(X|z, \theta)p(z|\theta) - \log p(X|\theta_n)$$

になります。

差の評価

この差を評価したいです。zについての確率の和は1で、log関数が凸関数であることに着目すると、イェンセンの不等式が使えそうです。

$$\begin{aligned} L(\theta) - L(\theta_n) &= \log \sum_z p(X|z, \theta)p(z|\theta) - \log p(X|\theta_n) \\ &= \log \sum_z p(z|X, \theta_n) \frac{p(X|z, \theta)p(z|\theta)}{p(z|X, \theta_n)} - \log p(X|\theta_n) \\ &\geq \sum_z p(z|X, \theta_n) \log \left( \frac{p(X|z, \theta)p(z|\theta)}{p(z|X, \theta_n)} \right) - \log p(X|\theta_n) \\ &= \sum_z p(z|X, \theta_n) \log \left( \frac{p(X|z, \theta)p(z|\theta)}{p(z|X, \theta_n)p(X|\theta_n)} \right) \\ &=: d(\theta, \theta_n) \end{aligned}$$

と評価できました。つまり以下のように書けます。

$$L(\theta) \geq L(\theta_n) + d(\theta|\theta_n)=:l(\theta|\theta_n)$$

つまり、$l(\theta|\theta_n)$は$L(\theta)$で上から抑えることができて、特に上に有界です。

また、

$$\begin{aligned} l(\theta_n|\theta_n) &= L(\theta_n) + d(\theta_n|\theta_n)\\ &= L(\theta_n) + \sum_z p(z|X, \theta_n) \log \left( \frac{p(X|z, \theta_n)p(z|\theta_n)}{p(z|X, \theta_n)p(X|\theta_n)} \right) \\ &= L(\theta_n) + \sum_z p(z|X, \theta_n) \log \left( \frac{p(X, z|\theta_n)}{p(X, z|\theta_n)} \right) \\ &= L(\theta_n) \end{aligned}$$

となるので、$l(\theta_n|\theta_n)$は尤度$L(\theta_n)$と等しくなります。
そのため、尤度関数の最大化の目的のためにこのような補助関数$l(\theta|\theta_n)$を使うことは有用そうです。
一応、私が思うこの補助関数の良い点を整理しておきます。

• $l(\theta_n|\theta_n)$は尤度関数の値$L(\theta_n)$と一致するので、良い$\theta_n$を見つけることさえできれば尤度の実現値となる
• $l(\theta|\theta_n)$は上に有界で、$L(\theta)$で抑えられているので、$\theta_n$の更新に都合がいい（後述）

$\theta_n$の更新

次に、$\theta_n$の更新はどうすればいいでしょうか？$l(\theta|\theta_n)$は、$L(\theta)$で上から抑えられているので、
$l(\theta|\theta_n)$が最大となるような、$\theta$を$\theta_{n+1}$として採用するとよさそうです。

実際に、計算してみましょう。

$$\begin{aligned} \theta_{n+1} &= \text{argmax}_\theta \ l(\theta | \theta_n)\\ &= \text{argmax}_\theta \ \Bigg(L(\theta_n) +\sum_z p(z|X, \theta_n) \log \left( \frac{p(X|z, \theta)p(z|\theta)}{p(z|X, \theta_n)p(X|\theta_n)} \right) \Bigg)\\ &= \text{argmax}_\theta \ \Bigg(\sum_z p(z|X, \theta_n) \log p(X|z, \theta)p(z|\theta) \Bigg)\\ &= \text{argmax}_\theta \ \Bigg(\sum_z p(z|X, \theta_n) \log \frac{p(X,z, \theta)}{p(z, \theta)}\frac{p(z, \theta)}{p(\theta)} \Bigg) \\ &= \text{argmax}_\theta \ \Bigg(\sum_z p(z|X, \theta_n) \log p(X,z|\theta) \Bigg) \\ &= \text{argmax}_\theta \ \Bigg(E_{z|X, \theta_n}[\log p(X,z|\theta)] \Bigg) \end{aligned}$$

よって、潜在変数ZとXの同時分布の対数尤度に対するzの条件付き期待値が最大となるような$\theta$を選べばよいということがわかりました。
これらの手順を以下のように分けて、EMアルゴリズムと呼びます。

• Eステップ: 条件付き期待値、$E_{z|X, \theta_n}[\log p(X,z|\theta)]=:Q(\theta|\theta_n)$を決定する。
• Mステップ: $Q(\theta|theta_n)$を$\theta$に関して、最大化して$\theta_{n+1}$を求める。

これらのステップを$\theta_n$が収束するまで交互に行うことに行うことによって、パラメータ推定を行うことができます。^[9]

混合正規分布モデルに対するEMアルゴリズム

一般論から混合正規分布モデルに戻って考えましょう。
$X = (X_1, \cdots, X_n)$はi.i.d.で次の混合正規分布の確率密度関数に従って生成されているとします。

$$f_x = \sum_{k=1}^K p_k \frac{\exp{-(x-\mu_k)^2/2\sigma^2_k}}{\sqrt{2 \pi \sigma_k^2}}$$

ここで、$\sum_k p_k = 1$とします。このまま、対数尤度を考えると難しいので、潜在変数zとして、次の確率変数を導入します。

$$p(z=k) = p_k \ k= 1, \cdots, K$$

zはXがどのクラスに属するかを決めるので、n個の確率変数$z_1, \cdots, z_n$をもち、$(z_1 ,\cdots, z_n) =: Z$と表します。
また、どのクラスに入るかが決まればどの確率密度関数に従うかがきまるので、以下のように条件付き確率密度関数を考えることができます。

$$f_{x|z}(x_n| z_n=k, \theta) = \frac{\exp(-(x-\mu_k)^2/2\sigma^2_k)}{\sqrt{2 \pi \sigma_k^2}}$$

ここで、$\theta = (p_1, \cdots, p_K, \mu_1, \cdots, \mu_K, \sigma_1, \cdots, \sigma_K)$として未知のパラメータをひとまとめに表します。

すると、対数尤度は、

$$L(\theta; X, Z) = \prod_n p_{z_n}\frac{1}{2\pi\sigma^2_{z_n}}\exp\big(-(x_n-\mu_{z_n})^2/2\sigma_{z_n}^2\big)$$

となります。潜在変数を導入することで、$\sum \log$の形をなくすことができました。EMアルゴリズムを導入して、Eステップと、Mステップについて考えます。^[10]

Eステップ

$$\begin{aligned} Q(\theta;\theta_{old}) &:= E_{z|X, \theta_{old}}[\log p(X, z|\theta)]\\ &= \sum_n \sum_k p(z_n= k| x_n,\theta_{old}) \log \big( f_{x|z}(x_n|z_n = k, \theta)p(z_n=k|\theta) \big) \end{aligned}$$

ここで、ベイズの定理を使うと$p(z_n=k| x_n, \theta_{old})$を以下のように整理できます。

$$\begin{aligned} \eta_{nk}&:=p(z_n=k| x_n, \theta_{old}) \\ &= \frac{p(z_n=k, X_n=x_n|\theta_{old})}{p(X_n=x_n|\theta_{old})} = \frac{f_{x|z}(x_n|z_n=k, \theta_{old})p(y_n=k|\theta_{old})}{\sum_k f_{x|y}(x_n|z_n=k, \theta_{old})p(z_n=k|\theta_{old})} \end{aligned}$$

を負担率と表現するそうです。

Mステップ

Eステップで求めた、$Q(\theta;\theta_n)$を$\theta$で微分して、最大となる点をもとめます。ただし、$\sum_k p_k = 1$という条件があるので、この条件に従って最大化する必要があります。等式制約の最大化なのでラグランジュの未定乗数法を使えばよさそうです。
あまり本質的ではないので、求めた結果を以下に載せます。^[11]

$$\begin{aligned} \mu_{k,new} &= \frac{\sum_n x_n \eta_{n,k}}{\sum_n \eta_{n,k}}\\ \sigma_{k, new}^2 &= \frac{\sum_n (x_n-\mu_k)^2 \eta_{n,k}}{\sum_n \eta_{n,k}} \\ p_{k, new} &= \frac{1}{N} \sum_n \eta_{n,k} \end{aligned}$$

シミュレーション

シミュレーション条件

• $K=3$
• $N=1000$
• $\hat{p}=(0.5, 0.2, 0.3)$
• $\hat{\mu}=(-1, 0, 1)$
• $\hat{\sigma}^2=(0.2, 1, 0.3)$
• イテレーション回数:1000^[12]

の設定でやってみました。データをサンプルした結果とサンプルを作るスクリプトは以下のようになります。

 import numpy as np

class GetMixtureGaussSample:

    def __init__(self, K=3, N=1000, p=(0.1, 0.7, 0.2), mu=(-1, 0, 1), v = (0.2, 1, 0.5)):
        # 混合正規分布のデータを取得するクラス

        # チェック
        assert K == len(p), "Kとpの長さが一致しません"
        assert K == len(mu), "Kとmuの長さが一致しません"
        assert K == len(v), "Kとvの長さが一致しません"
        assert np.isclose(sum(p), 1), "確率の合計が1になっていません"
        
        # p_i * N が整数になるように設定すること
        for i in range(K):
            assert np.isclose((p[i]*N) % 1, 0), f"p[{i}] * N が整数になるように設定されていません"
        
        self.K = K
        self.N = N
        self.p = np.array(p)
        self.mu = np.array(mu)
        self.v = np.array(v)
        
    def sample(self) -> np.ndarray:
        # 混合正規分布モデルからのサンプル
        data = np.concatenate([
            np.random.normal(loc=self.mu[i], scale=np.sqrt(self.v[i]), size=int(self.p[i] * self.N))
            for i in range(self.K)
        ])
        np.random.shuffle(data)  # データをシャッフル
        return data

EMアルゴリズムのソースコードは以下のようになります。
各パラメータの初期値は、一般的かわかりませんが、Kmeansによって決めました。^[13]

class EMAlgorithm:

    def __init__(self, N=1000, K=3, iteration=2000, history=True):
        self.K = K
        self.N = N
        self.iteration = iteration

        if history:
            # 履歴を保存するなら、履歴用のリストを作っておく
            self.mu_history = np.zeros((self.iteration+1, self.K))
            self.v_history = np.zeros((self.iteration+1, self.K))
            self.p_history = np.zeros((self.iteration+1, self.K))

    def init_em(self, data):
        # 初期化する
        self.data = data
        # K-means法でクラスタリング
        kmeans = KMeans(n_clusters=self.K, n_init=10)
        kmeans.fit(data.reshape(-1, 1))
        
        # 初期パラメータの設定
        self.mu = kmeans.cluster_centers_.flatten()
        self.v = np.array([np.var(data[kmeans.labels_ == k]) for k in range(self.K)])
        self.p = np.bincount(kmeans.labels_) / self.N
    
    def e_step(self):
        # 尤度は計算する必要はなく、負担率eta_nkさえわかればOK
        self.eta_nk = np.zeros((self.N, self.K))
        for k in range(self.K):
            self.eta_nk[:,k] = self.p[k] * norm.pdf(self.data, self.mu[k], np.sqrt(self.v[k]))
        self.eta_nk /= self.eta_nk.sum(axis=1, keepdims=True)

    def m_step(self):
        # 負担率をもとに、パラメータ更新
        sum_eta_n = self.eta_nk.sum(axis=0)
        self.mu = np.dot(self.eta_nk.T, self.data) / sum_eta_n
        self.p = sum_eta_n / self.N

        for k in range(self.K):
            self.v[k] = np.sum(self.eta_nk[:, k] * (self.data - self.mu[k])**2) / sum_eta_n[k]
    
    def simulation(self, data):
        # simulation
        self.init_em(data)
        self.save_history(0)
        for i in tqdm(range(self.iteration)):
            self.e_step()
            self.m_step()
            self.save_history(i+1)
        print('p:', self.p)
        print('m:', self.mu)
        print('v:', self.v)

    def save_history(self, i):
        # 履歴を保存する。
        self.mu_history[i, :] = self.mu
        self.v_history[i, :] = self.v
        self.p_history[i, :] = self.p

それぞれのパラメータの推移のグラフは以下です。

パラメータは収束していますが、iterationの1200回あたりの妙な動きが気になりました。およその収束値は以下のようになります。構成比率、平均、分散の順に記載します。

• 分布1: (0.04, 0.12, 0.03)
• 分布2: (0.584, -0.98, 0.20)
• 分布3: (0.38, 0.99, 0.29)

真の分布についても再掲します。

• 分布1: (0.2, 0, 1)
• 分布2: (0.5, -1, 0.2)
• 分布3: (0.3, 1, 0.3)

両端の分布の平均と分散はおおむね合致していますが、構成比率や真ん中の分布は正確に当てるのが難しいようです。真ん中の分布は両隣から大きな影響を受けるので、なかなか難しいというだと思いますが、数理的に何か言えるかはわかりません。

ヒストグラムに重ねがきしてみると以下のようになります。

分布推定としてみると、割と合っているような気もします。

まとめ

今回は、EMアルゴリズムを使って、混合正規分布のパラメータ推定をしてみました。本質的に最も重要なポイントは複雑な尤度関数が潜在変数を導入することによって、推定しやすい形に変形できることです。
最後になりますが、何かおかしいところがあればコメントもらえると助かります。

脚注

1. $\partial$とかくことも多いですが、$x_i$は定数と思っているので$d$でいいかなと思います。

2. 対数尤度関数の$\lambda$に関する2階微分は-$1/\lambda^2$なので常に負です（$\lambda>0$)。そのため、対数尤度関数は凹関数です。よって、微分が0になる点が対数尤度関数が最大となる$\lambda$です。

3. 指数分布に従う確率変数の期待値が$\lambda$になっていることから、納得しやすい結果ですね。

4. $\theta$はすべてのパラメータ$p_1, \cdots, p_K, \sigma^2_1, \cdots, \sigma_K^2, \mu_1, \cdots, \mu_K$をひとまとめにしたものです。

5. ベイズで混合分布の推定をしたことがある方はこの積と和の形が出てこないので、（ベイズでは尤度関数の最大化はしませんが）戸惑ったかもしれません。

6. 混合指数分布については一般の場合を導出してから考えます。

7. ここから、Lは対数尤度関数とします。

8. 連続の場合は、積分ですが測度論によってどちらでも正当化されるので、区別しません。

9. 収束性の議論はさぼります。

10. 今後、$\theta_n$を$\theta_{old}$、$\theta_{n+1}$を$\theta_{new}$に読み替えます。

11. 線形制約のため、凸制約になります。また凸制約のもとでの、凹関数の最大化問題は少なくとも１つの最適解をもちます。また凹関数の局所最大化は同時に大域的最大値でもあるので、最大値をもちます。$Q(\theta)$は対数関数の重ね合わせなので凹関数といっていいと思います。間違ってたら教えてください。

12. 賢い終了条件がありそうですが、今回はさぼります。

13. 初期値の決め方は、平均を各クラスタの平均値、分散を各クラスタ内分散、構成割合をクラス内に属するデータの比率としました。詳しくはコードを見てください。