EM 算法

EM 算法用于求解含隐变量的极大似然估计，先根据当前参数估计隐变量的分布，然后极大化该分布下似然函数的期望，以此得到新的参数估计。

• 1. Notations
• 2. Introduction
  • 2.1. 极大似然估计
  • 2.2. 隐变量
• 3. EM 算法
  • 3.1. 基本思想
  • 3.2. 理论推导
  • 3.3. 具体流程
  • 3.4. 收敛性
• 4. 思考
• 5. 补充更新

1. Notations

• 观测数据：$X=(x_1,x_2,\dots ,x_n)$​
• 参数：$\theta$
• 隐变量：$Z$
• 对数似然：$L(\theta;X)=\log P\left( X\mid \theta \right)$​​

2. Introduction

对于已知的观测数据 $X$，我们想要一个生成模型来描述 $X$ 的分布，这个分布 $P(X\mid \theta)$ 的结构通常是给定的，那么我们只需估计出一个合适的参数 $\theta$ 即可。

2.1. 极大似然估计

对于简单模型，最常用的方法是极大似然估计：寻找参数 $\theta$ 使得对数似然 $L(\theta;X)=\log P\left( X\mid \theta \right)$ 最大化。

\[\begin{aligned} L(\theta;X)&=\log P\left( X\mid \theta \right)\\ &=\log \prod_{i=1}^{n}P(x_i\mid \theta)\\ &=\sum_{i=1}^{n}\log P(x_i\mid \theta) \end{aligned}\] \[\theta^{\star}=\mathop{\arg \max}_{\theta}\sum_{i=1}^{n}\log P(x_i\mid \theta)\]

2.2. 隐变量

但是有的模型中会存在隐变量，这是一些不在观测数据中但非常重要的变量。

如果存在隐变量 $Z$，那么对数似然重新写成：

\[\begin{aligned} L(\theta;X)&=\log P\left( X\mid \theta \right)\\ &=\log \sum_{Z}P(X\mid Z,\theta)P(Z\mid \theta)\\ &=\log \sum_{Z}\prod_{i=1}^{n}P(x_i\mid Z,\theta)P(Z\mid \theta) \end{aligned}\]

这种形式很难直接求出极值，因为我们不仅要估计 $\theta$， 还要同时估计 $Z$ 的分布 $P(Z\mid \theta)$。

3. EM 算法

3.1. 基本思想

EM 算法的思想是：

• 在参数 $\theta$ 已知的情况下，我们可以推断出隐变量 $Z$ 的分布；
• 在 $Z$ 分布已知的情况下，我们可以通过极大似然估计得到参数 $\theta$。
• 两者交替迭代，最终收敛到局部最优解。

3.2. 理论推导

我们可以先假设 $Z$ 服从的分布为 $Z\sim q\left( Z\mid \theta \right)$，于是有：

\[\begin{aligned} \log P\left( X\mid \theta \right) &=\log P\left( X,Z\mid \theta \right) -\log P\left( Z\mid X,\theta \right)\\ &=\log \frac{P\left( X,Z\mid \theta \right)}{q\left( Z\mid \theta \right)}-\log \frac{P\left( Z\mid X,\theta \right)}{q\left( Z\mid \theta \right)}\\ \end{aligned}\]

两端关于 $Z\sim q\left( Z\mid \theta \right)$ 同时计算期望：

\[\begin{aligned} \log P\left( X\mid \theta \right) &=\sum_Z{q\left( Z\mid \theta \right) \log \frac{P\left( X,Z\mid \theta \right)}{q\left( Z\mid \theta \right)}}\color{blue}{ -\sum_Z{q\left( Z\mid \theta \right) \log \frac{P\left( Z\mid X,\theta \right)}{q\left( Z\mid \theta \right)}}}\\ &=\mathop {\mathbb{E}} \limits_{Z\sim q\left( Z\mid \theta \right)}\left[ \log P\left( X,Z\mid \theta \right) \right] \color{green}{ -\sum_Z{q\left( Z\mid \theta \right) \log q\left( Z\mid \theta \right) +\color{blue}{ \mathrm{KL}\left( q\left( Z\mid \theta \right) \parallel P\left( Z\mid X,\theta \right) \right) }}}\\ &=\mathop {\mathbb{E}} \limits_{Z\sim q\left( Z\mid \theta \right)}\left[ \log P\left( X,Z\mid \theta \right) \right] +\color{green}{ H\left( q\left( Z\mid \theta \right) \right) }+\color{blue}{ \mathrm{KL}\left( q\left( Z\mid \theta \right) \parallel P\left( Z\mid X,\theta \right) \right) }\\ &=\color{red}{ELBO\left( q,\theta \mid X \right) }+\color{blue}{\mathrm{KL}\left( q\left( Z\mid \theta \right) \parallel P\left( Z\mid X,\theta \right) \right) }\\ \end{aligned}\]

ELBO 是 evidence lower bound optimization(ELBO)，是 $L\left( \theta \right) $​​ 的一个下界。

3.3. 具体流程

EM 算法重复下面两步：

1. 寻找使得 KL 散度最小的 $q^{(t)}(Z) =P\left( Z\mid X,\theta ^{\left( t \right)} \right) $，使得 ELBO 进一步逼近 $L\left( \theta \right)$.

2. 寻找 $ELBO\left( \theta \mid q^{(t)},X \right) $ 的极大值点作为新参数 $\theta^{(t+1)}$.

3.3.1. E 步

\[L( \theta ) -ELBO( q ,\theta\mid X ) =\color{blue}{ \mathrm{KL}\left( q \parallel P\left( Z\mid X,\theta\right) \right) }\]

要使 ELBO 逼近 $L\left( \theta \right)$，就要让 KL 散度最小，先通过当前参数 $\theta^{(t)}$ 估计 $q^{(t)}$，得 $q^{(t)}(Z) =P\left( Z\mid X,\theta ^{\left( t \right)} \right)$，于是有：

\[\begin{aligned} L\left( \theta \right) &=\log P\left( X\mid \theta \right)\\ &=\mathop {\mathbb{E}} \limits_{Z\sim P\left( Z\mid X,\theta ^{\left( t \right)} \right)}\left[ \log P\left( X\mid \theta \right) \right]\\ &=Q\left( \theta ;\theta ^{\left( t \right)} \right) +\mathrm{KL}\left( P\left( Z\mid X,\theta ^{\left( t \right)} \right) \parallel P\left( Z\mid X,\theta \right) \right)\\ \end{aligned}\]

这里我们将 $ELBO\left( \theta \mid q^{(t)},X \right)$ 记为 $Q\left( \theta ;\theta ^{\left( t \right)} \right)$：

\[Q\left( \theta ;\theta ^{\left( t \right)} \right)=\mathop {\mathbb{E}} \limits_{Z\sim P\left( Z\mid X,\theta ^{\left( t \right)} \right)}\left[ \log P\left( X,Z\mid \theta \right) \right] +H\left( P\left( Z\mid X,\theta ^{\left( t \right)} \right) \right)\]

3.3.2. M 步

由于信息熵为常数项，因此最大化 $Q\left( \theta ;\theta ^{\left( t \right)} \right)$ 等价于将对数似然 $\log P ( X,Z\mid \theta )$ 的期望最大化。

\[\theta ^{\left( t+1 \right)}=\mathop {\mathrm{arg}\max} \limits_{\theta}\,Q\left( \theta ;\theta ^{\left( t \right)} \right) =\mathop {\mathrm{arg}\max} \limits_{\theta}\,\mathop {\mathbb{E}} \limits_{Z\sim P\left( Z\mid X,\theta ^{\left( t \right)} \right)}\left[ \log P\left( X,Z\mid \theta \right) \right]\]

3.4. 收敛性

EM 算法的解是逐步更优的。

Proof：

\[L\left( \theta \right) =Q\left( \theta ;\theta ^{\left( t \right)} \right) +\mathrm{KL}\left( P\left( Z\mid X,\theta ^{\left( t \right)} \right) \parallel P\left( Z\mid X,\theta \right) \right)\\\]

显然有 $L\left( \theta ^{\left( t+1 \right)} \right) \ge L\left( \theta ^{\left( t \right)} \right)$。

既然最优值是单调递增的，那么收敛性也不难理解了。

4. 思考

\[\begin{aligned} \log P\left( X\mid \theta \right) &=\color{red}{ELBO\left( q,\theta \mid X \right) }+\color{blue}{\mathrm{KL}\left( q\left( Z\mid \theta \right) \parallel P\left( Z\mid X,\theta \right) \right) }\\ &=\mathop {\mathbb{E}} \limits_{Z\sim q\left( Z\mid \theta \right)}\left[ \log P\left( X,Z\mid \theta \right) \right] +\color{green}{ H\left( q\left( Z\mid \theta \right) \right) }+\color{blue}{ \mathrm{KL}\left( q\left( Z\mid \theta \right) \parallel P\left( Z\mid X,\theta \right) \right) }\\ \end{aligned}\]

EM 算法中先极小化 $\mathrm{KL}\left( q \parallel P\left( Z\mid X,\theta^{(t)} \right) \right)$，将得到的 $q^{(t)}$ 带入到期望中，然后极大化期望。

注意到 $q^{(t)}$ 信息熵那一项被忽略掉了，因为 EM 算法认为 $q^{(t)}$ 确定后其信息熵就是一个常数项。

但是我认为信息熵作为一个正则化项，应当和 KL 散度一起考虑： $q(Z \mid \theta)$ 是 $Z$ 的真实分布，由于没法直接获得，所以只能暂时用 $P(Z\mid X,\theta^{(t)})$ 来近似代替， 但同时也应该考虑过拟合的风险，加上一个信息熵约束让 $q$ 不至于太复杂，这个想法应该是非常合理的。

所以我们需要的 $q^{(t)}$ 应该使 $\mathrm{KL}\left( q\parallel P\left( Z\mid X,\theta ^{(t)} \right) \right) +\color{red}{ H\left( q \right) }$ 最小。

E 步：

\[q^{(t)}=\mathop{\arg \max}_{q}\,\mathrm{KL}\left( q\parallel P\left( Z\mid X,\theta ^{(t)} \right) \right) + H\left( q \right)\]

M 步：

\[\theta ^{\left( t+1 \right)}=\mathop {\mathrm{arg}\max} \limits_{\theta}\,\mathop {\mathbb{E}} \limits_{Z\sim q^{(t)}} \left[ \log P\left( X,Z\mid \theta \right) \right]\]

但这样会带来两个问题：

1. 不能证明 $L(\theta^{(t+1)})$ 比 $L(\theta^{(t)})$ 更优。
2. E 步中的 $q^{(t)}$ 难以求解。

关于第一个问题，加了正则项后泛化性更强，但的确没办法保证下一步一定会更优，目前也没法证明一定能收敛，但是从经验角度来说应该效果会更好。

对于问题 2，如果我们要寻找使得 $\mathrm{KL}\left( q\parallel p\right) +{ H\left( q \right) }$ 最小的 $q$，首先设拉格朗日函数：

\[L(q)=\int q(x)\log \frac{q(x)}{p(x)}\,\mathrm{d}x-\int q(x)\log q(x)\,\mathrm{d}x +\lambda \left(1-\int q(x)\,\mathrm{d}x)\right)\]

计算偏导：

\[\frac{\partial L}{\partial q(x)}=-\log p(x)-\lambda\]

可见这种形式没有极值，那么我们转而寻找使得$\mathrm{KL}\left( p\parallel q\right) +{ H\left( q \right) }$ 最小的 $q$。

\[L(q)=\int p(x)\log \frac{p(x)}{q(x)}\,\mathrm{d}x-\int q(x)\log q(x)\,\mathrm{d}x +\lambda \left(1-\int q(x)\,\mathrm{d}x)\right)\] \[\frac{\partial L}{\partial q(x)}= -\frac{p(x)}{q(x)}-1-\log q(x)-\lambda\]

让偏导为0，得：

\[p(x)=-q(x)\left(\log q(x)+\lambda+1 \right)\]

两端积分得：

\[\lambda = H(q)-2\]

代入得：

\[p(x)=-q(x)\left(\log q(x)+H(q)-1 \right)\]

虽然等式关系得到了，但这个似乎解不出来，好吧，那就算了。。。

5. 补充更新

[Updated on 2021-08-30]

在 VAE 和 GAN 中一般都默认 $q(Z)$ 为正态分布，从某种角度来说，这也算是考虑了信息熵约束，因为在均值和方差确定的情况下，正态分布的信息熵是最小的。

[Updated on 2021-09-01]

看了李航的《统计学习方法》，里面提到了一个 F 函数：

\[F(q,\theta)=\mathop {\mathbb{E}} \limits_{Z\sim q\left( Z\right)}\left[ \log P\left( X,Z\mid \theta \right) \right] + H\left( q\left( Z\right) \right) \\\]

可以证明，对于固定的 $\theta$，存在唯一的 $q$ 使得 F 函数极大化。

\[L(q)=\mathop {\mathbb{E}} \limits_{Z\sim q\left( Z\right)}\left[ \log P\left( X,Z\mid \theta \right) \right] + H\left( q\left( Z\right) \right)+\lambda \left(1-\int q(z)\,\mathrm{d}z)\right)\] \[\frac{\partial L}{\partial q(Z)}= \log P\left( X,Z\mid \theta \right)-1-\log q(Z)-\lambda=0\] \[\frac{P\left( X,Z\mid \theta \right)}{q(Z)}=\mathrm{e}^{1+\lambda}\] \[q(Z)=\frac{P\left( X,Z\mid \theta \right)}{P\left( X\mid \theta \right)}=P(Z\mid X,\theta)\]

也就是说其实信息熵那一项其实考虑到里面去了，恰好就是 EM 的最优解，真是非常神奇。