这段时间研究了一下EM算法，发现之前对于EM算法的理解有些偏差，这一段时间的实践和学习让我对EM算法的理解又近了一步，所以这里写个文章记录一下。
背景
EM算法（Expectation Maximization）要解决的问题是：如何估计模型的参数。
存在一批观察数据，我们希望通过一个参数模型来对其进行描述，通过这批数据得到该模型的参数的估计。针对该问题，最常见的方法就是极大似然法：
\[\hat{\theta}=\argmax_{\theta}{\sum_{i=1}^N{L_{\theta}(x_i)}}= \argmax_{\theta}{\sum_{i=1}^N{\log p_{\theta}(x_i)}}\tag{1}\]
我们可以通过微积分的方式求解\(1\)：
\[\sum_{i=1}^N{\frac{\partial L_\theta(x_i)}{\partial\theta}}=0\]
很多情况下这种方式已经是足够用了。但是，另外的很多情况下，我们的参数模型是非常复杂的，上面这种直接的求解方法可能并不能用。比如，如果我们的参数模型是存在一个隐含变量\(z_i\)的话，就无法直接使用上面的方法了。
存在一个隐含变量的情况下，我们来计算一下对数似然函数的导数：
\[ \begin{aligned} \sum_{i=1}^N{\frac{\partial L_\theta(x_i)}{\partial\theta}} &=\sum_{i=1}^N{\frac{1}{p_{\theta}(x_i)}\frac{\partial p_{\theta}(x_i)}{\partial\theta}}\\ &\propto\sum_{i=1}^N{ \frac{1}{p_{\theta}(x_i)}\sum_{z} \frac{\partial p_{\theta}(x_i|z_i)}{\partial\theta}\qquad\text{这里假设先验是均匀的} } \end{aligned} \]
因为有隐变量的存在，我们会额外出现一个求和号\(\sum_z\)（或积分），这就会使得我们无法继续处理下去。比如若\(p_{\theta}\)是正态分布，则一般来说在进行对数似然的计算时可以化简计算。但这里因为有求和号的存在，我们无法再像之前那样处理，即便这个\(z\)是一个服从伯努利分布的二分类变量，更别提多分类或者连续分布变量了。
EM算法便善于解决上面的这种存在隐变量的问题，其思路是使用启发式的迭代方法：固定参数来估计隐含变量\(z_i\)（E step），固定隐含变量\(z_i\)估计参数（M step）。这样不断迭代下去，直至算法收敛。
其实，k-means算法就可以算作EM算法的一个特例：计算每个样本属于哪个簇的过程就是E step；重新更新每个簇的质心的过程就是M step。
推导
推导有两种思路：即通过Jensen不等式和KL散度，我对KL散度比较熟，所以以这条思路为主。
因为组成iid样本，我们先对单独一个样本进行分析：
\[\log p_{\theta}(x_i) =\log p_{\theta}(x_i,z_i)-\log p_{\theta}(z_i|x_i)\]
这里我们引入一个变分分布\(q_{\phi}(z_i)\)，将上式进行变换得到：
\[ \begin{aligned} \log p_{\theta}(x_i)&=\mathbb{E}_{q_{\phi}(z_i)}[\log p_{\theta}(x_i)] \\ &=\mathbb{E}_{q_{\phi}(z_i)}[ \frac{\log p_{\theta}(x_i,z_i)}{q_{\phi}(z_i)}] - \mathbb{E}_{q_{\phi}(z_i)}[ \frac{\log p_{\theta}(z_i|x_i)}{q_{\phi}(z_i)}] \\ &=\mathbb{E}_{q_{\phi}(z_i)}[ \frac{\log p_{\theta}(x_i,z_i)}{q_{\phi}(z_i)}] + KL(q_{\phi}(z_i)||p_{\theta}(z_i|x_i)) \\ &\ge \mathbb{E}_{q_{\phi}(z_i)}[ \frac{\log p_{\theta}(x_i,z_i)}{q_{\phi}(z_i)}] \\ &= ELBO(x_i, z_i, \theta, \phi) \end{aligned} \]
这时，我们可以看到，当\(KL(q_{\phi}(z_i)||p_{\theta}(z_i|x_i))=0\)，即\(q_{\phi}(z_i)=p_{\theta}(z_i|x_i)\)时，ELBO等于对数似然。
基于Jensen不等式的推导：
\[ \begin{aligned} \log p_{\theta}(x_i)&=\sum_{z_i}\log p_{\theta}(x_i,z_i) \\ &=\sum_{z_i}\log q_{\phi}(z_i)[ \frac{p_{\theta}(x_i,z_i)}{q_{\phi}(z_i)}] \\ &\ge \sum_{z_i}q_{\phi}(z_i)\log [ \frac{p_{\theta}(x_i,z_i)}{q_{\phi}(z_i)}] \\ &= \mathbb{E}_{q_{\phi}(z_i)}[\frac{\log p_{\theta}(x_i,z_i)}{q_{\phi}(z_i)}] \\ &= ELBO(x_i, z_i, \theta, \phi) \end{aligned} \]
这里使用的jensen不等式是：如果\(f\)是一个凸函数，则有\(f(\mathbb{E}(X))\le\mathbb{E}(f(X))\)，等号成立当且仅当随机变量\(X\)是个常数随机变量。对数函数是一个上凸函数，所以结论反过来即可。
根据上面的例子，我们知道当且仅当下面的条件成立的时候，ELBO和对数似然相等，即：
\[ \frac{p_{\theta}(x_i,z_i)}{q_{\phi}(z_i)}=c \]
又因为\(\sum_z{q_{\phi}(z)}=1\)，得到：
\[ \begin{aligned} q_{\phi}(z_i)&=\frac{p_{\theta}(x_i,z_i)}{c} \\ &=\frac{p_{\theta}(x_i,z_i)}{c\sum_z{q_{\phi}(z)}} \\ &=\frac{p_{\theta}(x_i,z_i)}{c\sum_z{\frac{p_\theta(x_i,z_i)}{c}}} \\ &=\frac{p_{\theta}(x_i,z_i)}{\sum_z{p_\theta(x_i,z_i)}} \\ &=\frac{p_{\theta}(x_i,z_i)}{p_\theta(x_i)} \\ &=p_{\theta}(z_i|x_i) \end{aligned} \]
我们可以看到，基于Jensen的过程和基于KL散度的过程是殊途同归的。
根据上面的公式，我们可以给出一个启发式的想法：
首先我们求出在当前\(\theta_1\)下的后验分布\(p_{\theta_1}(z_i|x_i)\)，将其作为变分分布\(q_{\phi}(z_i)\)，则此时\(ELBO_1\)就是当前\(\theta_1\)下的对数似然。
我们固定变分分布，改变\(\theta\)来最大化\(ELBO\)，得到新的\(\theta_2\)和\(ELBO_2\)。
注意，当前\(\theta_2\)下对应的对数似然一定是大于\(\theta_1\)的对数似然的。因为\(\theta_2\)的对数似然一定是大于等于\(ELBO_2\)（不管变分分布是什么，ELOB都小于等于对数似然）；而\(ELBO_2\)一定大于\(ELBO_1\)（最大化操作保证）；\(ELBO_1\)就是\(\theta_1\)下的对数似然。所以上述操作保证了\(\theta\)的变化一直在使其对应的对数似然提高的。
由此，我们得到了EM算法。
以上过程实际上证明了EM算法是一直在使对数似然函数值变大的。而很容易也知道，对数似然函数是有上界的（一定不会大于0），则可知EM算法一定收敛到一个稳定点。
这里最大化ELBO的步骤还可以进一步简化一下：
\[ \begin{aligned} ELBO(x_i, z_i, \theta, \phi) &=\mathbb{E}_{q_{\phi}(z_i)}[ \frac{\log p_{\theta}(x_i,z_i)}{q_{\phi}(z_i)}] \\ &=\mathbb{E}_{q_{\phi}(z_i)}[ \log p_{\theta}(x_i,z_i)]+H(q_{\phi}(z_i)) \end{aligned} \]
注意到，后面变分分布的熵不和\(\theta\)有关，所以我们寻找最大化\(\theta\)的ELBO的时候其没有贡献，可以去掉。另外将所有样本都考虑到：
\[ \hat{\theta}=\argmax_{\theta}\sum_i^N\mathbb{E}_{q_{\phi}(z_i)}[\log p_{\theta}(x_i,z_i)] \]
所以，我们求解的实际上是对数联合分布的期望，所以EM算法的第一个步骤是Expectation（期望）。求出来后，改变\(\theta\)来最大化这个期望，即Maximization（最大化）。
通过上面的推导，我们知道EM算法和VI（variational inference，变分推断）真的很像，其不同点在于：
VI将变分分布也参数化，也需要通过最大化ELBO来得到的，此时是固定\(\theta\)改变\(\phi\)来实现的，这时候使ELBO最大（达到对数似然）的变分分布正好就是后验分布。
EM则需要显式的知道后验分布怎么算，直接算出来。
另外，VI使用的ELBO，而EM使用更加简化的对数联合概率期望。
算法流程
- 输入数据：\(x=(x_1, x_2, \cdots, x_N)\)
- 隐变量：\(z_i\)
- 最大迭代次数：\(J\)
初始化模型参数\(\theta_0\)；
开始EM算法迭代：
- E step：计算后验分布作为变分分布，并由此计算对数联合分布的期望\(\sum_i^N\mathbb{E}_{p_{\theta_i}(z_i|x_i)}[\log p_{\theta_i}(x_i,z_i)]\)。
- M step：最大化上面的期望，得到新的参数估计\(\theta_{i+1}\)。
- 如果\(\theta_{i+1}\)对应的对数联合分布期望变化不大或达到设定的最大迭代次数\(J\)，则输出此参数，否则回到a。
在一些时候，我们也称得到后验分布的过程是E step，计算对数联合分布期望并最大化它的过程是M step。这时候更加注重强调隐变量\(z\)的作用。
注意事项
EM算法对于初值是非常敏感的，这一点上和Kmeans算法特别相似。为了解决这个问题，有下面的一些研究：
- Blömer J, Bujna K. Simple methods for initializing the EM algorithm for Gaussian mixture models[J]. CoRR, 2013.
- Chen F. An Improved EM algorithm[J]. arXiv preprint arXiv:1305.0626, 2013.
- Kwedlo W. (2013) A New Method for Random Initialization of the EM Algorithm for Multivariate Gaussian Mixture Learning. In: Burduk R., Jackowski K., Kurzynski M., Wozniak M., Zolnierek A. (eds) Proceedings of the 8th International Conference on Computer Recognition Systems CORES 2013. Advances in Intelligent Systems and Computing, vol 226. Springer, Heidelberg.
有兴趣的可以拜读一下。
高斯混合模型（GMM）
GMM是一个概率模型，其认为数据来自多个高斯分布的混合，其可以表示为：
\[p(x)=\sum_{i=1}^C{f(x|z=i)g(z)}\]
其中\(z\)服从的是一个Categorical distribution的随机变量，分类数量为\(C\)，表示有C个高斯分布。\(f(x|z=i)\)是第\(i\)个高斯分布的密度函数。我们的样本是\(x_i\)是\(x\)的采样。
这是一个典型的具有隐变量的问题，非常适合使用EM算法来解决。
公式推导
首先，我们需要估计的参数就是每个高斯分布的参数和分类分布的参数\(\theta=\{\mu_j,\sigma_j,\gamma_j|j=1,\dots,C\}\)：
\[g(z=j)=\gamma_j,\quad\sum_j\gamma_j=1\] \[f(x|z=j)=\frac{1}{\sqrt{2\pi}\sigma_j}\exp(-\frac{(x-\mu_j)^2}{2\sigma_j^2})\]
对于E step，关键在于知道\(p_{\theta}(z_i|x_i)\)。这个可以通过贝叶斯公式得到：
\[ \begin{aligned} w_{ij}=q_{\phi}(z_i=j|x_i)=p_{\theta}(z_i=j|x_i) &=\frac{p_{\theta}(x_i|z_i=j)\gamma_j} {\sum_{j}p_{\theta}(x_i|z_i=j)\gamma_j} \end{aligned} \]
之后，我们计算对数联合概率期望：
\[ \begin{aligned} &\sum_i^N\mathbb{E}_{p_{\theta}(z_i|x_i)}[\log p_{\theta}(x_i,z_i)] \\ =\quad&\sum_i^N\mathbb{E}_{p_{\theta}(z_i|x_i)} [\log p_{\theta}(x_i|z_i)g(z_i)] \\ =\quad&\sum_i^N\sum_{j=1}^C[\log (p_{\theta}(x_i|z_i=j)\gamma_j) p_{\theta}(z_i=j|x_i)] \\ =\quad&\sum_i^N\sum_{j=1}^Cw_{ij}\times(\log\gamma_j + [-\log(\sqrt{2\pi}\sigma_j)-\frac{(x_i-\mu_j)^2}{2\sigma^2_j}]) \end{aligned} \]
注意到，尽管\(p_{\theta}(z_i|x_i)\)也有参数下标，但这里其实表示的是变分分布的位置，所以在下面的M step中是不动的，使用\(w_{ij}\)来表示。
对于M step，我们可以通过微积分求极值的方式得到参数的估计。首先是分类分布的参数：
\[\gamma_j=\frac{\sum_{i=1}^Nw_{ij}}{\sum_{j=1}^C\sum_{i=1}^Nw_{ij}}\]
注意到，这里有约束条件\(\sum_j\gamma_j=1\)，利用\(\gamma_C=1-\sum_{j=1}^{C-1}\gamma_j\)或拉格朗日乘子法可以得到上面的结论。
令\(W_j=\sum_{i=1}^Nw_{ij}\)，则我们要最大化的其实是
\[\sum_{j=1}^CW_j\log\gamma_j\] 这是一个交叉熵。
我们首先使用\(\gamma_C=1-\sum_{j=1}^{C-1}\gamma_j\)，然后求导（这里不需加入\(\gt0\)的约束了，\(\log\)的定义域决定了不能\(\le0\)）：
\[ \begin{aligned} \frac{\partial \sum_{j=1}^{C-1}[W_j\log\gamma_j]+W_C\log(1-\sum_i^{C-1}\gamma_i)} {\partial\gamma_j}&=\frac{W_j}{\gamma_j}-\frac{W_C}{1-\sum_{i=1}^{C-1}\gamma_j} \end{aligned} \] 其中\(j=1,\cdots,C-1\)。然后令其等于0，得到： \[\frac{W_j}{W_C}=\frac{\gamma_j}{\gamma_C},\quad j=1,\cdots,C-1\] 其中\(\gamma_C=1-\sum_{i=1}^{C-1}\gamma_j\)。我们将以上\(C-1\)个式子加总，再加\(1\)，得到 \[\frac{\sum_jW_j}{W_C}=\frac{1}{\gamma_C}\] 得到\(\gamma_C=\frac{W_C}{\sum_jW_j}\)，进而得到\(\gamma_j=\frac{W_j}{\sum_jW_j}\)。
如果使用拉格朗日乘子法，则是最大化下面的式子
\[\sum_{j=1}^CW_j\log\gamma_j+\lambda(\sum_{j=1}^C\gamma_j-1)\]
求导得到：
\[ \begin{aligned} \frac{W_j}{\gamma_j}+\lambda&=0 \\ \sum_{j=1}^C\gamma_j&=1 \end{aligned} \]
继续
\[ \begin{aligned} \gamma_j=-\frac{W_j}{\lambda} \\ -\frac{\sum_jW_j}{\lambda}=1 \\ \lambda=-\sum_jW_j \\ \gamma_j=\frac{W_j}{\sum_jW_j} \end{aligned} \]
得到了相同的结果。
接下来是均值：
\[\hat{\mu_j}=\frac{\sum_i{w_{ij}x_i}}{\sum_i{w_{ij}}}\]
可以看到，这里的均值的估计就是以后验概率为权重的样本均值。
然后是方差的估计，这里的估计要分为两种情况：
组间方差不等：
\[\hat{\sigma_j}=\frac{\sum_i{w_{ij}(x_j-\hat{\mu_j})^2}} {\sum_i{w_{ij}}}\]
组间方差相等： \[\hat{\sigma}= \frac{\sum_j\sum_i{w_{ij}(x_j-\hat{\mu_j})^2}} {\sum_j\sum_i{w_{ij}}}\]
代码实现
# E step
def _e_step(self, X, gammas_, mus_, sigmas_):
"""
The Expectation step.
Compute the posterior distribution and expectation of joint
distribution.
"""
logp = ss.norm.logpdf(X[:, None], mus_, sigmas_) # p(x|z)
if self.gamma_type == "full":
logp = logp + np.log(gammas_)
logpsum = logsumexp(logp, axis=1)
logp_norm = logp - logpsum[:, None] # p(z|x)
return logp_norm, (logp * np.exp(logp_norm)).sum()# M step
def _m_step(self, X, logp):
"""
The Maximization step.
Maximize the expectation of joint distribution and compute the new
parameters.
"""
# 1. prepare
p = np.exp(logp)
# get rid of having NaN
psum = p.sum(axis=0) + 10 * np.finfo(p.dtype).eps
# 2. compute the parameters of categorical distribution
if self.gamma_type == "full":
gammas_ = psum / psum.sum()
else:
gammas_ = np.ones(self.n_components) / self.n_components
# 3. compute the mean of gaussian components
mus_ = (p * X[:, None]).sum(axis=0) / psum
# 4. compute the stdandard deviations of gaussian components
if self.sigma_type == "full":
sigmas2_ = ((X[:, None] - mus_) ** 2 * p).sum(axis=0) / psum
sigmas_ = np.sqrt(sigmas2_)
elif self.sigma_type == "equal":
sigmas2_ = ((X[:, None] - mus_) ** 2 * p).sum() / psum.sum()
sigmas1_ = np.sqrt(sigmas2_)
sigmas_ = np.full(self.n_components, sigmas1_)
else:
sigmas_ = self.sigmas
return gammas_, mus_, sigmas_EM算法的收敛性
- 在上面推导的部分，我其实已经证明了，EM算法一定会收敛到一个稳定点。刘建平Pinard的博客也有一个证明，大家可以一起看一下。
- EM算法并不能保证收敛到极大似然解，其很有可能陷入局部最优值。这很大程度上取决了初始值的设置。
