反向传播计算顺序详解：从计算图到PyTorch自动微分

反向传播难的不是链式法则难的是搞明白每一步该算什么、先算什么、后算什么。很多初学者对着公式能看懂一让自己从头推到尾就卡壳或者写自定义算子时梯度死活不对。这篇就用最朴素的方式把“反向传播的计算顺序”拆开讲清楚适合正在啃反向传播的学生、准备面试的候选人以及写 PyTorch 自定义扩展时对梯度流向拿不准的工程师。我先把结论放在前面反向传播的计算顺序本质上就是计算图的“反向拓扑序”。也就是说从输出节点开始沿着计算图逆着正向传播的方向一步一步往输入端走。每一步计算梯度时它依赖的上游梯度必须已经算好。这个顺序是唯一的不能乱。理解了这一点后面所有细节都会变得很顺。1. 反向传播到底在算什么1.1 先用一个最简单的计算图理解顺序假设有一个复合函数 y f(g(h(x)))。正向传播时你会先算 a h(x)再算 b g(a)最后算 y f(b)。反向传播时目标是求出 dy/dx还有中间那些参数可能对应的梯度。反推的顺序其实很固定第一步从输出开始令 dy/dy 1这是反向传播的起点。第二步算 dy/db f(b)。这里只用到局部导数 f(b) 和刚才的 dy/dy。第三步算 dy/da (dy/db) * g(a)。第四步算 dy/dx (dy/da) * h(x)。注意这个顺序你不可能先算 dy/da因为 dy/da 要用到 dy/db而 dy/db 还没算出来。每一步都依赖上一步的结果这就是“反向拓扑序”的含义。打个比方正向传播像从起点把一串链子往前甩出去反向传播则像从链子的末端开始一节一节把力传回来。如果你从中间开始往回拉中间那段可能会松脱根本传不到起点。1.2 链式法则与梯度累积顺序问题的起点当计算图里出现分支时顺序会更严格因为还要考虑梯度累加。举个例子某个中间节点 a 同时被 b 和 c 两个下游节点使用而 b、c 最终都影响 y。那根据全导数公式dy/da (dy/db) * (db/da) (dy/dc) * (dc/da)也就是说梯度要从两条路径分别传回来然后在节点 a 上相加。顺序上你必须等两条路径的上游梯度都算完再求和。这也是 PyTorch 里多次调用 backward 时梯度会累加的原因之一更是网络里参数共享、残差连接、多分支结构梯度计算的底层逻辑。很多人第一次看自动微分框架的实现会被“梯度累积”这个概念绕晕。其实它就是链式法则里“求和”二字的工程化表达。只要计算图有分支反向梯度就一定会在某个节点汇聚汇聚的方式就是相加。2. 为什么计算顺序不能乱反向拓扑序2.1 依赖关系决定了一切计算图本身是一个有向无环图DAG。正向传播时数据从输入节点流向输出节点每个节点的输出被后续节点消费。反向传播时梯度从输出节点流回输入节点每个节点要计算它对上游节点的梯度必须是它的下游梯度已经就绪。如果把计算顺序打乱比如某个节点的下游梯度还没算就直接去算上上游梯度那结果一定是错的甚至根本没法算。现实中你不会在手工推导时犯这种错但框架的自动微分引擎必须保证这点。PyTorch 里每次 forward 都会记录一个 grad_fn 链backward 时就是按照这个依赖关系反向调用的。这里有个容易误解的点反向计算顺序是“从后往前”但并不是简单地把链式法则倒过来念一遍。因为计算图可能很复杂有分支、有合并、有共享节点你得严格按依赖关系走。这个依赖关系就是反向拓扑排序。反向拓扑排序的规则很朴素每次选择一个入度为 0 的节点在反向图中入度表示“上游梯度还没来”的依赖数量计算它的梯度然后把它从图里移除更新相邻节点的入度。重复这个过程直到所有节点都处理完。2.2 从输出到输入逐层推进一个通用推导套路我平时手推梯度基本按下面四步走这套路可以用在任何网络上包括 Transformer画出计算图标清每个节点是什么运算记录正向传播时每个中间节点的值。从最终的损失节点开始令损失对自身的梯度为 1。沿着计算图逆序处理每个节点用链式法则计算它对其直接上游节点的梯度。如果某个节点有多条下游路径把它接受到的所有上游梯度求和得到总梯度再继续向上传。第 3 步里那句“用链式法则”拆开就是当前节点的上游梯度 × 当前运算对输入的局部梯度。注意这里“上游梯度”特指损失对当前节点输出的梯度而不是对输入。很多公式写成 dy/dx dy/du * du/dx其中 dy/du 就是上游梯度du/dx 是局部梯度。这套路看起来简单但真正动手时容易在形状和转置上卡住。我的建议是每算一步先确认要算的梯度是什么形状。比如输入是 (N, D) 的矩阵参数是 (H, D)那梯度一定也是 (H, D)。形状能帮你快速判断是不是转置放错了。2.3 中间变量的保存与释放显存和顺序的纠缠反向传播要用到正向传播时的中间变量这些变量也叫激活值。比如 ReLU 需要记住哪些位置是正数softmax 需要记住归一化后的概率矩阵乘法需要记住输入矩阵。框架会在 forward 时把必要的东西保存下来等 backward 时读取。这些中间变量占显存而且占得非常多。于是就有各种节省显存的技术核心思路就是“重新计算”。PyTorch 的梯度检查点checkpoint会把一段计算图的中间激活丢掉反向传播时再重新算一次前向来恢复这些值。恢复之后再按原来的反向顺序继续推梯度。代价是前向计算会多跑一遍时间变长但显存降下来了。我训练大模型时会用 checkpoint_wrapper 包住 transformer 的每一层配合混合精度能把 7B 模型塞进单卡训练。理解了计算顺序你才会明白它为什么是“用时间换显存”反传本身需要的中间值没有消失只是被延后到反传时重新生成。3. 手推一个两层全连接网络3.1 网络结构与符号约定空谈顺序没有感觉直接手推一个实际例子。我们构造一个非常小的网络输入 x 是二维向量假设一个样本 x [1.0, 0.5]。隐藏层z1 W1 x b1激活函数用 ReLU得到 h1。输出层z2 W2 h1 b2接 Softmax 得到预测概率 p用交叉熵损失。参数尺寸如下W1 是 2×2b1 是二维。W2 是 2×2b2 是二维。我随便给一组具体的值方便你跟着手算验证W1 [[0.2, -0.1], [0.3, 0.4]]b1 [0.1, -0.2]W2 [[0.5, -0.2], [0.1, 0.6]]b2 [0.0, 0.1]标签 y 的 one-hot 是 [1, 0]也就是正确类别是第 0 类。这个例子很小你可以拿笔算一遍也可以直接用 NumPy 验算。后面我列出的每一步都是按照反向传播应有的计算顺序排的。3.2 正向传播各节点的值先算正向顺便把中间值都记录下来因为反传都要用。z1 W1 x b1z1[0] 0.2×1.0 (-0.1)×0.5 0.1 0.25 z1[1] 0.3×1.0 0.4×0.5 (-0.2) 0.30h1 ReLU(z1) [0.25, 0.30]这里恰好全是正数mask 是 [1, 1]z2 W2 h1 b2z2[0] 0.5×0.25 (-0.2)×0.30 0.0 0.065 z2[1] 0.1×0.25 0.6×0.30 0.1 0.305softmax 概率exp(0.065)≈1.0672exp(0.305)≈1.3566总和≈2.4238p[0] ≈ 0.4403p[1] ≈ 0.5597交叉熵损失L -log(p[0]) ≈ 0.8206记住这些数字后面会用到。3.3 反传每一步的计算顺序反向传播的计算顺序从损失开始往前推第一步算损失对 z2 的梯度。Softmax 加上交叉熵有一个非常简洁的梯度形式dz2 p - y所以 dz2 [0.4403 - 1, 0.5597 - 0] [-0.5597, 0.5597]。这个“p 减 one-hot”的结果直接就是损失对 logits 的梯度不需要先算 softmax 的雅可比矩阵再乘交叉熵梯度。这一技巧值得记住很多框架实现里也是这么做的。第二步有了 dz2先算输出层参数的梯度并向上一层传递。dW2 outer(dz2, h1)也就是 dz2 的每个分量乘以 h1 的每个分量得到形状 (2,2)。dW2 [[-0.5597×0.25, -0.5597×0.30], [0.5597×0.25, 0.5597×0.30]] [[-0.1399, -0.1679], [0.1399, 0.1679]]db2 dz2 [-0.5597, 0.5597]dh1 W2^T dz2dh1[0] 0.5×(-0.5597) 0.1×0.5597 ≈ -0.2239 dh1[1] -0.2×(-0.5597) 0.6×0.5597 ≈ 0.4478注意这里计算顺序很重要dh1 是为继续往更上游传播而算的中间梯度必须在一层内先算好因为你下一步要用它来计算隐藏层的梯度。第三步计算损失对 z1 的梯度。因为 h1 ReLU(z1)所以 dz1 dh1 * mask其中 mask 是 ReLU 前向时记录的指示向量正数的位置为 1非正数为 0。这个例子里 mask 是 [1, 1]所以 dz1 [-0.2239, 0.4478]。如果某个神经元输出为负数反传梯度会直接变成 0这就是 ReLU 死亡问题的本质也解释了为什么反向传播顺序里需要保存 ReLU 的掩码。第四步继续算隐藏层参数的梯度以及损失对输入 x 的梯度。dW1 outer(dz1, x)dW1 [[-0.2239×1.0, -0.2239×0.5], [0.4478×1.0, 0.4478×0.5]] [[-0.2239, -0.1120], [0.4478, 0.2239]]db1 dz1 [-0.2239, 0.4478]dx W1^T dz1dx[0] 0.2×(-0.2239) 0.3×0.4478 ≈ 0.0896 dx[1] -0.1×(-0.2239) 0.4×0.4478 ≈ 0.2015到这里所有参数的梯度就算完了。你可以看到整个流程是“输出层 logits 梯度 → 输出层参数梯度 → 隐藏层输出梯度 → 隐藏层预激活梯度 → 隐藏层参数梯度 → 输入梯度”这个顺序像剥洋葱一样一层一层往回走。3.4 权重梯度的累积与代码对应上面是单个样本的梯度。实际训练时一个 batch 里有多个样本梯度会在 batch 维度上求和或求平均。PyTorch 的 loss.backward() 默认把梯度累积到 .grad 里所以每步优化前要调用 optimizer.zero_grad() 清零否则上一轮的梯度会累加进来。此外如果同一个权重被多个分支使用梯度还要跨路径累加。比如一个共享的 W 同时作用在两个输入上那么 dW 应该等于两条路径各自贡献的梯度之和。计算顺序上框架会先把所有路径的上游梯度求出来再统一累加到这个参数上。这也解释了为什么手动实现多分支网络时不能简单地在第一个分支算完就更新参数要等所有分支的梯度都回来再做优化。4. 自动微分框架如何决定计算顺序4.1 动态图与静态图的差异PyTorch 的动态图机制每次 forward 都会动态构建一张计算图。Tensor 上的 grad_fn 记录了它是通过什么运算产生的这些 grad_fn 之间通过 next_functions 连成一张反向图。调用 backward() 时PyTorch 会从输出节点的 grad_fn 出发做一次反向拓扑遍历依次调用每个 grad_fn 的 backward 方法把梯度传给输入。TensorFlow 的静态图机制则是先建立完整的计算图再通过自动微分生成对应的反向子图。由于图是固定的TensorFlow 可以对反向子图做更多编译优化比如算子融合、内存复用等。但用户侧看到的效果是一致的反向计算顺序永远是从 loss 往输入方向走。这里有一个我经常跟人强调的点动态图虽然灵活但每次迭代的图都可能不一样所以 backprop 的具体顺序依赖本次 forward 路径。PyTorch 的 Debug 工具 traceback 经常看到 grad_fn 链就是为了让你理解当前这一次反传的顺序。4.2 以 PyTorch 为例一次 forward/backward 中发生了什么下面这段代码很典型import torch x torch.tensor([1.0, 0.5], requires_gradTrue) w1 torch.tensor([[0.2, -0.1], [0.3, 0.4]], requires_gradTrue) b1 torch.tensor([0.1, -0.2], requires_gradTrue) w2 torch.tensor([[0.5, -0.2], [0.1, 0.6]], requires_gradTrue) b2 torch.tensor([0.0, 0.1], requires_gradTrue) z1 x w1.T b1 h1 torch.relu(z1) z2 h1 w2.T b2 loss torch.nn.functional.cross_entropy(z2.unsqueeze(0), torch.tensor([0])) loss.backward()调用 loss.backward() 时PyTorch 内部从 loss 的 grad_fn可能是 NllLossBackward出发沿着 next_functions 一直往前遍历。每经过一个节点就会调用该节点对应的反向函数。比如经过 z2 的 grad_fnAddBackward它会计算对 z2 输入的梯度经过 h1 的 grad_fnReluBackward它会用到前向保存的 mask计算 dz1。最终所有 requires_gradTrue 的叶子节点如 w1、b1、w2、b2、x都会收到梯度。用代码检查梯度也很简单print(x.grad) print(w1.grad) print(w2.grad)如果你手推的结果和框架给出的梯度对不上那一定是某一步转置或顺序错了。我经常这样自检。4.3 中间值保存、释放与 checkpoint 的影响自动微分框架里默认前向会保存反向需要的中间张量。比如上面的 h1它是计算 w2 梯度时需要的所以 ReLU 会把输出 h1 保存下来。什么时候释放在对应反向节点计算完毕后如果这个张量不再被任何后续反向节点使用框架才可以释放。这个释放顺序直接决定了显存峰值。你用 torch.cuda.max_memory_allocated() 观察一次大模型训练会发现显存占用在前向结束时达到峰值因为前向把所有中间值都保存了。反向开始后随着各层梯度算完中间值逐渐释放显存才慢慢降下来。gradient checkpointing 的思路是故意丢弃一部分中间值反传到该段时再按顺序重算一次前向得到中间值。这样前向需要保存的中间值数量大大减少代价是计算时间变长。我在训练长序列模型时经常用 checkpoint因为它能把显存需求砍掉一半甚至更多。4.4 自定义 autograd.Function 的梯度顺序如果你需要自己实现一个算子比如自定义一个带参数的模块就必须非常小心 backward 的返回顺序。PyTorch 规定backward 返回的梯度元组必须与 forward 的输入参数顺序一一对应。class MyLinear(torch.autograd.Function): staticmethod def forward(ctx, x, w, b): ctx.save_for_backward(x, w) return x w.T b staticmethod def backward(ctx, grad_output): x, w ctx.saved_tensors grad_x grad_output w # 注意形状匹配 grad_w grad_output.T x grad_b grad_output.sum(0) return grad_x, grad_w, grad_b这里 forward 的输入顺序是 x, w, b所以 backward 必须返回三个梯度顺序也是 grad_x, grad_w, grad_b。如果你把顺序写反了PyTorch 不会报错但梯度会错误地传给别的参数训练结果会莫名其妙地发散。这种 bug 特别隐蔽所以我的习惯是先写一个单步测试用 torch.autograd.gradcheck 验证自定义算子的梯度是否正确。5. 常见问题与排查技巧实录5.1 梯度为 NaN、为 0、不更新的典型原因很多人问为什么我的 loss 一开始就是 NaN第一个要查的就是反向传播这条链上的数值稳定性。Softmax 里如果直接算 log(softmax(z))z 很大时 exp(z) 会上溢log 里会变成 inf梯度自然就是 NaN。正确做法是用 log_softmax 再套 NLLLoss或者在交叉熵计算里做好数值修正。这是计算图表达方式带来的顺序问题先做一次稳定的前向反向梯度顺序才能正确且数值可靠。梯度为 0 的常见原因有三个ReLU 输出负区间的神经元永远不激活初始化导致所有激活值过大softmax 饱和梯度极接近 0或者某个中间节点被 detach 了梯度被硬生生切断。梯度不更新先确认三件事requires_grad 是否打开loss 是否含有梯度路径以及 optimizer.zero_grad() 是否调用了。如果忘了 zero_grad梯度会跨 batch 累积表现就是 loss 经常突然跳一下。这和反向传播本身的“梯度累积顺序”是直接相关的你连续 backward 了两次第一次的梯度还没清零第二次又加上去optimizer 拿到的是一个混合了多个 batch 的梯度。5.2 用 gradcheck 验证梯度计算顺序是否正确手推梯度容易翻车尤其是形状复杂的算子。PyTorch 提供一个数值梯度检查工具from torch.autograd import gradcheck input torch.randn(1, 2, dtypetorch.float64, requires_gradTrue) w torch.randn(2, 2, dtypetorch.float64, requires_gradTrue) b torch.randn(2, dtypetorch.float64, requires_gradTrue) test gradcheck(MyLinear.apply, (input, w, b), eps1e-6, atol1e-4) print(test)gradcheck 会用中心差分算数值梯度再和你实现的 backward 输出的梯度做对比。如果顺序错了数值梯度能正确反映“输入 e_i 对 loss 的影响”而你的解析梯度会张冠李戴两者对不上gradcheck 就会返回 False。这个工具对我排查自定义算子帮助很大。注意 gradcheck 对输入类型有要求通常要用 float64因为 float32 下数值差分误差会比较大可能误报。5.3 共享参数与多分支的梯度累积顺序权重共享在实际模型里很常见比如 Siamese 网络或者某些结构里同一个矩阵在不同时间步被反复使用。这样的参数一次 forward 会经历多次正向路径反传时它的梯度是各路径梯度之和。PyTorch 会自动累加但复杂度在于如果路径的深度不同梯度到达参数的顺序不同累加结果可能受到浮点误差影响但不会影响最终值。真正危险的是在有分支的情况下提前修改参数这会破坏后续路径的反向计算顺序导致梯度算错。所以正确做法永远是先完成一轮完整的 forward再完整 backward不要在一个层刚算完梯度就立刻更新参数。等所有梯度都累积到 .grad 里调用 optimizer.step()最后 zero_grad。这个顺序是深度学习训练流程的黄金法则。5.4 混合精度下的顺序与数值稳定性混合精度训练时前向用 FP16梯度通常用 FP32 累积。如果框架没有做好梯度累加的精度保持比如在 FP16 里连续加很多个小梯度可能因为数值范围不够出现精度损失甚至下溢。这也是为什么 AMPAutomatic Mixed Precision在反向传播的梯度累加处特别小心梯度暂时存在 FP32 buffer只在需要时才转成 FP16。另外在不同设备或不同 reduction 顺序下floating point 的求和顺序不同结果也可能差 1e-7。如果你复现论文时发现 loss 曲线和官方不完全一致先别急着怀疑自己代码看看是否是计算顺序带来的浮点误差。只要在合理的误差范围内就不用管。6. BPTT随时间反向传播的计算顺序6.1 把时间维展开成计算图顺序就清楚了循环神经网络 RNN包括 LSTM、GRU训练时的核心算法叫随时间反向传播简称 BPTT。它本质上就是把时间维展开成一个很深的计算图然后在展开后的图上做标准的反向传播。所以 BPTT 的计算顺序依然是“从后往前”只不过多了一个时间方向。展开到第三步的简单 RNNh1 tanh(W_h h0 W_x x1 b)h2 tanh(W_h h1 W_x x2 b)h3 tanh(W_h h2 W_x x3 b)每个时刻都会输出一个预测损失可能是各时刻损失之和。反传时先算最后一个时刻 h3 的梯度然后通过 W_h 回传到 h2在 h2 处还要加上第 2 时刻的输出损失带来的梯度再继续回传到 h1。整个过程就像一个深度为 T 的前馈网络只是隐藏层之间共享 W_h。6.2 一个三步展开的 RNN 手推思路用 T3 的例子如果损失 L L1 L2 L3那么反向顺序可以描述为计算 L3 对 h3 的梯度设为 g3。把 g3 通过 tanh 的局部导数传到 h3 的输入再通过 W_h 传到 h2得到一部分 g2。把 L2 对 h2 的梯度加到 g2 上这就是 h2 的总梯度。再通过同样的方式回传到 h1同时加上 L1 对 h1 的梯度。最后对 W_h 的梯度等于每个时间步贡献之和对 W_x 的梯度也类似。可以看到BPTT 的顺序和前馈网络的逐层回传没有本质区别只是把“层”换成了“时间步”。在实际实现中PyTorch 的 LSTM 会为每个时间步创建一次反向节点所以反向时会逐时间步回溯。6.3 长序列中的梯度爆炸与截断 BPTTBPTT 一个臭名昭著的问题是梯度爆炸或梯度消失。原因在于损失对较早时间步的梯度需要在时间方向上连乘一长串转移矩阵的导数。如果矩阵的特征值大于 1梯度指数增长很快就变成 NaN如果小于 1梯度指数衰减早期时间步几乎学不到东西。工程上常用的缓解手段有三个梯度裁剪对所有参数梯度的范数设一个上限比如 max_norm1.0。如果梯度范数超过这个值就整体缩放。这个操作要放在 optimizer.step() 之前顺序不能反。截断 BPTT不把整个长序列一次性反传到第一个时间步而是每 k 步截断一次只反传最近 k 步。PyTorch 的 LSTM 里通常用 detach 切断 h 的图这相当于人为控制反传深度。我处理上千长度序列时会设计类似策略否则显存和时间都吃不消。门控结构LSTM 和 GRU 通过遗忘门、输入门控制梯度流能在一定程度上缓解梯度消失。这也是为什么实际任务中很少用朴素 RNN。从顺序角度看梯度裁剪的位置也有讲究。应该在整个反向传播完成后对所有参数梯度统一做范数裁剪再更新参数。如果边反传边裁剪等价于改变了梯度量级可能影响收敛。最后分享一个我自己的习惯每次实现新模型我都会先手动推一遍至少一层网络的正向和反向顺序用一个小矩阵把每步梯度形状写下来再和框架输出对照。过程很枯燥但真的能帮你建立对“反向传播计算顺序”的直觉。哪怕后来用惯了 PyTorch 的自动微分一旦遇到梯度异常我还是会回到这个最朴素的手推流程沿着计算图一步步查通常很快就能定位问题。