EM算法的推導、證明和例子

EM

推導 ¹

$\mathcal{X}=\{x_1, x_2, \cdots, x_N\}$: 觀察數據

z: 潛在變量

似然估計函數
$\log P(\mathcal{X};\theta)=\log \prod_i^N P(x_i;\theta) = \sum_i^N \log P(x_i;\theta) \tag{1}$
我們的目標是找到最大化似然估計函數的參數$\theta$
$\underset{\theta}{\arg \max} \log P(\mathcal{X};\theta) \tag{2}$

現有
$P(x,z;\theta) = P(x;\theta)P(z|x;\theta) \tag{3}$
因此
$\log P(x;\theta) = \log \frac{P(x,z;\theta)}{P(z|x;\theta)} = \log P(x,z;\theta) - \log P(z|x;\theta) \tag{4}$
假設z屬於$Q(z;\phi)$分佈
$\begin{aligned} \log P(x;\theta) & = \log P(x,z;\theta) - \log P(z|x;\theta) - \log Q(z;\phi) - \log Q(z;\phi) \\ & = (\log P(x,z;\theta) - \log Q(z;\phi)) - (\log P(z|x;\theta) - \log Q(z;\phi) ) \\ & = \log \frac{P(x,z;\theta)}{Q(z;\phi)} - \log \frac{\log P(z|x;\theta)}{Q(z;\phi)} \end{aligned} \tag{5}$
公式5左右兩邊求關於z的期望
$\text{left} = \int_z Q(z;\phi) \log P(x;\theta) dz = \log P(x;\theta) \int_z Q(z;\phi) dz = \log P(x;\theta) \tag{6}$
因爲$\int_z Q(z;\phi) dz = 1$。

$\begin{aligned} \text{right} & = \int_z Q(z;\phi) \log \frac{P(x,z;\theta)}{Q(z;\phi)} dz - \int_z Q(z;\phi) \log \frac{\log P(z|x;\theta)}{Q(z;\phi)} dz \\ & = \text{ELBO} + \text{KL}(Q(z;\phi)||P(z|x;\theta)) \end{aligned} \tag{7}$
公式7第一項稱爲ELBO(evidence lower bound)，第二項是KL散度。
因此，我們得到
$\log P(x;\theta) = \text{ELBO} + \text{KL}(Q(z;\phi)||P(z|x;\theta)) \tag{8}$
因爲$\text{KL}(\cdot) \ge 0$，因此$\log P(x;\theta) \ge \text{ELBO}$，當且僅當$Q(z;\phi) = P(z|x;\theta)$時，等號成立。ELBO相當於一個下界，不斷地提高ELBO，就能不斷提高$\log P(x;\theta)$，達到我們的目的——最大化似然估計函數。

假設我們有$\theta^{(t)}$，我們想要最大化ELBO，即最小化$\text{KL}(Q(z;\phi)||P(z|x;\theta))$:
$\phi^{(t)} = \underset{\phi}{\arg \min} \text{KL}(Q(z;\phi)||P(z|x;\theta^{(t)})) = \underset{\phi}{\arg \max} \text{ELBO}(\phi, \theta^{(t)}) \tag{9}$
當得到最優的$\phi^{(t)}$，有$Q(z;\phi^{(t)}) = P(z|x;\theta^{t})$。實際情況下很難得到最優的$\phi^{(t)}$，我們的目的時儘可能最大化ELBO。當我們計算出ELBO後，我們可以反過來求$\theta^{(t+1)} = \underset{\theta}{\arg \max}\text{ELBO}(\phi^{(t)}, \theta)$。

通過不斷重複這個兩個最大化的過程，我們就可以近似地求出最大化似然估計函數的參數$\theta$。

繼續看看$\text{ELBO}(\phi^{(t)}, \theta)$

$\begin{aligned} \text{ELBO}(\phi^{(t)}, \theta) &= \int_z Q(z;\phi^{(t)}) \log \frac{P(x,z;\theta)}{Q(z;\phi^{(t)})} dz \\ &= \int_z Q(z;\phi^{(t)}) \log P(x,z;\theta) dz - \int_z Q(z;\phi^{(t)}) \log Q(z;\phi^{(t)} dz \\ &= E_{z\backsim Q(z;\phi^{(t)})}[\log P(x,z;\theta)] - E_{z\backsim Q(z;\phi^{(t)})}[\log Q(z;\phi^{(t)}] \end{aligned} \tag{10}$
公式10的最後一行的第一項是關於z的期望，第二項是一個常數（$\phi^{(t)}$已知），所以
$\begin{aligned} \theta^{(t+1)} &= \underset{\theta}{\arg \max}\text{ELBO}(\phi^{(t)}, \theta) \\ &= \underset{\theta}{\arg \max} E_{z\backsim Q(z;\phi^{(t)})}[\log P(x,z;\theta)] \end{aligned} \tag{11}$
因此EM算法叫做期望最大化算法。

總結，EM算法的迭代過程如下

• E-step: 固定$\theta^{(t)}$，$\phi^{(t)}=\underset{\phi}{\arg \max} E_{z\backsim Q(z;\phi^{(t)})}[\log P(x,z;\theta)]$；
• M-step: 固定$\phi^{(t)}$，$\theta^{(t+1)} = \underset{\theta}{\arg \max}E_{z\backsim Q(z;\phi^{(t)})}[\log P(x,z;\theta)]$。

E-step和M-step的順序可以互換。

EM算法收斂性證明¹

簡單的證明：只要$\theta^{(t)} \to \theta^{(t+1)}, \log P(x;\theta^{(t)}) \le \log P(x;\theta^{(t+1)})$，就能保證算法收斂。

從公式4出發，兩邊求關於z的期望
$\begin{aligned} \text{left} &= \int_z Q(z;\phi^{(t)}) \log P(x;\theta) dz \\ &= \log P(x;\theta) \int_z Q(z; \phi^{(t)}) dz \\ &= \log P(x;\theta) \end{aligned} \tag{12}$

$\text{right} = \int_z Q(z;\phi^{(t)}) \log P(x,z;\theta)dz - \int_z Q(z;\phi^{(t)}) \log P(z|x;\theta)dz \tag{13}$

因爲$\phi^{(t)}$根據公式9求解得到的，假設我們得到的是最優解，則$Q(z;\phi^{(t)}) = P(z|x;\theta^{t})$，代入公式13得
$\begin{aligned} \text{right} &= \int_z P(z|x;\theta^{(t)}) \log P(x,z;\theta)dz - \int_z P(z|x;\theta^{(t)}) \log P(z|x;\theta)dz \\ &= H_1 (\theta, \theta^{(t)}) - H_2 (\theta, \theta^{(t)}) \end{aligned} \tag{14}$
我們分別用$H_1$和$H_2$指代公式14的兩項。

由公式12和14得
$\log P(x;\theta) = H_1 (\theta, \theta^{(t)}) - H_2 (\theta, \theta^{(t)}) \tag{15}$

因爲$\theta^{(t+1)} = \underset{\theta}{\arg \max}E_{z\backsim Q(z;\phi^{(t)})}[\log P(x,z;\theta)]$，所以有$H_1 (\theta^{(t+1)}, \theta^{(t)}) \ge H_1 (\theta^{(t)}, \theta^{(t)})$。接下來，只要證明$-H_2(\theta^{(t+1)}, \theta^{(t)}) \ge -H_2(\theta^{(t)}, \theta^{(t)})$，就能證明$\log P(x;\theta^{(t+1)}) \ge \log P(x;\theta^{(t+1)})$。

現在
$\begin{aligned} & H_2(\theta^{(t+1)}, \theta^{(t)}) - H_2(\theta^{(t)}, \theta^{(t)}) \\ =& \int_z P(z|x;\theta^{(t)}) \log P(z|x; \theta^{(t+1)}) dz - \int_z P(z|x;\theta^{(t)}) \log P(z|x; \theta^{(t)}) dz \\ =& \int_z P(z|x;\theta^{(t)}) \log \frac{P(z|x; \theta^{(t+1)})}{P(z|x; \theta^{(t)})} dz \end{aligned} \tag{16}$
證明公式16小於等於0：

方法1：公式16是負KL散度$-KL(P(z|x;\theta^{(t)})||P(z|x;\theta^{(t+1)})) \le 0$。

方法2：公式16等於$E_{z \backsim P(z|x,\theta^{(t)})}[\log \frac{P(z|x; \theta^{(t+1)})}{P(z|x; \theta^{(t)})}]$。根據Jensen不等式$E[\log (x)] \le \log E[x]$，因此
$\begin{aligned} & E_{z \backsim P(z|x,\theta^{(t)})}[\log \frac{P(z|x; \theta^{(t+1)})}{P(z|x; \theta^{(t)})}] \\ \le & \log E_{z \backsim P(z|x,\theta^{(t)})}[\frac{P(z|x; \theta^{(t+1)})}{P(z|x; \theta^{(t)})}] \\ = & \log \int_z P(z|x;\theta^{(t)}) \frac{P(z|x; \theta^{(t+1)})}{P(z|x; \theta^{(t)})} dz \\ =& \log \int_z P(z|x; \theta^{(t+1)}) dz \\ =& \log 1 = 0 \end{aligned} \tag{17}$

例子 ²

拋硬幣，有兩個硬幣，但是兩個硬幣的材質不同導致其出現正反面的概率不一樣，目前我們只有一組觀測數據，要求出每一種硬幣投擲時正面向上的概率。總共投了五輪，每輪投擲五次。假設我們不知道每一次投擲用的是哪一種硬幣，等於是現在的問題加上了一個隱變量，就是每一次選取的硬幣的種類。

(圖片來自https://blog.csdn.net/u010834867/article/details/90762296)

設兩個硬幣分別是AB，P(正|A)=$x_1$，P(反|A)=$1-x_1$，P(正|B)=$x_2$，P(正|B)=$x_2$。

假設第i次實驗選擇硬幣A的概率是$P(z_{i}=A)=y_i$，選擇硬幣B的概率是$P(z_i=B)=1-y_i$。

看實驗i的數據j，用$x_{ij}$表示，似然估計函數爲
$\log P(x) = \log \prod_i \prod_j P(x_{ij}) = \sum_i \sum_j \log P(x_{ij})$
$P(x_{ij})=\frac{P(x,z)}{P(z|x)}=\frac{P(z)P(x|z)}{P(z|x)}$

其中$P(z|x)$不好求出來。我們使用EM算法來解$x_1$和$x_2$。

首先，我們求出期望，
$\begin{aligned} & E_{z\backsim Q(z;\phi^{(t)})}[\log P(x,z;\theta)] \\ =& y_1 \log y_1(x_1 x_1 (1-x_1) x_1 (1-x_1)) &+& (1-y_1)\log (1-y_1)(x_2 x_2 (1-x_2) x_2 (1-x_2)) \\ +& y_2 \log y_2((1-x_1)(1-x_1)x_1 x_1(1-x_1)) &+& (1-y_2)\log (1-y_2)((1-x_2)(1-x_2)x_2 x_2(1-x_2)) \\ +& y_3 \log y_3(x_1 (1-x_1)(1-x_1)(1-x_1)(1-x_1)) &+& (1-y_3)\log (1-y_3)(x_2 (1-x_2)(1-x_2)(1-x_2)(1-x_2)) \\ +& y_4 \log y_4(x_1 (1-x_1)(1-x_1) x_1 x_1) &+& (1-y_4) \log (1-y_4)(x_2 (1-x_2)(1-x_2) x_2 x_2) \\ +& y_5 \log y_5((1-x_1)x_1 x_1 (1-x_1) (1-x_1)) &+& (1-y_5) \log (1-y_5)((1-x_2)x_2 x_2 (1-x_2) (1-x_2)) \end{aligned}$

假設$x_1=0.2,x_2=0.7$，代入上式得
$\begin{aligned} & E_{z\backsim Q(z;\phi^{(t)})}[\log P(x,z;\theta)] \\ =& y_1 \log 0.00512 y_1 &+& (1-y_1)\log 0.03087(1-y_1) \\ +& y_2 \log 0.02048 y_2 &+& (1-y_2)\log 0.01323(1-y_2) \\ +& y_3 \log 0.08192 y_3 &+& (1-y_3)\log 0.00567 (1-y_3) \\ +& y_4 \log 0.00512 y_4 &+& (1-y_4) \log 0.03087 (1-y_4) \\ +& y_5 \log 0.02048 y_5 &+& (1-y_5) \log 0.01323(1-y_5) \end{aligned}$

(圖片來自https://blog.csdn.net/u010834867/article/details/90762296)

現在求
$\max E_{z\backsim Q(z;\phi^{(t)})}[\log P(x,z;\theta)]$

爲了簡單運行，我們取$Q(z;\phi^{(t)})$爲$y=\{0,1,1,0,1\}$，即$z=\{B,A,A,B,A\}$。雖然求出來的期望不是最大的，但不影響算法的收斂。因爲z的結果已經固定了，可以直接計算$\theta^{(t+1)}$：
$x_1 = (2+1+2) / 15 = 0.33, x_2 = (3+3) / 10 = 0.6$

接着不斷迭代，直到z或者$x_1$和$x_2$收斂。

若有不恰當之處，請指正。

---

1. https://www.bilibili.com/video/av31906558?p=1 ↩︎ ↩︎

2. https://blog.csdn.net/u010834867/article/details/90762296 ↩︎