逐样本累加与全批量梯度计算结果差异的原因探究

从数学理论上，逐个遍历MNIST样本计算单个梯度后累加平均，与直接用全量数据集计算梯度的结果应该完全一致，但实际运行两段代码后，梯度范数平方差异显著，核心原因如下：

1. 代码中的直接bug #

代码片段2中计算梯度范数平方时，未初始化norm_sq为0，直接执行norm_sq += torch.norm(param.grad) ** 2会导致变量初始值为未定义的垃圾值，这是最直观的差异来源。

2. 浮点数运算的累积精度误差 #

• 样本顺序影响：代码片段2使用shuffle=True打乱了数据集顺序，而片段1按原始顺序遍历。虽然数学上梯度的线性求和与顺序无关，但浮点数的累加存在精度限制，60000个小梯度的不同累加顺序会产生累积误差，最终导致平均后的梯度出现细微偏差。
• 梯度计算路径差异：逐个样本计算时，每次反向传播的是单个样本的梯度，经过60000次累加再平均；而全量计算时，直接基于整个批次的损失计算梯度，两种路径的浮点数舍入误差累积方式不同，当样本量极大时，这种差异会被放大。

3. ReLU激活的状态偏差 #

在逐个样本前向传播时，每个样本的ReLU激活是独立计算的；而全量数据前向传播时，是批量计算激活。由于浮点数精度的细微差异，部分神经元的激活状态（是否被ReLU置0）可能在两种计算方式下出现不一致：比如某个神经元在单个样本计算时输出刚好略大于0，而批量计算时由于数值精度累积变成略小于0，反之亦然。这种激活状态的差异会反向传播到梯度，最终导致整体梯度的差异。

代码片段1：逐个样本计算梯度并累加 #

import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch import nn

# 下载MNIST数据集
train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)

# 定义LeNet-300-100-10模型
class LeNet(torch.nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.fc1 = torch.nn.Linear(28 * 28, 300)
        self.fc2 = torch.nn.Linear(300, 100)
        self.fc3 = torch.nn.Linear(100, 10)
        
    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 固定随机种子初始化模型
torch.manual_seed(0)
model = LeNet()

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

# 初始化梯度累加容器
grad_sum = [torch.zeros_like(param) for param in model.parameters()]

criterion = nn.CrossEntropyLoss()

for i, (images, labels) in enumerate(train_dataset):
    # 前向传播
    images, labels = images.to(device), labels.to(device)
    outputs = model(images)
    # 计算损失
    loss = criterion(outputs, labels.unsqueeze(0))   
    # 清空梯度
    model.zero_grad()
    # 单个样本反向传播
    loss.backward()
    # 累加梯度
    for i, param in enumerate(model.parameters()):
        grad_sum[i] += param.grad

# 计算平均梯度
grad = [grad_sum[i] / len(train_dataset) for i in range(len(grad_sum))]

# 计算梯度范数平方
norm_sq = 0
for p_grad in grad:
      norm_sq += torch.norm(p_grad) ** 2

print(norm_sq)

代码片段2：全量数据集直接计算梯度 #

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor

# 定义LeNet-300-100-10模型
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 300)
        self.fc2 = nn.Linear(300, 100)
        self.fc3 = nn.Linear(100, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 加载MNIST数据集
train_dataset = MNIST(root='./data', train=True, download=True, transform=ToTensor())

# 创建全量数据加载器（注意这里shuffle=True会打乱顺序）
train_loader = DataLoader(train_dataset, batch_size=len(train_dataset), shuffle=True)

# 选择设备（GPU优先）
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 固定随机种子初始化模型
torch.manual_seed(0)
model = LeNet().to(device)

# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 前向传播与反向传播
for images, labels in train_loader:
    images, labels = images.to(device), labels.to(device)
    
    # 计算输出与损失
    outputs = model(images)
    loss = criterion(outputs, labels)

    # 清空梯度
    model.zero_grad()
    
    # 反向传播计算全量梯度
    loss.backward()

# 计算梯度范数平方（已修正原代码未初始化的问题）
norm_sq = 0
for param in model.parameters():
    norm_sq += torch.norm(param.grad) ** 2

print(norm_sq)

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