期望最大化算法(EM算法)举例加推导详解

期望最大化算法(EM算法)

1. 背景

极大似然估计 (MLE) 是一种非常有效的参数估计方法，但当分布中有多余参数或数据为截尾或缺失时，其 MLE 的求取是比较困难的。Dempster 等人于 1977 年提出了 EM 算法，它是寻找具有潜变量的概率模型的最大似然解的一种通用方法。其出发点是把求 MLE 的过程分两步走，第一步求期望，以便把多余的部分去掉，第二部分求极大值。

2. 一个例子引出 EM 算法

假设我们投掷一个不均匀的四面体骰子，将四个面分别标记为 $A$、$B$、$C$、$D$，其朝上发生的概率分别为$\frac{1}{2}-\frac{\theta }{4}$，$\frac{1-\theta }{4}$，$\frac{1+\theta }{4}$，$\frac{\theta }{4}$，其中$\theta \in \left(0,1\right)$ ，现进行了 197 次实验，四种结果发生的次数分别为 75，18，70，34。试求$\theta$的 MLE。

2.1 采用极大似然估计的方法

以$x_1$，$x_2$，$x_3$，$x_4$表示四种结果发生的次数，它们的集合记为$\mathbf{X}$，记参数的集合为$\theta \theta$，由于该例子中参数只有一个$\theta$，因此 θ = { θ } {\pmb{\theta}} = \left\{ \theta \right\} θθ={θ}。此时总体分布为多项分布，故其对数似然函数
$$\ln p(\mathbf{X}\mid \theta \theta )\propto {\left(\frac{1}{2}-\frac{\theta }{4}\right)}^{x_1}{\left(\frac{1-\theta }{4}\right)}^{x_2}{\left(\frac{1+\theta }{4}\right)}^{x_3}{\left(\frac{\theta }{4}\right)}^{x_4}\propto (2-\theta {)}^{x_1}(1-\theta {)}^{x_2}(1+\theta {)}^{x_3}{\theta }^{x_4}$$
其对数似然方程是一个三次多项式。

2.2 采用 EM 算法

我们可以通过引入 2 个变量$z_1,z_2$后，使得求解变得比较容易，$z_1,z_2$的集合记为$\mathbf{Z}$。现假设第一种结果可以分成两部分，其发生的概率分别为$\frac{1-\theta }{4}$和$\frac{1}{4}$，令$z_1$和$x_1-z_1$分别表示落入这两部分的次数；再假设第三种结果分成两部分，其发生的概率分别为$\frac{\theta }{4}$和$\frac{1}{4}$，令$z_2$和$x_3-z_2$分别表示落入这两部分的次数。

那为什么可以这样分呢，我们可以想象将不均匀的四面体骰子的$A$ 面和 $C$ 面又分别分成了两个面 $A1$，$A2$ 和 $C1$，$C2$，于是我们可以将上述不均匀的四面体骰子假想成一个六面体骰子，六个面分别记为 $A1$、$A2$、$B$、$C1$、$C2$、$D$，其朝上发生的概率分别为$\frac{1-\theta }{4}$，$\frac{1}{4}$，$\frac{1-\theta }{4}$，$\frac{\theta }{4}，\frac{1}{4}$，$\frac{\theta }{4}$，其中$\theta \in \left(0,1\right)$，那么四面体骰子 $A$ 面朝上的概率等价于六面体骰子 $A1$ 面朝上或 $A2$ 面朝上的概率，四面体骰子 $B$ 面朝上的概率等价于六面体骰子 $B1$ 面朝上或 $B2$ 面朝上的概率。

$\mathbf{Z}$在文献中被称为 $latent$ $variable$, 即潜变量。称数据 { X , Z } \{\mathbf{X},\mathbf{Z}\} {X,Z}为完全数据，而观测到的数据$\mathbf{X}$称为不完全数据。此时，完全数据的对数似然函数为

$$\ln p(\mathbf{X},\mathbf{Z}\mid \theta \theta )=\left(z_2+x_4\right)\ln \theta +\left(z_1+x_2\right)\ln (1-\theta )$$

可以看到上式的形式更加简单，然而，这样的六面体骰子是人为假想出来的，在真实的观测中，我们无法观测到 $A1$ 面朝上的次数$z_1$和 $C1$ 面朝上的次数$z_2$，也就是说，我们没有完整数据集 { X , Z } \{\mathbf{X},\mathbf{Z}\} {X,Z}，只有不完整数据$\mathbf{X}$。我们关于潜在变量 的取值的知识仅仅来源于后验分布$p(\mathbf{Z}\mid \mathbf{X},\theta \theta )$。由于我们不能使⽤完整数据的对数似然函数，因此我们反过来考虑在潜在变量的后验概率分布下，它的期望值，这对应于 EM 算法中的 E 步骤（稍后会看到）。在接下来的 M 步骤中，我们最⼤化这个期望。当$\mathbf{X}$及$\theta \theta$已知时，潜在变量的后验分布$p(\mathbf{Z}\mid \mathbf{X},\theta \theta )$为$$z_1\sim b\left(y_1,\frac{1-\theta }{2-\theta }\right),z_2\sim b\left(y_3,\frac{\theta }{1+\theta }\right)$$，于是，Dempster 等人建议分两步进行迭代求解。

E 步：在已有的观测数据$\mathbf{X}$及上一步估计值$\theta \theta ={\theta \theta }^{old}$的条件下，求基于完全数据的对数似然函数的期望（即把其中与$\mathbf{Z}$有关的部分积分掉）：
$$Q\left(\theta \theta \mid \mathbf{X},{\theta \theta }^{old}\right)=E_{\mathbf{Z}}\left(\ln p(\mathbf{X},\mathbf{Z}\mid \theta \theta )\right)=\sum _{\mathbf{Z}}\left(p\left(\mathbf{Z}\mid \mathbf{X},{\theta \theta }^{old}\right)\ln p(\mathbf{X},\mathbf{Z}\mid \theta \theta )\right)$$
M 步：求$Q\left(\theta \theta \mid \mathbf{X},{\theta \theta }^{old}\right)$关于$\theta \theta$的最大值${\theta \theta }^{new}$，即找${\theta \theta }^{new}$使得
$${\theta \theta }^{new}={\max }_{\theta \theta }Q\left(\theta \theta \mid \mathbf{X},{\theta \theta }^{old}\right)$$
检查是否收敛，如果不收敛，则令${\theta \theta }^{old}={\theta \theta }^{new}$，重复上述步骤直到收敛即可得到$\theta$的 MLE。

对于本例，其 E 步为
$\begin{array}{rl}Q\left(\theta \theta \mid \mathbf{X},{\theta \theta }^{old}\right)& =\left(E\left(z_2\mid \mathbf{X},{\theta \theta }^{old}\right)+x_4\right)\ln \theta +\left(E\left(z_1\mid \mathbf{X},{\theta \theta }^{old}\right)+x_2\right)\ln (1-\theta )\\ & =\left(\frac{{\theta }^{old}}{1+{\theta }^{old}}{x}_3+x_4\right)\ln \theta +\left(\frac{1-{\theta }^{old}}{2-{\theta }^{old}}{x}_1+x_2\right)\ln (1-\theta )\end{array}$
其 M 步即为上式对于$\theta$求导，并令其等于 0，即
$$\frac{\frac{{\theta }^{old}}{1+{\theta }^{old}}{x}_3+x_4}{{\theta }^{new}}-\frac{\frac{1-{\theta }^{old}}{2-{\theta }^{old}}{x}_1+x_2}{1-{\theta }^{new}}=0$$
解之可得如下迭代公式
$${\theta }^{new}=\frac{\frac{{\theta }^{old}}{1+{\theta }^{old}}{x}_3+x_4}{\frac{{\theta }^{old}}{1+{\theta }^{old}}{x}_3+x_4+\frac{1-{\theta }^{old}}{2-{\theta }^{old}}{x}_1+x_2}$$

那么，对完整数据的对数似然函数，我们考虑其在潜在变量的后验概率分布下的期望值，这个做法的合理性是什么呢，有没有更加合理能够令人信服的解释，EM 算法在迭代过程中究竟是怎么发挥作用的。

3. EM 算法的再思考

考虑一个概率模型，其中我们将所有的观测变量联合起来记作$\mathbf{X}$，将所有的隐含变量记作 $\mathbf{Z}$。联合概率分布 $p(\mathbf{X},\mathbf{Z}\mid \theta \theta )$ 由一组参数控制，记作$\theta \theta$。我们的目标是最大化似然函数
$$\begin{array}{rl}\ln p(\mathbf{X}\mid \theta \theta )& =\sum _{\mathbf{Z}}q(\mathbf{Z})\ln p(\mathbf{X}\mid \theta \theta )\\ & =\sum _{\mathbf{Z}}q(\mathbf{Z})\ln p(\mathbf{X}\mid \theta \theta )+\sum _{\mathbf{Z}}q(\mathbf{Z})\ln \left\{\frac{p(\mathbf{Z}\mid \mathbf{X},\theta \theta )}{q(\mathbf{Z})}\right\}-\sum _{\mathbf{Z}}q(\mathbf{Z})\ln \left\{\frac{p(\mathbf{Z}\mid \mathbf{X},\theta \theta )}{q(\mathbf{Z})}\right\}\\ & =\sum _{\mathbf{Z}}q(\mathbf{Z})\ln \left\{\frac{p(\mathbf{Z}\mid \mathbf{X},\theta \theta )p(\mathbf{X}\mid \theta \theta )}{q(\mathbf{Z})}\right\}-\sum _{\mathbf{Z}}q(\mathbf{Z})\ln \left\{\frac{p(\mathbf{Z}\mid \mathbf{X},\theta \theta )}{q(\mathbf{Z})}\right\}\\ & =\sum _{\mathbf{Z}}q(\mathbf{Z})\ln \left\{\frac{p(\mathbf{X},\mathbf{Z}\mid \theta \theta )}{q(\mathbf{Z})}\right\}-\sum _{\mathbf{Z}}q(\mathbf{Z})\ln \left\{\frac{p(\mathbf{Z}\mid \mathbf{X},\theta \theta )}{q(\mathbf{Z})}\right\}\\ & =\mathcal{L}(q,\theta \theta )+\text{KL}(q\Vert p)\end{array}$$
令
$$\begin{array}{rl}\mathcal{L}(q,\theta \theta )& =\sum _{\mathbf{Z}}q(\mathbf{Z})\ln \left\{\frac{p(\mathbf{X},\mathbf{Z}\mid \theta \theta )}{q(\mathbf{Z})}\right\}\\ \text{KL}(q\Vert p)& =-\sum _{\mathbf{Z}}q(\mathbf{Z})\ln \left\{\frac{p(\mathbf{Z}\mid \mathbf{X},\theta \theta )}{q(\mathbf{Z})}\right\}\end{array}$$
则
$$\ln p(\mathbf{X}\mid \theta \theta )=\mathcal{L}(q,\theta \theta )+\text{KL}(q\Vert p)\text{\hspace{1em}(1)}$$
$\text{KL}(q\Vert p)$是$q(\mathbf{Z})$和后验概率分布$p(\mathbf{Z}\mid \mathbf{X},\theta \theta )$之间的 Kullback-Leibler 散度，Kullback-Leibler 散度满足$\text{KL}(q\Vert p)\ge 0$，当且仅当$p(\mathbf{Z}\mid \mathbf{X},\theta \theta )$时等号成立。因此，根据公式（1），$\mathcal{L}(q,\theta \theta )\le \ln p(\mathbf{X}\mid \theta \theta )$（利用 Jensen 不等式也可得到该结果），也即$\mathcal{L}(q,\theta \theta )$是$\ln p(\mathbf{X}\mid \theta \theta )$的一个下界。图 1 说明了公式(1)的分解

图 1 由公式(1)给出的分解的说明，它对于分布$q(\mathbf{Z})$的任意选择都成立。由于 Kullback-Leibler 散度满足$\text{KL}(q\Vert p)\ge 0$，因此我们看到$\mathcal{L}(q,\theta \theta )$是对数似然函数$\ln p(\mathbf{X}\mid \theta \theta )$的一个下界。

EM 算法是一个两阶段的迭代优化算法，用于寻找最大似然解。我们可以使用公式(1)来定义EM算法，证明它确实最大化了对数似然函数。
假设参数向量的当前值为${\theta \theta }^{old}$。在 E 步骤中，$\ln p(\mathbf{X}\mid {\theta \theta }^{old})$与$q(\mathbf{Z})$无关，因此任意改变$q(\mathbf{Z})$不影响$\ln p(\mathbf{X}\mid {\theta \theta }^{old})$的值。所以根据公式 (1)，$\ln p(\mathbf{X}\mid {\theta \theta }^{old})$不变，对$\mathcal{L}(q,{\theta \theta }^{old})$关于$q(\mathbf{Z})$的最大化，只需要使得$\text{KL}(q\Vert p)$最小，也就是$\text{KL}(q\Vert p)=0$，也即${\theta \theta }^{old}$。此时，下界$\mathcal{L}(q,{\theta \theta }^{old})$等于对数似然函数$\ln p(\mathbf{X}\mid {\theta \theta }^{old})$，如图 2 所示。

图 2 EM 算法的 E 步骤的说明。q 分布被设置为当前参数值${\theta \theta }^{old}$下的后验概率分布，这使得下界上移到与对数似然函数值相同的位置，此时 KL 散度为零。

在接下来的 M 步骤中，固定${\theta \theta }^{old}$后，下界的形式为
$\begin{array}{rl}\mathcal{L}(q,\theta \theta )& =\sum _{\mathbf{Z}}q(\mathbf{Z})\ln \left\{\frac{p(\mathbf{X},\mathbf{Z}\mid \theta \theta )}{q(\mathbf{Z})}\right\}\\ & =\sum _{\mathbf{Z}}p(\mathbf{Z}\mid \mathbf{X},{\theta \theta }^{old})\ln \left\{\frac{p(\mathbf{X},\mathbf{Z}\mid \theta \theta )}{p(\mathbf{Z}\mid \mathbf{X},{\theta \theta }^{old})}\right\}\\ & =\sum _{\mathbf{Z}}p\left(\mathbf{Z}\mid \mathbf{X},{\theta \theta }^{old}\right)\ln p(\mathbf{X},\mathbf{Z}\mid \theta \theta )-\sum _{\mathbf{Z}}p\left(\mathbf{Z}\mid \mathbf{X},{\theta \theta }^{old}\right)\ln p\left(\mathbf{Z}\mid \mathbf{X},{\theta \theta }^{old}\right)\\ & =\mathcal{Q}\left(\theta \theta ,{\theta \theta }^{old}\right)+\text{const}\end{array}$
下界$\mathcal{L}(q,\theta \theta )$关于$\theta \theta$进行最大化等价于$\mathcal{Q}\left(\theta \theta ,{\theta \theta }^{old}\right)=\sum _{\mathbf{Z}}p\left(\mathbf{Z}\mid \mathbf{X},{\theta \theta }^{old}\right)\ln p(\mathbf{X},\mathbf{Z}\mid \theta \theta )$关于$\theta \theta$进行最大化，由此可得到某个新值${\theta \theta }^{new}$。这会使得下界$\mathcal{L}(q,\theta \theta )$增大，也即$\mathcal{L}(q,{\theta \theta }^{new})>\mathcal{L}(q,{\theta \theta }^{old})$（除非已经达到了极大值）。由于概率分布${\theta \theta }^{old}$由旧的参数值${\theta \theta }^{old}$确定，并且在 M 步骤中保持固定，因此它不会等于新的后验概率分布$p(\mathbf{Z}\mid \mathbf{X},{\theta \theta }^{new})$，从而 KL 散度非零。于是，对数似然函数$\ln p(\mathbf{X}\mid \theta \theta )$的增加量包括$\mathcal{L}(q,\theta \theta )$的增量和 的增量两部分，因此$\ln p(\mathbf{X}\mid \theta \theta )$大于下界$\mathcal{L}(q,\theta \theta )$的增加量，如图 3 所示。

图 3: EM 算法的 M 步骤的说明。分布${\theta \theta }^{old}$保持固定，下界$\mathcal{L}(q,\theta \theta )$关于参数向量$\theta \theta$最大化，得到修正值${\theta \theta }^{new}$。由于 KL 散度非负，因此这使得对数似然函数$\ln p(\mathbf{X}\mid \theta \theta )$的增量至少与下界的增量相等。

EM 算法的计算也可以被看做参数空间中的运算，如图 4 所示。这里，红色曲线表示（不完整数据）对数似然函数$\ln p(\mathbf{X}\mid \theta \theta )$，我们希望找到一个参数$\theta \theta$使其最大。我们首先选择某个初始的参数值${\theta \theta }^{old}$，然后在第一个 E 步骤中，得到了下界$\mathcal{L}(q,\theta \theta )$（事实上是其等价的优化函数$\mathcal{Q}\left(\theta \theta ,{\theta \theta }^{old}\right)$，但用$\mathcal{L}(q,\theta \theta )$更方便描述整个过程），用蓝色曲线表示。注意，下界$\mathcal{L}(q,\theta \theta )$与不完全数据的对数似然函数$\ln p(\mathbf{X}\mid \theta \theta )$在${\theta \theta }^{old}$处以切线的方式连接，因此两条曲线的梯度相同。这个下界$\mathcal{L}(q,\theta \theta )$是一个凹函数，对于指数族分布的混合分布来说，有唯一的最大值。在 M 步骤中，下界$\mathcal{L}(q,\theta \theta )$被最大化，得到了新的值${\theta \theta }^{new}$，这个值给出了比${\theta \theta }^{old}$处更大的对数似然函数值。接下来的 E 步骤又通过${\theta \theta }^{new}$构建了一个新的下界$\mathcal{L}(q,\theta \theta )$，它在${\theta \theta }^{new}$处与对数似然函数切线连接，用绿色曲线表示。如此循环往复直到收敛。

图 4 EM 算法涉及到交替计算当前参数值下的对数似然函数的下界以及最大化下界的值得到新的参数值。