PyTorch多轴方差计算的原理、分步实现及参数解析

一、底层数学原理 #

方差的核心是衡量数据离散程度，分两种核心计算方式：

• 有偏方差（总体方差）：先计算所有目标数据的均值$\mu$，再求每个数据与$\mu$差的平方的平均值，公式为：
  $$\sigma^2 = \frac{1}{N}\sum_{i=1}^N (x_i - \mu)^2$$
• 无偏方差（样本方差）：为修正样本估计的偏差，分母用$N-1$代替$N$，公式为：
  $$s^2 = \frac{1}{N-1}\sum_{i=1}^N (x_i - \mu)^2$$

当在PyTorch中对多个轴同时计算方差时，本质是把指定轴上的所有元素视为一个独立集合：比如你对形状(2,2,3)的张量的轴0和轴2计算方差，就是把每个轴1位置对应的所有元素（轴0的2个元素×轴2的3个元素=6个元素）单独提取出来，分别计算它们的方差，最终得到形状为(2,)的结果。

二、多轴计算时的张量形态变化 #

拿你的示例张量input_tensor（形状(2,2,3)）来说：

1. 计算input_tensor.mean(axis=(0,2), keepdims=True)时，keepdims=True会保留被缩减轴的维度（设为1），结果形状为(1,2,1)。这样做是为了让均值张量能和原张量(2,2,3)进行广播运算——原张量的轴0和轴2会自动匹配均值张量的对应维度（1扩展为2，1扩展为3），从而完成每个元素减均值的操作。
2. 如果不用keepdims=True，均值结果形状是(2,)，PyTorch也会自动广播，但分步计算时容易混淆维度对应关系，所以推荐用keepdims=True保持维度一致。

三、unbiased参数的影响 #

PyTorch的torch.var()默认unbiased=True，计算的是无偏方差；当设置unbiased=False时，计算的是有偏方差：

• 你的示例中，input_tensor.var(axis=(0,2), unbiased=False)等价于分步计算的((input_tensor - mean)**2).mean(axis=(0,2))，因为两者都是除以元素总数N（每个轴1位置对应6个元素）。
• 如果要和默认的torch.var(axis=(0,2))（unbiased=True）对应，分步计算时最后一步不能用mean，而是要除以$(N-1)$，比如：

  n_elements = input_tensor.shape[0] * input_tensor.shape[2]  # 2*3=6
  ((input_tensor - mean)**2).sum(axis=(0,2)) / (n_elements - 1)
  

四、分步实现的详细拆解 #

把你的分步代码逐行解释：

1. 计算均值：

   mean = input_tensor.mean(axis=(0,2), keepdims=True)
   

   这一步对轴0和轴2求均值，得到形状(1,2,1)的张量：

   tensor([[[ 5.],
            [ 8.]]])
   

   对应每个轴1位置的6个元素的均值：(1+2+3+7+8+9)/6=5，(4+5+6+10+11+12)/6=8。

2. 计算平方差：

   squared_diff = (input_tensor - mean)**2
   

   由于广播机制，原张量每个元素都减去对应轴1位置的均值，然后平方，结果形状仍为(2,2,3)，比如第一个元素(1-5)²=16，第二个(2-5)²=9，以此类推。

3. 计算平方差的均值：

   result = squared_diff.mean(axis=(0,2))
   

   对轴0和轴2求均值，得到形状(2,)的结果：

   tensor([8.6667, 8.6667])
   

   这个结果和input_tensor.var(axis=(0,2), unbiased=False)完全一致，因为都是除以6得到的有偏方差。

如果要验证无偏方差的分步实现，代码如下：

import torch

input_tensor = torch.tensor([
    [[1.0, 2.0, 3.0],
     [4.0, 5.0, 6.0]],
    [[7.0, 8.0, 9.0],
     [10.0, 11.0, 12.0]]
])

# 直接调用无偏方差
direct_unbiased = input_tensor.var(axis=(0,2))
print("直接计算无偏方差:", direct_unbiased)

# 分步实现无偏方差
mean = input_tensor.mean(axis=(0,2), keepdims=True)
squared_diff = (input_tensor - mean)**2
n_elements = input_tensor.shape[0] * input_tensor.shape[2]
step_unbiased = squared_diff.sum(axis=(0,2)) / (n_elements - 1)
print("分步计算无偏方差:", step_unbiased)

运行结果会显示两者都是tensor([10.4000, 10.4000])，因为除以的是6-1=5。

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