EM算法精解

用大众能看懂的文字来详细剖析EM(Expectation-Maximum) 算法。

一、问题描述

有A，B，C三枚硬币，正面朝上的概率分别是a, b, c，是未知参数，也是EM算法要估的参数。现在对三枚硬币进行试验，先抛A硬币，若正面朝上，则抛B硬币，否则抛C硬币。如果抛B硬币或者C硬币时，正面朝上记为1，反面朝上记为0。此为一次试验。现重复进行10次如上的试验，记录观测值为：1, 0, 1, 1, 0, 1, 0, 1, 1, 1。问：a, b, c分别是多少？令$\theta$=(a, b, c)。

二、变量定义

定义随机变量Xi为第i次试验的观测值，其中 X i ∈ { 0 , 1 } , i = 1 , 2 , … , 10 X_i \in \{0,1 \}, i=1,2,\ldots,10 Xi​∈{0,1},i=1,2,…,10。定义隐含随机变量$Z_i$为第i次试验中A硬币正面朝上与否，其中 Z i ∈ { 0 , 1 } , i = 1 , 2 , … , 10 Z_i\in \{0,1\}, i=1,2,\ldots,10 Zi​∈{0,1},i=1,2,…,10，当A硬币正面朝上，$Z_i$为1，否则为0。其中$X_i$是已知，$Z_i$是未知。10次试验的观测值为 x i ∈ { 0 ,   1 } , i = 1 , 2 , … , 10 x_i\in\{0,\ 1\},i=1,2,\ldots,10 xi​∈{0, 1},i=1,2,…,10。

三、似然函数

现有n=10个观测值，它们之间相互独立，故似然函数为：
$$L\left(\theta \right)=\prod _{i=1}^{10}p\left(X_i=x_i;\theta \right)=\prod _{i=1}^{10}[a{b}^{x_i}(1-b{)}^{1-x_i}+(1-a)c^{x_i}(1-c{)}^{1-x_i}]$$
对数似然函数为：
$$l(\theta )=logL(\theta )=\sum _{i=1}^{10}log[a{b}^{x_i}(1-b{)}^{1-x_i}+(1-a)c^{x_i}(1-c{)}^{1-x_i}]$$
我们的目标是求$\theta$，也即a, b, c三个介于0-1的数字使得对数似然函数最大，即：
$${\theta }^{\ast }=\arg \underset{\theta }{\max }l(\theta )$$
这种多个对数里面求和的累加是很难通过求偏导来算解析极值的。于是用Jensen不等式来进行转换，把对数里面的求和提到外面。

$$\begin{gathered} l(\theta )=logL(\theta )=log(\prod _{i=1}^{10}p(X_i=x_i;\theta )) \\ =log(\prod _{i=1}^{10}\sum _{z=0}^{1}p(X_i=x_i,Z_i=z;\theta ))全概率公式 \\ =\sum _{i=1}^{10}log(\sum _{z=0}^{1}p(X_i=x_i,Z_i=z;\theta )) \\ =\sum _{i=1}^{10}log(\sum _{z=0}^{1}B(Z_i=z;\theta )\frac{p(X_i=x_i,Z_i=z;\theta )}{B(Z_i=z;\theta )}) \\ 分子分母同乘B(Z_i=z;\theta )，其中\sum _{z=0}^{1}B(Z_i=z;\theta )=1 \\ =\sum _{i=1}^{10}log[E_Z(\frac{p(X_i=x_i,Z_i;\theta )}{B(Z_i;\theta )})] \\ \ge \sum _{i=1}^{10}{E}_Z(log(\frac{p(X_i=x_i,Z_i;\theta )}{B(Z_i;\theta )}))根据Jensen不等式 \\ =\sum _{i=1}^{10}\sum _{z=0}^{1}B(Z_i=z;\theta )\log [\frac{p(X_i=x_i,Z_i=z;\theta )}{B(Z_i=z;\theta )}] \end{gathered}$$

若令$B(Z_i=z;\theta )=p(Z_i=z|X_i=x_i;\theta ),Z_i=0,1,i=1,2,\dots ,10$则上面等号成立（证明见附录一），即：
$$l(\theta )=\sum _{i=1}^{10}\sum _{z=0}^{1}p(Z_i=z|X_i=x_i;\theta )\log [\frac{P(X_i=x_i,Z_i=z;\theta )}{p(Z_i=z|X_i=x_i;\theta )}]$$
也就是说，我们的目标变成了求$\theta$，使得上述$l(\theta )$ 最大化，即：
$${\theta }^{\ast }=\arg \underset{\theta }{\max }\sum _{i=1}^{10}\sum _{z=0}^{1}p(Z_i=z|X_i=x_i;\theta )log[\frac{p(X_i=x_i,Z_i=z;\theta )}{p(Z_i=z|X_i=x_i;\theta )}]$$

四、迭代求参数

接下来用迭代的方法来求参数$\theta$使得$l(\theta )$最大化。

1. 首先初始化${\theta }_j={\theta }_0$，例如$a=0.4,b=0.5,c=0.6$。
2. 则第$j+1$次迭代的$\theta$如下：
   $$\begin{gathered} {\theta }_{j+1}=\arg \underset{\theta }{\max }\sum _{i=1}^{10}\sum _{z=0}^{1}p(Z_i=z|X_i=x_i;{\theta }_j)log[\frac{p(X_i=x_i,Z_i=z;\theta )}{p(Z_i=z|X_i=x_i;\theta )}] \\ =\arg \underset{\theta }{\max }\sum _{i=1}^{10}\sum _{z=0}^1[p(Z_i=z_i|X_i=x_i;{\theta }_j)\log p(X_i=x_i,Z_i=z;\theta )-p(Z_i=z_i|X_i=x_i;{\theta }_j)\log p(Z_I=z|X_i=x_i;{\theta }_j)] \end{gathered}$$
   去掉和$\theta$无关的项，得：
   $$\arg \underset{\theta }{\max }\sum _{i=1}^{10}\sum _{z=0}^{1}p(Z_i=z_i|X_i=x_i;{\theta }_j)\log p(X_i=x_i,Z_i=z;\theta )$$
   记${\sum }_{i=1}^{10}{\sum }_{z=0}^{1}p(Z_i=z_i|X_i=x_i;{\theta }_j)\log p(X_i=x_i,Z_i=z;\theta )$为$Q$函数，记作$Q(\theta ,{\theta }_j)$。因为$Q(\theta ,{\theta }_j)={\sum }_{i=1}^{10}{\sum }_{z=0}^{1}p(Z_i=z_i|X_i=x_i;{\theta }_j)\log p(X_i=x_i,Z_i=z;\theta )={\sum }_{i=1}^{10}{E}_Z[\log p(X_i=x_i,Z_i=z;\theta )]$，这就是EM算法中E的由来。
3. 继而，求使得$Q(\theta ,{\theta }_j)$最大化的${\theta }_{j+1}$，即${\theta }_{j+1}=\arg \underset{\theta }{\max }Q(\theta ,{\theta }_j)$，这就是M的由来。
   $$\begin{gathered} Q(\theta ,{\theta }_j)=\sum _{i=1}^{10}\sum _{z=0}^1[p(Z_i=z|X_i=x_i;{\theta }_j)\log p(X_i=x_i,Z_i=z;\theta )] \\ =\sum _{i=1}^{10}[p(Z_i=0|X_i=x_i;{\theta }_j)\log p(X_i=x_i,Z_i=0;\theta )+p(Z_i=1|X_i=x_i;{\theta }_j)\log p(X_i=x_i,Z_i=1;\theta )] \\ =\sum _{i=1}^{10}[\frac{p(Z_i=0,X_i=x_i;{\theta }_j)}{p(X_i=x_i;{\theta }_j)}\log p(X_i=x_i,Z_i=0;\theta )+\frac{p(Z_i=1,X_i=x_i;{\theta }_j)}{p(X_i=x_i;{\theta }_j)}\log p(X_i=x_i,Z_i=1;\theta )]条件概率 \\ =\sum _{i=1}^{10}[\frac{p(Z_i=0,X_i=x_i;{\theta }_j)}{{\sum }_{z=0}^{1}p(X_i=x_i,Z_i=z;{\theta }_j)}\log p(X_i=x_i,Z_i=0;\theta )+\frac{p(Z_i=1,X_i=x_i;{\theta }_j)}{{\sum }_{z=0}^{1}p(X_i=x_i,Z_i=z;{\theta }_j)}\log p(X_i=x_i,Z_i=1;\theta )] \\ =\sum _{i=1}^{10}\frac{(1-a_j)c_j^{x_i}(1-c_j{)}^{1-x_i}}{(1-a_j)c_j^{x_i}(1-c_j{)}^{1-x_i}+a_j{b}_j^{x_i}(1-b_j{)}^{1-x_i}}\log [(1-a_{j+1})c_{j+1}^{x_i}(1-c_{j+1}{)}^{1-x_i}]+ \\ \frac{a_j{b}_j^{x_i}(1-b_j{)}^{1-x_i}}{(1-a_j)c_j^{x_i}(1-c_j{)}^{1-x_i}+a_j{b}_j^{x_i}(1-b_j{)}^{1-x_i}}\log [a_{j+1}{b}_{j+1}^{x_i}(1-b_{j+1}{)}^{1-x_i}] \end{gathered}$$
   令$m_i=\frac{(1-a_j)c_j^{x_i}(1-c_j{)}^{1-x_i}}{(1-a_j)c_j^{x_i}(1-c_j{)}^{1-x_i}+a_j{b}_j^{x_i}(1-b_j{)}^{1-x_i}},i=1,2,\dots ,10$，令$n_i=\frac{a_j{b}_j^{x_i}(1-b_j{)}^{1-x_i}}{(1-a_j)c_j^{x_i}(1-c_j{)}^{1-x_i}+a_j{b}_j^{x_i}(1-b_j{)}^{1-x_i}},i=1,2,\dots ,10$，有$m_i+n_i=1,i=1,2,\dots ,10$

$Q$函数中$a_j,b_j,c_j$是上一轮迭代的结果，是已知的。其中$a_{j+1},b_{j+1},c_{j+1}$是未知参数，现在分别对它们求偏导并等于$0$来求使$Q$函数最大的$a_{j+1},b_{j+1},c_{j+1}$。
$$\frac{\partial Q(\theta ,{\theta }_j)}{\partial a_{j+1}}=\sum _{i=1}^{10}(\frac{m_i}{a_{j+1}-1}+\frac{n_i}{a_{j+1}})=\sum _{i=1}^{10}\frac{a_{j+1}{m}_i+a_{j+1}{n}_i-n_i}{(a_{j+1}-1)a_{j+1}}$$
令$\frac{\partial Q(\theta ,{\theta }_j)}{\partial a_{j+1}}=0$得：$a_{j+1}=\frac{1}{10}{\sum }_{i=1}^{10}{n}_i$。
又：
$$\begin{gathered} \frac{\partial Q(\theta ,{\theta }_j)}{\partial b_{j+1}}=\sum _{i=1}^{10}{n}_i\frac{a_{j+1}}{a_{j+1}{b}_{j+1}^{x_i}(1-b_{j+1}{)}^{1-x_i}}(b_{j+1}^{x_i}(1-b_{j+1}{)}^{1-x_i})\prime \\ =\sum _{i=1}^{10}{n}_i(\frac{x_i}{b_{j+1}}-\frac{1-x_i}{1-b_{j+1}})根据求导乘法法则并化简 \end{gathered}$$
令$\frac{\partial Q(\theta ,{\theta }_j)}{\partial b_{j+1}}=0$得：$b_{j+1}=\frac{{\sum }_{i=1}^{10}{n}_i{x}_i}{{\sum }_{i=1}^{10}{n}_i}$
同理可得：$c_{j+1}=\frac{{\sum }_{i=1}^{10}{m}_i{x}_i}{{\sum }_{i=1}^{10}{m}_i}$
4. 重复2，3两步， 直至收敛，或者达到预设标准，最后的$\theta$即为所求。

五、代码示例

import random
import math

a_true = 0.9 # A硬币正面朝上的概率是0.9
b_true = 0.1 # B硬币正面朝上的概率是0.1
c_true = 0.8 # C硬币正面朝上的概率是0.8
n = 5000 # 一共执行5000次试验

obs = []
for i in range(n):
    if random.random() <= a_true:
        obs.append(1) if random.random() <= b_true else obs.append(0)
    else:
        obs.append(1) if random.random() <= c_true else obs.append(0)
print(obs)
# 输出5000次试验的观察结果，1是正面朝上，0是 反面朝上
[1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, ...] 共5000个

# 初始化第一次的参数
thetaj = [0.45, 0.2, 0.98] # a0, b0, c0 = 0.45, 0.2, 0.98

# thetaj(s) are from the previous iteration or initialization
def em(thetaj, obs):
    m = []
    n = []
    for i in range(len(obs)):
        m_now = ( (1-thetaj[0])*math.pow(thetaj[2],obs[i])*math.pow(1-thetaj[2],1-obs[i]) )/( (1-thetaj[0])*math.pow(thetaj[2],obs[i])*math.pow(1-thetaj[2],1-obs[i]) + thetaj[0]*math.pow(thetaj[1],obs[i])*math.pow(1-thetaj[1],1-obs[i]) )
        m.append(m_now)
        n_now = 1 - m_now
        n.append(n_now)
  
    m_total, n_total, mx_total, nx_total = 0, 0, 0, 0
    for i in range(len(obs)):
        m_total += m[i]
        n_total += n[i]
        mx_total += m[i]*obs[i]
        nx_total += n[i]*obs[i]
  
    thetaj[0] = n_total/len(obs)
    thetaj[1] = nx_total/n_total
    thetaj[2] = mx_total/m_total

# 一共迭代1000次来估算参数
for j in range(1000):
    em(thetaj, obs)

print("Type: a,b,c")
print("True: " + str(a_true) + "," + str(b_true) + "," + str(c_true))
print("Initialization: " + str(a0) + "," + str(b0) + "," + str(c0))
print("Estimated: " + ",".join([str(i) for i in thetaj]))
# 输出
Type: a,b,c
True: 0.9,0.1,0.8
Initialization: 0.45,0.2,0.98
Estimated: 0.8250824695060178,0.030452224013034972,0.8602591368880216

六、扩展

类似的，可以扩展到其他分布类型，例如混合高斯。只需要根据分布来修改$Q(\theta ,{\theta }_j)$中的$p(Z_i,X_i;{\theta }_j)$表达式，再分别对要估算的参数求偏导数就可以。例如现有某个年级100名学生的身高的测量值，要求估算出男生和女生身高的均值和方差，也即4个参数需要估计。假设男生和女生身高呈现正态分布。这里的$Z$随机变量表示该观测值是来自男生还是女生，是一个$0-1$分布。$X$随机变量服从混合高斯分布。在给定$Z$的前提下$X$服从一维正态分布。

七、附录一

由Jensen不等式可知，若要上面的等号成立，$\frac{p(X_i=x_i,Z_i=z;\theta )}{B(Z_i=z;\theta )}$必须是一个常数。现在证明当$B(Z_i=z;\theta )=p(Z_i=z|X_i=x_i;\theta ),Z_i=0,1,i=1,2,\dots ,10,且{\sum }_{z=0}^{1}B(Z_i=z;\theta )=1$时，$\frac{p(X_i=x_i,Z_i=z;\theta )}{B(Z_i=z;\theta )}$是个常数。
设$\frac{p(X_i=x_i,Z_i=z;\theta )}{B(Z_i=z;\theta )}=c$，$c$是任意一个常数，则：
$$\begin{gathered} B(Z_i=z;\theta )=\frac{p(X_i=x_i,Z_i=z;\theta )}{c}=\frac{p(X_i=x_i,Z_i=z;\theta )}{{\sum }_{z=0}^{1}cB(Z_i=z;\theta )} \\ =\frac{p(X_i=x_i,Z_i=z;\theta )}{{\sum }_{z=0}^{1}p(X_i=x_i,Z_i=z;\theta )} \\ =\frac{p(X_i=x_i,Z_i=z;\theta )}{p(X_i=x_i;\theta )} \\ =p(Z_i=z|X_i=x_i;\theta ) \end{gathered}$$
所以，当$B(Z_i=z;\theta )=p(Z_i=z|X_i=x_i;\theta ),Z_i=0,1,i=1,2,\dots ,10$时上述等号成立。