计算建模之EM算法

EM(Expectation-Maximization)算法于1977年提出，被用于解决如下的参数估计问题：
已知（或者至少我们的模型如此假定）随机变量$X$服从概率密度函数$p(X,Z;\theta ),$其中$\theta$为待确定的参数，$Z$为隐变量(latent variable)，是由于某些原因我们无法直接观测得到的随机变量。EM算法可以在我们只能观测到$X$的情况下对参数$\theta$进行估计。下面我们举两个例子来说明该算法。

例一

这个经典的例子出自¹,问题描述如下：
我们有两枚硬币A和B，它们投出正面的概率分别为${\theta }_A$和${\theta }_B,$是待估计的参数。假如我们现在投掷$5$轮，每轮先以等概率选择其中一枚，再投掷这枚硬币$10$次，得到如下结果：

轮次	硬币	结果
1	B	5正5反
2	A	9正1反
3	A	8正2反
4	B	4正6反
5	A	7正3反

这个模型中，每轮的结果（即观测变量）可以正面出现的次数唯一确定，记为$y$;此外，对结果有影响的变量还有每轮选择的硬币类型(A/B)，我们记为$z$,当$z=1$时，表示我们选择了$A$硬币；否则当$z=0$时，表示此时选择$B$硬币。
如果这时让我们估算未知的参数${\theta }_A$，可以直接用轮次2,3,5的正面次数除以总次数($\widehat{{\theta }_A}=\frac{9+8+7}{30}=\frac{4}{5}$);但是为了与后面的EM算法的思想保持一致，我们稍微绕一下远：可以写出如下的概率分布函数

$$P(y,z;\theta \theta )=P(y|z;\theta \theta )P(z)=\frac{1}{2}P(y|z)=\{\begin{array}{r}\frac{1}{2}{C}_{10}^y{\theta }_A^y(1-{\theta }_A{)}^{10-y},z=1\\ \frac{1}{2}{C}_{10}^y{\theta }_B^y(1-{\theta }_B{)}^{10-y},z=0\end{array}$$

容易知道它等价于下面的式子：

$$P(y,z;\theta \theta )=\frac{1}{2}z{C}_{10}^y{\theta }_A^y(1-{\theta }_A{)}^{10-y}+\frac{1}{2}(1-z)C_{10}^y{\theta }_B^y(1-{\theta }_B{)}^{10-y}$$
其中$z\in \{0,1\},y\in \{0,1,...,n\},\theta \theta =({\theta }_A,{\theta }_B).$

为了估算式子中的$\theta \theta$的值，可以用概率论中的极大似然估计构造对数似然函数：

$$\begin{array}{rl}LL(\theta \theta |y,z)& =ln\prod _{i=1}^{5}P(y_i,z_i;\theta \theta )\\ & =ln\prod _{i=1}^5[\frac{1}{2}{z}_i{C}_{10}^{y_i}{\theta }_A^{y_i}(1-{\theta }_A{)}^{10-y_i}+\frac{1}{2}(1-z_i)C_{10}^{y_i}{\theta }_B^{y_i}(1-{\theta }_B{)}^{10-y_i}]\\ & =\sum _{i=1}^{5}ln[\frac{1}{2}{z}_i{C}_{10}^{y_i}{\theta }_A^{y_i}(1-{\theta }_A{)}^{10-y_i}+\frac{1}{2}(1-z_i)C_{10}^{y_i}{\theta }_B^{y_i}(1-{\theta }_B{)}^{10-y_i}]\end{array}$$

（注：对数似然函数$LL(\theta \theta |y,z)$的记号表示$y,z$是已知的，此时$\theta \theta$才是该函数的自变量）

之后可以令该函数对$\theta \theta$偏导为0，以求出$\theta \theta$的值使得$LL(\theta \theta |y,z)$最大，即令$\partial (LL)/\partial \theta \theta =0,$解出$\theta \theta$的值即可。

这个问题的难度可以升一下级：当我们不知道每轮取到的硬币时（见下表），能否估计参数$\theta \theta$的值呢？EM算法就可以用来解决这个问题。

轮次	硬币	结果
1	？	5正5反
2	？	9正1反
3	？	8正2反
4	？	4正6反
5	？	7正3反

不知道每轮取到的硬币，也就是上面的概率分布函数$P(y,z)$中，$z$的值无法确定了，从而在求解对数似然函数偏导数为$0$时，我们没法解出$\theta \theta$的确定值（因为$z_i$是未知的）。EM算法的解决方案是，初始化一个参数${\theta \theta }^0,$于是对于每个样本$y_i$,概率分布函数

$$P(z|y,{\theta \theta }^0)=\frac{1}{2}z{C}_{10}^y{\theta }_A^{0y}(1-{\theta }_A^0{)}^{10-y}+\frac{1}{2}(1-z)C_{10}^y{\theta }_B^{0y}(1-{\theta }_B^0{)}^{10-y}$$

中，只有$z_i$是不确定的了（这里的符号有一些变化，即观测变量$y_i$变成了条件，已知），也就是说，我们可以求出$z_i$的概率分布。那么，我们就可以求出对数似然函数$LL(\theta \theta |y,z)$的期望${\mathbb{E}}_z[LL(\theta \theta |y,z)]$，来代替原来的对数似然函数（取期望这个操作消去了不确定的$z$，使其变为仅与$\theta \theta$有关的函数）。这时我们再求出使得${\mathbb{E}}_z[LL(\theta \theta |y,z)]$最大的${\theta \theta }^1$，作为新一轮的$\theta \theta$的估计值，重复上面的操作直到参数$\theta \theta$收敛（收敛性在此不证明，分析可参见²），算法停止。

总结一下EM算法的步骤：

(0)：用随机数初始化位置参数${\theta \theta }^0$;

(1)（E步）将观测变量$y_i$代入$LL(\theta \theta |y,z)=ln\prod _{i=1}^{N}P(y_i,z_i;\theta \theta )$,求出对数似然函数关于$z$的期望${\mathbb{E}}_z[LL(\theta \theta |y,z)]$;

(2)（M步）求出使得(1)中期望最大的参数${\theta \theta }^{t+1}.$这里，$z$服从的概率分布由上一轮的参数${\theta \theta }^t$确定。

(3) 若迭代次数到达指定上限或$\theta \theta$已经收敛（即${\theta \theta }^{t+1}$和${\theta \theta }^t$的差距已经非常小），结束算法；否则回到(1)继续进行迭代。

说明：这里出现的$\theta \theta$和${\theta \theta }^t$容易让人弄混。我们在EM算法中的目标函数是${\mathbb{E}}_z[LL(\theta \theta |y,z)]$，这是一个关于未知参数$\theta \theta$的函数；估计出来的${\theta \theta }^t$用于确定隐变量$z$的分布，以计算$LL(\theta \theta |y,z)$关于$z$的期望。

下面我们对上述硬币问题应用EM算法，来熟悉一下这个过程。我们这里把轮次记为$N$，每轮投掷次数为$n$（上面的例子中，$N=5,n=10$）.设每轮观测到正面朝上的次数为$y_i,$每轮选出的硬币为$z_i\in \{0,1\}$,是一个隐变量。

（1）初始化${\theta \theta }^0,$可以调用随机函数生成，这里不指定具体数值。

（2）写出概率分布函数：

$$P(y,z)=\frac{1}{2}z{C}_n^y{\theta }_A^y(1-{\theta }_A{)}^{n-y}+\frac{1}{2}(1-z)C_n^y{\theta }_B^y(1-{\theta }_B{)}^{n-y}$$

（3）表示对数似然函数关于隐变量$z$的期望：
$$\begin{array}{rl}{\mathbb{E}}_z[LL(\theta \theta |y,z)]& ={\mathbb{E}}_z\{\sum _{i=1}^{N}ln[\frac{1}{2}{z}_i{C}_n^{y_i}{\theta }_A^{y_i}(1-{\theta }_A{)}^{n-y_i}+\frac{1}{2}(1-z_i)C_n^{y_i}{\theta }_B^y(1-{\theta }_B{)}^{n-y_i}]\}\\ & =\sum _{i=1}^N{\mathbb{E}}_z\{ln[\frac{1}{2}{z}_i{C}_n^{y_i}{\theta }_A^{y_i}(1-{\theta }_A{)}^{n-y_i}+\frac{1}{2}(1-z_i)C_n^{y_i}{\theta }_B^y(1-{\theta }_B{)}^{n-y_i}]\}(把期望号放进求和号)\\ & =\sum _{i=1}^N\{P(z=1|y_i,{\theta \theta }^t)ln[\frac{1}{2}{C}_n^{y_i}{\theta }_A^{y_i}(1-{\theta }_A{)}^{n-y_i}]+P(z=0|y_i,{\theta \theta }^t)ln[C_n^{y_i}{\theta }_B^y(1-{\theta }_B{)}^{n-y_i}]\}(离散型随机变量期望的定义)\end{array}$$

（4）求使得${\mathbb{E}}_z[LL(\theta \theta |y,z)]$最大的$\theta \theta$，需要对${\theta }_A,{\theta }_B$求偏导,这里以${\theta }_A$为例：
$$\begin{array}{rl}\frac{\partial {\mathbb{E}}_z}{\partial {\theta }_A}& =\sum _{i=1}^{N}P(z=1|y_i,{\theta \theta }^t)\frac{\frac{1}{2}{C}_n^{y_i}{\theta }_A^{y_i-1}(1-{\theta }_A{)}^{n-y_i-1}(y_i-n{\theta }_A)}{\frac{1}{2}{C}_n^{y_i}{\theta }_A^{y_i}(1-{\theta }_A{)}^{n-y_i}}\\ & =\sum _{i=1}^{N}P(z=1|y_i,{\theta \theta }^t)\frac{y_i-n{\theta }_A}{{\theta }_A(1-{\theta }_A)}\end{array}$$

令$$\frac{\partial {\mathbb{E}}_z}{\partial {\theta }_A}=0,$$
有
$$\begin{gathered} \sum _{i=1}^{N}P(z=1|y_i,{\theta \theta }^t)(y_i-n{\theta }_A)=0, \\ \sum _{i=1}^{N}n{\theta }_{A}P(z=1|y_i,{\theta \theta }^t)=\sum _{i=1}^N{y}_{i}P(z=1|y_i,{\theta \theta }^t) \end{gathered}$$
解得
$${\theta }_A^{t+1}=\frac{1}{n}\cdot \frac{\sum _{i=1}^N{y}_i\cdot P(z=1|y_i,{\theta \theta }^t)}{\sum _{i=1}^{N}P(z=1|y_i,{\theta \theta }^t)}(1)$$
同理
$${\theta }_B^{t+1}=\frac{1}{n}\cdot \frac{\sum _{i=1}^N{y}_i\cdot P(z=0|y_i,{\theta \theta }^t)}{\sum _{i=1}^{N}P(z=0|y_i,{\theta \theta }^t)}(2)$$
其中(注意${\theta }_A,{\theta }_B$上面的角标$t$)
$$\begin{array}{rl}P(z=1|y_i,{\theta \theta }^t)& =\frac{C_n^{y_i}{\theta }_A^{t{y}_i}(1-{\theta }_A^t{)}^{n-y_i}}{C_n^{y_i}{\theta }_A^{t{y}_i}(1-{\theta }_A^t{)}^{n-y_i}+C_n^{y_i}{\theta }_B^{t{y}_i}(1-{\theta }_B^t{)}^{n-y_i}}(已知y_i次正面朝上，投的硬币是A的概率)\\ & =\frac{{\theta }_A^{t{y}_i}(1-{\theta }_A^t{)}^{n-y_i}}{{\theta }_A^{t{y}_i}(1-{\theta }_A^t{)}^{n-y_i}+{\theta }_B^{t{y}_i}(1-{\theta }_B^t{)}^{n-y_i}}.\end{array}$$
那么
$$\begin{array}{rl}P(z=0|y_i,{\theta \theta }^t)& =1-P(z=1|y_i,{\theta \theta }^t)(对立事件)\\ & =\frac{{\theta }_B^{t{y}_i}(1-{\theta }_B^t{)}^{n-y_i}}{{\theta }_A^{t{y}_i}(1-{\theta }_A^t{)}^{n-y_i}+{\theta }_B^{t{y}_i}(1-{\theta }_B^t{)}^{n-y_i}}.\end{array}$$

因此我们可以用式$(1)(2)$迭代地更新参数直到收敛。下面我们用python模拟一下这个过程，验证式子的正确性以及算法的收敛性：

import random
import numpy as np

# 每轮的投掷次数n
n = 10
thetaA, thetaB = 0.5, 0.7

# 模拟概率为p，次数为n的二项分布(投硬币)，返回正面朝上的次数
def observe(n, p):
    sum = 0
    for i in range(n):
        t = random.random()
        if t < p:
            sum += 1
    return sum

# N 为数据集大小，返回 N 轮后的结果
def get_dataset(N):
    dataset = []
    for i in range(N):
        if random.random() < 0.5:
            dataset.append(observe(n, thetaA))
        else:
            dataset.append(observe(n, thetaB))
    return dataset
    
# 对dataset应用EM算法
def EM(dataset)
	# y是N维向量，表示观测结果之集
    y = np.array(dataset, dtype = 'float')
    # 初始化估计参数。random.random()返回0~1中的小数
    thetaA_hat, thetaB_hat = random.random(), random.random()
    
    iteration = 0
    while True:
        iteration += 1
        A = ((thetaA_hat ** y) * (1 - thetaA_hat) ** (n - y))
        B = ((thetaB_hat ** y) * (1 - thetaB_hat) ** (n - y))
        # f,g 分别表示上文的P(z=1|y,theta),P(z=0|y,theta),是N维向量
        # 向量的第 i 维代表P(z|y_i,theta)
        f = A / (A + B)
        g = B / (A + B)
        
        # 这个更新与上文(1)(2)本质相同，不过用了numpy中的点乘简化运算
        
        # 即 [y_1, y_2, ..., y_N] 与 [P(z=1|y_1,theta), ..., P(z=1|y_N,theta)]的点乘
        new_thetaA_hat = np.dot(y, f) / n / np.sum(f)

		# 即 [y_1, y_2, ..., y_N] 与 [P(z=0|y_1,theta), ..., P(z=0|y_N,theta)]的点乘
        new_thetaB_hat = np.dot(y, g) / n / np.sum(g)
        
        # 计算一下两次参数之间的变化量(欧几里得距离平方)
        conv = (thetaA_hat - new_thetaA_hat) ** 2 + (thetaB_hat - new_thetaB_hat) ** 2
        if conv < 1e-10:
            break
        
        thetaA_hat = new_thetaA_hat
        thetaB_hat = new_thetaB_hat
    print(f'iteration = {iteration}')
    return thetaA_hat, thetaB_hat
# 做10000轮实验
dataset = get_dataset(10000)
# 重复调用20次EM算法，查看每次是否收敛到一个结果，以及迭代次数
for i in range(20):
    print(EM(dataset))

运行结果：

iteration = 15
(0.4977663316100947, 0.6962688553151227)
iteration = 15
(0.4977649567622754, 0.6962702416641224)
iteration = 16
(0.497762899611069, 0.6962724001030222)
iteration = 15
(0.6962733823360499, 0.4977619579412015)
iteration = 14
(0.49776333638894465, 0.6962718926393595)
# ...(后面省略)

可以看到概率值与真实值$(0.5,0.7)$非常接近了，平均只需要迭代十几次即可收敛。但我们发现有的时候结果会收敛于$(0.7,0.5)$，与我们初始设置的参数相反。这其实可以理解：我们已经“屏蔽”了每次投硬币的编号，所以我们至多只能得出“有两枚硬币，其中一枚有0.5的几率朝上，另一枚有0.7的几率朝上”的结论。至于它们的编号是无关紧要也无从得知的。

例二

这个例子出自³,也是哈工大课件中的例子，问题描述如下：

在一个模式识别任务中，物体大致上可分为两类：暗色物体和亮色物体。其中暗色物体还可以分为方形和圆形。一个特征提取器可以将物体分为暗色物体$(y_1)$和亮色物体$(y_2)$，却无法区分方形和圆形的物体。此外，我们已知这三种物体（暗圆$x_1$，暗方$x_2$，亮色$x_3$）服从如下的三项分布，它是二项分布的推广：

$$P(x_1,x_2,x_3;\theta )=\frac{(x_1+x_2+x_3)!}{x_1!x_2!x_3!}(\frac{1}{4}{)}^{x_1}(\frac{1+\theta }{4}{)}^{x_2}(\frac{2-\theta }{4}{)}^{x_3}\theta \in (-1,2).$$

在这个问题中，观测变量为$Y=(y_1,y_2),$全部变量为$X=(x_1,x_2,x_3)$，有
$$\begin{gathered} y_1=x_1+x_2 \\ {y}_2=x_3 \end{gathered}$$
隐变量为$x_1,x_2$（$x_3=y_2$，因此是可以被观测的）。我们尝试对这个问题应用EM算法，以估计概率密度函数中的未知参数$\theta$。设$X_i=(x_{i1},x_{i2},x_{i3})$为第$i$个样本的全部变量，$Y_i=(y_{i1},y_{i2})$为第$i$个样本的观测变量，EM算法流程如下：
（1）初始化${\theta \theta }^0$;
（2）表示对数似然函数关于隐变量$X_h$的期望(由于我们可以求出$x_2=y_1-x_1$，可以省略它)：
$$\begin{array}{rl}{\mathbb{E}}_{X_h}[LL(\theta \theta |X,Y)]& ={\mathbb{E}}_{X_h}[\sum _{i=1}^{N}lnP(x_{i1},x_{i2},x_{i3}|y_{i1},y_{i2};\theta )]\\ & =\sum _{i=1}^N{\mathbb{E}}_{X_h}[lnP(x_{i1},x_{i2},x_{i3}|y_{i1},y_{i2};\theta )]\end{array}$$

其中

$$\begin{array}{rl}{\mathbb{E}}_{X_h}[lnP(x_{i1},x_{i2},x_{i3}|y_{i1},y_{i2};\theta )]& ={\mathbb{E}}_{X_h}[ln\frac{(x_1+x_2+x_3)!}{x_1!x_2!x_3!}(\frac{1}{4}{)}^{x_1}(\frac{1+\theta }{4}{)}^{x_2}(\frac{2-\theta }{4}{)}^{x_3}]\\ & ={\mathbb{E}}_{X_h}[ln\frac{(x_1+x_2+x_3)!}{x_1!x_2!x_3!}+x_{1}ln\frac{1}{4}+x_{2}ln(\frac{1+\theta }{4})+x_{3}ln(\frac{2-\theta }{4})]\\ & ={\mathbb{E}}_{X_h}[ln\frac{(x_1+x_2+x_3)!}{x_1!x_2!x_3!}]+ln\frac{1}{4}{\mathbb{E}}_{X_h}[x_1]+ln(\frac{1+\theta }{4}){\mathbb{E}}_{X_h}[x_2]+ln(\frac{2-\theta }{4}){\mathbb{E}}_{X_h}[x_3]\end{array}.$$
（3）求得使(2)中期望最大的$\theta$值,为此需要对期望求$\theta$的导数：
$$\begin{array}{rl}\sum _{i=1}^N\frac{d{\mathbb{E}}_{X_h}}{d\theta }& =\sum _{i=1}^N(\frac{1}{\theta +1}{\mathbb{E}}_{X_h}[x_{i2}]-\frac{1}{2-\theta }{\mathbb{E}}_{X_h}[x_{i3}])=0\end{array}$$
得
$${\theta }^{t+1}=\frac{\sum _{i=1}^N(2{\mathbb{E}}_{X_h}[x_{i2}]-{\mathbb{E}}_{X_h}[x_{i3}])}{\sum _{i=1}^N({\mathbb{E}}_{X_h}[x_{i2}]+{\mathbb{E}}_{X_h}[x_{i3}])}.$$
由于${\mathbb{E}}_{X_h}[\cdot ]$表示已知${\theta }^t,y_{i1},y_{i2}$的对$x_1$的期望，因此
$$\begin{gathered} {\mathbb{E}}_{X_h}[x_{i3}|y_{i1},y_{i2}]=y_{i2}(x_{i3}=y_{i2}); \\ {\mathbb{E}}_{X_h}[x_{i2}|y_{i1},y_{i2}]=y_{i1}\frac{\frac{1+{\theta }^t}{4}}{\frac{1}{4}+\frac{1+{\theta }^t}{4}}=y_{i1}\frac{1+{\theta }^t}{2+{\theta }^t}(见下方说明). \end{gathered}$$
对于上式的第二行，原论文的第53页附加了一个简单的证明。这里我们不作严格的证明，仅从三项分布的意义和性质来理解：三项分布的一般形式如
$$P(x_1,x_2,x_3)=\frac{(x_1+x_2+x_3)!}{x_1!x_2!x_3!}{p}^{x_1}{q}^{x_2}{r}^{x_3}$$
其中$p+q+r=1,$分别代表在总体中抽出第一/二/三类样本的概率为$p/q/r$（放回抽样，和二项分布很类似）；整个式子代表抽取$n=x_1+x_2+x_3$次，分别抽得一/二/三类样本$x_1/x_2/x_3$个的概率。回到上面第二行的问题：当$x_{i1}+x_{i2}$固定为$y_{i1}$时,我们想求得${\mathbb{E}}_{X_h}[x_{i2}|y_{i1},y_{i2}].$这时相当于一个总数为$y_{i1},$概率为$\frac{\frac{1+{\theta }^t}{4}}{\frac{1}{4}+\frac{1+{\theta }^t}{4}}$(将$x_1$和$x_2$的概率进行“加权”分配，即归一化)的二项分布，即给定$y_{i1}$的条件下$x_{i2}\sim B(y_{i1},\frac{1+{\theta }^t}{2+{\theta }^t}).$由二项分布的性质得
$${\mathbb{E}}_{X_h}[x_{i2}|y_{i1},y_{i2}]=y_{i1}\frac{1+{\theta }^t}{2+{\theta }^t}.$$

综上所述，参数更新表达式为
$${\theta }^{t+1}=\frac{\sum _{i=1}^N(2y_{i1}\frac{1+{\theta }^t}{2+{\theta }^t}-y_{i2})}{\sum _{i=1}^N(y_{i1}\frac{1+{\theta }^t}{2+{\theta }^t}+y_{i2})}.$$

下面我们也是用python来验证一下这个结果：

import random
import numpy as np

# 每次三项分布总共取出20个样本
n = 20
# 参数theta
theta = 1.5

# 返回观测值y_1, y_2
def observe(n, p1, p2, p3):
    sum1, sum2, sum3 = 0, 0, 0
    for i in range(n):
        t = random.random()
        if t < p1:
            sum1 += 1
        elif t < p1 + p2:
            sum2 += 1
        else:
            sum3 += 1
    return [sum1 + sum2, sum3]

# N 为数据集大小
def get_dataset(N):
    dataset = []
    for i in range(N):
        dataset.append(observe(n, 1 / 4, (1 + theta) / 4, (2 - theta) / 4))
    return dataset

def EM(dataset):
    # theta在(-1, 2)之间
    theta_hat = random.random() * 3 - 1
    y = np.array(dataset, dtype = 'float')
    
    # A,B 分别为y_i1,y_i2之和
    A = np.sum(y[:, 0])
    B = np.sum(y[:, 1])
    
    iteration = 0
    while True:
        iteration += 1
        
        tmp = (1 + theta_hat) / (2 + theta_hat)
        new_theta_hat = (2 * A * tmp - B) / (A * tmp + B)
        
        conv = abs(new_theta_hat - theta_hat)
        if conv < 1e-10 or iteration > 1000:
            break
            
        theta_hat = new_theta_hat
            
    print(f'iteration = {iteration}')
    return theta_hat

dataset = get_dataset(10000)
for i in range(20):
	print(EM(dataset))

运行结果：

iteration = 8
1.5024600000467414
iteration = 9
1.5024599999731247
iteration = 9
1.5024600000072472
iteration = 10
1.502459999983291
iteration = 9
1.5024600000063015
# ...(后面省略)

结果普遍收敛到初始定下的$\theta =1.5,$且收敛速度很快。

参考文献

---

1. Do, C., Batzoglou, S. What is the expectation maximization algorithm?. Nat Biotechnol 26, 897–899 (2008). ↩︎

2. A. O. Hero, “On the Convergence of the Em Algorithm,” Proceedings. IEEE International Symposium on Information Theory, 1993, pp. 187-187, doi: 10.1109/ISIT.1993.748501. ↩︎

3. Molenberghs G , Verbeke G . The Expectation-Maximization Algorithm[M]. 2000. ↩︎