EM算法理论及实战

🏷NLP 🏷动态规划

概率图模型中，除可观测的 显变量 \(X\) 外, 如果还依赖于无法观测的 隐变量 \(Z\) ，则可考虑通过EM算法对 参数 \(\theta\) 进行最大似然估计(或MAP估计)。

1 简单例子：抛硬币

建模过程(硬币结果生成的过程)

两枚硬币a，b，抛掷后正面朝上的概率为 \(\theta_a, \theta_b\)

  • for i in range(N):

    • 随机选择a或者b:

      • 如果是a，连续抛10次硬币a，记录结果
      • 如果是b，连续抛10次硬币b，记录结果

符号表示如下:

  • N为样本个数, 比如5
  • 10为特征数
  • \(\mathbf{X} \in \mathcal{R}^{(N, d)}\) 为观测到的硬币结果, \(\mathbf{X}_{i,j}\) 表示第i组第j次抛硬币的结果, 1为正面朝上, 0为正面朝下

目标: 预估 \(\theta: = (\theta_a, \theta_b)\)

  • 隐变量 \(\mathbf{Z} = (z_1, z_2, ..., z_N)\) 不可观测，即不知道每次抛的是硬币A还是硬币B

    • \(z_i = 0\) : 第i组抛的是硬币a
    • \(z_i = 1\) : 第i组抛的是硬币b

鸡生蛋、蛋生鸡问题

  • 估计 $Z=(z_1, z_2, ..., z_5)\(（即每轮抛的是A还是B），需要已知 $\theta=(\theta_a, \theta_b)\)
  • 估计 \(\theta_a,\theta_b\) ，需要已知 \(\mathbf{Z}\)

解决方案：反复迭代

  • E: 随机初始化 \(\theta_a, \theta_b\) ，估计z (准确说是 \(p(z|x, \theta_a, \theta_b)\) )
  • M: 基于z，重新估计 \(\theta_a, \theta_b\)

2 理论

令 \(\mathcal{X}\) 表示可观察的显变量，比如人脸图像，\(\mathcal{Z}\) 表示不可观测的隐变量，比如人脸图像中人的性别、种族、姿势等。

  • Complete Case: 如果对于每个观察到的 \(\mathbf{X}\), 我们同时也知道对应的 \(\mathbf{Z}\) 的值，即

    • \(\mathcal{Z}\) 可观测到（或者通过人工标注得到）
    • \(\hat{\theta}_{MLE} = \mathop{\arg \max} \limits_{\theta} \log p(\mathbf{X}, \mathbf{Z}|\theta)\)

  • *Incomplete Case*，是下文的重点。

    • \(\mathcal{Z}\) 不可观测，也不知道对应的值
    • \(\hat{\theta}_{MLE} = \log p(\mathbf{X}|\theta)\)

2.1 算法

输入

  • 显变量状态空间\(\mathcal{X}\) 上观测到的数据集\(\mathbf{X} \in \mathcal{R}^{N \times d}\), N个观测数据样本
  • 隐变量状态空间\(\mathcal{Z}\) 上的数据集 \(\mathbf{Z}\) 不可观测，但是知道：\(p(\mathbf{Z}|\mathbf{X}, \mathbf{\theta})\)
  • \(p(\mathbf{X, Z}|\mathbf{\theta})\)

输出

  • \(\hat{\theta}_{MLE} = \mathop{\arg \max} \limits_{\theta} p(\mathbf{X}|\theta)\)

过程

  • 初始化:i=0; \(\theta^{i}\)
  • E: 计算隐状态 \(\mathbf{Z}\) 的分布(pdf)：\(\color{blue}{p(\mathbf{Z}|\mathbf{X}, \mathbf{\theta}^{i})}\) (注： \(\mathbf{X}\) 可观测，\(\theta^{i}\) 为上一步的预估值)
  • M: 估计参数 \(\mathbf{\theta}^{i+1} = \mathop{\arg \max}\limits_{\theta} Q(\mathbf{\theta}, \mathbf{\theta}^{i})\); i+=1
  • 如果未收敛，返回步骤E

2.2 理论基础

\begin{aligned} \log p(\mathbf{X}|\theta) &= \log \sum_{\mathbf{Z}} p(\mathcal{X} = \mathbf{X}, \mathcal{Z}=\mathbf{Z}|\mathcal{\Theta}=\theta)\\ &:= \log \sum_{\mathbf{Z}} \color{red}{p(\mathbf{X, Z}|\theta)} \qquad \text{很难优化} \\ &= Q(\mathbf{\theta}, \theta^{i}) + const +KL(p||q) \\ &\geq Q(\mathbf{\theta}, \theta^{i}) + const \qquad \text{下界, 容易优化} \end{aligned}

其中 \(Q(\mathbf{\theta}, \theta^{i}) = \sum_{\mathbf{Z}} \color{blue}{p(\mathbf{Z}|\mathbf{X}, \mathbf{\theta}^{i})} \log \color{red}{p(\mathbf{X}, \mathbf{Z}|\mathbf{\theta})}\) (注：\(\theta^{i}\) 为上一步的预估值)

  • 优化目标函数的一个下界，来得到近似解。
  • 可证明, \(p(\mathbf{X}|\theta^{i+1}) \geq p(\mathbf{X}|\theta^{i})\), 即每次迭代都会比上一次更好，最终收敛到某一个值

3 应用

  • K-means是 EM 的一个特殊例子
  • HMM 的参数估计，也是用的EM算法

4 实战

将抛硬币套用EM算法，重写如下:

假设 \(p(z_i=0|\theta) = p(z_i=1|\theta) = 0.5\),

目标: 预估 \(\theta: = (\theta_a, \theta_b)\)

输入

  • X 硬币结果

    第i组 1=Head, 0=Tail
    1 1, 0, 0, 0, 1, 1, 0, 1, 0, 1
    2 1, 1, 1, 1, 0, 1, 1, 1, 1, 1
    3 1, 0, 1, 1, 1, 1, 1, 0, 1, 1
    4 1, 0, 1, 0, 0, 0, 1, 1, 0, 0
    5 0, 1, 1, 1, 0, 1, 1, 1, 0, 1

  • p(Z, X|θ)

    \[ p(Z, X| \theta) = p(X|Z, \theta) p(Z|\theta) = p(X|Z, \theta) 0.5 \]

  • p(Z|X, θ)
\begin{aligned} p(Z|X, \theta) &= \frac{ p(X|\theta, Z)p(Z|\theta) }{P(X|\theta)} \\ &= \frac{ p(X|\theta, Z)p(Z|\theta) }{\sum_{Z} p(X|Z, \theta)p(Z|\theta) } \\ &= \frac{ p(X|\theta, Z) }{\sum_{Z} p(X|Z, \theta) } \end{aligned}

\[ p(z_i=0|\theta_a, \theta_b, X_i) = \frac{ \theta_a^{\sum_j X_{i,j}} (1-\theta_a)^{5-\sum_jX_{i,j}}} { \theta_a^{\sum_j X_{i,j}} (1-\theta_a)^{5-\sum_jX_{i,j}}+ \theta_b^{\sum_j X_{i,j}} (1-\theta_b)^{5-\sum_jX_{i,j}} } \]

\[ p(z_i=1|\theta_a, \theta_b, X_i) = \frac{ \theta_b^{\sum_j X_{i,j}} (1-\theta_b)^{5-\sum_jX_{i,j}}} { \theta_a^{\sum_j X_{i,j}} (1-\theta_a)^{5-\sum_jX_{i,j}}+ \theta_b^{\sum_j X_{i,j}} (1-\theta_b)^{5-\sum_jX_{i,j}} } \]

em_coin.png

def main():

    # 做实验
    # n_time = 10  # 每组实验抛几次
    # n_group = 100  # 实验组数
    # result = toss_coins(n_time, n_group, p_a=0.8, p_b=0.5)

    n_time = 10  # 每组实验抛几次
    n_group = 5  # 实验组数
    result = [
        [1, 0, 0, 0, 1, 1, 0, 1, 0, 1],
        [1, 1, 1, 1, 0, 1, 1, 1, 1, 1],
        [1, 0, 1, 1, 1, 1, 1, 0, 1, 1],
        [1, 0, 1, 0, 0, 0, 1, 1, 0, 0],
        [0, 1, 1, 1, 0, 1, 1, 1, 0, 1]
    ]

    # 初始化:
    i_iter = 0
    theta_a, theta_b = 0.6, 0.5

    for i_iter in range(1, 11):
        # E: P(Z|X, theta)
        pz0 = []  # P(z=0|theta_a, theta_b, X) ~ 抛硬币a的概率
        pz1 = []  # P(z=1|theta_a, theta_b, X) ~ 抛硬币b的概率
        for i in range(n_group):
            data = result[i]
            score_a = theta_a ** sum(data) * (1 - theta_a) ** (n_time - sum(data))
            score_b = theta_b ** sum(data) * (1 - theta_b) ** (n_time - sum(data))
            pz0.append((score_a * 0.5) / (score_a * 0.5 + score_b * 0.5))
            pz1.append((score_b * 0.5) / (score_a * 0.5 + score_b * 0.5))

        # M: update theta to max Q(theta, theta_old)
        n_head_from_a, n_head_from_b = 0, 0
        n_tail_from_a, n_tail_from_b = 0, 0
        for i in range(n_group):
            data = result[i]
            n_head_from_a += sum(data) * pz0[i]
            n_tail_from_a += (n_time - sum(data)) * pz0[i]
            n_head_from_b += sum(data) * pz1[i]
            n_tail_from_b += (n_time - sum(data)) * pz1[i]
        old_theta_a = theta_a
        old_theta_b = theta_b
        theta_a = n_head_from_a / (n_head_from_a + n_tail_from_a)
        theta_b = n_head_from_b / (n_head_from_b + n_tail_from_b)

        print(f"""iter={i_iter:02d}, (theta_a, theta_b) = ({theta_a:.3f}, {theta_b:.3f})""")

        # 退出循环
        if abs(old_theta_a - theta_a) < 0.001 and abs(old_theta_b - theta_b) < 0.001:
            break

5 参考文献