如何在NumPy中向量化EM算法的聚类协方差矩阵计算？

当然可以实现完全向量化的计算，彻底摆脱双重循环，利用NumPy的广播机制和批量矩阵运算来大幅提升效率。以下是具体实现思路和代码：

核心思路拆解 #

1. 计算数据点与各聚类均值的差值：通过广播将X（形状(n,2)）和mu（形状(k,2)）扩展为可广播的维度，得到每个数据点对应所有聚类的差值矩阵，形状为(n,k,2)。
2. 批量计算加权外积：将权重矩阵Z（形状(n,k)）扩展维度后，与差值的外积进行逐元素相乘，得到每个数据点对应每个聚类的加权外积矩阵，形状为(n,k,2,2)。
3. 求和得到分子：对所有数据点维度（n维度）求和，得到每个聚类的分子部分，形状为(k,2,2)。
4. 归一化处理：计算每个聚类的权重总和（Z的列和，形状(k,)），扩展维度后与分子做除法，得到最终的协方差矩阵估计。

完整向量化代码 #

import numpy as np

# 假设已有X(n,2), Z(n,k), mu(k,2)
n, k = X.shape[0], mu.shape[0]

# 步骤1：计算每个数据点与各聚类均值的差值 (n,k,2)
diff = X[:, np.newaxis, :] - mu[np.newaxis, :, :]
# 等价简化写法：diff = X[:, None, :] - mu[None, :, :]

# 步骤2：计算加权外积 (n,k,2,2)
weighted_outer = Z[:, :, np.newaxis, np.newaxis] * np.matmul(diff[:, :, :, np.newaxis], diff[:, :, np.newaxis, :])

# 步骤3：对数据点维度求和得到分子 (k,2,2)
num_sum = np.sum(weighted_outer, axis=0)

# 步骤4：计算分母并归一化
denominator = np.sum(Z, axis=0)[:, np.newaxis, np.newaxis]
sigma = num_sum / denominator

正确性验证 #

你可以通过对比向量化代码和原循环代码的输出结果来验证正确性：

# 原循环实现的sigma（假设已初始化sigma为(k,2,2)）
sigma_loop = np.zeros((k,2,2))
for kk in range(k):
    numsum = 0
    for ii in range(X.shape[0]):
        diff = (X[ii, :]-mu[kk, :]).reshape(-1, 1)
        numsum = numsum + Z[ii, kk]*np.matmul(diff, diff.T)
    sigma_loop[kk] = numsum / np.sum(Z[:, kk])

# 验证结果是否一致
print(np.allclose(sigma, sigma_loop))  # 应输出True

效率优势 #

向量化实现避免了Python层面的循环，所有运算都在NumPy的底层C代码中执行，对于大规模数据（比如n=10000以上），效率提升会非常明显。

内容的提问来源于stack exchange，提问作者Branden Keck