EM算法|机器学习推导系列（十二）

一、概述

假设有如下数据：

$X$:observed data
$Z$:unobserved data(latent variable)
$(X,Z)$:complete data
$\theta$:parameter

EM算法的目的是解决具有隐变量的参数估计（MLE、MAP）问题。EM算法是一种迭代更新的算法，其计算公式为：

$$\begin{gathered} {\theta }^{t+1}=E_{z|x,{\theta }^t}[logp(x,z|\theta )] \\ =\underset{\theta }{\mathrm{argmax}}{\int }_{z}logp(x,z|\theta )\cdot p(z|x,{\theta }^t)\mathrm{d}z \end{gathered}$$

这个公式包含了迭代的两步：
①E step：计算$p(x,z|\theta )$在概率分布$p(z|x,{\theta }^t)$下的期望
②S step：计算使这个期望最大化的参数得到下一个EM步骤的输入

二、EM的算法收敛性

现在要证明迭代求得的${\theta }^t$序列会使得对应的$p(x|{\theta }^t)$是单调递增的，也就是说要证明$p(x|{\theta }^t)\le p(x|{\theta }^{t+1})$。首先我们有：

$$logp(x|\theta )=logp(x,z|\theta )-logp(z|x,\theta )$$

接下来等式两边同时求关于$p(z|x,{\theta }^t)$的期望：

$$\begin{gathered} 左边={\int }_{z}p(z|x,{\theta }^t)\cdot logp(x|\theta )\mathrm{d}z \\ =logp(x|\theta ){\int }_{z}p(z|x,{\theta }^t)\mathrm{d}z \\ =logp(x|\theta ) \\ 右边=\underset{Q(\theta ,{\theta }^t)}{\underbrace{{\int }_{z}p(z|x,{\theta }^t)\cdot p(x,z|\theta )\mathrm{d}z}}-\underset{H(\theta ,{\theta }^t)}{\underbrace{{\int }_{z}p(z|x,{\theta }^t)\cdot logp(z|x,\theta )\mathrm{d}z}} \\ 因此有logp(x|\theta )={\int }_{z}p(z|x,{\theta }^t)\cdot p(x,z|\theta )\mathrm{d}z-{\int }_{z}p(z|x,{\theta }^t)\cdot logp(z|x,\theta )\mathrm{d}z \end{gathered}$$

这里我们定义了$Q(\theta ,{\theta }^t)$，称为Q函数（Q function），这个函数也就是上面的概述中迭代公式里用到的函数，因此满足$Q({\theta }^{t+1},{\theta }^t)\ge Q({\theta }^t,{\theta }^t)$。

接下来将上面的等式两边$\theta$分别取${\theta }^t$和${\theta }^{t+1}$并相减：

$$logp(x|{\theta }^{t+1})-logp(x|{\theta }^t)=[Q({\theta }^{t+1},{\theta }^t)-Q({\theta }^t,{\theta }^t)]-[H({\theta }^{t+1},{\theta }^t)-H({\theta }^t,{\theta }^t)]$$

我们需要证明$logp(x|{\theta }^{t+1})-logp(x|{\theta }^t)\ge 0$，同时已知$Q({\theta }^{t+1},{\theta }^t)-Q({\theta }^t,{\theta }^t)\ge 0$，现在来观察$H({\theta }^{t+1},{\theta }^t)-H({\theta }^t,{\theta }^t)$：

$$\begin{gathered} H({\theta }^{t+1},{\theta }^t)-H({\theta }^t,{\theta }^t) \\ ={\int }_{z}p(z|x,{\theta }^t)\cdot logp(z|x,{\theta }^{t+1})\mathrm{d}z-{\int }_{z}p(z|x,{\theta }^t)\cdot logp(z|x,{\theta }^t)\mathrm{d}z \\ ={\int }_{z}p(z|x,{\theta }^t)\cdot log\frac{p(z|x,{\theta }^{t+1})}{p(z|x,{\theta }^t)}\mathrm{d}z \\ \le log{\int }_{z}p(z|x,{\theta }^t)\frac{p(z|x,{\theta }^{t+1})}{p(z|x,{\theta }^t)}\mathrm{d}z \\ =log{\int }_{z}p(z|x,{\theta }^{t+1})\mathrm{d}z \\ =log1 \\ =0 \\ 这里应用了Jensen不等式：log\sum _j{\lambda }_j{y}_j\ge \sum _j{\lambda }_{j}log{y}_j,其中{\lambda }_j\ge 0，log\sum _j{\lambda }_j=1 \\ 也可以使用KL散度来证明{\int }_{z}p(z|x,{\theta }^t)\cdot log\frac{p(z|x,{\theta }^{t+1})}{p(z|x,{\theta }^t)}\mathrm{d}z\le 0： \\ 两个概率分布P(x)和Q(x)的KL散度的定义为 \\ {D}_{KL}(P||Q)=E_{x\sim P}[log\frac{P(x)}{Q(x)}]，KL散度是恒\ge 0的。 \\ 因此{\int }_{z}p(z|x,{\theta }^t)\cdot log\frac{p(z|x,{\theta }^{t+1})}{p(z|x,{\theta }^t)}\mathrm{d}z=-KL(p(z|x,{\theta }^t)||p(z|x,{\theta }^{t+1}))\le 0 \end{gathered}$$

因此得证$logp(x|{\theta }^{t+1})-logp(x|{\theta }^t)\ge 0$。这说明使用EM算法迭代更新参数可以使得$logp(x|\theta )$逐步增大。

另外还有其他定理保证了EM的算法收敛性。首先对于${\theta }^i(i=1,2,\cdots )$序列和其对应的对数似然序列$L({\theta }^t)=logp(x|{\theta }^t)(t=1,2,\cdots )$有如下定理：
①如果$p(x|\theta )$有上界，则$L({\theta }^t)=logp(x|{\theta }^t)$收敛到某一值$L^{\ast }$；
②在函数$Q(\theta ,{\theta }^{{}^{\prime }})$与$L(\theta )$满足一定条件下，由EM算法得到的参数估计序列${\theta }^t$的收敛值${\theta }^{\ast }$是$L(\theta )$的稳定点。

三、EM的算法的导出

1. ELBO+KL散度的方法

$$\begin{gathered} logp(x|\theta )=logp(x,z|\theta )-logp(z|x,\theta ) \\ =log\frac{p(x,z|\theta )}{q(z)}-log\frac{p(z|x,\theta )}{q(z)}q(z)\ne 0 \\ 以上引入一个关于z的概率分布q(z)，然后式子两边同时求对q(z)的期望 \\ 左边={\int }_{z}q(z)\cdot logp(x|\theta )\mathrm{d}z=logp(x|\theta ){\int }_{z}q(z)\mathrm{d}z=logp(x|\theta ) \\ 右边=\underset{ELBO(evidencelowerbound)}{\underbrace{{\int }_{z}q(z)log\frac{p(x,z|\theta )}{q(z)}\mathrm{d}z}}\underset{KL(q(z)||p(z|x,\theta ))}{\underbrace{-{\int }_{z}q(z)log\frac{p(z|x,\theta )}{q(z)}\mathrm{d}z}} \end{gathered}$$

因此我们得出$logp(x|\theta )=ELBO+KL(q||p)$，由于KL散度恒$\ge 0$，因此$logp(x|\theta )\ge ELBO$，则$ELBO$就是似然函数$logp(x|\theta )$的下界。

使得$logp(x|\theta )=ELBO$时，就必须有$KL(q||p)=0$，也就是$q(z)=p(z|x,\theta )$时。

在每次迭代中我们取$q(z)=p(z|x,{\theta }^t)$，就可以保证$logp(x|{\theta }^t)$与$ELBO$相等，也就是：

$$\begin{gathered} logp(x|\theta )=\underset{ELBO}{\underbrace{{\int }_{z}p(z|x,{\theta }^t)log\frac{p(x,z|\theta )}{p(z|x,{\theta }^t)}\mathrm{d}z}}\underset{KL(p(z|x,{\theta }^t)||p(z|x,\theta ))}{\underbrace{-{\int }_{z}p(z|x,{\theta }^t)log\frac{p(z|x,\theta )}{p(z|x,{\theta }^t)}\mathrm{d}z}} \\ 当\theta ={\theta }^t时，logp(x|{\theta }^t)取ELBO，即 \\ logp(x|{\theta }^t)=\underset{ELBO}{\underbrace{{\int }_{z}p(z|x,{\theta }^t)log\frac{p(x,z|{\theta }^t)}{p(z|x,{\theta }^t)}\mathrm{d}z}}\underset{=0}{\underbrace{-{\int }_{z}p(z|x,{\theta }^t)log\frac{p(z|x,{\theta }^t)}{p(z|x,{\theta }^t)}\mathrm{d}z}}=ELBO \end{gathered}$$

也就是说$logp(x|\theta )$与$ELBO$都是关于$\theta$的函数，且满足$logp(x|\theta )\ge ELBO$，也就是说$logp(x|\theta )$的图像总是在$ELBO$的图像的上面。对于$q(z)$，我们取$q(z)=p(z|x,{\theta }^t)$，这也就保证了只有在$\theta ={\theta }^t$时$logp(x|\theta )$与$ELBO$才会相等，因此使$ELBO$取极大值的${\theta }^{t+1}$一定能使得$logp(x|{\theta }^{t+1})\ge logp(x|{\theta }^t)$。该过程如下图所示：

然后我们观察一下取$ELBO$取极大值的过程：

$$\begin{gathered} {\theta }^{t+1}=\underset{\theta }{\mathrm{argmax}}ELBO \\ =\underset{\theta }{\mathrm{argmax}}{\int }_{z}p(z|x,{\theta }^t)log\frac{p(x,z|\theta )}{p(z|x,{\theta }^t)}\mathrm{d}z \\ =\underset{\theta }{\mathrm{argmax}}{\int }_{z}p(z|x,{\theta }^t)logp(x,z|\theta )\mathrm{d}z-\underset{与\theta 无关}{\underbrace{\underset{\theta }{\mathrm{argmax}}{\int }_{z}p(z|x,{\theta }^t)p(z|x,{\theta }^t)\mathrm{d}z}} \\ =\underset{\theta }{\mathrm{argmax}}{\int }_{z}p(z|x,{\theta }^t)logp(x,z|\theta )\mathrm{d}z \\ =\underset{\theta }{\mathrm{argmax}}{E}_{z|x,{\theta }^t}[logp(x,z|\theta )] \end{gathered}$$

由此我们就导出了EM算法的迭代公式。

2. ELBO+Jensen不等式的方法

首先要具体介绍一下Jensen不等式：对于一个凹函数 $f(x)$（国内外对凹凸函数的定义恰好相反，这里的凹函数指的是国外定义的凹函数），我们查看其图像如下：

$$\begin{gathered} t\in [0,1] \\ c=ta+(1-t)b \\ \varphi =tf(a)+(1-t)f(b) \\ 凹函数恒有f(c)\ge \varphi \\ 也就是f(ta+(1-t)b)\ge tf(a)+(1-t)f(b) \\ 当t=\frac{1}{2}时有f(\frac{a}{2}+\frac{b}{2})\ge \frac{f(a)}{2}+\frac{f(b)}{2} \\ 可以理解为先求期望再求函数恒\ge 先求函数再求期望 \\ 即f(E)\ge E[f] \end{gathered}$$

接下来应用Jensen不等式来导出EM算法：

$$\begin{gathered} logp(x|\theta )=log{\int }_{z}p(x,z|\theta )\mathrm{d}z \\ =log{\int }_z\frac{p(x,z|\theta )}{q(z)}\cdot q(z)\mathrm{d}z \\ =log{E}_{q(z)}[\frac{p(x,z|\theta )}{q(z)}] \\ \ge \underset{ELBO}{\underbrace{E_{q(z)}[log\frac{p(x,z|\theta )}{q(z)}]}} \end{gathered}$$

这里应用了Jensen不等式得到了上面出现过的$ELBO$，这里的$f(x)$函数也就是$log$函数，显然这是一个凹函数。当$log\frac{P(x,z|\theta )}{q(z)}$这个函数是一个常数时会取得等号：

$$\begin{gathered} \frac{p(x,z|\theta )}{q(z)}=C \\ \Rightarrow q(z)=\frac{p(x,z|\theta )}{C} \\ \Rightarrow {\int }_{z}q(z)\mathrm{d}z={\int }_z\frac{1}{C}p(x,z|\theta )\mathrm{d}z \\ \Rightarrow 1=\frac{1}{C}{\int }_{z}p(x,z|\theta )\mathrm{d}z \\ \Rightarrow C=p(x|\theta ) \\ 将C代入q(z)=\frac{p(x,z|\theta )}{C}得 \\ q(z)=\frac{p(x,z|\theta )}{p(x|\theta )}=p(z|x,\theta ) \end{gathered}$$

这种方法到这里就和上面的方法一样了，总结来说就是：

$$\begin{gathered} logp(x|\theta )\ge \underset{ELBO}{\underbrace{E_{q(z)}[log\frac{p(x,z|\theta )}{q(z)}]}} \\ 当q(z)=p(z|x|\theta )时取等号，因此在迭代更新过程中取q(z)=p(z|x,{\theta }^t)。 \\ 接下来的推导过程就和第1种方法一样了。 \end{gathered}$$

四、广义EM

上面介绍的EM算法属于狭义的EM算法，它是广义EM的一个特例。在上面介绍的EM算法的E步中我们假定$q(z)=p(z|x,{\theta }^t)$，但是如果这个后验$p(z|x,{\theta }^t)$无法求解，那么必须使⽤采样（MCMC）或者变分推断等⽅法来近似推断这个后验。前面我们得出了以下关系：

$$logp(x|\theta )={\int }_{z}q(z)log\frac{p(x,z|\theta )}{q(z)}\mathrm{d}z-{\int }_{z}q(z)log\frac{p(z|x,\theta )}{q(z)}\mathrm{d}z=ELBO+KL(q||p)$$

当我们对于固定的$\theta$，我们希望$KL(q||p)$越小越好，这样就能使得$ELBO$更大：

$$固定\theta ，\hat{q}=\underset{q}{\mathrm{argmin}}KL(q||p)=\underset{q}{\mathrm{argmax}}ELBO$$

$ELBO$是关于$q$和$\theta$的函数，写作$L(q,\theta )$。以下是广义EM算法的基本思路：

$$\begin{gathered} Estep：q^{t+1}=\underset{q}{\mathrm{argmax}}L(q,{\theta }^t) \\ Mstep：{\theta }^{t+1}=\underset{q}{\mathrm{argmax}}L(q^{t+1},\theta ) \end{gathered}$$

再次观察一下$ELBO$：

$$\begin{gathered} ELBO=L(q,\theta )=E_q[logp(x,z)-logq] \\ =E_q[logp(x,z)]\underset{H[q]}{\underbrace{-E_q[logq]}} \end{gathered}$$

因此，我们看到，⼴义 EM 相当于在原来的式⼦中加⼊熵这⼀项。

五、EM的变种

EM 算法类似于坐标上升法，固定部分坐标，优化其他坐标，再⼀遍⼀遍的迭代。如果在 EM 框架中，⽆法求解$z$后验概率，那么需要采⽤⼀些变种的 EM 来估算这个后验：
①基于平均场的变分推断，VBEM/VEM
②基于蒙特卡洛的EM，MCEM