PyTorch自定义backward实现MSE损失 线性回归示例出现损失爆炸

核心规则要求 #

PyTorch自定义torch.autograd.Function的反向传播逻辑必须严格遵循以下要求，否则会出现梯度计算错误：

• 返回梯度的顺序必须与forward方法的输入参数顺序完全一一对应，不需要计算梯度的输入需要返回None占位
• 每个返回的梯度张量形状，必须与对应的forward输入形状完全一致，不能存在维度缩放、拼接、重复等改变原形状的操作
• 必须使用backward方法入参grad_output完成链式法则计算：grad_output是上游计算节点传递回来的、对当前函数输出的梯度，需要与你计算得到的局部偏导逐元素相乘，不能省略这一步
• 梯度数值必须严格对应数学上的偏导关系，与PyTorch自动求导的结果完全对齐

原代码的错误定位 #

当前自定义backward的实现存在两处核心问题，直接导致梯度计算失真、训练崩溃：

1. 未使用传入的grad_output参与计算，违反链式法则要求
2. 梯度计算逻辑完全错误：
   标准MSE损失的数学定义为 $\text{MSE} = \frac{1}{N}\sum_{i=1}^N (y_{pred,i} - y_i)^2$，对输入y_pred的逐元素梯度为 $\frac{\partial \text{MSE}}{\partial y_{pred}} = \frac{2(y_{pred} - y)}{N}$。
   原代码中先对所有样本的梯度取全局均值（等价于把所有样本的梯度求总和再除以N），再通过repeat复制到100个位置，等效把每个样本位置的梯度值放大了100倍——原设定的学习率被隐式放大100倍，参数更新步长远超收敛范围，最终导致损失快速溢出为inf、nan值。
   同时代码中未除以样本数N，即使去掉repeat操作，梯度也会比正确值大N倍，依旧无法正常收敛。

修正后的可运行代码 #

from torch.autograd import Function

class MyMSELoss(Function):
    @staticmethod
    def forward(ctx, y_pred, y):    
        ctx.save_for_backward(y_pred, y)
        return ((y - y_pred)**2).mean()
    
    @staticmethod
    def backward(ctx, grad_output):
        y_pred, y = ctx.saved_tensors
        n_samples = y_pred.numel()
        # 逐元素计算局部梯度，乘以上游传来的梯度，形状与y_pred完全对齐
        grad_y_pred = 2 * (y_pred - y) / n_samples * grad_output
        # 标签y不需要计算梯度，返回None
        return grad_y_pred, None

替换原有自定义损失类后，训练收敛曲线会与PyTorch内置MSE损失、无自定义backward的MSE实现完全一致，最终收敛到0.0824左右。

内容的提问来源于stack exchange，提问作者Björn