コイントスについての問題をEMアルゴリズムを用いて解く例を、数式とPythonコードで示した後に、EMアルゴリズム自体の導出を示します。

例題 : コイントスで表が出る確率の推定

問題 : ふたつのコインA, Bがあります。コイントスをすると、AのほうがBより表が出やすいことが分かっています。どちらかのコインを使って20回コイントスを行うという試行を10回繰り返した結果、表が出た回数は次のようになりました。A、Bそれぞれのコインの表が出る確率を推定しなさい。

表が出た回数	15	4	6	13	3	5	4	4	7	2

解答 : A、Bそれぞれの表が出る確率を $\theta_A,\theta_B$ 、$N$ 回の試行のうち、$i$ 回目の試行で $M$ 回のコイントスの結果、表が出た回数を $x_i$ 、使ったコインを $z_i \in \{A, B\}$ とおく。

EMアルゴリズムとよばれる、次のアルゴリズムによって $\theta = (\theta_A, \theta_B)$ の推定値 $\hat{\theta}$ を求める。

$\theta$ の推定値の初期値 $\theta^{(0)}$ を決める
推定値 $\theta^{(t)}$ を使って、各試行 $i$ における$z_i$ の事後分布 $p(z_i|x_i,\theta^{(t)})$ を計算する
$\sum_i \sum_{z_i} p(z_i|x_i,\theta^{(t)}) \log p(x_i|z_i,\theta)$ を最大にする $\theta$ を計算し、 $\theta^{(t+1)}$ とする
$\theta^{(t)}$ が収束するまで $\theta^{(t)}$ に $\theta^{(t+1)}$ を入れてステップ1に戻り計算を繰り返す

各ステップのの具体的な計算式とPythonによるコードを以下に示す。

ステップ1

$\theta_A^{(0)} = 0.5, \theta_B^{(0)} = 0.4$ とする。

import numpy as np
import scipy.stats as st

np.random.seed(0)

# Problem setting
N = 10
M = 20
theta_true = np.array([0.8, 0.2])
z_true = np.random.randint(2, size=N)
x = np.zeros(N, dtype=int)
for i in range(0, N):
    x[i] = st.binom.rvs(M, theta_true[z_true[i]])
print(x)

# Initial estimates
t = 0
t_max = 50
thetas = np.zeros([2, t_max])
thetas[:,t] = np.array([0.5, 0.4])

ステップ2.

$\theta^{(t)}$ が与えられたときの $z_i$ の事後分布は、

\begin{eqnarray}
p(z_i|x_i,\theta^{(t)})&=& p(z_i|x_i,\theta^{(t)})\\
 &=&  \frac{p(x_i|z_i, \theta^{(t)}) p(z_i)}{\sum_{z_i} P(x_i|z_i,\theta^{(t)})p(z_i)}\\
 &=&  \frac{p(x_i|z_i, \theta^{(t)})}{\sum_{z_i} P(x_i|z_i,\theta^{(t)})}\\
 &=&  \frac{B(x_i|M,\theta_{z_i}^{(t)})}{\sum_{z_i} B(x_i|M,\theta_{z_i}^{(t)})} \ .
\end{eqnarray}

2行目の等式はベイズの定理を用い、3行目の等式はどちらのコインを使ったかについての情報はないことから事前分布を $p(z_i=A)=p(z_i=B)$ として計算した。また、$B(x|n,p)$ は二項分布を表す。二項分布 $B(x|n,p)$ は事象$U,V$のうち、事象$U$ が確率 $p$ で選択される試行を $n$ 回独立に行った時に $x$ 回 $U$ が選択される確率を表す。

def get_post_z_i(z_i, x_i, M, theta):
    norm_c = 0 # Normalization term in denominator
    for j in range(0,2):
        p_j = st.binom.pmf(x_i, M, theta[j])
        norm_c = norm_c + p_j
        if z_i == j:
            ll = p_j
    return ll / norm_c

ステップ3

ステップ1での計算結果も使い、勾配法によって最大になる $\theta$ を計算する。

def neg_expect(theta_next, theta_cur, x):
    N = x.size
    e = 0
    for i in range(0, N): # for trial i
        for j in range(0, 2): # used coin j
            post_z = get_post_z_i(j, x[i], M, theta_cur)
            prob_x = st.binom.logpmf(x[i], M, theta_next[j])
            e = e + post_z * prob_x
    return -e

# Sample calculation
bnds = [(0,0.99), (0,0.99)]
res = minimize(neg_expect, thetas[:,t], args=(thetas[:,t], x),
          bounds=bnds, method='SLSQP', options={'maxiter': 1000})
res.x

ステップ4

収束するまで繰り返す。

t = 0
while t < t_max-1:
    if t % 10 == 0:
        print(str(t) + "/" + str(t_max))
    res = minimize(neg_expect, thetas[:,t], args=(thetas[:,t], x),
         bounds=bnds, method='SLSQP', options={'maxiter': 1000})
    thetas[:,t+1] = res.x
    t = t + 1

結果

import matplotlib.pyplot as plt
import seaborn as sns
# result
plt.plot(np.arange(0,t_max), thetas[0,:])
plt.plot(np.arange(0,t_max), thetas[1,:])
# true value
plt.plot([0, t_max], np.ones(2) * theta_true[0])
plt.plot([0, t_max], np.ones(2) * theta_true[1])
plt.ylim([0, 1])

今回の問題では、データ数が少なかったため、そこまで正確には推定できませんでした。

EMアルゴリズムの導出

N回の試行で得られた確率変数の標本を $x=\{x_0, \cdots, x_N\}$ 、観測できない確率変数を $z$ 、推定したいパラメーターを $\theta$ とします。$\theta$ の対数尤度関数 $l(\theta)$ は次にように書けます。（以下では、今回の問題にそって $z$ が離散変数であるとして $\sum_z$ を考えますが、 $z$ が連続変数の場合は積分を考えれば同じです。）

l(\theta) := \log p(x|\theta) = \log \sum_z p(x,z|\theta) \ . (1)

実際の計算では、

l(\theta) = \sum_i \log \sum_{z_i} p(x_i,z_i|\theta)\ ,

としたほうが計算しやすいですが、数式が見づらくなるので、ここでは(1)式のまま進めます。ここで、$z$ についての任意の確率分布 $Q(z)$ を考え、Jensenの不等式を使うと、

l(\theta) = \log \sum_z Q(z) \frac{p(x,z|\theta)}{Q(z)} \geq \sum_z Q(z) \log \frac{p(x,z|\theta)}{Q(z)}\ ,

が成り立ちます。上式で不等式で等号が成り立つのは、定数を $C$ として、

\frac{p(x,z|\theta)}{Q(z)} = C \ ,

であるときだけです。これを満たすような $Q(z)$ は、規格化定数を考えれば

\begin{eqnarray}
Q(z)&=\frac{p(x,z|\theta) }{\sum_zp(x,z|\theta)} \\
&=\frac{p(x,z|\theta)}{p(x|\theta)}\\
&=p(z|x,\theta)  \ ,
\end{eqnarray}

です。すなわち、ある $\theta^{(t)}$ を与えると、

\log \sum_z p(z|x,\theta^{(t)}) \frac{p(x,z|\theta^{(t)})}{p(z|x,\theta^{(t)})} = \sum_z p(z|x,\theta^{(t)}) \log \frac{p(x,z|\theta^{(t)})}{p(z|x,\theta^{(t)})}\ ,

が成り立ち、どのような $\theta^{(*)}\neq\theta^{(t)}$ を与えても、

\log \sum_z p(z|x,\theta^{(t)}) \frac{p(x,z|\theta^{(*)})}{p(z|x,\theta^{(t)})} \geq \sum_z p(z|x,\theta^{(t)}) \log \frac{p(x,z|\theta^{(*)})}{p(z|x,\theta^{(t)})}\ ,  \ \ (2)

が成り立つということです。そこで、不等式の下限、すなわち、 $\log p(x,z|\theta^{(*)})/p(z|x,\theta^{(t)})$ の $z$ についての期待値を最大化(Maximization of Expextation) するような $\theta^{(*)}$ を $\theta^{(t+1)}$ とします ;

\begin{eqnarray}
\theta^{(t+1)} &=& \mathrm{arg \ max}_{\theta^{(*)}}  \sum_z p(z|x,\theta^{(t)}) \log \frac{p(x,z|\theta^{(*)})}{p(z|x,\theta^{(t)})} \  \ (3)\\
&=& \mathrm{arg \ max}_{\theta^{(*)}}  \sum_z p(z|x,\theta^{(t)}) \log p(x,z|\theta^{(*)}) \ .  \  \ (4)
\end{eqnarray}

すると、次の不等式が成り立つことが分かります。

\begin{eqnarray}
l(\theta^{(t+1)})
&\geq &\sum_z p(z|x,\theta^{(t)})\log  \frac{p(x,z|\theta^{(t+1)})}{p(z|x,\theta^{(t)})} \\
&\geq& \sum_z p(z|x,\theta^{(t)})\log  \frac{p(x,z|\theta^{(t)})}{p(z|x,\theta^{(t)})} \\
&=& l(\theta^{(t)}) \ .
\end{eqnarray}

ここで最初の不等式には (2)の関係を、次の不等式には(3)の関係を使いました。

これでEMアルゴリズムの導出ができました！これまでのことをまとめると、あるパラメーター $\theta^{(t)}$ が与えられたとき、次の2ステップの計算によって得られる $\theta^{(t+1)}$ は $\theta$ の尤度関数 $l(\theta)$ について、必ず $l(\theta^{(t)})\leq l(\theta^{(t+1)})$ が成り立つことが分かりました。

隠れ変数である $z$ について、$\theta^{(t)}$ を使って事後分布 $p(z|x,\theta^{(t)})$ を計算する
$\sum_z p(z|x,\theta^{(t)}) \log p(x,z|\theta)$ を最大にする $\theta$ を計算し、 $\theta^{(t+1)}$ とする

つまり、適当に $\theta^{(0)}$ を決めて、上記の手順を反復的に上繰り返せば、局所最適な推定値 $\hat{\theta}$ が得られるということです。これがEMアルゴリズムです。

参考資料

CS229 Lecture notes, by Andrew NG
Computational Statistics in Python, Expectation Maximization (EM) Algorithm