EMアルゴリズムで混合正規分布(MoG)のパラメータを導出

EMアルゴリズムで混合正規分布(MoG)のパラメータを求めるメモ．

参考：

混合正規分布（MoG）の定義

MoGやGMM（gaussian mixture model）と呼ばれる

正規分布に重みづけ和する．重みは足して1にする．

$Pr(\vec{x} | \vec{\theta}) = \sum_{k=1}^{K} \lambda_{k} \text{Norm}_{\vec{x}}(\vec{\mu}_{k}, \mathbf{\Sigma}_{k}) \\\\ \sum_{k=1}^{K} \lambda_{k} = 1$

この重み・正規分布のパラメータを直接，MLで最大化することが難しいので，EMアルゴリズムが使われる．

pdfのプロット

pdfも重みづけ和にすればいい．確率の公理で全体が1を満たさなければならないので， 1つの正規分布の重みが$\alpha$ならば，他の正規分布が密度を$1-\alpha$で負担することになる．

正規分布が2つで，均等の重みの混合正規分布であれば，以下の式になる．

$Pr(\vec{x} | \vec{\theta}) = \frac{1}{2} \text{Norm}_{\vec{x}}(\vec{\mu}_{1}, \mathbf{\Sigma}_{1}) + \frac{1}{2}\text{Norm}_{\vec{x}}(\vec{\mu}_{2}, \mathbf{\Sigma}_{2})$

1次元の混合正規分布の密度は以下のように計算できる．

code:

import scipy as sp
from scipy import stats

def calc_mix_pdf(k_size, x, lambda_vec, mu_vec, sigma2_vec):
    pdf = sp.zeros_like(x)

    for k in sp.arange(k_size):
        norm_k = stats.norm(loc=mu_vec[k], scale=sp.sqrt(sigma2_vec[k]))
        pdf += lambda_vec[k] * norm_k.pdf(x)
    return pdf

1次元正規分布の混合正規分布のpdfをプロットする．

import scipy as sp
from scipy import stats
import matplotlib.pyplot as plt

%matplotlib inline

def calc_mix_pdf(k_size, x, lambda_vec, mu_vec, sigma2_vec):
    pdf = sp.zeros_like(x)

    for k in sp.arange(k_size):
        norm_k = stats.norm(loc=mu_vec[k], scale=sp.sqrt(sigma2_vec[k]))
        pdf += lambda_vec[k] * norm_k.pdf(x)
    return pdf

x = sp.linspace(-6, 6, 100)
N1 = stats.norm(loc=2, scale=1)
N2 = stats.norm(loc=-2, scale=1)

pdf_mix = calc_mix_pdf(2, x, lambda_vec=[0.5, 0.5], mu_vec=[2, -2], sigma2_vec=[1, 1])
plt.plot(x,N1.pdf(x)/2, linestyle="--", label="N1")
plt.plot(x,N2.pdf(x)/2, linestyle="--", label="N2")
plt.plot(x, pdf_mix, label="mix norm")

plt.legend()
plt.show()

MLでパラメータ導出を考えると

最尤推定（ML）で考えると，以下の最大化を考えることにより，微分によるパラメータ導出が困難になる，

Q.Can we learn this model with maximum likelihood ?

$\begin{eqnarray} \hat{\vec{\theta}} &=& \mathop{argmax}_{\vec{\theta}} \left[ \sum_{i=1}^{I} \log{Pr(\vec{x}_{i} | \vec{\theta})} \right] \\\\ &=& \mathop{argmax}_{\vec{\theta}} \left[ \sum_{i=1}^{I} \log{ \left[ \sum_{k=1}^{K} \lambda_{k} \text{Norm}_{\vec{x}_{i}}\left[ \vec{\mu}_{k}, \mathbf{\Sigma}_{k} \right] \right] } \right] \\ \end{eqnarray}$

対数の中の積であれば，$\log{AB} = \log{A} + \log{B}$のように和で変換して微分を容易にできるが，対数の中の和の微分は難しい．

また，パラメータにも条件がある．

共分散行列は正定値
重み$\lambda$は足して1.

EMアルゴリズム

特徴として，以下の問題が解け欠損値に強いが，初期値依存性が高く局所最適解に陥りやすい．

潜在変数の導入

Key idea:

represent density $Pr(x)$ as marginalization of joint density with another variable h that we do not see

観測データ$x$があったとしてこの分布$Pr(\vec{x})$ を求めるのが困難だとする．ただし，隠れ変数$h$との同時分布は簡単だとする．ただし，hはわからない．ただ，周辺化すればhは消せて$Pr(x)$を計算できる．同時分布が簡単であれば周辺化で隠れ変数の部分は消すことができる．

$$ Pr( \vec{x} ) = \int Pr(\vec{x}, \vec{h}) d \vec{h} $$

パラメータに依存していても，パラメータに依存した隠れ変数との同時分布を考える．周辺化すれば$Pr(\vec{x} | \vec{\theta})$を求められる．

$$ Pr(\vec{x} | \vec{\theta}) = \int Pr(\vec{x}, \vec{h} | \vec{\theta}) d \vec{h} $$

期待値最大化

この問題を解くのに，うまいやり方として，期待値最大化法（EMアルゴリズム）がある．

次の式の右辺を最大化するパラメータを求める．

$$ \hat{\vec{\theta}} = \mathop{\text{argmax}}_{\vec{\theta}} \left[ \sum_{ i =1 }^{I} \log{ \left[ \int Pr(\vec{x}_{i}, \vec{h}_{i} | \vec{\theta}) d\vec{h}_{i} \right] } \right] $$

積分の部分は周辺化に該当し$x$の分布$Pr(x|\theta)$になり，観測値$x_{i}$を与えられた（固定した）とき，パラメータ$\vec{\theta}$を変数としたときの密度より尤度になっている．
尤度の対数をとっているので対数尤度になる．
全ての観測データは独立として，尤度の積は対数尤度では和の計算になる．

つまり，対数尤度の最大化の式になっている．

このまま最大化を考えるのは難しいので次の下界を考える．

Lower Bound 下界

次の下界（lower bound）の関数$\cal{B}$を定義する．この関数は，常に先ほどの最大化したい対数尤度より小さい関係（下界）にある．

なので，この関数の値を最大化していけば元の対数尤度関数が最大化していくことと同じになる．

$\begin{eqnarray} \cal{B}\left[ \left\{ q_{i}(\vec{h}_{i}) \right\}, \vec{\theta} \right] &=& \sum_{i=1}^{I} \int q_{i}(\vec{h}_{i}) \log{ \left[ \frac{Pr(\vec{x}_{i}, \vec{h}_{i} | \vec{\theta})}{q_{i}(\vec{h}_{i})} \right] }d\vec{h}_{i} \\ &\leq& \sum_{i=1}^{I} \log{ \left[ \int q_{i}(\vec{h_{i}}) \frac{Pr(\vec{x}_{i}, \vec{h}_{i} | \vec{\theta})}{ q_{i}(\vec{h_{i}}) } d\vec{h}_{i} \right] } = \sum_{i=1}^{I} \log{ \left[ \int Pr(\vec{x}_{i}, \vec{h}_{i} | \vec{\theta}) d\vec{h}_{i} \right] } \\ \end{eqnarray}$

E-step & M-step

Lower bound is a function of parameters $\vec{\theta}$ and a set of probailitiy ditributions $q_{i}(\vec{h}_{i})$

$q_{i}$は隠れ変数$\vec{h}_{i}$の確率分布であり，関数$\cal{B}$はこれにも依存している．

$q_{i}$は各データごとに計算していることに注意．

EMアルゴリズムでは，パラメータ$\vec{\theta}$と$q_{i}(\vec{h}_{i})$の2種類を求めなければならない． EMアルゴリズムではこれらをE-stap, M-stepで交互に求める．

E-Step - Maximize bound w.r.t.(with reference to) distributions $q(\vec{h}_{i})$：前回のM-stepで推定したパラメータ$\vec{\theta}$で固定して，$q$をよりいいものに変えて下界を上げる（最大化する）．これは，前回推定したM-stepのパラメータを使った隠れ変数の分布を事前確率として観測値の尤度で更新した事後確率を求めることに該当する．
M-step - Maximize bound w.r.t. parameters ：隠れ変数hの分布$q_{i}$を固定してパラメータ$\vec{\theta}$を変えて最大化する．

イメージとしては以下の図になる．

E-stepとM-stepを繰り返す（引用："Computer vision: models, learning and inference" by Simon Prince, p.129 Figure7.24）

E-step & M-stepでの最大化の数式

E-step

$q_{i}$は隠れ変数の確率分布だった．

前回のM-stepで推定したパラメータ$\vec{\theta}$で固定して，隠れ変数$h$の確率を事前分布とし観測データ$x$の尤度でベイズ更新して $h$の事後分布を計算すると最大化に相当する．

$\hat{q}_{ i }( \vec{h}_{i} ) = Pr(\vec{h}_{i} | \vec{x}_{i}, \vec{\theta}^{[ t ]}) \\\\ = \frac{Pr(\vec{x}_{i} | \vec{h}_{i} , \vec{\theta}^{ [ t ] } ) Pr( \vec{h}_{i} | \vec{\theta}^{ [ t ] })}{Pr(\vec{x}_{i})} \\\\$

$\hat{q}_{ i }( \vec{h}_{i} )$は，その観測データがどの隠れ変数の値の場合が妥当であるかを表すので， responsibility(負担率)とも呼ばれる．

M-step

E-stepで推定した事後確率で$q_{i}$を固定して，パラメータ$\vec{\theta}$を変えて最大化する．

分母に$\theta$は無いので省略できる

$\begin{eqnarray} \hat{\vec{\theta}}^{[ t+1 ]} &=& \mathop{argmax}_{\vec{\theta}} \left[ \sum_{i=1}^{I} \int \hat{q}_{i}(\vec{h}_{i}) \log{ \left[ \frac{Pr(\vec{x}_{i}, \vec{h}_{i} | \vec{\theta})}{\hat{q}_{i}(\vec{h}_{i})} \right] }d\vec{h}_{i} \right] \\ &=& \mathop{argmax}_{\vec{\theta}} \left[ \sum_{i=1}^{I} \int \hat{q}_{i}(\vec{h}_{i}) \log{ \left[ Pr(\vec{x}_{i}, \vec{h}_{i} | \vec{\theta}) \right] }d\vec{h}_{i} \right] \\ \end{eqnarray}$

右辺の積分は$q$についてのパラメータ$\theta$を条件とした条件付期待値になっている．

この最大化は，制約条件を入れてラグランジュの未定乗数法などを使った目的関数の微分=0と置いた式から導出できる．

混合正規分布におけるEMアルゴリズムの適用

EMアルゴリズムを混合正規分布に適用していく．

重み$\lambda_{k}$を（合計1より）隠れ変数$h=k$の確率として拡張する． $h=1$のとき$\lambda_{1}$, $h=2$のとき$\lambda_{2}$, $h=K$のとき$\lambda_{K}$と考え，これはcategorical分布になる．

また，各正規分布を隠れ変数の値によって切り替わるように考え，条件付分布として拡張する．

式にすると以下になる．

$\begin{eqnarray} Pr(\vec{x} | h, \vec{\theta}) &=& \text{Norm}_{\vec{x}}( \vec{\mu}_{h}, \mathbf{\Sigma}_{h}) \\\\ Pr(h | \vec{\theta}) &=& \text{Cat}_{h}( \vec{\lambda} ) \end{eqnarray}$

同時分布を周辺化してxの分布を計算したいので次を定義する．

$\begin{eqnarray} Pr(\vec{x} | \vec{\theta}) &=& \sum_{k=1}^{K} Pr(\vec{x}, h=k | \vec{\theta}) \\\\ &=& \sum_{k=1}^{K} Pr(\vec{x} | h=k | \vec{\theta}) Pr(h=k | \vec{\theta}) \\\\ &=& \sum_{k=1}^{K} \lambda_{k} \text{Norm}_{\vec{x}}\left[ \vec{\mu}_{k}, \mathbf{\Sigma}_{k} \right] \end{eqnarray}$

同時分布を隠れ変数$h$で周辺化した分布が目的の混合正規分布になる．

例：

隠れ変数$h = 1$について正規分布
隠れ変数$h = 2$ について正規分布
隠れ変数$h = 3$について正規分布とすれば

3つの正規分布の混合正規分布になる．

隠れ変数との同時分布の周辺分布が混合正規分布になる（引用："Computer vision: models, learning and inference" by Simon Prince, Figure7.7, p.107 ）

混合正規分布からのサンプリング

categorical分布を導入したことで，

categorical分布からkをサンプル，
サンプルした値kで固定した，つまりk番目の正規分布から値をサンプリング

を繰り返すことでサンプリングが可能である．

混合1次元正規分布を例にcodeにすると，

import scipy as sp
from scipy import stats

def sampling_mix(k_size, lambda_vec, mu_vec, sigma2_vec, n=10):
    cat = stats.multinomial(n=1, p=lambda_vec)
    k_arr = cat.rvs(size=n)[:, 0]

    sampling = sp.empty(n)
    for k in sp.arange(k_size):
        N = stats.norm(loc=mu_vec[k], scale=sp.sqrt(sigma2_vec[k]))
        sampling[k_arr==k] = N.rvs(size=len(k_arr[k_arr==k]))

    return sampling

sampling_mix(2, [0.5, 0.5], [2, -2], [1, 1], n=10)

array([-0.45735961,  1.99554384, -1.26413578, -2.20297206, -3.02347406,
       -1.44119158,  2.08046362, -1.55275303,  2.61056233, -3.05617989])

E-step

まず，各データ$x_{i}$に対して隠れ変数$h_{i}$の事後確率を求める．分母は$h_{i}$で周辺化して，正規化する．

$\begin{eqnarray} Pr(h_{i}=k | \vec{x}_{i}, \vec{\theta}^{[ t ] }) &=& \frac{Pr(\vec{x}_{i} | h_{i} , \vec{\theta}^{ [ t ] }) Pr(h_{i} | \vec{\theta}^{[ t ]})}{Pr(\vec{x}_{i})} \\\\ &=& \frac{Pr(\vec{x}_{i} | h_{i}=k , \vec{\theta}^{[ t ] }) Pr(h_{i}=k | \vec{\theta}^{[ t] })}{ \sum_{ j =1}^{K} Pr(\vec{x}_{i} | h_{i}=j , \vec{\theta}^{ [ t ] }) Pr(h_{i}=j | \vec{\theta}^{ [ t ] })} \\\\ &=& \frac{\text{Norm}_{\vec{x}_{i}}(\vec{\mu}_{k}, \mathbf{\Sigma}_{k}) \cdot \lambda_{k}}{ \sum_{ j = 1 }^{K} \lambda_{j} \text{Norm}_{\vec{x}_{i}}(\vec{\mu}_{j}, \mathbf{\Sigma}_{j}) } \\\\ &=& r_{i k} \end{eqnarray}$

We'll call this the responsibility of the k-th Gaussian for the i-th data point.

この$r_{ik}$はresponsibilityと呼ばれる． k番目の正規分布がi番目のデータに対してどのくらい責任をもっているのか，生成したことになるかを表している．

h=1,2の混合正規分布であれば，正規分布が2つあるので，responsibilityも2つ計算している．

Repeat this procedure for every datapoint!

このresponsibilityをすべてのデータに対して求める．

M-step

E-step計算後，M-Stepを計算する．

次の下界の最大化より，

$\begin{eqnarray} \hat{\vec{\theta}}^{[ t+1 ] } &=& \mathop{argmax}_{\vec{\theta}} \left[ \sum_{i=1}^{I} \int \hat{q}_{i}(\vec{h}_{i}) \log{ \left[ Pr(\vec{x}_{i}, \vec{h}_{i} | \vec{\theta}) \right] }d\vec{h}_{i} \right] \\ \end{eqnarray}$

隠れ変数が離散で，1次元なので次のように書き直せる．

$\begin{eqnarray} \hat{ \vec{\theta}}^{ [ t +1 ] } &=& \mathop{ \text{argmax}}_{\vec{\theta}} \left[ \sum_{i=1}^{I} \sum_{k=1}^{K} \hat{q}_{i}( h_{i}=k ) \log{ \left[ Pr(\vec{x}_{i}, h_{i}=k | \vec{\theta}) \right] } \right] \\\\ &=& \mathop{\text{argmax}}_{\vec{\theta}} \left[ \sum_{i=1}^{I} \sum_{k=1}^{K} r_{ i k } \log{ \left[ Pr(\vec{x}_{i}, \vec{h}_{i} | \vec{\theta}) \right] } \right] \\\\ &=& \mathop{\text{argmax}}_{\vec{\theta}} \left[ \sum_{i=1}^{I} \sum_{k=1}^{K} r_{ik} \log{ \left[ \lambda_{k} \text{Norm}_{\vec{x}_{i}}( \vec{\mu}_{k}, \mathbf{\Sigma}_{k}) \right] } \right] \end{eqnarray}$

これをラグランジュ乗数法で解いて，パラメータを導出する．

ラグランジュで解いてパラメータを導出

パラメータでの微分，つまり，平均ベクトル$\vec{\mu}$での微分，共分散行列 ${\Sigma}$ での微分が必要になる．

まず，目的関数を定める．等式制約$\sum_{k=1}^{K} \lambda_{k} = 1$を未定乗数$\nu$として入れると，以下になる．

$\begin{eqnarray} L &=& \sum_{i=1}^{I} \sum_{k=1}^{K} r_{ik} \log{ \left[ \lambda_{k} \text{Norm}_{\vec{x}_{i}}(\vec{\mu}_{k}, \mathbf{\Sigma}_{k}) \right] } - \nu \left(\sum_{j=1}^{K} {\lambda_{j}} - 1\right) \\\\ &=& \sum_{i=1}^{I} \sum_{k=1}^{K} r_{ik} (\log{\lambda_{k}} + \log{\left[ \text{Norm}_{\vec{x}_{i}}(\vec{\mu}_{k}, \mathbf{\Sigma}_{k}) \right] }) - \nu \left(\sum_{j=1}^{K} {\lambda_{j}} - 1\right) \end{eqnarray}$

このラグランジュ関数を各パラメータで微分する．

$\lambda_{k}$での微分

$\begin{eqnarray} \frac{\partial L}{\partial \lambda_{k}} = \sum_{i=1}^{I} r_{ik} \frac{1}{ \lambda_{k} } - \nu &=& 0 \\ \nu\lambda_{k} &=& \sum_{i=1}^{I} r_{ik} \\ \lambda_{k}^{[ t+1 ] } &=& \frac{\sum_{i=1}^{I} r_{ik}}{\nu} \\ \end{eqnarray}$

等式制約より，未定乗数の値が導出される．

$\sum_{k=1}^{K}\lambda_{k} = 1\\ \sum_{k=1}^{K} \frac{\sum_{i=1}^{I} r_{ik}}{\nu} = 1 \\ \nu = \sum_{k=1}^{K} \sum_{i=1}^{I} r_{ik}$

これを代入すると重み（categorical分布の確率）$\lambda_{k}$が決まる．

$\begin{eqnarray} \lambda_{k}^{ [ t+1 ] } &=& \frac{ \sum_{i=1}^{I} r_{ik} }{ \sum_{k=1}^{K} \sum_{i=1}^{I} r_{ik}} \\ \end{eqnarray}$

平均ベクトル$\mu_{k}$での微分

これは二次形式をベクトルでの微分になる．

参考：ベクトル・行列の微分メモ - はしくれエンジニアもどきのメモ

$\begin{eqnarray} \frac{\partial L}{\partial \vec{\mu}_{k}} = \sum_{i=1}^{I} r_{ik} \cdot -\frac{1}{2} \left\{ - 2( \mathbf{\Sigma}_{k}^{-1} (\vec{x}_{i} - \vec{\mu}_{k})) \right\} &=& 0 \\\\ \sum_{i=1}^{I} r_{ i k } (\vec{ x }_{ i } - \vec{\mu}_{ k } ) &=& 0 \\\\ \vec{\mu}_{k} \sum_{ i = 1 }^{ I } r_{ i k } &=& \sum_{i=1}^{I} r_{ik} \vec{x}_{i} \\\\ \vec{\mu}_{k}^{ [ t+1 ] } &=& \frac{\sum_{i=1}^{I} r_{ik} \vec{x}_{i}}{\sum_{i=1}^{ I } r_{ik}} \end{eqnarray}$

共分散行列 $\Sigma_{k}$ での微分

$\begin{eqnarray} \frac{\partial L}{\partial \mathbf{\Sigma}_{k}} = \frac{\partial}{\partial \mathbf{\Sigma}_{k}} \sum_{i=1}^{I} r_{ik} \cdot (- \frac{1}{2} \log{|\mathbf{\Sigma}_{k}|}) + \frac{\partial }{\partial \mathbf{\Sigma}_{k}} \sum_{i=1}^{I} r_{ik} (-\frac{1}{2}(\vec{x}_{i} - \vec{\mu}_{k})^{T}\mathbf{\Sigma}_{k}^{-1} (\vec{x}_{i} - \vec{\mu}_{k})) &=& 0 \end{eqnarray}$

この微分では"対数行列式の微分"と"逆行列の入った二次形式"の微分が必要になる．

分けて考える．

対数行列式の微分：

参考：ベクトル・行列の微分メモ - はしくれエンジニアもどきのメモ

$$ \frac{d \log{|\mathbf{A}|}}{d \mathbf{A}} = (\mathbf{A}^{-1})^{T} $$

より，

$\frac{\partial}{\partial \mathbf{\Sigma}_{k}} \sum_{i=1}^{I} r_{ik} \cdot (- \frac{1}{2} \log{|\mathbf{\Sigma}_{k}|}) = - \frac{1}{2} \sum_{i=1}^{I} r_{ik} (\mathbf{\Sigma}_{k}^{-1})^{T} = - \frac{1}{2} \sum_{i=1}^{I} r_{ik} \mathbf{\Sigma}_{k}^{-1}$

逆行列の入った二次形式を行列で微分：

参考：ベクトル・行列の微分メモ - はしくれエンジニアもどきのメモ

$\begin{eqnarray} \frac{d \vec{x}^{T} \mathbf{A}^{-1} \vec{x} }{d \mathbf{A}} &=& - \mathbf{A}^{-1} \vec{x} \vec{x}^{T} \mathbf{A}^{-1} \end{eqnarray}$

$\begin{eqnarray} \frac{\partial }{\partial \mathbf{\Sigma}_{k}} \sum_{i=1}^{I} r_{ik} (-\frac{1}{2}(\vec{x}_{i} - \vec{\mu}_{k})^{T}\mathbf{\Sigma}_{k}^{-1} (\vec{x}_{i} - \vec{\mu}_{k})) &=& - \frac{1}{2}\sum_{i=1}^{I} r_{ik}\mathbf{\Sigma}_{k}^{-1} \{- (\vec{x}_{i} - \vec{\mu}_{k}) (\vec{x}_{i} - \vec{\mu}_{k})^{T}\mathbf{\Sigma}_{k}^{-1}\} \\\\ &=& \frac{1}{2}\sum_{i=1}^{I} r_{ik}\mathbf{\Sigma}_{k}^{-1} (\vec{x}_{i} - \vec{\mu}_{k}) (\vec{x}_{i} - \vec{\mu}_{k})^{T}\mathbf{\Sigma}_{k}^{-1} \end{eqnarray}$

1次元の混合正規分布に対してのEMアルゴリズムを実装・描画

この場合の，E-step，M-stepの式を載せておく．

E-Step：

$\begin{eqnarray} Pr(h_{i}=k | x_{i}, \vec{\theta}^{[ t ] }) &=& \frac{Pr(x_{i} | h_{i} , \vec{\theta}^{[ t ]}) Pr( h_{i} | \vec{\theta}^{[ t ]})}{Pr( x_{i})} \\\\ &=& \frac{Pr( x_{i} | h_{i}=k , \vec{\theta}^{[ t ] }) Pr(h_{i}=k | \vec{\theta}^{[ t ]})}{ \sum_{j=1}^{K} Pr(\vec{x}_{i} | h_{i}=j , \vec{\theta}^{ [ t ] }) Pr( h_{i} = j | \vec{\theta}^{ [ t ] })} \\\\ &=& \frac{\text{Norm}_{x_{i}}(\mu_{k}, \sigma^{2}_{k}) \cdot \lambda_{k}}{ \sum_{j=1}^{K} \lambda_{j} \text{Norm}_{x_{i}}(\mu_{j}, \sigma^{2}_{j}) } \\ &=& r_{i k} \end{eqnarray}$

M-Step:

$\begin{eqnarray} \hat{\vec{\theta}}^{[ t+1 ] } &=& \mathop{ \text{argmax} }_{\vec{\theta}} \left[ \sum_{i=1}^{I} \sum_{k=1}^{K} r_{ik} \log{ \left[ \lambda_{k} \text{Norm}_{x_{i}} \mu_{k}, \sigma^{2}_{k}) \right] } \right] \\\\ \end{eqnarray}$

$\begin{eqnarray} \lambda_{k}^{ [ t+1] } &=& \frac{\sum_{i=1}^{I} r_{ik}}{\sum_{k=1}^{K} \sum_{i=1}^{I} r_{ik}} \\ \end{eqnarray}$

$\begin{eqnarray} \mu_{k}^{[t+1]} &=& \frac{\sum_{i=1}^{I} r_{ik} x_{i}}{\sum_{i=1}^{I} r_{ik}} \\ \end{eqnarray}$

$\begin{eqnarray} (\sigma^{2}_{k})^{[ t+1 ] } &=& \frac{\sum_{i=1}^{I} r_{ik}( x_{i} - \mu_{k}^{[ t+1 ] })^{2} }{\sum_{i=1}^{I} r_{ik}} \\ \end{eqnarray}$

真の混合正規分布を$\frac{2}{3}N(2, 1) + \frac{1}{3}N(-2, 1)$として，この分布から1000個の観測データが得られたとする．初期分布を$\frac{1}{2}N(1, 1) + \frac{1}{2}N(-1, 1)$として， EMstepを50回行って，後のパラメータの値を出力し，そのパラメータの分布を描画する．

import scipy as sp
from scipy import stats
import matplotlib.pyplot as plt

%matplotlib inline

# E-step
def EStep(k_size, data_x, lambda_vec, mu_vec, sigma2_vec):
    '''
    data_x[i]
    lambda[k]
    mu[k]
    sigma[k]
    '''
    I = len(data_x)
    responsibility = sp.empty((k_size, I))
    for k in sp.arange(k_size):
        norm = stats.norm(loc=mu_vec[k], scale=sp.sqrt(sigma2_vec[k]))
        responsibility[k] = norm.pdf(data_x)

    responsibility = responsibility / sp.sum(responsibility, axis=0)
    return responsibility

def MStep(k_size, responsibility, data_x):

    lambda_vec = sp.empty(k_size)
    mu_vec = sp.empty(k_size)
    sigma2_vec = sp.empty(k_size)

    for k in sp.arange(k_size):
        r_k = responsibility[k]
        lambda_vec[k] = sp.sum(r_k) / sp.sum(responsibility)
        mu_vec[k] = sp.sum(r_k * data_x) / sp.sum(r_k)
        sigma2_vec[k] = sp.sum(r_k * (data_x - mu_vec[k])**2) / sp.sum(r_k)

    return lambda_vec, mu_vec, sigma2_vec

def calc_mix_pdf(k_size, x, lambda_vec, mu_vec, sigma2_vec):
    pdf = sp.zeros_like(x)

    for k in sp.arange(k_size):
        norm_k = stats.norm(loc=mu_vec[k], scale=sp.sqrt(sigma2_vec[k]))
        pdf += lambda_vec[k] * norm_k.pdf(x)
    return pdf

N1 = stats.norm(loc=2, scale=sp.sqrt(1))
N2 = stats.norm(loc=-2, scale=sp.sqrt(1))
s1 = N1.rvs(size=int(sp.around(1000/3*2)))
s2 = N2.rvs(size=int(sp.around(1000/3*1)))
data_x = sp.hstack((s1, s2))
data_x

# print("True:", 2/3, 1/3)
plt.hist(data_x, bins=30, density=True)

#init
k_size = 2
lambda_vec=[0.5, 0.5]
mu_vec=[1, -1]
sigma2_vec=[1, 1, 1, 1]

x = sp.linspace(-6, 6, 200)
plt.plot(x, calc_mix_pdf(k_size, x, lambda_vec, mu_vec, sigma2_vec), label="ini")

for i in sp.arange(50+1):
    responsibility = EStep(k_size, data_x, lambda_vec, mu_vec, sigma2_vec)
    lambda_vec, mu_vec, sigma2_vec = MStep(k_size, responsibility, data_x)
    if i % 5 == 0:
        print("i:", i, "lambda", lambda_vec, "mu", mu_vec, "sigma2", sigma2_vec)
        plt.plot(x, calc_mix_pdf(k_size, x, lambda_vec, mu_vec, sigma2_vec), label=f"i:{i}", linestyle="--")

plt.plot(x, calc_mix_pdf(k_size, x, lambda_vec, mu_vec, sigma2_vec), label=f"last{i}", linestyle="-")

plt.legend()
plt.tight_layout()
plt.show()

i: 0 lambda [0.64534364 0.35465636] mu [ 1.99636031 -1.7754273 ] sigma2 [1.10508121 1.79948778]
i: 5 lambda [0.59236344 0.40763656] mu [ 2.16143918 -1.52509804] sigma2 [0.78561226 2.1105988 ]
i: 10 lambda [0.56973801 0.43026199] mu [ 2.19380078 -1.37409267] sigma2 [0.76310397 2.4452504 ]
i: 15 lambda [0.55812694 0.44187306] mu [ 2.20703843 -1.29705978] sigma2 [0.75407164 2.62149057]
i: 20 lambda [0.55298589 0.44701411] mu [ 2.21208021 -1.26299661] sigma2 [0.75040063 2.70086877]
i: 25 lambda [0.55087658 0.44912342] mu [ 2.2139888 -1.2490169] sigma2 [0.74893837 2.73374747]
i: 30 lambda [0.55003897 0.44996103] mu [ 2.21471953 -1.24346374] sigma2 [0.74836388 2.74686075]
i: 35 lambda [0.54971069 0.45028931] mu [ 2.21500161 -1.24128695] sigma2 [0.74813964 2.75200957]
i: 40 lambda [0.54958269 0.45041731] mu [ 2.21511094 -1.24043812] sigma2 [0.74805233 2.75401863]
i: 45 lambda [0.54953288 0.45046712] mu [ 2.21515338 -1.24010781] sigma2 [0.74801838 2.75480063]
i: 50 lambda [0.54951352 0.45048648] mu [ 2.21516986 -1.23997938] sigma2 [0.74800519 2.75510473]

オレンジの実線：初期値の混合分布