EM不详解

为什么要EM

EM是极大似然估计的强化版，用于估计当模型中包含不可观测节点（即隐变量）时的似然函数，准确地说是估计似然函数的下界。

举例说明EM的必要性：

假设一个由 $K$ 个正态分布组成的高斯混合模型（GMM），$N$ 个样本，都不知道属于哪个正态分布，令变量 $z$ 代表所属的正态分布，即 $z$ 是隐变量。

则似然函数为 $\prod_{i=1}^Np(x_i)$，一般转为对数似然，即 $L=\sum_{i=1}^N\log p(x_i)$，由于隐变量 $z$ 无法被观测，因此不能将 $p(x)$ 写成某个正态分布形式，只能展开为 $L=\sum_{i=1}^N\log \sum_{k=1}^K\pi_{z_k}p(x_i|z_k)$，在极大似然估计中，令似然函数对各参数求导等于零，联立方程组，解得各参数。假设我们对 $p(x_i|z_1)$ 即第一个正态分布中的均值 $\mu_1$ 求导并令导数等于零：

$$\begin{align} \frac{\partial L}{\partial \mu_1}=\sum_{i=1}^N \frac{\pi_{z_1}}{\sum_{k=1}^K\pi_{z_k}p(x_i|z_k)}\frac{\partial p(x_i|z_1)}{\partial \mu_1}=0 \end{align}$$

类似求出其他参数对应的方程，求解这个方程组。

再来看如果用EM，则需要最大化的函数为 $ELBO=\sum_{i=1}^N\sum_{k=1}^K q(z_k)\log p(x_i|z_k)$，同样对 $\mu_1$ 求导并令导数等于零：

$$\begin{align} \frac{\partial ELBO}{\partial \mu_1}=\sum_{i=1}^N \frac{q(z_k)}{p(x_i|z_k)}\frac{\partial p(x_i|z_1)}{\partial \mu_1}=0 \end{align}$$

在 $(1)$ 中每个方程都包含所有正态分布的参数（在分母上，注意这里分母是无法通过右边的零去掉的），而 $(2)$ 中只包含自己这个正态分布的参数，求解难度差距可想而知。

并且这个方程组是存在许多约束的，例如 $\sum_z \pi_z=1$，协方差矩阵半正定等等。在多数情况下，$(1)$ 完全无法求解。

什么是EM

在上一节中已经给出了EM中需要最大化的函数 $ELBO=\sum_{i=1}^N\sum_{k=1}^K q(z_k)\log p(x_i|z_k)$，这一节将进行推导。

在上一节已经说过，EM与极大似然估计一样，希望最大化对数似然 $L=\sum_{i=1}^N\log p(x_i)$，由于不可观测的隐变量导致无法写成单一的分布，故只能展开为 $L=\sum_{i=1}^N\log \sum_{k=1}^Kp(x_i,z_k)$。注意这里进一步没有展开成 $\sum_{i=1}^N\log \sum_{k=1}^K\pi_{z_k}p(x_i|z_k)$，是因为我们要引入一个变分分布 $q(z_k)$ 以使Jensen不等式能够取到等号。引入 $q(z_k)$，则可以转化为

$$L=\sum_{i=1}^N\log \sum_{k=1}^Kq(z_k)\frac{p(x_i,z_k)}{q(z_k)}$$

上一节说过，$\log\sum$ 这个形式导致方程组难解，所以有人提出通过Jensen不等式能将其转化成 $\sum\log$。

Jensen不等式：

如果 $X$ 是随机变量，$g$ 是凸函数，则

$$g(\mathbb E[X])\le \mathbb E[g(X)]$$

等号当且仅当 $X$ 是一个常数或 $g$ 是线性时成立。

在我们这里有 $g=\log$，$X=\frac{p(x_i,z_k)}{q(z_k)}$，由于 $\log$ 是凹函数，所以由Jensen不等式可知

$$\begin{align} \log \sum_{k=1}^Kq(z_k)\frac{p(x_i,z_k)}{q(z_k)}\ge \sum_{k=1}^K q(z_k)\log\frac{p(x_i,z_k)}{q(z_k)} \end{align}$$

等号当且仅当 $\frac{p(x_i,z_k)}{q(z_k)}$ 是一个常数时成立。

则

$$\begin{align} L&=\sum_{i=1}^N\log p(x_i)\\ &=\sum_{i=1}^N\log \sum_{k=1}^Kq(z_k)\frac{p(x_i,z_k)}{q(z_k)}\\ &\ge \sum_{i=1}^N\sum_{k=1}^K q(z_k)\log\frac{p(x_i,z_k)}{q(z_k)}\\ &\triangleq ELBO(q,x) \end{align}$$

最后一行我们定义 $\sum_{i=1}^N\sum_{k=1}^K q(z_k)\log\frac{p(x_i,z_k)}{q(z_k)}$ 为证据下界（Evidence Lower BOund，ELBO），即对数似然函数 $L$（证据）的下界。

接下来就是正式讲解EM算法了。

EM算法分为E步和M步：

1. 在E步中确定我们引入的变分分布 $q(z_k)$ 以使上面的等号成立，即找到使 $ELBO=L$ 的 $q(z_k)$；
2. 在M步中寻找参数 $\Theta$ 以最大化 $ELBO$，这里的 $\Theta$ 是模型的参数，例如GMM中所有正态分布的均值和方差。

所以EM实际上是通过优化 $ELBO$ 来间接优化对数似然，而 $ELBO$ 是 $\log\frac{p(x_i,z_k)}{q(z_k)}$ 的期望（Expectation），所以得名EM（Expectation Maximum）算法。

E-step

在E步中确定我们引入的变分分布 $q(z_k)$ 以使上面的等号成立，即找到使 $ELBO=L$ 的 $q(z_k)$；

上面说过 $(3)$ 式中等号当且仅当 $\frac{p(x_i,z_k)}{q(z_k)}$ 是一个常数时成立。则我们设这个常数为 $c$，即 $\frac{p(x_i,z_k)}{q(z_k)}=c$，进行以下变换：

$$\begin{align} &\frac{p(x_i,z_k)}{q(z_k)}=c\\ \implies& \frac{p(x_i,z_k)}{c}=q(z_k)\\ \implies& \sum_{k=1}^K\frac{p(x_i,z_k)}{c}=\sum_{k=1}^Kq(z_k)=1\\ \implies& c=\sum_{k=1}^K p(x_i,z_k) \end{align}$$

则

$$\begin{align} q(z_k)=\frac{p(x_i,z_k)}{c}=\frac{p(x_i,z_k)}{\sum_{k'=1}^K p(x_i,z_{k'})}=\frac{p(x_i,z_k)}{p(x_i)}=p(z_k|x_i) \end{align}$$

换句话说，我们找到的那个 $q(z_k)$ 就是 $z_k$ 的后验分布 $p(z_k|x_i)$，而这个后验分布可以通过贝叶斯公式算出来。（插一句，用贝叶斯公式算的时候里面会用到 $z_k$ 的先验分布 $\pi_{z_k}$，要把它和我们引入的变分分布 $q(z_k)$ 区分开。）在GMM中算后验分布可以用贝叶斯公式比较简单，但其他情况下也可能会比较复杂，例如涉及到变分推断等，不在本文讨论范围内。

M-step

在M步中寻找参数 $\Theta$ 以最大化 $ELBO$，这里的 $\Theta$ 是模型的参数，例如GMM中所有正态分布的均值和方差。

这一步与不包含隐变量的极大似然估计的区别在于：极大似然估计中最大化对数似然，而EM中最大化的是对数似然的下界 $ELBO$。将 $ELBO$ 对 $\Theta$ 中每个参数求导，令导数为零，联立方程组求解出参数。（插一句，尽管经过E步后已经有 $ELBO(q,x)=L(x)$，但由于Jensen不等式的存在，所以在新的参数$\hat{\Theta}$下 $ELBO$ 与对数似然 $L$ 不一定仍然相等，所以直接优化对数似然与优化 $ELBO$ 还是不同的，这也是EM算法必须要迭代的原因。）

M步是一个约束优化问题，其中的约束如在第一节最后提到的 $\sum_z \pi_z=1$ 以及协方差矩阵半正定等等，需要用到拉格朗日乘数法等，这里不再展开。

当然，如果觉得求解方程组太繁琐，也可以用梯度上升来优化 $ELBO$，但如果能直接求出解析解，为什么还要在EM已经有迭代的情况下再套一层梯度上升迭代呢？况且用梯度上升为什么不直接用神经网络呢？

EM的收敛性

参考 邱锡鹏，神经网络与深度学习，机械工业出版社，https://nndl.github.io/, 2020.

假设在第 $t$ 次迭代时模型参数为 $\Theta_t$，在E步时找到了分布 $q_{t}(z_k)$ 使得 $p(x;\Theta_t)=ELBO(q,x;\Theta_t)$. 在M步时固定 $q_{t}(z_k)$，找到一组参数 $\Theta_{t+1}$ 使得 $ELBO(q,x;\Theta_{t+1})\ge ELBO(q,x;\Theta_t)$. 因此有

$$\begin{align} \log p(x;\Theta_{t+1})\ge ELBO(q_{t},x;\Theta_{t+1})\ge ELBO(q_{t},x;\Theta_t)=\log p(x;\Theta_t) \end{align}$$

即经过每一次迭代，对数似然增加，即 $\log p(x;\Theta_{t+1})\ge \log p(x;\Theta_t)$.