PyTorch中float32张量std/log_prob计算精度不一致的解决方法

问题重现 #

代码：

import torch
import numpy as np

a = torch.tensor(np.repeat(3, 10))
print(a)
print(a.log().std())
b = torch.tensor(np.repeat(3, 5))
print(b)
print(b.log().std())

输出：

tensor([3, 3, 3, 3, 3, 3, 3, 3, 3, 3])
tensor(1.2566e-07)
tensor([3, 3, 3, 3, 3])
tensor(0.)

原因分析 #

这是float32浮点精度的固有问题：虽然所有元素的理论log值完全相同，但torch.log的浮点运算可能引入极其微小的舍入误差；计算标准差时，不同长度的张量会导致误差累积方式不同，部分场景下误差被后续计算舍入为0，部分场景下保留了极小的非零值。

解决方案 #

不需要手动重写log_prob函数，有几种高效的处理方式：

1. 提前判断元素一致性（推荐） #

既然全相同元素的理论标准差为0，可先检查张量的log值是否全部相等（兼容浮点误差），直接返回0：

def stable_std(tensor):
    log_tensor = tensor.log()
    # 检查所有元素是否与第一个元素在浮点精度下一致
    if torch.allclose(log_tensor, log_tensor[0].expand_as(log_tensor)):
        return torch.tensor(0.0, dtype=tensor.dtype, device=tensor.device)
    return log_tensor.std()

# 测试
print(stable_std(a))  # tensor(0.)
print(stable_std(b))  # tensor(0.)

2. 截断极小值为0 #

如果不想提前判断，也可以在计算标准差后，将小于浮点精度阈值的值强制设为0：

log_a = a.log()
std_a = log_a.std()
# 以float32机器epsilon的5倍作为阈值，避免误判有效小值
threshold = 5 * torch.finfo(torch.float32).eps
std_a = torch.where(std_a < threshold, torch.tensor(0.0, dtype=torch.float32), std_a)
print(std_a)  # tensor(0.)

3. 临时提升中间计算精度（可选） #

如果允许临时提升计算精度（最终返回仍为float32），可以先将张量转为float64计算，再转回float32：

std_a = a.log().double().std().float()
print(std_a)  # tensor(0.)

这种方法能大幅降低浮点误差，但会增加少量计算开销，适合对精度敏感的场景。

关于log_prob的扩展 #

如果是处理分布log_prob输出的标准差，逻辑完全一致：只需检查log_prob的结果是否全部一致，或对结果进行极小值截断，不需要重写整个log_prob函数。

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