PyTorch批量2x2协方差矩阵正定检查异常问题求助

PyTorch的MultivariateNormal在验证协方差矩阵正定时，依赖torch.linalg.cholesky的默认检查逻辑，对数值误差非常敏感。你遇到的情况是批量构造的协方差矩阵中，个别矩阵因数值舍入导致特征值接近0（~3.17e-10），批量检查时触发失败，但单独检查时因计算精度的细微差异通过。直接修改检查逻辑为半正定并不稳妥，因为半正定矩阵可能无法完成Cholesky分解，导致后续采样、log概率计算出错。以下是更优的解决方案：

1. 添加极小对角扰动（最常用方案） #

给协方差矩阵的对角元素添加一个极小的正数，确保所有特征值严格大于0，同时不显著改变原矩阵的统计性质。

# 构造匹配批量维度的单位矩阵扰动
epsilon = 1e-6  # 根据数据尺度调整，比如原协方差量级是1e-3则用1e-9
eye = torch.eye(2, device=cov_matrix.device)
# 扩展维度匹配批量形状(117,100,20,2,2)
perturbation = epsilon * eye[None, None, None, :, :]
# 应用扰动
cov_matrix = cov_matrix + perturbation

注意：epsilon的选择要平衡数值稳定性和结果准确性，避免过大扰动扭曲原数据分布。

2. 带容错的Cholesky分解（精准处理异常矩阵） #

手动执行Cholesky分解，仅在分解失败时对对应矩阵添加扰动，而非全局修改：

def safe_cholesky(cov, epsilon=1e-6):
    # 尝试正常分解
    try:
        return torch.linalg.cholesky(cov)
    except RuntimeError:
        # 对分解失败的矩阵添加对角扰动后重试
        cov_shape = cov.shape
        eye = torch.eye(cov_shape[-1], device=cov.device)
        # 扩展维度匹配批量
        perturbation = epsilon * eye.view((1,) * (len(cov_shape)-2) + cov_shape[-2:])
        cov = cov + perturbation
        return torch.linalg.cholesky(cov)

# 用处理后的Cholesky因子初始化MultivariateNormal
scale_tril = safe_cholesky(cov_matrix)
mvn = torch.distributions.MultivariateNormal(mean, scale_tril=scale_tril)

这种方式仅修正确实存在数值问题的矩阵，最小化对整体数据的影响。

3. 提升数值精度 #

将张量从float32转为float64，更高的精度可以减少数值舍入误差，可能让接近零的特征值在计算时满足正定要求：

# 转换协方差矩阵和均值为double精度
cov_matrix = cov_matrix.double()
mean = mean.double()
# 初始化MultivariateNormal
mvn = torch.distributions.MultivariateNormal(mean, covariance_matrix=cov_matrix)

注意：该方案会增加内存占用和计算时间，需根据硬件资源权衡。

4. 精准定位并修正异常矩阵 #

先批量计算所有协方差矩阵的特征值，定位特征值过小的矩阵并单独修正：

# 计算对称矩阵的特征值（eigvalsh针对对称矩阵更高效）
eigvals = torch.linalg.eigvalsh(cov_matrix)
# 找出任意特征值小于阈值的矩阵位置
threshold = 1e-8
mask = (eigvals < threshold).any(dim=-1)
# 对异常矩阵添加对角扰动
eye = torch.eye(2, device=cov_matrix.device)
cov_matrix[mask] = cov_matrix[mask] + threshold * eye

该方案精准修正问题矩阵，完全不影响正常矩阵的分布。

不建议修改MultivariateNormal的检查逻辑为半正定，因为半正定矩阵无法保证Cholesky分解的可行性，会导致后续计算出现错误。优先尝试添加对角扰动或带容错的Cholesky分解，这两种方案在数值稳定性和结果准确性之间的平衡最优。

内容的提问来源于stack exchange，提问作者Thuan Nguyen