反向传播

一个把相机图像映射到关节指令的策略网络，参数 \(\boldsymbol\theta\) 可能有几千万个。训练它，就是用梯度下降（6.3 节，配合 6.4 节的小批量）去最小化损失 \(L(\boldsymbol\theta)\)，所以每一步都需要梯度 \(\nabla_{\boldsymbol\theta}L\)：每个参数一个数，几千万个数，整个训练要算上几十万次。有限差分每个参数至少要多做一次前向计算（6.11 节算过：\(d\) 维要 \(2d\) 次求值），根本付不起。反向传播只需要再花大约两次前向计算的工作量，就把所有偏导一次算完。它没有新的数学：就是 2.2 节的链式法则，按聪明的顺序执行，并且从不显式地构造 3.3 节的雅可比矩阵。

机器人里到处都在用它：

• 训练每一个学习出来的部件：抓取成功率分类器（16.1 节）、视觉运动策略、动力学模型（24.1 节）、价值函数。它们的训练循环里，"算梯度"这一步就是反向传播。
• 可微的运动学与动力学：把"末端离目标多远"对关节角求导，就能用梯度下降做逆运动学；轨迹优化（6.13 节）需要动力学对状态和控制的导数。
• 可微渲染：12.4 节的 NeRF 与高斯泼溅，是把像素误差穿过渲染器反向传播回场景参数。
• 对输入求导，而不只是对参数：GAN 的生成器更新（19.2 节）、能量模型的朗之万采样（19.5 节），都要损失对网络输入的梯度，反向传播顺便就给了。

这一节依次讲：(1) 一个能看到每个节点的数值与伴随量的计算图；(2) 把反向传播看成向量–雅可比积，以及反向模式为什么只花前向计算的一个小倍数；(3) 用梯度检验确认你的反向传播没写错；(4) 深层网络里梯度为什么会消失或爆炸，四种标准的修法；(5) PyTorch 的 autograd 到底做了什么。

把一段计算想成一条流水线：每个工位（节点）拿几个数，算出一个数。前向传播从输入跑到损失。现在对每个工位问一个问题：如果它的输出被轻轻推动一个很小的 \(\varepsilon\)，损失会动多少？这个灵敏度叫这个节点的伴随量（adjoint），沿用 2.2 节的记号 \(\bar v=\partial L/\partial v\)。损失自己的伴随量是 1。

要把"责任"往回传，一个工位只需要知道两件事：自己的伴随量，以及自己对每个输入的局部导数。后者只依赖它在前向传播时见过的那几个数，自己就能算。两者相乘，就是每个输入应分担的那一份；如果一个输入同时喂给好几个工位，它从各处分到的份额相加：

其中 \(i\to j\) 表示 \(v_j\) 直接用到了 \(v_i\)，求和遍历所有这样的节点。2.2 节在一条标量链上一步步走过这个递推。这一节有三件新东西：(1) 真实的计算图会分叉与汇合，所以那个求和是关键；(2) 节点是向量和矩阵，"局部导数"变成一个我们从不写出来的雅可比（第 4 节）；(3) 计算图是成千上万层浮点运算，这改变了会出什么错（第 5、6 节）。

我们的第一个计算图根本不是神经网络，而是一台平面两连杆臂的正运动学（\(l_1=1\) m，\(l_2=0.8\) m，与 8.3 节同一台臂），后面接上"末端到目标 \(\mathbf t\) 的距离平方"。这里的"参数"是关节角 \(\mathbf q=(q_1,q_2)\)，损失是 \(L=\tfrac12\|\mathbf p(\mathbf q)-\mathbf t\|^2\)。在这张图上做反向传播，得到的 \(\partial L/\partial\mathbf q\) 必须等于 \(J^\top\mathbf e\)，即雅可比的转置乘以末端误差 \(\mathbf e=\mathbf p-\mathbf t\)（3.3 节的 \(J\)）。训练神经网络用的也正是同一个算法。

上图是计算图：每个方框是一个基本运算，青色数字是前向传播算出的值，琥珀色数字是反向传播算出的伴随量 \(\bar v\)。按"单步"一步一步走：前 12 步是前向（从左到右），后 12 步是反向（从右到左），正在处理的边会标出紫色的局部导数。拖动机械臂（关节圆圈）或琥珀色的目标，整张图立刻重算。"梯度一步"执行 \(\mathbf q\leftarrow\mathbf q-\eta\,\bar{\mathbf q}\)，"持续下降"让它一直走下去。

默认姿态 \(\mathbf q=(0.5,1.1)\) rad、目标 \(\mathbf t=(-0.6,1.2)\) m 时，\(L=1.061\)，反向传播给出 \(\bar{\mathbf q}=(-1.793,-1.165)\)，与 \(J^\top\mathbf e\) 逐位相同。注意 \(q_1\)：它同时喂给 \(c_1=\cos q_1\)、\(s_1=\sin q_1\) 和 \(a=q_1+q_2\) 三个节点，所以它的伴随量是三份之和：\(-0.697+0.069-1.165=-1.793\)。漏掉任何一份，梯度就错了。

从默认姿态出发，\(\eta=0.3\) 走 20 步就让 \(L<10^{-4}\)。把 \(\eta\) 调到 1 左右，下降就不再收敛，而是来回振荡（6.3 节的步长问题）。目标在可达范围外时，手臂最后伸直指向目标，\(L\) 停在 \(\tfrac12(\|\mathbf t\|-1.8)^2\)：这是这个问题能做到的最好结果，梯度为零只说明"附近没有更好的"。目标离基座超过约 3 m 时，\(\eta=0.3\) 又会来回振荡：损失的曲率随距离增大，需要把 \(\eta\) 调小。

把一整层看成一个节点：\(\mathbf y=f(\mathbf x)\)，\(\mathbf x\in\mathbb R^n\)，\(\mathbf y\in\mathbb R^m\)，雅可比 \(J=\partial\mathbf y/\partial\mathbf x\) 是 \(m\times n\) 矩阵（3.3 节）。它有两种用法：

• 向前推一个切向量（JVP，Jacobian–vector product）：输入动一点 \(\delta\mathbf x\)，输出动 \(\delta\mathbf y\approx J\,\delta\mathbf x\)。
• 向后拉一个伴随量（VJP，vector–Jacobian product）：已知输出的伴随量 \(\bar{\mathbf y}=\partial L/\partial\mathbf y\)，输入的伴随量是 \(\bar{\mathbf x}=J^\top\bar{\mathbf y}\)。理由只有一行：\(\delta L=\bar{\mathbf y}^\top\delta\mathbf y=\bar{\mathbf y}^\top J\,\delta\mathbf x=(J^\top\bar{\mathbf y})^\top\delta\mathbf x\)。

对一层 \(\mathbf y=\sigma(W\mathbf x+\mathbf b)\)，记 \(\mathbf z=W\mathbf x+\mathbf b\)，雅可比是 \(J=\mathrm{diag}(\sigma'(\mathbf z))\,W\)，于是

一次逐元素乘法，加一次乘 \(W^\top\)。宽 4096 的一层，\(J\) 有一千六百多万个元素，可 VJP 的代价和前向的那次矩阵乘法一样。下图是一层 \(2\to2\) 的网络（\(W=\begin{bmatrix}1.3&-0.7\\0.5&1.2\end{bmatrix}\)，\(\mathbf b=(0.1,-0.2)\)）。左边是输入平面，右边是输出平面。拖动输入点 \(\mathbf x\)、青色的切向量 \(\delta\mathbf x\) 和右边琥珀色的伴随量 \(\bar{\mathbf y}\)。

读数里两个内积 \(\langle\bar{\mathbf y},J\delta\mathbf x\rangle\) 与 \(\langle\bar{\mathbf x},\delta\mathbf x\rangle\) 永远相等（默认设置下都是 0.831）：这就是"\(J^\top\) 是 \(J\) 的伴随"这句话的全部含义。把 \(\mathbf x\) 拖进两个单元都饱和的角落，\(\sigma'(\mathbf z)\) 趋于 0，不管 \(\bar{\mathbf y}\) 多大，\(\bar{\mathbf x}\) 都几乎为零：梯度消失在单独一层里就已经出现了。

整个网络的雅可比是各层雅可比之积 \(J=J_D\cdots J_2J_1\)（3.3 节的链式法则）。2.2 节说过，这个乘积可以从右往左乘，也可以从左往右乘：

• 前向模式：从输入端出发，每一遍推一个切向量 \(\dot{\mathbf x}\) 穿过整个网络（每层一次 JVP），得到 \(J\) 的一列。要整个 \(m\times n\) 的 \(J\)，跑 \(n\) 遍。
• 反向模式：先跑一次前向并存下每层的中间值，再从输出端拉一个伴随量 \(\bar{\mathbf y}\) 回来（每层一次 VJP），得到 \(J\) 的一行。要整个 \(J\)，跑 \(m\) 遍。

训练时损失是标量，\(m=1\)：一遍反向就得到对所有参数的梯度。下图在一个 \(n\to32\to32\to m\) 的 tanh 网络上真的把两种模式都跑一遍，逐个数乘加和逐元素运算，并检查两种方法得到的 \(J\) 是否一致。

\(m=1\) 时反向模式便宜 \(n\) 倍左右（\(n=48\) 时正好 48 倍）。前向模式的一遍 JVP 约是两次前向（原值和切向量各算一遍），反向模式的一遍 VJP 约是一次前向，所以当 \(2n<m+1\) 左右时前向模式更便宜，例如 \(n=2\)、\(m=8\)：4.1 对 9.4。读数的最后一行是对全部参数求一个标量损失的梯度：反向模式要 3.0 次前向的工作量（一次前向，加上 VJP 和 \(\bar W=\bar{\mathbf z}\mathbf h^\top\) 各约一次），前向模式则要对每个参数推一遍切向量。

手写的反向传播很容易错得很隐蔽：程序照常运行，损失甚至还在下降，只是更慢，或者停在更差的地方。验证方法是把它和数值差分逐项对比。对第 \(k\) 个参数，中心差分是

其中 \(\mathbf g\) 是反向传播给出的梯度，\(\tilde{\mathbf g}\) 是数值梯度，\(E\) 是相对误差（分母让它与梯度的尺度无关）。前向差分 \(\hl{fdd}{(L(\boldsymbol\theta+h\mathbf e_k)-L(\boldsymbol\theta))/h}\) 少算一次，但误差大得多。下图检验一个 \([2,8,8,1]\) 的 tanh 网络（105 个参数），它在 16 个样本上拟合 \(y=\cos q_1+0.5\cos(q_1+q_2)\)（一个两连杆臂关节 1 重力力矩形状的函数）。表格按参数组列出相对误差；你可以往反向传播代码里植入四种常见 bug，看检验能不能抓住它们，还可以把损失的计算切到 float32。

float64 下中心差分的最佳 \(h\) 约是 \(10^{-5.25}\approx6\times10^{-6}\)，误差 \(3\times10^{-11}\)；前向差分的最佳 \(h\) 约是 \(3\times10^{-8}\)，误差也只有 \(2\times10^{-8}\)。float32 下中心差分最好只能做到约 \(2\times10^{-5}\)（在 \(h\approx6\times10^{-3}\)），而 \(h=10^{-6}\) 时误差约 10%：一个完全正确的实现会被误判成有 bug。再把四种 bug 各试一遍，看哪些组出问题：

• 忘了乘 \(\sigma'\)：前两层（W₁、b₁、W₂、b₂）全错，线性输出层（W₃、b₃）仍然对，因为它没有激活函数。
• \(W_2\) 用错了转置：只有 \(W_2\) 下面的组（W₁、b₁）错，因为错的是从第 2 层往回传的 \(\bar{\mathbf h}\)。
• 偏置没在批上求和（只留了最后一个样本）：只有偏置组错。
• 损失的导数多乘了 2：所有组的误差都正好是 \(|2g-g|/(2g+g)=1/3\)。

出错的位置本身就是线索：检验结果告诉你 bug 在哪一层、哪一步。

深层网络是一条很长的链。第 \(l\) 层输入处的伴随量是 \(\bar{\mathbf h}_l=J_l^\top J_{l+1}^\top\cdots J_{D-1}^\top\,\bar{\mathbf h}_D\)：每过一层，乘一次雅可比的转置。2.2 节的 FIG 3 在标量链上见过这件事：每层的放大倍数相乘，略小于 1 就指数地消失，略大于 1 就指数地爆炸。对一层 \(\mathbf h_{l+1}=\phi(W_l\mathbf h_l)\)，这个倍数由三样东西决定：激活函数的斜率 \(\phi'\)（sigmoid 处处 \(\le1/4\)）、权重的尺度（初始化），以及有没有恒等通路。于是有四种标准修法：

1. 不饱和的激活函数：ReLU 在正半轴上斜率恒为 1。
2. 合适的初始化：Xavier/Glorot 初始化 \(\mathrm{Var}(w)=2/(n+m)\)（tanh，\(n,m\) 为输入、输出宽度；方阵层即 \(1/n\)），He 初始化 \(\mathrm{Var}(w)=2/n\)（ReLU），让每层的信号方差大致不变（第 7 节推导）。
3. 归一化：LayerNorm 把每个样本的 \(\mathbf u=W\mathbf h\) 在特征维上标准化，BatchNorm 把每个特征在一个批上标准化，让激活函数的输入一直待在有斜率的区间里。
4. 残差连接：\(\mathbf h_{l+1}=\mathbf h_l+F(\mathbf h_l)\)，雅可比变成 \(I+\partial F/\partial\mathbf h\)，梯度有一条直通的路（17.1 节讲 ResNet 结构本身）。

下图在一个宽 32 的深层网络上做一次前向与反向（32 个随机输入、随机目标），画出每一层的伴随量大小 \(\|\bar{\mathbf h}_l\|\) 与激活值的均方根。读数里的 \(\rho=\|\bar{\mathbf h}_0\|/\|\bar{\mathbf h}_D\|\) 是"到达第一层的梯度"与"最后一层的梯度"之比，绿色带是 \(0.1\le\rho\le10\)。权重按 \(\mathcal N(0,\,g^2c/n)\) 初始化，\(c=2\)（ReLU，He）或 \(c=1\)（tanh、sigmoid，Xavier），\(g\) 是增益滑块。

默认设置（tanh、20 层、Xavier）下 \(\rho\approx0.11\)，已经有些衰减。sigmoid 的斜率不超过 1/4，这里每层大约损失 4 倍：10 层 \(\rho\approx6\times10^{-7}\)，17 层起跌破 \(10^{-10}\)，20 层 \(\approx2\times10^{-13}\)；把增益调到 \(g=2\) 也只到 \(\rho\approx10^{-8}\)：更大的权重确实放大了反向的信号，但也把 sigmoid 推向饱和、斜率变小，每层的增益远小于 2 倍。ReLU 加 He 初始化（\(g=1\)）在 40 层时 \(\rho\approx0.16\)，仍在绿色带附近；增益只要偏一点，深度就会把它放大：\(g=1.5\)、30 层时 \(\rho\approx2.6\times10^5\)。只加残差的 ReLU 网络，40 层时 \(\rho\approx10^9\)：每一块都往"残差流"里加一份正的输出，激活值越积越大（青色线一路上升），梯度也跟着涨；再加上 LayerNorm，\(\rho\approx1.3\)。一个出人意料的结果：不加残差、只加 BatchNorm 的 60 层 ReLU 网络，\(\rho\approx2.4\times10^5\)，归一化并不自动保证梯度健康。

\(\rho\) 只是一次前向反向的诊断。下图真的训练三个同样深、同样宽（16）的网络去拟合一个机器人信号：一个关节的重力力矩加上电机的齿槽转矩纹波，\(y=\cos q+0.3\sin6q\)，64 个点，Adam，学习率 \(3\times10^{-3}\)，全批量。三者只差结构：普通（紫）、残差（青）、残差 + LayerNorm（琥珀）。

默认（sigmoid、16 块）下，普通网络的损失在 3000 步内一直停在 0.27，正是"预测平均值"的水平：前向信号同样被冲掉了：到最后一块时，隐藏状态在 64 个 \(q\) 之间的差别只有 \(10^{-11}\) 量级，输出层看到的几乎是一个常数，只能学平均值；反向时前几层几乎收不到梯度（\(\rho\) 在 \(10^{-10}\) 量级），而 Adam 按 \(\sqrt{v}+\epsilon\) 归一化步长，这么小的梯度被 \(\epsilon=10^{-8}\) 压住，参数几乎不动。残差网络在 500 步时约 0.02、残差 + LayerNorm 约 0.013，1000 步时分别约 0.019 与 0.0094。残差网络先在 0.0225 附近停留很久：它学会了平滑的 \(\cos q\)，纹波却还没学（约 1000 步后才开始明显往下走）；残差 + LayerNorm 更早学到一部分纹波。先学平滑部分、高频部分学得慢，就是 12.4 节提到的谱偏置。换成 ReLU（He 初始化），16 块的普通网络在 700 步左右也降到 \(10^{-4}\) 以下：激活函数和初始化选对了，这个深度还用不着残差。

机器人学习的代码几乎都建在 PyTorch 或 JAX 上（例如 Hugging Face 的 LeRobot 用 PyTorch，MuJoCo 的 MJX 与 Brax 用 JAX）。PyTorch 的 autograd 就是本节的反向模式，外加一点工程：

• 边算边建图（define-by-run）：对 requires_grad=True 的张量每做一次运算，就在图里记下一个节点（grad_fn），并保存反向需要的张量。Python 的 if 和循环都照常工作，因为图就是这一次实际执行的记录。
• loss.backward() 从 \(\bar L=1\) 出发，按逆拓扑顺序调用每个节点的 VJP，把结果累加到叶子张量的 .grad 里。因为是累加，每一步训练前都要 optimizer.zero_grad()。
• 反向完成后图默认被释放（要再反向一次需 retain_graph=True）；torch.no_grad() 里不建图，推理时省内存；.detach() 把一个张量从图中切开。
• 自定义运算写成 torch.autograd.Function，自己给出 forward 和 backward（即 VJP），用 torch.autograd.gradcheck 以 double 精度做第 5 节的检验。
• torch.func 提供 vjp、jvp、jacrev、jacfwd：前向模式与反向模式都有，按 FIG 3 的规则选。JAX 的 jax.grad、jax.vjp、jax.jvp、jacrev、jacfwd 是同样的分工。

FIG 1 那张图用 PyTorch 写只有几行，打印出的梯度就是图里的 \(\bar{\mathbf q}\)：

去掉张量与 GPU，autograd 的核心不过二十行。下面是一个标量版本（与 Karpathy 的 micrograd 同一个思路）：每个节点记住自己的父节点和局部导数，backward 先做拓扑排序，再按逆序把伴随量"推"给父节点。