如何在numpy中向量化矩阵和的分数（期望最大化）？

Question

我正在尝试使用numpy对一个关于二维高斯分布的期望最大化/聚类方程进行向量化处理。下面我会在问题的最后附上我简单的做法：

为了让大家更好理解，下面是一些变量和维度的定义：

• = 数据点的索引（比如说从1到1000）
• = 聚类的索引（比如说1到3）
• = 条件概率，表示数据点 属于聚类 的概率（范围在0到1之间）
• = 数据点 的值（形状为(2,)）
• = 当前聚类 的多变量均值估计（形状为(2,)）

最终的结果是一个分子，它是形状为(2, 2)的矩阵的和，而分母是一个标量。最后的值是一个(2, 2)的协方差矩阵估计。这个过程需要对每个“k”的值（1, 2, 3）都进行一次。

我已经为其他值实现了向量化的方法，定义了以下numpy数组：

• = 每个数据点和聚类的估计概率值
• = 多变量数据矩阵
• = 聚类均值的估计

我的简单代码如下：

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[kk] = numsum / np.sum(Z[:, kk])

长话短说，有没有更好的方法来实现这个？

Answer 1

下面的代码应该可以正常工作：

diff = X[np.newaxis, :, :] - mu[:, np.newaxis, :]  # kxnx2
numsum = np.matmul(Z.T[:, np.newaxis, :] * diff.transpose(0, 2, 1), diff)  # kx2x2
sigma_proposed = numsum / Z.sum(axis=0)[:, np.newaxis, np.newaxis]  # kx2x2

总的来说，我用以下代码进行了检查：

import numpy as np

n, k = 1000, 3

# Create some data
rand = np.random.default_rng(seed=0xC0FFEE)  # For reproducibility
Z = rand.uniform(size=(n, k))
X = rand.normal(size=(n, 2))
mu = rand.normal(size=(k, 2))
sigma = np.zeros((k, 2, 2))

# Code from question
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[kk] = numsum / np.sum(Z[:, kk])
    
# Proposed
diff = X[np.newaxis, :, :] - mu[:, np.newaxis, :]  # kxnx2
numsum = np.matmul(Z.T[:, np.newaxis, :] * diff.transpose(0, 2, 1), diff)  # kx2x2
sigma_proposed = numsum / Z.sum(axis=0)[:, np.newaxis, np.newaxis]  # kx2x2

assert np.allclose(sigma, sigma_proposed)

Answer 2

你可以使用 np.einsum 这个函数：

d = X - mu[:,None]
np.einsum('ijk,ijm,ji->imk', d, d, Z/Z.sum(0, keepdims=True))

---