EM算法理解的第一层境界：期望E和最大化M（一）

前言：学生时代入门机器学习的时候接触的EM算法，当时感觉这后面应该有一套数学逻辑来约束EM算法的可行性。最近偶然在知乎上拜读了史博大佬的《EM算法理解的九层境界》^[1]，顿时感觉自己还是局限了。重新学习思考了一段时间，对EM算法有了更深的理解。

一、EM算法的形式

通常印象中EM算法一般应用于有隐变量的极大似然估计中。对于没有隐变量的极大似然来说，我们需要最大化似然估计$p(X|\theta )$。但是当问题中有了隐变量的时候，我们就需要把隐变量给积分掉（遍历隐变量的所有可能性），这个时候的似然估计为：
$$L(\theta |X)=p(X|\theta )=\int p(X,Z|\theta )dZ$$
一般而言，因为式子中需要将隐变量给积分掉，直接求解这个式子会非常复杂，这个时候EM算法就派上用场了。EM算法是个迭代算法，它由交替进行的两个部分组成：E-step和M-step。在迭代过程中，待估计参数$\theta$会逐步接近、到达最优解。
形式上说，EM算法的E-step就是利用上一步迭代得到的待估计参数${\theta }^{(t)}$来“估计”隐变量$Z$的“近似”分布，借由$Z$的“近似”分布，将隐变量$Z$给积分掉，从而得到待估计变量$\theta$的“似然”期望值：
$$Q(\theta |{\theta }^{(t)})=E_{Z|X,{\theta }^{(t)}}[\log L(\theta ;X,Z)]=\int p(Z|X,{\theta }^{(t)})\log L(\theta ;X,Z)dZ$$
EM算法的M-step就比较直接了，最大化这个变量$\theta$的“似然”，得到这一轮迭代变量$\theta$的估计值${\theta }^{(t+1)}$:
$${\theta }^{(t+1)}=\underset{\theta }{\mathrm{argmax}}Q(\theta |{\theta }^{(t)})$$

二、经典例子

最为经典简单的隐变量求解问题就是三硬币问题。给定A、B、C三枚硬币，它们抛出正面的概率分别为${\theta }_A$、${\theta }_B$、${\theta }_C$。对于每一个轮次，先抛硬币C来决定使用A、B中的哪枚硬币，正面（${\theta }_C$）使用硬币A，反面（$1-{\theta }_C$）则使用硬币B，接下来将硬币连续抛$\delta$次，记录每次正反面情况。然后将上述轮次进行$n$轮，得出如图1中所示的结果。

图1：EM算法三硬币例子 ^[2]

如果我们知道每个轮次使用的是哪一枚硬币，那么可以直接使用极大似然（maximum likelihood）来求解A、B两枚硬币的正面概率${\theta }_A$、${\theta }_B$，如图1中情景a所示。

但是如果我们并不知道每个轮次使用的是哪一枚硬币，那么就必须引入隐变量来求解这个问题。

这里我们将例子中的问题形式化一下，方面后面理解EM算法的形式。一共进行了$n$论次抛硬币游戏，每个轮次抛$\delta$次选定的硬币，观测结果记为 X = [ x 1 , x 2 , . . . , x n ] , x i ∈ { 0 , 1 , 2 , . . . , δ } X=[x_1, x_2, ..., x_n], x_i\in\{0, 1, 2, ..., \delta\} X=[x1​,x2​,...,xn​],xi​∈{0,1,2,...,δ}，即轮次$i$抛硬币观察到有$x_i$次正面，$\delta -x_i$次反面。所有的待求解参数$\theta =[{\theta }_A,{\theta }_B,{\theta }_C],{\theta }_A\in [0,1],{\theta }_B\in [0,1],{\theta }_C\in [0,1]$。

接下来就是隐变量的表示，这里我们可以引入硬币C的正面概率${\theta }_C$作为隐变量，但是为了与图1中的例子对应起来，我们引入隐变量 Z = [ z 1 , z 2 , . . . , z n ] , z i ∈ { 0 , 1 } Z=[z_1, z_2,...,z_n], z_i\in\{0,1\} Z=[z1​,z2​,...,zn​],zi​∈{0,1}表示每一轮次（$i$）是硬币A（$z_i=1$）还是硬币B（$z_i=0$）。

三、直觉上理解EM算法的形式

下面我们直觉上来套用EM算法的形式来求解图1中的三硬币问题。
首先是EM算法的E-step，我们需要依据上一轮的参数估计$[{\hat{\theta }}_A^{(t)},{\hat{\theta }}_B^{(t)}]$来“估计”隐变量$Z$。这里我们依据直觉来，对于轮次$i$，隐变量$z_i$有$P(z_i=1|x_i)+P(z_i=0|x_i)=1$，而依据上一轮的参数估计，我们有在第$i$轮是硬币A时出现$x_i$次正面和$\delta -x_i$次反面的概率：$P(x_i|z_i=1)=({\hat{\theta }}_A^{(t)}{)}^{x_i}\ast (1-{\hat{\theta }}_A^{(t)}{)}^{\delta -x_i}$，在第$i$轮是硬币B时出现$x_i$次正面和$\delta -x_i$次反面的概率$P(x_i|z_i=0)=({\hat{\theta }}_B^{(t)}{)}^{x_i}\ast (1-{\hat{\theta }}_B^{(t)}{)}^{\delta -x_i}$，这样我们就可以"估计"出轮次$i$中$z_i$的概率分布：
$$\begin{array}{rl}P(z_i=1|\ast )& =P(x_i|z_i=1)/(P(x_i|z_i=1)+P(x_i|z_i=0))\\ & =\frac{({\hat{\theta }}_A^{(t)}{)}^{x_i}\ast (1-{\hat{\theta }}_A^{(t)}{)}^{\delta -x_i}}{({\hat{\theta }}_A^{(t)}{)}^{x_i}\ast (1-{\hat{\theta }}_A^{(t)}{)}^{\delta -x_i}+({\hat{\theta }}_B^{(t)}{)}^{x_i}\ast (1-{\hat{\theta }}_B^{(t)}{)}^{\delta -x_i}}\end{array}$$
$$\begin{array}{rl}P(z_i=0|\ast )& =P(x_i|z_i=0)/(P(x_i|z_i=1)+P(x_i|z_i=0))\\ & =\frac{({\hat{\theta }}_B^{(t)}{)}^{x_i}\ast (1-{\hat{\theta }}_B^{(t)}{)}^{\delta -x_i}}{({\hat{\theta }}_A^{(t)}{)}^{x_i}\ast (1-{\hat{\theta }}_A^{(t)}{)}^{\delta -x_i}+({\hat{\theta }}_B^{(t)}{)}^{x_i}\ast (1-{\hat{\theta }}_B^{(t)}{)}^{\delta -x_i}}\end{array}$$
接下来是EM算法的M-setp，在上面得到的隐变量$Z$的分布之后，我们可以用来估计参数${\theta }_A$和${\theta }_B$。对于硬币A而言，它在轮次$i$中抛硬币的次数期望为$\delta \ast P(z_i=1)$，抛硬币为正面的次数期望为$x_i\ast P(z_i=1)$；对于硬币B而言，它在轮次$i$中抛硬币的次数期望为$\delta \ast P(z_i=0)$，抛硬币为正面的次数期望为$x_i\ast P(z_i=0)$。将所有轮次的总的次数期望和跑正面的次数期望加起来，我们有：
$$\begin{array}{rl}{\hat{\theta }}_A^{(t+1)}& =\frac{{\sum }_{i=1}^n{x}_i\ast P(z_i=1|\ast )}{{\sum }_{i=1}^n\delta \ast P(z_i=1|\ast )}\end{array}$$
$$\begin{array}{rl}{\hat{\theta }}_B^{(t+1)}& =\frac{{\sum }_{i=1}^n{x}_i\ast P(z_i=0|\ast )}{{\sum }_{i=1}^n\delta \ast P(z_i=0|\ast )}\end{array}$$
最后硬币$C$，我们将其估计为所有轮次为硬币$A$的"期望"，即：
$$\begin{array}{rl}{\hat{\theta }}_C^{(t+1)}& =\frac{{\sum }_{i=1}^{n}P(z_i=1|\ast )}{{\sum }_{i=1}^n[P(z_i=0|\ast )+P(z_i=1|\ast )]}=\frac{1}{n}\sum _{i=1}^{n}P(z_i=1|\ast )\end{array}$$
我们凭借直觉理解，将EM算法的形式套用在三硬币问题上面，得出了每个迭代轮次E-step和M-step的计算过程，两个计算过程都可以与图1中情形b中的例子过程对应起来。

References

[1] https://www.zhihu.com/question/40797593/answer/275171156
[2] Do C B, Batzoglou S. What is the expectation maximization algorithm?[J]. Nature biotechnology, 2008, 26(8): 897-899.