PyTorch反向传播梯度维度实操：线性层梯度形状匹配问题

所有梯度计算的底层遵循统一约束：损失对任意参数的梯度，形状永远和参数本身的形状完全一致。你遇到的形状不匹配问题，本质是1维特例下隐去了矩阵维度变换步骤，换到多维场景后没有补全链式法则对应的矩阵运算维度对齐操作，不需要修改已验证正确的mse_grad函数。 #

无偏置线性层梯度计算逻辑（单样本场景） #

先明确无偏置线性层的前向传播定义：

• 权重W形状固定为(out_features, in_features)
• 层输入x形状为(in_features,)
• 层输出y = W @ x，形状为(out_features,)，和你计算得到的l2层输出梯度(2,)形状完全匹配，这部分计算没有问题。

梯度计算按以下步骤做维度对齐即可：

• 计算本层权重梯度：损失对W的梯度dL/dW形状必须和W一致，计算公式为dL/dW = dL/dy.reshape(-1, 1) @ x.reshape(1, -1)。
  以你使用的Linear(2,2)层为例：形状为(2,)的输出梯度dL/dyreshape为(2,1)列向量，本层输入（也就是l1层的输出）形状为(2,)，reshape为(1,2)行向量，两个矩阵做乘法得到的结果形状恰好为(2,2)，和权重形状完全匹配。
• 计算向上一层传递的输出梯度：损失对本层输入x的梯度dL/dx形状和x一致，计算公式为dL/dx = W.T @ dL/dy。
  同样以Linear(2,2)层为例：形状为(2,2)的权重转置后为(2,2)，乘以形状为(2,)的dL/dy，得到的梯度形状为(2,)，和上一层输出形状完全匹配，可以直接用于上一层的梯度计算。

为什么Linear(1,1)场景没有出现形状问题 #

Linear(1,1)是多维线性层的特例：1维向量reshape为列/行向量后相乘结果仍为1x1的标量，和权重的(1,1)形状天然对齐，直接做标量乘法就能得到正确结果，感知不到显式的维度变换步骤，因此切换到多维输入输出场景时容易遗漏这步操作。 #

验证代码 #

你可以运行以下代码核对手动计算结果和PyTorch原生反向传播结果的一致性：

import torch
import torch.nn as nn

# 固定随机种子保证可复现
torch.manual_seed(42)
# 构造两层无偏置Linear(2,2)网络
l1 = nn.Linear(2, 2, bias=False)
l2 = nn.Linear(2, 2, bias=False)
# 随机生成单样本输入和标签
x = torch.randn(2)
target = torch.randn(2)

# PyTorch原生前向+反向传播
h1 = l1(x)
y = l2(h1)
loss = nn.functional.mse_loss(y, target)
loss.backward()

# 手动实现反向传播
# MSE对l2输出的梯度：dL/dy = 2*(y-target)/N，N为输出元素个数
dL_dy = 2 * (y - target) / 2
# 计算l2权重梯度
dL_dW2 = dL_dy.reshape(-1, 1) @ h1.reshape(1, -1)
# 计算传递到l1输出的梯度
dL_dh1 = l2.weight.T @ dL_dy
# 计算l1权重梯度
dL_dW1 = dL_dh1.reshape(-1, 1) @ x.reshape(1, -1)

# 核对结果
print(f"l2权重梯度手动计算与PyTorch结果一致：{torch.allclose(dL_dW2, l2.weight.grad)}")
print(f"l1权重梯度手动计算与PyTorch结果一致：{torch.allclose(dL_dW1, l1.weight.grad)}")

运行后两个判断都会输出True，全程不需要修改MSE梯度的计算逻辑。

如果是批量输入场景（输入形状为(batch_size, in_features)），只需要在计算权重梯度时对批量维度做平均归约即可，最终权重梯度形状依然和权重本身完全一致。

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