EM算法入门：三硬币模型推导与Python实现

1. 从一个让人抓狂的硬币游戏说起三枚硬币A、B、C扔十次记录正面朝上的次数。规则是这样的先扔A如果A是正面就用B扔十次如果A是反面就用C扔十次。你只看到了每次正面朝上的总数但不知道这一轮到底用的是B还是C。现在问题来了让你估计A、B、C三枚硬币正面朝上的概率你怎么估我第一次碰到这个问题的时候脑子里第一反应是这还不简单数一数不就完了。但仔细一想根本数不了——因为你压根不知道每一轮用的是哪枚硬币。这就是EM算法要解决的核心困境观测数据不完整存在隐变量。三硬币模型是EM算法最经典的入门案例没有之一。它足够简单简单到你可以拿张纸手算又足够典型典型到把EM算法的核心思想——E步和M步的交替迭代——展现得淋漓尽致。如果你正在学EM算法或者被Jensen不等式、隐变量、极大似然估计这些词搞得头大那这篇文章就是写给你的。我会从零开始把三硬币模型的EM推导、代码实现、踩坑经验全部摊开讲清楚保证你看完能自己手推一遍也能写出一份能跑的代码。2. 为什么这个问题不能直接用极大似然估计2.1 极大似然估计的“理想条件”在这里失效了极大似然估计MLE是我们最熟悉的参数估计方法。它的逻辑很直白找一组参数使得当前观测数据出现的概率最大。如果三硬币问题里我们知道每一轮用的是B还是C那MLE的答案就是简单的频率统计——用B的那些轮次里正面朝上的比例就是B的正面概率用C的同理。但现实是我们只知道每轮正面朝上的总数不知道用的是哪枚硬币。这就导致似然函数里出现了一个“求和再取对数”的结构假设观测数据是 ( y (y_1, y_2, ..., y_n) )每轮扔了10次隐变量 ( z_i ) 表示第 ( i ) 轮用的是B还是C。似然函数长这样[ L(\theta) \prod_{i1}^{n} [\pi \cdot C(10, y_i) \cdot p_B^{y_i}(1-p_B)^{10-y_i} (1-\pi) \cdot C(10, y_i) \cdot p_C^{y_i}(1-p_C)^{10-y_i}] ]取对数之后[ \log L(\theta) \sum_{i1}^{n} \log [\pi \cdot f_B(y_i) (1-\pi) \cdot f_C(y_i)] ]问题就出在这个 (\log) 里面套了一个求和。如果是 (\log(\pi \cdot f_B) \log((1-\pi) \cdot f_C))那求导很轻松。但现在是 (\log(\pi \cdot f_B (1-\pi) \cdot f_C))求导之后各项纠缠在一起没有解析解。这就是隐变量带来的核心困难对数似然函数中出现了“对数-求和”结构导致无法直接求导得到闭式解。2.2 EM算法的破局思路先猜后迭代既然直接求不了那就换个思路。EM算法的核心思想可以用一句话概括既然不知道隐变量那我就先猜一组参数用这组参数去猜隐变量的分布再用猜出来的隐变量去更新参数反复迭代直到收敛。这就像鸡生蛋蛋生鸡的问题——你不知道参数就不知道隐变量不知道隐变量就估不了参数。EM算法的做法是随便从某个地方开始先假定一组参数然后交替优化每次都比上一次更好最终收敛到一个局部最优。具体到三硬币模型E步固定当前参数计算每轮使用B硬币的概率这就是隐变量的后验概率M步用E步算出的概率作为权重加权更新三枚硬币的正面概率听起来简单但背后的数学保证是什么为什么这样迭代一定能收敛这就不得不提Jensen不等式和Q函数了。2.3 Jensen不等式EM算法收敛性的数学基石Jensen不等式是EM算法推导中最重要的数学工具。它的表述很简洁对于凸函数 ( f )有 ( f(E[X]) \leq E[f(X)] )对于凹函数不等号反向。(\log) 函数是凹函数所以[ \log E[X] \geq E[\log X] ]EM算法利用这个不等式构造了一个对数似然函数的下界称为Q函数然后通过最大化这个下界来间接提升对数似然。每次迭代分两步E步在当前参数 ( \theta^{(t)} ) 下计算隐变量的后验分布 ( P(z|y, \theta^{(t)}) )构造Q函数 ( Q(\theta|\theta^{(t)}) E_{z|y,\theta^{(t)}}[\log P(y,z|\theta)] )M步最大化Q函数得到新参数 ( \theta^{(t1)} \arg\max_\theta Q(\theta|\theta^{(t)}) )可以证明每次迭代后对数似然函数单调不减。这个性质保证了EM算法的收敛性虽然可能收敛到局部最优而非全局最优但至少不会越迭代越差。3. 三硬币模型的完整数学推导3.1 符号定义与模型设定在开始推导之前先把符号理清楚。这是很多人看EM算法推导时第一个卡住的地方——符号一多就晕。符号含义取值范围( \pi )A硬币正面朝上的概率[0, 1]( p )B硬币正面朝上的概率[0, 1]( q )C硬币正面朝上的概率[0, 1]( y_i )第i轮观测到的正面次数0~10的整数( z_i )第i轮使用的硬币B或C离散隐变量( n )总轮数正整数模型生成过程每轮先扔A以 ( \pi ) 的概率选B以 ( 1-\pi ) 的概率选C然后用选中的硬币扔10次记录正面次数 ( y_i )。我们观测到 ( y_1, ..., y_n )但不知道每轮的 ( z_i )。3.2 E步计算隐变量的后验概率假设当前参数估计为 ( \theta^{(t)} (\pi^{(t)}, p^{(t)}, q^{(t)}) )。对于第 ( i ) 轮隐变量 ( z_i ) 的后验概率为[ P(z_i B | y_i, \theta^{(t)}) \frac{\pi^{(t)} \cdot C(10, y_i) \cdot (p^{(t)})^{y_i}(1-p^{(t)})^{10-y_i}}{\pi^{(t)} \cdot C(10, y_i) \cdot (p^{(t)})^{y_i}(1-p^{(t)})^{10-y_i} (1-\pi^{(t)}) \cdot C(10, y_i) \cdot (q^{(t)})^{y_i}(1-q^{(t)})^{10-y_i}} ]组合数 ( C(10, y_i) ) 在分子分母中约掉简化为[ \gamma_i P(z_i B | y_i, \theta^{(t)}) \frac{\pi^{(t)} \cdot (p^{(t)})^{y_i}(1-p^{(t)})^{10-y_i}}{\pi^{(t)} \cdot (p^{(t)})^{y_i}(1-p^{(t)})^{10-y_i} (1-\pi^{(t)}) \cdot (q^{(t)})^{y_i}(1-q^{(t)})^{10-y_i}} ]这个 ( \gamma_i ) 就是E步的核心输出它表示第 ( i ) 轮使用B硬币的“责任度”或“权重”。3.3 M步加权更新参数有了 ( \gamma_i ) 之后M步就是最大化Q函数。Q函数的形式为[ Q(\theta|\theta^{(t)}) \sum_{i1}^{n} [\gamma_i \log(\pi \cdot f_B(y_i)) (1-\gamma_i) \log((1-\pi) \cdot f_C(y_i))] ]分别对 ( \pi )、( p )、( q ) 求导并令导数为零得到更新公式更新 ( \pi )[ \pi^{(t1)} \frac{1}{n} \sum_{i1}^{n} \gamma_i ]这个很直观——A硬币正面概率的更新值就是所有轮次中使用B的平均责任度。更新 ( p )[ p^{(t1)} \frac{\sum_{i1}^{n} \gamma_i \cdot y_i}{\sum_{i1}^{n} \gamma_i \cdot 10} ]分子是加权正面次数分母是加权总次数比值就是B硬币的正面概率估计。更新 ( q )[ q^{(t1)} \frac{\sum_{i1}^{n} (1-\gamma_i) \cdot y_i}{\sum_{i1}^{n} (1-\gamma_i) \cdot 10} ]同理这是C硬币的正面概率估计。注意M步的更新公式本质上就是加权版的频率统计。如果没有隐变量( \gamma_i ) 就是0或1退化为普通的频率估计。EM算法的精妙之处就在于用软分配概率权重代替了硬分配0/1标签。3.4 收敛条件与迭代终止EM算法的迭代终止条件通常有三种对数似然函数的变化小于某个阈值如 ( 10^{-6} )参数的变化小于某个阈值达到最大迭代次数实际使用中我一般同时设置对数似然变化阈值和最大迭代次数双保险。因为有些情况下参数变化很小但对数似然还在缓慢爬升只看参数变化可能会过早停止。4. 手把手代码实现与逐行解析4.1 环境准备与数据生成我用Python来实现依赖只有numpy。先造一组模拟数据这样我们知道真实参数可以验证算法是否能收敛到正确答案附近。import numpy as np np.random.seed(42) # 真实参数 true_pi 0.4 # A硬币正面概率 true_p 0.7 # B硬币正面概率 true_q 0.3 # C硬币正面概率 n_rounds 100 # 轮数 n_tosses 10 # 每轮扔10次 # 生成数据 z_true np.random.binomial(1, true_pi, n_rounds) # 1表示用B0表示用C y np.zeros(n_rounds, dtypeint) for i in range(n_rounds): if z_true[i] 1: y[i] np.random.binomial(n_tosses, true_p) else: y[i] np.random.binomial(n_tosses, true_q) print(观测数据前20个:, y[:20]) print(使用B的轮数:, z_true.sum(), 使用C的轮数:, n_rounds - z_true.sum())这段代码生成了100轮数据每轮扔10次。真实情况下我们只能看到 ( y )看不到 ( z_{true} )。这里保留 ( z_{true} ) 只是为了最后验证算法效果。4.2 EM算法核心实现def em_algorithm(y, n_tosses10, max_iter1000, tol1e-6): n len(y) # 初始化参数随机初始化注意避开0和1 pi 0.5 p 0.6 q 0.4 log_likelihood_history [] for iteration in range(max_iter): # E步 # 计算每轮使用B的后验概率 gamma np.zeros(n) for i in range(n): # B硬币产生y[i]次正面的概率 prob_b (p ** y[i]) * ((1 - p) ** (n_tosses - y[i])) # C硬币产生y[i]次正面的概率 prob_c (q ** y[i]) * ((1 - q) ** (n_tosses - y[i])) # 后验概率贝叶斯公式 gamma[i] (pi * prob_b) / (pi * prob_b (1 - pi) * prob_c) # M步 # 更新pi pi_new np.mean(gamma) # 更新p p_new np.sum(gamma * y) / np.sum(gamma * n_tosses) # 更新q q_new np.sum((1 - gamma) * y) / np.sum((1 - gamma) * n_tosses) # 计算对数似然 log_likelihood 0 for i in range(n): prob_b (p ** y[i]) * ((1 - p) ** (n_tosses - y[i])) prob_c (q ** y[i]) * ((1 - q) ** (n_tosses - y[i])) log_likelihood np.log(pi * prob_b (1 - pi) * prob_c) log_likelihood_history.append(log_likelihood) # 检查收敛 if iteration 0 and abs(log_likelihood_history[-1] - log_likelihood_history[-2]) tol: print(f在第 {iteration} 次迭代收敛) break # 更新参数 pi, p, q pi_new, p_new, q_new return pi, p, q, log_likelihood_history # 运行EM算法 pi_est, p_est, q_est, ll_history em_algorithm(y) print(f\n估计结果:) print(fpi {pi_est:.4f} (真实值: {true_pi})) print(fp {p_est:.4f} (真实值: {true_p})) print(fq {q_est:.4f} (真实值: {true_q}))跑完这段代码你会看到估计值在真实值附近波动。由于EM算法只保证局部最优加上数据量有限估计值不会完全等于真实值但应该比较接近。4.3 对数似然曲线的可视化import matplotlib.pyplot as plt plt.figure(figsize(10, 5)) plt.plot(ll_history, linewidth2) plt.xlabel(Iteration) plt.ylabel(Log-Likelihood) plt.title(EM Algorithm Convergence Curve) plt.grid(True, alpha0.3) plt.tight_layout() plt.show()这条曲线应该是单调递增的前期上升快后期逐渐平缓。如果你看到曲线有下降那一定是代码写错了——EM算法的理论保证就是对数似然单调不减。4.4 初始化敏感性实验EM算法对初始值敏感这是它的一个显著缺点。我做了个实验用不同的初始值跑100次看看结果分布def run_multiple_inits(y, n_trials100): results [] for _ in range(n_trials): pi_init np.random.uniform(0.2, 0.8) p_init np.random.uniform(0.2, 0.8) q_init np.random.uniform(0.2, 0.8) pi, p, q, _ em_algorithm_with_init(y, pi_init, p_init, q_init) results.append((pi, p, q)) results np.array(results) return results # 统计不同初始化下的收敛结果 results run_multiple_inits(y) print(pi估计的均值:, results[:, 0].mean(), 标准差:, results[:, 0].std()) print(p估计的均值:, results[:, 1].mean(), 标准差:, results[:, 1].std()) print(q估计的均值:, results[:, 2].mean(), 标准差:, results[:, 2].std())实测下来大部分初始化都能收敛到差不多的结果但偶尔会陷入一个“交换解”——把B和C的角色对调了。这是因为三硬币模型本身存在标签交换对称性如果把 ( \pi ) 换成 ( 1-\pi )( p ) 和 ( q ) 互换似然函数值不变。所以EM可能收敛到 ( (\pi, p, q) ) 或 ( (1-\pi, q, p) )两者都合理。实操心得遇到标签交换问题可以通过固定 ( p q ) 或者固定 ( \pi 0.5 ) 来消除对称性。或者在多次随机初始化中选对数似然最大的那组解。5. 实操中踩过的坑与排查技巧5.1 数值下溢概率连乘的经典陷阱三硬币模型里每轮扔10次概率连乘还好。但如果把每轮次数改成100次、1000次( p^{y_i} ) 很快就会下溢到0。我一开始写代码的时候没注意跑大数实验直接得到一堆NaN。解决方案是取对数后再计算。E步的后验概率可以写成[ \gamma_i \frac{1}{1 \exp(\log(1-\pi) \log f_C(y_i) - \log \pi - \log f_B(y_i))} ]用log-sum-exp技巧避免指数溢出。具体实现def safe_gamma(y_i, pi, p, q, n_tosses): log_prob_b y_i * np.log(p 1e-300) (n_tosses - y_i) * np.log(1 - p 1e-300) log_prob_c y_i * np.log(q 1e-300) (n_tosses - y_i) * np.log(1 - q 1e-300) log_alpha np.log(pi 1e-300) log_prob_b log_beta np.log(1 - pi 1e-300) log_prob_c # log-sum-exp max_log max(log_alpha, log_beta) gamma np.exp(log_alpha - max_log) / (np.exp(log_alpha - max_log) np.exp(log_beta - max_log)) return gamma加 ( 10^{-300} ) 是为了防止 ( \log(0) )。这个技巧在处理大规模数据时是必须的。5.2 参数初始化别小看这个步骤我试过用全0.5初始化结果算法收敛到一个鞍点附近对数似然几乎不涨。后来改成随机初始化跑多次取最优效果稳定很多。初始化策略优点缺点推荐场景固定值(0.5, 0.5, 0.5)简单可复现容易陷入鞍点快速验证随机均匀探索性强结果不稳定多次运行取最优K-means预热起点较好实现复杂对精度要求高网格搜索覆盖全面计算量大小规模问题我的建议是至少跑10次随机初始化选对数似然最大的结果。多花几秒钟换来更可靠的估计。5.3 收敛判断别只看参数变化前面提过参数变化小不代表真的收敛了。我遇到过一次参数几乎不变了但对数似然还在缓慢上升。后来把收敛条件改成“对数似然变化小于阈值”才正确停止。另外最大迭代次数一定要设。有些病态初始化会导致算法在鞍点附近震荡不设上限会死循环。5.4 常见问题速查表问题现象可能原因排查方法解决方案对数似然下降代码逻辑错误检查E步和M步公式重新推导更新公式参数收敛到0或1数据极端或初始化不当检查数据分布加平滑项或换初始化结果每次不一样随机初始化固定随机种子多次运行取最优收敛极慢初始值离最优解远打印迭代曲线换更好的初始化标签交换模型对称性检查p和q是否互换固定pq或多次运行NaN出现数值下溢检查概率连乘改用对数域计算6. 从三硬币到更复杂的EM应用三硬币模型虽然简单但它包含了EM算法的全部核心要素隐变量、E步后验计算、M步加权更新、Jensen不等式保证收敛。把这套逻辑吃透之后再看高斯混合模型GMM、隐马尔可夫模型HMM的参数估计会发现套路完全一样。GMM的E步是计算每个样本属于每个高斯分量的后验概率M步是用这些后验概率加权更新均值、协方差和混合系数。HMM的Baum-Welch算法本质上也是EME步是前向-后向算法计算状态的后验概率M步是加权更新转移矩阵和发射矩阵。我个人的学习路径建议是先手推三硬币模型的EM再手推GMM的EM最后看HMM的Baum-Welch。每一步都自己写代码实现一遍不要直接调库。调库虽然方便但调完之后你对EM的理解还是停留在表面。还有一个值得注意的点EM算法得到的是局部最优解不是全局最优。在实际应用中如果问题规模不大可以多跑几次不同初始化如果问题规模很大可以考虑用随机初始化加早停策略或者用变分推断等替代方法。最后分享一个我在实际项目中用到的小技巧把EM算法的每次迭代参数和对数似然都记录下来画成动态图。这样不仅能直观看到收敛过程还能在结果异常时快速定位是哪一步出了问题。这个习惯帮我省了很多调试时间。