如何高效对多组均值向量与协方差矩阵应用np.random.multivariate_normal？

场景：
我拥有一个形状为(N, M)的numpy数组mean，包含N个M维均值向量；同时有一个形状为(N, M, M)的numpy数组covs，对应N个M×M的协方差矩阵。需要高效生成N个分别符合对应均值与协方差的随机向量，替代以下循环实现：

import numpy as np

def GenRandVar(mean, covs):
    rand_result = []
    for m, c in zip(mean, covs):
        rand_result.append(np.random.multivariate_normal(m, c))
    return np.array(rand_result)

高效实现方案 #

原循环的问题在于每次调用np.random.multivariate_normal都有Python层面的开销，当N较大时效率低下。我们可以利用多元正态分布的数学性质结合numpy的向量化操作来优化：

多元正态分布的随机向量可以表示为：X = μ + L @ Z，其中：

• μ是均值向量
• L是协方差矩阵Σ的Cholesky分解结果（满足Σ = L @ L.T）
• Z是标准正态分布（均值0，方差1）的随机向量

基于这个性质，我们可以实现完全向量化的版本：

import numpy as np

def GenRandVar(mean, covs):
    # 批量对协方差矩阵做Cholesky分解，得到下三角矩阵L
    L = np.linalg.cholesky(covs)
    # 生成标准正态随机数，形状与mean一致
    z = np.random.normal(size=mean.shape)
    # 批量计算L @ z，再加上均值得到结果
    return mean + np.einsum('nij,nj->ni', L, z)

关键优化点说明 #

• 批量Cholesky分解：np.linalg.cholesky原生支持对(N, M, M)形状的数组进行批量分解，无需循环，效率远高于逐次处理单个矩阵
• 批量矩阵-向量乘法：np.einsum通过爱因斯坦求和约定高效处理批量的矩阵与向量乘法，避免了Python循环的开销
• 完全向量化：整个流程在numpy的C底层执行，大幅提升大N场景下的运行速度

注意事项 #

• 所有协方差矩阵必须是正定矩阵，否则Cholesky分解会抛出错误（这和np.random.multivariate_normal的要求一致）
• 如果需要为每个分布生成K个样本，只需调整随机数形状为(N, K, M)，并修改einsum参数为'nij,nkj->nki'即可

内容的提问来源于stack exchange，提问作者Tommy Tsang