EM算法

EM算法原理

概率模型中常常既含有可观测变量，有含有不可被观测的隐变量(latent variable)

记可观测变量为$X$，隐变量为$Z$，模型参数为$\theta$

若概率模型的变量都可被观测，则可以直接用极大似然$\ln L(\theta)=\sum_{X}\ln P(X|\theta)$来估计参数$\theta$

然而若模型含有不可观测的隐变量，则只能通过边缘概率求似然，即似然函数变为$\ln L(\theta)=\sum{X}\ln P(X|\theta)=\sum{X}\ln \sum_Z P(X,Z|\theta)$，该式无法解析求解

EM(Expectation-Maximization)算法就是对含有隐变量的概率模型应用极大似然估计或极大后验概率的方法

E-step

首先引入Jensen不等式

设$X$是随机变量，$f(x)$为(下)凸函数，则有
$E[f(X)]\geq f(E[X])$
当且仅当$f(x)$为常数时等号成立。特别的，若$f(x)$为凹函数，则不等号相反。

因为 $\ln$ 是凹函数，所以根据Jensen不等式有$E[\ln f(X)]\leq\ln E[f(X)]$

设$Z$的概率分布为$Q(Z|\theta)$，则

$\begin{aligned} \sum_{X}\ln \sum_Z P(X,Z|\theta) &=\sum_{X}\ln \sum_Z P(Z|\theta)\frac{P(X,Z|\theta)}{Q(Z|\theta)} \\ &=\sum_{X}\ln E_Q[\frac{P(X,Z|\theta)}{Q(Z|\theta)}] \\ &\geq \sum_{X}E_Q[\ln \frac{P(X,Z|\theta)}{Q(Z|\theta)}] \\ &=\sum_{X}\sum_Z Q(Z|\theta)\ln \frac{P(X,Z|\theta)}{Q(Z|\theta)} \end{aligned}$

为了使$\ln P(X|\theta)$最大化，只需最大化下界$E_Q[\ln \frac{P(X,Z|\theta)}{Q(Z|\theta)}]$，称其为证据下界(evidence lower bound, ELBO)

由此我们得到EM算法的E-step：固定参数$\theta$，寻找一个$Q(Z|\theta)$的近似分布使得ELBO最大

显然若$\theta$固定，则$\ln P(X|\theta)$为常数，只需令$\frac{P(X,Z|\theta)}{Q(Z|\theta)}$等于常数$C$使得上式等号成立，ELBO就取到了最大值，此时

$\begin{aligned} P(X,Z|\theta)&=C\cdot Q(Z|\theta)\\ \sum_ZP(X,Z|\theta)&=C\sum_Z Q(Z|\theta)\\ P(X|\theta)&=C\\ P(X|\theta)&=\frac{P(X,Z|\theta)}{Q(Z|\theta)}\\ Q(Z|\theta)&=P(Z|X,\theta) \end{aligned}$

也即EM算法的E-step用$P(Z|X,\theta)$来近似$Q(Z|\theta)$，这里隐含了$P(Z|X,\theta)$易处理的假设

M-step

获得$Q(Z|\theta)$的近似后，似然函数可变为

$\begin{aligned} E_Q[\ln \frac{P(X,Z|\theta)}{Q(Z|\theta)}] &=\sum_{X}\sum_Z P(Z|X,\theta)\ln \frac{P(X,Z|\theta)}{P(Z|X,\theta)}\\ &=\sum_{X}\sum_Z (P(Z|X,\theta)\ln P(X,Z|\theta) - P(Z|X,\theta)\ln P(Z|X,\theta))\\ &=-H(P(Z|X,\theta))+\sum_{X}\sum_Z P(Z|X,\theta)\ln P(X,Z|\theta) \end{aligned}$

其中信息熵$H(P(Z|X,\theta))$是定值，因此最大似然即为

$\arg \max_\theta L(\theta)=\arg \max_\theta \sum_{X}\sum_Z P(Z|X,\theta)\ln P(X,Z|\theta)=E_Z[\ln P(X,Z|\theta)|X,\theta]$

也即EM算法的M-step：固定分布$Q(Z|\theta)$，寻找最优参数$\theta$使得ELBO最大

综上，EM算法按如下步骤迭代地计算$\theta$

选定初始参数$\theta^{(0)}$
E步：记$\theta^{(i)}$为第$i$次迭代得到的$\theta$估计值，则第$i+1$次计算期望
$Q(\theta|\theta^{(i)})=E_Z[\ln P(X,Z|\theta)|X,\theta^{(i)}]=\sum_X\sum_ZP(Z|X,\theta^{(i)})\ln P(X,Z|\theta)$
M步：求使得$Q(\theta|\theta^{(i)})$最大化的$\theta$，即有
$\theta^{(i+1)}=\arg \max_{\theta} Q(\theta|\theta^{(i)})$
重复2，3直到$\theta$收敛

需要注意的是，EM算法对初值$\theta^{(0)}$十分敏感

三硬币模型

这里引用《统计学习方法》（李航著）里面的三硬币模型来解释EM算法

三硬币模型：假设有3枚硬币A、B、C，他们出现正面的概率分别是$\pi,p,q$，同时用1表示正面，0表示反面。进行如下掷硬币试验：先掷硬币A，结果记为$Z$，若$Z=0$则选硬币B，否则选硬币C，然后掷选出的硬币，结果记为$X$，独立地重复n次试验

该模型中$Y$为可观测变量，$Z$为不可观测变量，$\theta=(\pi,p,q)$为模型参数，首先有

$P(Z=0|x,\theta)=\frac{P(Z,x|\theta)}{P(x|\theta)}=\frac{\pi p^x(1-p)^{1-x}}{\pi p^x(1-p)^{1-x}+(1-\pi)q^x(1-q)^{1-x}}$

记$P(Z=0|x,\theta^{(i)})=\mu^{(i)},\ P(Z=1|x,\theta^{(i)})=1-\mu^{(i)}$，则对于E-step有

$\begin{aligned} Q(\theta|\theta^{(i)}) &=\sum_X\sum_ZP(Z|X,\theta^{(i)})\ln P(X,Z|\theta)\\ &=\sum_{x\in X}P(Z=0|x,\theta^{(i)})\ln P(x,Z=0|\theta)+P(Z=1|x,\theta^{(i)})\ln P(x,Z=1|\theta)\\ &=\sum_{x\in X}\mu^{(i)}\ln \pi p^x(1-p)^{1-x} + (1-\mu^{(i)})\ln (1-\pi)q^x(1-q)^{1-x} \end{aligned}$

对于M-step有

$\begin{aligned} \frac{\partial Q(\theta|\theta^{(i)})}{\partial \pi} &=\sum_{x\in X}\mu^{(i)}\frac{p^x(1-p)^{1-x}}{\pi p^x(1-p)^{1-x}}-(1-\mu^{(i)})\frac{q^x(1-q)^{1-x}}{(1-\pi)q^x(1-q)^{1-x}}\\ &=\sum_{x\in X}\frac{\mu^{(i)}}{\pi}-\frac{1-\mu^{(i)}}{1-\pi}\\ &=\frac{1}{\pi(1-\pi)}\sum_{x\in X}\mu^{(i)}-\pi \end{aligned}$

令$\frac{\partial Q(\theta|\theta^{(i)})}{\partial \pi}=0$得

$\pi^{(i+1)}=\frac{1}{n}\sum_{x\in X}\mu^{(i)}$

同理可解得

$p^{(i+1)}=\frac{\sum_{x\in X}\mu^{(i+1)}x}{\sum_{x\in X}\mu^{(i+1)}}\\ q^{(i+1)}=\frac{\sum_{x\in X}(1-\mu^{(i+1)})x}{\sum_{x\in X}(1-\mu^{(i+1)})}$

EM算法的收敛性

我们先证似然函数序列$P(X|\theta^{(i)})$是单调递增的，即$P(X|\theta^{(i+1)})\geq P(X|\theta^{(i)})$

首先有

$P(X|\theta)=\frac{P(X,Z|\theta)}{P(Z|X,\theta)}$

两边取对数得

$\ln P(X|\theta)=\ln P(X,Z|\theta)-\ln P(Z|X,\theta)$

两边同时取关于$P(Z|X,\theta^{(i)})$的期望得

$E_Z[\ln P(X|\theta)|X,\theta^{(i)}]=E_Z[\ln P(X,Z|\theta)|X,\theta^{(i)}]-E_Z[\ln P(Z|X,\theta)|X,\theta^{(i)}]$

因为$\ln P(X|\theta)$与$P(Z)$无关，因此左式$E_Z[\ln P(X|\theta)|X,\theta^{(i)}]=\ln P(X|\theta)$

此外右式第一项$E_Z[\ln P(X,Z|\theta)|X,\theta^{(i)}]=Q(\theta|\theta^{(i)})$，因此有

$\ln P(X|\theta)=Q(\theta|\theta^{(i)}) - E_Z[\ln P(Z|X,\theta)|X,\theta^{(i)}]$

由于$\theta^{(i+1)}=\arg \max_{\theta} Q(\theta|\theta^{(i)})$，即$\theta^{(i+1)}$是极大值，因此有

$Q(\theta^{(i+1)}|\theta^{(i)})\geq Q(\theta^{(i)}|\theta^{(i)})$

又有KL散度满足

$\begin{aligned} D_{KL}(P(Z|X,\theta^{(i)})\|P(Z|X,\theta^{(i+1)})) &=E_Z[\ln\frac{P(Z|X,\theta^{(i)})}{P(Z|X,\theta^{(i+1)})}|X,\theta^{(i)}]\\ &=E_Z[\ln P(Z|X,\theta^{(i)})|X,\theta^{(i)}]-E_Z[\ln P(Z|X,\theta^{(i+1)})|X,\theta^{(i)}]\\ &\geq 0 \end{aligned}$

因此

$\begin{aligned} \ln P(X|\theta^{i+1})-\ln P(X|\theta^{i})=&Q(\theta^{(i+1)}|\theta^{(i)})-Q(\theta^{(i)}|\theta^{(i)})+\\ &E_Z[\ln P(Z|X,\theta^{(i)})|X,\theta^{(i)}]-E_Z[\ln P(Z|X,\theta^{(i+1)})|X,\theta^{(i)}] \\ \geq& 0 \end{aligned}$

综上所述，有$P(X|\theta^{i+1})>P(X|\theta^{i})$，得证

假设似然函数$P(X|\theta)$有上界，那么根据极限的性质易知EM算法收敛