PyTorch计算聚类经验方差出现负值的问题排查与优化

负值出现的核心原因 #

• 浮点精度误差：使用float32等低精度浮点类型时，减法操作（如x²均值 - 均值平方）会因精度损失出现极小负值，尤其当方差本身接近0时。
• 空聚类/样本数为1的处理失误：
  • 原代码给n_K加1e-10后，若聚类无样本（n_K原本为0），计算分母n_K-1会得到-1e-10，除以负数会让非负的分子变为负值。
  • 当聚类仅含1个样本时，分母n_K-1=0，直接计算会得到无穷大或NaN，后续处理易产生异常值。
• 潜在的维度错误：你描述enc维度为(N,D)，但K是聚类数，正常独热编码enc应是(N,K)，若D≠K会导致矩阵乘法逻辑错误，间接引发数值异常。

优化实现方案 #

以下是修正后的代码，同时优化了计算效率与数值稳定性：

方案一：中心化后平方求和（直观版） #

import torch

def get_mean_var(x, enc):
    # 确认维度：x(N,D)，enc(N,K)
    K = enc.size(1)
    device = x.device

    # 统计每个聚类的样本数，形状(K,1)
    n_K = enc.sum(dim=0, keepdim=True).t()
    # 处理空聚类（设样本数为1避免除零）和单样本聚类
    n_K = torch.where(n_K == 0, torch.tensor(1.0, device=device), n_K)
    
    # 计算聚类均值 (K,D)
    mu_e = enc.t() @ x / n_K
    # 每个样本对应的聚类均值 (N,D)
    encoded_mu = enc @ mu_e
    
    # 计算中心化平方项并求和
    centered_sq = (x - encoded_mu) ** 2
    var_e = enc.t() @ centered_sq
    
    # 计算方差：单样本聚类方差为0，否则除以n_K-1
    var_e = torch.where(n_K == 1, torch.tensor(0.0, device=device), var_e / (n_K - 1))
    # 强制非负并添加极小值避免后续计算问题
    var_e = torch.clamp(var_e, min=1e-10)
    
    return mu_e, var_e

方案二：平方和减均值平方（高效版） #

这种方式减少了一次矩阵乘法，计算更高效，数值稳定性也更好：

import torch

def get_mean_var(x, enc):
    device = x.device
    # 统计每个聚类的样本数 (K,1)
    n_K = enc.sum(dim=0, keepdim=True).t()
    n_K = torch.where(n_K == 0, torch.tensor(1.0, device=device), n_K)
    
    # 计算每个聚类的x总和与x²总和
    sum_x = enc.t() @ x
    sum_x_sq = enc.t() @ (x ** 2)
    
    # 计算均值与方差
    mu_e = sum_x / n_K
    var_e = (sum_x_sq - sum_x * mu_e)  # 等价于 sum((x - mu)^2)
    # 处理单样本聚类，避免除以0
    var_e = torch.where(n_K == 1, torch.tensor(0.0, device=device), var_e / (n_K - 1))
    var_e = torch.clamp(var_e, min=1e-10)
    
    return mu_e, var_e

关键优化点说明 #

1. 特殊聚类处理：对空聚类（样本数0）设为1避免除零，对单样本聚类直接设方差为0，符合统计逻辑。
2. 数值稳定性保障：用torch.clamp强制方差非负，解决浮点精度导致的极小负值；优先使用float64类型计算可进一步降低精度损失。
3. 效率提升：方案二避免了encoded_mu的计算，减少了矩阵乘法操作，适合大数据量场景。

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