机器学习-白板推导系列笔记（十）-EM算法

此文章主要是结合哔站shuhuai008大佬的白板推导视频：EM_100min

全部笔记的汇总贴：机器学习-白板推导系列笔记

一、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步骤的输入

二、收敛性证明

现在要证明迭代求得的${\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)}{{\int }_{z}p(z|x,{\theta }^t)\cdot p(x,z|\theta )\mathrm{d}z}-\underset{H(\theta ,{\theta }^t)}{{\int }_{z}p(z|x,{\theta }^t)\cdot logp(z|x,\theta )\mathrm{d}z} \end{gathered}$$

因此有：

$$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$$

这里我们定义了$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 \end{gathered}$$

这里应用了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$$

因此得证$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的算法的导出

（一）ELBO+KL散度的方法

由EM算法公式可得：

$$\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 \end{gathered}$$

以上引入一个关于z的概率分布$q(z)$，然后式子两边同时求对$q(z)$的期望

$$\begin{gathered} 左边={\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)}{{\int }_{z}q(z)log\frac{p(x,z|\theta )}{q(z)}\mathrm{d}z}\underset{KL(q(z)||p(z|x,\theta ))}{-{\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$相等，也就是：

$$logp(x|\theta )=\underset{ELBO}{{\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 ))}{-{\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}{{\int }_{z}p(z|x,{\theta }^t)log\frac{p(x,z|{\theta }^t)}{p(z|x,{\theta }^t)}\mathrm{d}z}\underset{=0}{-{\int }_{z}p(z|x,{\theta }^t)log\frac{p(z|x,{\theta }^t)}{p(z|x,{\theta }^t)}\mathrm{d}z}=ELBO$$

也就是说$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 无关}{\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算法的迭代公式。

（二）ELBO+Jensen不等式的方法

Jensen不等式可以理解为先求期望再求函数恒$\ge$先求函数再求期望，即$f(E)\ge E[f]$。
应用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}{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 ) \end{gathered}$$

将$C$代入$q(z)=\frac{p(x,z|\theta )}{C}$得，

$$q(z)=\frac{p(x,z|\theta )}{p(x|\theta )}=p(z|x,\theta )$$

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

$$logp(x|\theta )\ge \underset{ELBO}{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种方法一样了。

四、广义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$更大：

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

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

$$\begin{gathered} E-Step：q^{t+1}=\underset{q}{\mathrm{argmax}}L(q,{\theta }^t) \\ M-Step：{\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]}{-E_q[logq]} \end{gathered}$$

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

五、EM的变种

EM 算法类似于坐标上升法，固定部分坐标，优化其他坐标，再⼀遍⼀遍的迭代。如果在 EM 框架中，⽆法求解$z$后验概率，那么需要采⽤⼀些变种的 EM 来估算这个后验：

①基于平均场的变分推断，VBEM/VEM
②基于蒙特卡洛的EM，MCEM

下一章传送门：白板推导系列笔记（十一）-高斯混合模型

参考文章
EM算法|机器学习推导系列