Expectation Maximization

22 Mar 2014 Comments

Introduction

给定数据 $X$，一种常用的估计模型参数 $\theta$ 的方法是 maximum likelihood：

$$\theta = \underset{\theta}{\arg\max} \;\; p(X|\theta)$$

这里的 $X$ 是我们可见的，即 visible variable。如果给定的问题涉及到 hidden variable，ML 就需要引入新变量：

$$\theta = \underset{\theta}{\arg\max} \;\; p(X|\theta) = \underset{\theta}{\arg\max} \;\; \int p(X, z|\theta) dz$$

其中 $z$ 对应 hidden variable。在很多情况下，引入积分后问题会变得不可求解，这也是为什么需要 EM 算法的原因。EM 尝试用迭代的方式去求解 $\theta$，从而避免掉积分操作。

另外，为了计算方便，我们通常并不直接操作 $p(X

\theta)$，而是用 $\log$ 形式代替：

$$\theta = \underset{\theta}{\arg\max} \; \log p(X|\theta)$$

为了使用方便，下面记 $L(\theta) = \log p(X

\theta)$

EM Iteration

假设第 $n$ 轮迭代得到的参数是 $\theta_n$，由于我们的目标是最大化 $L(\theta)$，因此我们希望下一轮迭代得到的 $\theta$ 可以进一步提升 $L(\theta)$。换句话说，在第 $n+1$ 轮迭代时，希望能找到 $\theta$ 使得 $L(\theta) > L(\theta_n)$。 EM 的做法是利用 $\theta_n$ 得到 $L(\theta)$ 的一个下界，然后去优化这个下界得到一个更优的 $\theta$ (当然这个下界通常是易于优化的)。

下面的推导首先给出 $L(\theta)$ 的下界并证明可以通过优化这个下界得到更优的 $\theta$。

推导过程中用到了 concave function 的一个性质：

If $f(x)$ is a concave function, $\forall \sum_i a_i = 1 \;\; f(\sum_{i}a_i x_i) \geq \sum_i a_i f(x_i)$

特别的，如果 $f(x) = \log x$，则有 $\log \sum_i a_i x_i \geq \sum_i a_i \log x_i \;\; \forall \sum_i a_i = 1$

得到下界

假设 $X = \{\b{x}_1, \b{x}_2, …, \b{x}_n\}$

$$\begin{align*} L(\theta) = & \sum\_i \log p(\b{x}^{(i)}|\theta) \\\\ = & \sum\_i \log \sum\_\b{z} p(\b{x}^{(i)}, \b{z}|\theta) \\\\ = & \sum\_i \log \sum\_\b{z} p(\b{x}^{(i)}, \b{z}|\theta) \frac{p(\b{z}|\b{x}^{(i)}, \theta\_n)}{p(\b{z}|\b{x}^{(i)}, \theta\_n)} \\\\ \geq & \sum\_i \sum\_\b{z} p(\b{z}|\b{x}^{(i)}, \theta\_n) \log \frac{p(\b{x}^{(i)}, \b{z}|\theta)}{p(\b{z}|\b{x}^{(i)}, \theta\_n)} \\\\ = & l(\theta) \end{align*}$$

这样我们就得到了 $L(\theta)$ 的下界 $l(\theta)$。推导的核心的是 **第三个等式引入了 hidden variable distribution $p(\b{z}

\b{x}^{(i)}, \theta_n)$**。

证明优化下界提升 likelihood

当 $\theta = \theta_n$ 时，我们有：

$$\begin{align*} l(\theta\_n) = & \sum\_i \sum\_\b{z} p(\b{z}|\b{x}^{(i)}, \theta\_n) \log \frac{p(\b{x}^{(i)}, \b{z}|\theta\_n)}{p(\b{z}|\b{x}^{(i)}, \theta\_n)} \\\\ = & \sum\_i \sum\_\b{z} p(\b{z}|\b{x}^{(i)}, \theta\_n) \log {p(\b{x}^{(i)}|\theta\_n)} \\\\ = & \sum\_i \log {p(\b{x}^{(i)}|\theta\_n)} \\\\ = & L(\theta\_n) \end{align*}$$

因此，如果 $\theta_{n+1} = \underset{\theta}{\arg\max}\; l(\theta)$ 则有 $L(\theta_{n+1}) \geq l(\theta_{n+1}) \geq l(\theta_{n}) = L(\theta_n)$，也就是通过优化下界我们可以得到一个 likelihood 更高的 $\theta$。

E-step & M-step

下面进一步精简 $\arg\max$：

$$\begin{align*} \theta\_{n+1} = & \underset{\theta}{\arg\max}\; l(\theta) \\\\ = & \underset{\theta}{\arg\max}\; \sum\_i \sum\_\b{z} p(\b{z}|\b{x}^{(i)}, \theta\_n) \log \frac{p(\b{x}^{(i)}, \b{z}|\theta)}{p(\b{z}|\b{x}^{(i)}, \theta\_n)} \\\\ = & \underset{\theta}{\arg\max}\; \sum\_i \sum\_\b{z} p(\b{z}|\b{x}^{(i)}, \theta\_n) \log p(\b{x}^{(i)}, \b{z}|\theta) \\\\ \end{align*}$$

分母相对于 $\theta$ 是个常数，所以可以直接去掉而不影响 $\arg\max$ 的结果。

这样就可以得到 E-step 和 M-step:

• E-step: get distribution of $\b{z}: p(\b{z}

  \b{x}^{(i)}, \theta_n)$

• M-step: get updated $\theta$: $\theta_{n+1} = \arg\max\; l(\theta)$

下面给出两个具体的例子，分别是 Hidden Markov Model (HMM) 和 Probabilistic Latent Semantic Analysis (PLSA) 的训练过程。

PLSA

PLSA 用于浅层语义分析，其 likelihood 可以表示为

$$L(\theta) = \sum\_{d} \sum\_{w} n(w, d) \log p(d, w)$$

其中 $d$ 表示文档，$w$ 表示文档中的词 (其实 $d, w$ 可以表示所有有关联的事物)，$n(w, d)$ 表示 $d, w$ 共现的次数，$p(w, d)$ 表示 $w, d$ 的联合概率。

引入 $z$ 后，likelihood 表示为

$$\begin{align*} L(\theta) = & \sum\_{d} \sum\_{w} n(w, d) \log \sum\_{z} p(d, w, z) \\\\ = & \sum\_{d} \sum\_{w} n(w, d) \log \sum\_{z} p(z)p(w|z)p(d|z) \end{align*}$$

所以对于 PLSA，$\theta$ 包括 3 个部分，$p(z), p(w

z), p(d

z)$

---

假设 EM 第 n 轮迭代得到的参数为 $\theta_n$，则有

$$\begin{align*} L(\theta) = & \sum\_{d} \sum\_{w} n(w, d) \log \sum\_{z} p(z)p(w|z)p(d|z) \frac{p(z|w, d, \theta\_n)}{p(z|w, d, \theta\_n)} \\\\ \geq & \sum\_{d} \sum\_{w} n(w, d) \sum\_{z} p(z|w, d, \theta\_n) \log p(z)p(w|z)p(d|z) \frac{1}{p(z|w, d, \theta\_n)} \\\\ = & l(\theta) \end{align*}$$

其中 $p(z

w, d, \theta_n)$ 表示由 $\theta_n$ 得到的 hidden variable distribution，在 E-step 计算。

---

M-step 最大化 $l(\theta)$

$$\begin{align*} \underset{\theta}{\arg\max} \;\; l(\theta) = & \underset{\theta}{\arg\max} \sum\_{d} \sum\_{w} n(w, d) \sum\_{z} p(z|w, d, \theta\_n) \log p(z)p(w|z)p(d|z) \\\\ s.t. & \sum\_z p(z) = 1, \; \sum\_d p(d|z) = 1, \; \sum\_w p(w|z) = 1 \end{align*}$$

令

$$\begin{align*} f(x) = & \sum\_{d} \sum\_{w} n(w, d) \sum\_{z} p(z|w, d, \theta\_n) (\log p(z) + \log p(w|z) + \log p(d|z)) \\\\ & - \alpha(\sum\_z p(z) - 1) - \beta(\sum\_d p(d|z) - 1) - \gamma(\sum\_w p(w|z) - 1) \end{align*}$$

根据 $\frac{\partial f(x)}{\partial p(z)} = 0, \frac{\partial f(x)}{\partial p(d

z)} = 0, \frac{\partial f(x)}{\partial p(w

z)} = 0$ 计算出 $p(z), p(d

z), p(w

z)$。以 $p(z)$ 为例

$$\begin{align*} & \frac{\partial f(x)}{\partial p(z)} = \sum\_{d} \sum\_{w} n(w, d) p(z|w, d, \theta\_n) \frac{1}{p(z)} - \alpha = 0 \\\\ \Rightarrow & p(z) = \frac{\sum\_{d} \sum\_{w} n(w, d) p(z|w, d, \theta\_n)}{\alpha} \end{align*}$$

根据 $\sum_z p(z) = 1$ 推出 $\alpha = \sum_{z} \sum_{d} \sum_{w} n(w, d) p(z

w, d, \theta_n)$，因此

$$p(z) = \frac{\sum\_{d} \sum\_{w} n(w, d) p(z|w, d, \theta\_n)}{\sum\_{z} \sum\_{d} \sum\_{w} n(w, d) p(z|w, d, \theta\_n)}$$

$p(d

z), p(w

z)$ 同理可得。

---

PLSA 迭代总结如下：

• E-step

  $$p(z|w, d, \theta\_n) = \frac{p(z)p(w|z)p(d|z)}{\sum\_z p(z)p(w|z)p(d|z)}$$

• M-step

  $$ p(z) = \frac{\sum_{d} \sum_{w} n(w, d) p(z	w, d, \theta_n)}{\sum_{z} \sum_{d} \sum_{w} n(w, d) p(z	w, d, \theta_n)} $$
  $$ p(w	z) = \frac{\sum_{d} n(w, d) p(z	w, d, \theta_n)}{\sum_{d} \sum_{w} n(w, d) p(z	w, d, \theta_n)} $$
  $$ p(d	z) = \frac{\sum_{w} n(w, d) p(z	w, d, \theta_n)}{\sum_{d} \sum_{w} n(w, d) p(z	w, d, \theta_n)} $$

HMM

HMM 包含一个观测序列和一个状态序列，如果训练样本中的状态序列未知，则需要使用 EM 进行模型的参数估计。假设训练数据中包含 n 个观测序列，$D = \{O_1, O_2, \cdots, O_n\}$，则数据的 likelihood 可表示为

$$L(\theta) = \sum\_{i=1}^n \log P(O\_i|\theta)$$

令 $S$ 表示状态序列，则有

$$L(\theta) = \sum\_{i=1}^n \log P(O\_i|\theta) = \sum\_{i=1}^n \log \sum\_{S} P(O\_i, S|\theta)$$

其中 $\theta$ 包含 3 个部分：

• $\pi_k$: 表示第一个状态为 $k$ 的概率 $\pi_k = P(S_1 = k)$

• $a_{kl}$: 表示从状态 $k$ 转移到状态 $l$ 的概率 $a_{kl} = P(S_{t+1} = l

  S_{t} = k)$

• $b_k(u)$: 表示状态 $k$ 发射出观测 $u$ 的概率 $b_k(u) = P(O_t = u

  S_t = k)$

这样一个观测序列和一个状态序列的联合概率就可以表示成 (假设序列长度为 $l$)

$$P(O, S|\theta) = \pi\_{S\_1} \prod\_{t=1}^{l-1} a\_{S\_{t}S\_{t+1}} \prod\_{t=1}^{l} b\_{S\_t}(O\_t)$$

---

假设 EM 第 n 轮迭代得到的参数为 $\theta_n$，则有

$$\begin{align*} L(\theta) = & \sum\_{i=1}^n \log \sum\_{S} P(O\_i, S|\theta) \frac{P(S|O\_i, \theta\_n)}{P(S|O\_i, \theta\_n)} \\\\ \geq & \sum\_{i=1}^n \sum\_{S} P(S|O\_i, \theta\_n) \log P(O\_i, S|\theta) \frac{1}{P(S|O\_i, \theta\_n)} \\\\ = & l(\theta) \end{align*}$$

---

M-step 最大化 $l(\theta)$

$\begin{align*} \underset{\theta}{\arg\max} \; l(\theta) = & \underset{\theta}{\arg\max} \sum\_{i=1}^n \sum\_{S} P(S|O\_i, \theta\_n) \log P(O\_i, S|\theta) \\\\ = & \underset{\theta}{\arg\max} \sum\_{i=1}^n \sum\_{S} P(S|O\_i, \theta\_n) \log \pi\_{S\_1} \prod\_{t=1}^{l-1} a\_{S\_{t}S\_{t+1}} \prod\_{t=1}^{l} b\_{S\_t}(O\_{it}) \\\\ \end{align*}$ $s.t. \sum\_k \pi\_k = 1, \;\; \sum\_l a\_{kl} = 1, \;\; \sum\_u b\_k(u) = 1$

令

$$\begin{align*} f(x) = & \sum\_{i=1}^n \sum\_{S} P(S|O\_i, \theta\_n) \log \pi\_{S\_1} \prod\_{t=1}^{l-1} a\_{S\_{t}S\_{t+1}} \prod\_{t=1}^{l} b\_{S\_t}(O\_{it}) + \\\\ & \alpha(\sum\_k \pi\_k - 1) + \beta(\sum\_l a\_{kl} - 1) + \gamma(\sum\_u b\_k(u) - 1) \end{align*}$$

根据 $\frac{\partial f(x)}{\partial \pi_k} = 0, \frac{\partial f(x)}{\partial a_{kl}} = 0, \frac{\partial f(x)}{\partial b_k(u)} = 0$ 计算出 $\pi_k, a_{kl}, b_k(u)$。以 $\pi_k$ 为例

$$\begin{align*} \frac{\partial f(x)}{\partial \pi\_k} = & \frac{\partial \sum\_{i=1}^n \sum\_{S} P(S|O\_i, \theta\_n) \log \pi\_{S\_1}}{\partial \pi\_k} - \alpha \\\\ = & \frac{\partial \sum\_{i=1}^n \sum\_{S} P(S|O\_i, \theta\_n) \sum\_k I(S\_1 = k) \log \pi\_{k}}{\partial \pi\_k} - \alpha \\\\ = & \frac{\partial \sum\_{i=1}^n \sum\_k \log \pi\_{k} \sum\_{S} P(S|O\_i, \theta\_n) I(S\_1 = k)}{\partial \pi\_k} - \alpha \\\\ = & \frac{\partial \sum\_k \log \pi\_{k} \sum\_{i=1}^n P(S\_1 = k|O\_i, \theta\_n)}{\partial \pi\_k} - \alpha \\\\ = & \frac{1}{\pi\_k} \sum\_{i=1}^n P(S\_1 = k|O\_i, \theta\_n) - \alpha \\\\ = & 0 \end{align*}$$

由此推出

$$\pi\_k = \frac{\sum\_{i=1}^n P(S\_1 = k|O\_i, \theta\_n)}{n}$$

注意到上面的推导有一个 trick，即引入了 indicator function $I$，在推导 $a_{kl}$ 和 $b_{k}(u)$ 的过程中也会用到类似的 trick，具体过程这里就不给出了。

---

HMM 迭代总结如下：

• E-step

  $$\text{calculate} \;\; P(S\_t = k|O\_i, \theta\_n) \;\; \text{and} \;\; P(S\_t = k, S\_{t+1} = l|O\_i, \theta\_n) \;\; \forall \; t, i, k, l$$

• M-step

  $$\pi_k = \frac{\sum_{i=1}^n P(S_1 = k	O_i, \theta_n)}{n}$$
  $$a_{kl} = \frac{\sum_{i=1}^n\sum_{t=1}^{l-1} P(S_t = k, S_{t+1} = l	O_i, \theta_n)}{\sum_{i=1}^n\sum_{t=1}^{l-1} P(S_t = k	O_i, \theta_n)}$$
  $$b_k(u) = \frac{\sum_{i=1}^n\sum_{t=1}^{l} P(S_t = k	O_i, \theta_n) I(O_{it} = u)}{\sum_{i=1}^n\sum_{t=1}^{l} P(S_t = k	O_i, \theta_n)}$$

这里没有给出计算 $P(S_t = k

O_i, \theta_n)$ 和 $P(S_t = k, S_{t+1} = l

O_i, \theta_n)$ 的具体细节。