反向传播与自动微分

训练神经网络时，真正想得到的不是某层输出本身的导数，而是损失函数对每个可训练参数的梯度：

$\dfrac{\partial L}{\partial W},\qquad\dfrac{\partial L}{\partial b}$

梯度告诉优化器：参数稍微增大时，损失会朝哪个方向变化、变化多快。反向传播（backpropagation）就是高效计算这些梯度的方法。

本文不把它当作一组需要背诵的规则，而是从一个可手算的计算图出发，走到一个三层小网络，解释链式法则、Jacobian 和自动微分分别在做什么。

1
2
前向传播：输入 -> 中间结果 -> 预测 -> 损失
反向传播：损失 -> 梯度 -> 每层参数的梯度

1. 先区分前向和反向：训练到底在算什么？

考虑最简单的模型：

$\hat y=wx$

设输入 x=3，当前权重 w=2，正确答案 y=10。使用平方误差：

$L=(\hat y-y)^2$

前向计算为：

$\hat y=2\times3=6$

$L=(6-10)^2=16$

此时不是直接凭感觉把 w 改大，而是求：

$\dfrac{\partial L}{\partial w}=2(wx-y)x=2(6-10)\times3=-24$

梯度为负，说明增大 w 会让损失下降。若学习率为 0.01：

$w_{new}=w-0.01\dfrac{\partial L}{\partial w}=2.24$

神经网络只是在这里多了大量中间层和参数。反向传播处理的就是：怎样不遗漏任何路径地算出每个参数对应的梯度。

2. 一个计算图：梯度如何从输出倒着传回输入？

先看一个标量例子：

$y=\ln(x_1)+x_1x_2-\sin(x_2)$

令：

$x_1=2,\qquad x_2=5$

为了让每一步只有一个简单运算，将它拆成中间变量：

1
2
3
4
5
6
v1 = ln(x1)
v2 = x1 * x2
v3 = sin(x2)
v4 = v1 + v2
v5 = v4 - v3
y  = v5

前向计算的数值如下：

2.1 从 $\partial y/\partial y=1$ 开始

输出对自己的导数总是 1：

$\dfrac{\partial y}{\partial v_5}=\dfrac{\partial y}{\partial y}=1$

又因为：

$v_5=v_4-v_3$

所以：

$\dfrac{\partial y}{\partial v_4}=1\times\dfrac{\partial v_5}{\partial v_4}=1\times1=1$

这里的 1 不是 v4=10.693 这个数值，而是“只把 v4 增加一点时，y 也增加同样一点”的变化比例。

同理：

$\dfrac{\partial y}{\partial v_3}=1\times\dfrac{\partial(v_4-v_3)}{\partial v_3}=-1$

继续穿过加法节点：

$\dfrac{\partial y}{\partial v_1}=\dfrac{\partial y}{\partial v_4}\dfrac{\partial v_4}{\partial v_1}=1\times1=1$

$\dfrac{\partial y}{\partial v_2}=\dfrac{\partial y}{\partial v_4}\dfrac{\partial v_4}{\partial v_2}=1\times1=1$

2.2 分叉处为什么必须相加？

x1 同时影响 v1=ln(x1) 和 v2=x1*x2，所以它到输出有两条路径：

1
2
x1 -> v1 -> v4 -> v5 -> y
 \-> v2 -> v4 -> v5 -> y

两条路径的贡献必须相加：

$\dfrac{\partial y}{\partial x_1}=\dfrac{\partial y}{\partial v_1}\dfrac{\partial v_1}{\partial x_1}+\dfrac{\partial y}{\partial v_2}\dfrac{\partial v_2}{\partial x_1}$

$=1\times\dfrac12+1\times5=5.5$

x2 也有两条路径：

$\dfrac{\partial y}{\partial x_2}=\dfrac{\partial y}{\partial v_2}\dfrac{\partial v_2}{\partial x_2}+\dfrac{\partial y}{\partial v_3}\dfrac{\partial v_3}{\partial x_2}$

$=1\times2+(-1)\times\cos5\approx1.716$

这就是反向传播最基本的工作方式：从输出开始，沿边应用局部导数；在一个变量的多条后续路径汇合时，将梯度贡献相加。

3. 为什么向量层要提 Jacobian？

上面的每个变量都是一个数，普通链式法则就足够了。但神经网络的一层通常是向量到向量：

$y=f(x),\qquad x\in\mathbb R^n,\quad y\in\mathbb R^m$

此时“$\partial y/\partial x$”不是一个数，因为每个输入分量都可能影响每个输出分量。Jacobian 把所有局部影响放进一张表：

$J_{y,x}=\dfrac{\partial y}{\partial x}=\begin{bmatrix}\dfrac{\partial y_1}{\partial x_1}&\cdots&\dfrac{\partial y_1}{\partial x_n}\\\vdots&\ddots&\vdots\\\dfrac{\partial y_m}{\partial x_1}&\cdots&\dfrac{\partial y_m}{\partial x_n}\end{bmatrix}$

它表达的是局部近似：

$\Delta y\approx J_{y,x}\Delta x$

也就是输入出现一个很小的改动时，输出会怎样变化。

训练里最终的 (L) 是标量。若已知后一层传回的梯度：

$g_y=\dfrac{\partial L}{\partial y}$

则前一层的梯度为：

$g_x=\dfrac{\partial L}{\partial x}=J_{y,x}^Tg_y$

这就是向量形式的链式法则。

4. 三层数值网络：从前向到参数梯度

下面用一个没有激活函数黑箱的三层例子。第二层只做逐元素平方：

$x\xrightarrow{W_1}z\xrightarrow{\text{square}}h\xrightarrow{W_2}q\xrightarrow{\text{squared error}}L$

定义：

$z=W_1x,\qquad h=z\odot z,\qquad q=W_2h,\qquad L=\dfrac12(q-t)^2$

取具体数值：

$x=\begin{bmatrix}1\\2\end{bmatrix},\quad W_1=\begin{bmatrix}1&1\\2&-1\end{bmatrix},\quad W_2=\begin{bmatrix}1&-2\end{bmatrix},\quad t=5$

4.1 前向传播

第一层：

$z=W_1x=\begin{bmatrix}1&1\\2&-1\end{bmatrix}\begin{bmatrix}1\\2\end{bmatrix}=\begin{bmatrix}3\\0\end{bmatrix}$

第二层：

$h=z\odot z=\begin{bmatrix}9\\0\end{bmatrix}$

第三层：

$q=W_2h=\begin{bmatrix}1&-2\end{bmatrix}\begin{bmatrix}9\\0\end{bmatrix}=9$

损失：

$L=\dfrac12(9-5)^2=8$

4.2 先从损失传给预测值

$g_q=\dfrac{\partial L}{\partial q}=q-t=4$

当前预测值 9 大于目标 5，所以继续增大 (q) 会让损失变大，梯度为正。

4.3 穿过线性层 $q=W_2h$

逐元素写为：

$q=1\cdot h_1-2\cdot h_2$

因此：

$g_h=\dfrac{\partial L}{\partial h}=W_2^Tg_q=\begin{bmatrix}1\\-2\end{bmatrix}\times4=\begin{bmatrix}4\\-8\end{bmatrix}$

同一层中，参数 (W_2) 的梯度为：

$\dfrac{\partial L}{\partial W_2}=g_qh^T=4\begin{bmatrix}9&0\end{bmatrix}=\begin{bmatrix}36&0\end{bmatrix}$

4.4 穿过逐元素平方 $h=z\odot z$

对每一项有：

$h_i=z_i^2,\qquad\dfrac{\partial h_i}{\partial z_i}=2z_i$

所以：

$g_z=g_h\odot(2z)=\begin{bmatrix}4\\-8\end{bmatrix}\odot\begin{bmatrix}6\\0\end{bmatrix}=\begin{bmatrix}24\\0\end{bmatrix}$

4.5 穿过第一层 $z=W_1x$

输入梯度为：

$g_x=W_1^Tg_z=\begin{bmatrix}1&2\\1&-1\end{bmatrix}\begin{bmatrix}24\\0\end{bmatrix}=\begin{bmatrix}24\\24\end{bmatrix}$

第一层参数梯度为：

$\dfrac{\partial L}{\partial W_1}=g_zx^T=\begin{bmatrix}24\\0\end{bmatrix}\begin{bmatrix}1&2\end{bmatrix}=\begin{bmatrix}24&48\\0&0\end{bmatrix}$

至此，优化器已经拿到了两个可训练矩阵的梯度：

$\dfrac{\partial L}{\partial W_1}=\begin{bmatrix}24&48\\0&0\end{bmatrix},\qquad\dfrac{\partial L}{\partial W_2}=\begin{bmatrix}36&0\end{bmatrix}$

5. 这些反向规则怎样对应 Jacobian？

以平方层为例：

$h=\begin{bmatrix}z_1^2\\z_2^2\end{bmatrix}$

它的完整 Jacobian 是：

$J_{h,z}=\begin{bmatrix}2z_1&0\\0&2z_2\end{bmatrix}=\begin{bmatrix}6&0\\0&0\end{bmatrix}$

如果显式使用 Jacobian，反向计算是：

$g_z=J_{h,z}^Tg_h=\begin{bmatrix}6&0\\0&0\end{bmatrix}\begin{bmatrix}4\\-8\end{bmatrix}=\begin{bmatrix}24\\0\end{bmatrix}$

但这张矩阵只是对角线可能非零。直接计算：

$g_z=g_h\odot(2z)$

得到完全相同的结果，却不需要创建一个 (H\times H) 的巨大对角矩阵。

对线性层 (z=W_1x)，Jacobian 恰好就是 (W_1)：

$J_{z,x}=\dfrac{\partial z}{\partial x}=W_1$

因此：

$g_x=J_{z,x}^Tg_z=W_1^Tg_z$

这说明“线性层反向规则”并不是替代了 Jacobian，而是直接计算了 (J^Tg)。对矩阵参数 (W) 的梯度，也能逐元素从链式法则推出：

$z_i=\sum_jW_{ij}x_j$

$\dfrac{\partial L}{\partial W_{ij}}=\dfrac{\partial L}{\partial z_i}\dfrac{\partial z_i}{\partial W_{ij}}=g_{z_i}x_j$

把所有 (i,j) 排成矩阵，就是：

$\boxed{\dfrac{\partial L}{\partial W}=g_zx^T}$

6. 为什么不计算整个网络的 Jacobian？

数学上，可以先构造三层网络从 (x) 到 (q) 的完整 Jacobian：

$\dfrac{\partial q}{\partial x}=J_{q,h}J_{h,z}J_{z,x}$

再乘上损失的梯度：

$\dfrac{\partial L}{\partial x}=\dfrac{\partial L}{\partial q}\left(\dfrac{\partial q}{\partial x}\right)^T$

在上面的玩具网络里，这当然可行。但 Transformer 隐藏维常为 4096；单个逐元素算子的 Jacobian 就可能有：

$4096\times4096\approx1678\text{ 万个元素}$

而且网络中有很多层。构造、保存、相乘这些矩阵，会造成巨大的显存占用和额外计算。

反向传播改为逐层计算：

1
2
3
4
gq = 4
gh = W2.T @ gq
gz = gh * (2 * z)
gx = W1.T @ gz

这等价于：

$J_{z,x}^T\left(J_{h,z}^T\left(J_{q,h}^Tg_q\right)\right)$

但不会构造或相乘整个 Jacobian 链。实际中，这类计算被称为 vector-Jacobian product（VJP）：只求 Jacobian 对当前上游梯度的作用。

7. 常见算子的局部反向规则

令 (g_y=\partial L/\partial y) 表示某个算子输出接收到的上游梯度。常见规则如下：

加法规则中的 += 代表累加，而不是简单赋值。只要一个变量流向多个后续节点，反向传播就会把每条路径传回的梯度加到同一位置。

Softmax、交叉熵、LayerNorm 等高层算子通常也有专门的融合反向实现。数学含义仍然是链式法则，但融合 Kernel 可以减少中间张量、访存和数值问题。

8. PyTorch 如何做这件事？

下面的代码与本文的三层数值例子一致：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import torch

x = torch.tensor([1.0, 2.0], requires_grad=True)
W1 = torch.tensor([[1.0, 1.0], [2.0, -1.0]], requires_grad=True)
W2 = torch.tensor([[1.0, -2.0]], requires_grad=True)
target = torch.tensor(5.0)

z = W1 @ x
h = z.square()
q = (W2 @ h).squeeze()
loss = 0.5 * (q - target).square()

loss.backward()

print(loss.item())  # 8.0
print(x.grad)       # tensor([24., 24.])
print(W1.grad)      # tensor([[24., 48.], [ 0.,  0.]])
print(W2.grad)      # tensor([[36.,  0.]])

调用 loss.backward() 后，框架沿着 loss -> q -> h -> z -> x 的反方向遍历计算图。它保存前向阶段反向所需的值，例如平方层需要 z，矩阵乘法需要对应的输入；然后调用各算子已经实现好的局部 VJP 规则。

9. 一句话总结

反向传播不是神秘的全局公式，而是一种组织链式法则的方式：

$\boxed{\text{从标量损失开始，倒序通过每个算子的局部导数，并在分叉处累加梯度}}$

Jacobian 用来准确描述向量到向量的局部导数；但训练时只需要 (J^Tg)，因此框架会利用每种算子的结构直接计算这个结果，而不会显式构造巨大的 Jacobian 矩阵。

• Previous
  Softmax、交叉熵与 KL 散度