loss.backward() 是训练代码里最短的一行，也是被理解得最少的一行。它做的事在 1986 年就已经写清楚了：沿着计算图反向应用链式法则。但只有自己推过一遍、写过一遍、用有限差分验证过一遍，才会真的知道三件后面每一篇都要用的事——梯度的形状与被求导的量相同、反向的计算量是前向的两倍、前向的中间结果必须保留到反向。第一件决定了怎么读任何一个梯度公式，第二件是训练 FLOPs 等于 \(6ND\) 的来源，第三件是激活显存与激活重算的全部原因。

本篇用一个两层 MLP（Linear → ReLU → Linear → softmax → 交叉熵）把这三件事推到底。选它是因为它足够小——每一步的形状能写在一行里——又足够完整：Transformer 里除了 attention 之外的每个部件（Linear、激活函数、norm、loss）的反向都是同一套规则。全篇的核心问题是：

不用框架，能不能手推并手写一个两层网络的反向传播，用有限差分验证到 \(10^{-6}\) 以内？¹ 能不能由此说出为什么训练 FLOPs 是 \(6ND\)、为什么激活要存？²

一、总览：从链式法则到三条结论

1. 本文的对象：一个两层 MLP

MLP（多层感知机，multi-layer perceptron）是最基础的前馈神经网络，四个特征把它定义清楚：

1. 层级结构：一个输入层（习惯上不计入层数）、一个或多个隐藏层、一个输出层。”两层 MLP”指两个带参数的层——一个隐藏层加一个输出层。
2. 全连接（fully connected）：每一层的每个神经元与上一层的所有神经元相连，所以又叫 FC 层、Dense 层；数学上一层就是一次矩阵乘加偏置 \(xW + b\)。
3. 非线性激活：隐藏层的输出要过一个非线性函数（ReLU、sigmoid、tanh）。没有它，两个线性层叠起来 \((xW_1)W_2 = x(W_1 W_2)\) 仍是一个线性层，再深也只能拟合线性函数。
4. 单向传播：信号从输入逐层向前到输出，内部没有循环（第六篇的 RNN 有）。

本文的 MLP 把 784 维输入（一张 28×28 的手写数字图拉直）分成 10 类：隐藏层 256 个神经元、激活 ReLU，输出层 10 个神经元、过 softmax 变成 10 个类别的概率：

按 batch 一次算 \(m\) 个样本，每一步的形状是：

参数是 \(W_1 \in \mathbb{R}^{784 \times 256}, b_1 \in \mathbb{R}^{256}, W_2 \in \mathbb{R}^{256 \times 10}, b_2 \in \mathbb{R}^{10}\)。个数就是这四个数组的元素数之和——图 1 里每条边一个权重、每个隐藏 / 输出神经元一个偏置：

全文要做的是：给出 \(L\) 对这四个参数的梯度、说出每一步的计算量与需要保存的量、用数值方法验证公式没错、在 MNIST 上训到 97%。

2. 本文的章节安排

3. 来龙去脉：一个算法，被发明了三次

三次发明说明一件事：反向传播不是神经网络专属的，它是”复合函数求导”的通用高效算法——任何能写成计算图的东西都能用。1986 年之前没人用它训网络，不是数学没有，而是没人相信多层网络能学到东西；1986 年之后它成了唯一的训练方式，至今没有替代者（所有”不用反向传播”的尝试——目标传播、前向-前向——都没有在大模型上站住）。今天手推它的意义，不是为了自己写框架，而是第五、六章那两笔账：反向为什么是前向的两倍算量、激活为什么要占那么多显存——这两笔账决定了大模型训练的成本结构，而框架把它们藏起来了。

二、链式法则与计算图

1. 三种导数的形状

深度学习里只出现三种求导，各自的结果形状要先记住：

前两种是训练最终要的东西——loss 对每个参数的梯度，形状与参数一样，所以可以直接 W -= lr * dW。第三种是中间量：每一层把输入变成输出，层的局部导数是一个 Jacobian。

2. 链式法则的向量形式

复合函数 \(L = \ell(y), \; y = f(x)\)，\(L\) 对 \(x\) 的梯度是

也就是”上游传来的梯度，左乘本层 Jacobian 的转置”。把网络看成一串函数 \(x_0 \to x_1 \to \dots \to x_l \to L\)，从 \(L\) 出发反向逐层套用这一行，就是反向传播的全部：

每一层只需要知道两件事：自己的 Jacobian 是什么、上游传来的梯度是什么。它不需要知道网络其他部分长什么样——这就是为什么框架可以把任意算子拼起来自动求导。

3. 为什么从不构造 Jacobian

上面的公式里 Jacobian 只以”乘一个向量”的形式出现：\(J^T v\)，叫 vector-Jacobian product（VJP）。反向传播每一步算的是 VJP，从不把 \(J\) 本身写出来。原因是尺寸：本文第一层 Linear 在 batch 128 下，输入 \(X \in \mathbb{R}^{128 \times 784}\)、输出 \(H \in \mathbb{R}^{128 \times 256}\)，把它们拉直后 Jacobian 是 \((128 \times 256) \times (128 \times 784) = 32768 \times 100352\)，约 \(3.3 \times 10^9\) 个元素，fp32 下 13 GB——一个 20 万参数的小网络的一层。而 VJP \(J^T v\) 的结果只有 \(X\) 那么大（100352 个数，392 KiB），并且对 Linear 这类结构化的层，它可以直接写成一个矩阵乘法，下一章推。

所以”反向传播”这个词精确的含义是：按计算图的反向拓扑序，对每个算子算一次 VJP。每个算子的 VJP 有自己的闭式公式，框架里每个 Function 的 backward 就是那条公式。

三、矩阵求导的形状规则

1. Linear 层的三条公式

\(Y = XW + b\)，\(X \in \mathbb{R}^{m \times k}\)（\(m\) 个样本、每个 \(k\) 维），\(W \in \mathbb{R}^{k \times n}\)，\(b \in \mathbb{R}^n\) 广播到每一行，\(Y \in \mathbb{R}^{m \times n}\)。记上游梯度 \(G = \partial L / \partial Y \in \mathbb{R}^{m \times n}\)。三条公式：

推导。 按元素写：\(Y_{ij} = \sum_r X_{ir} W_{rj} + b_j\)。\(L\) 对 \(W_{rj}\) 的导数，按链式法则对所有依赖 \(W_{rj}\) 的 \(Y_{ij}\) 求和：

对 \(X_{ir}\) 同理：\(\partial L / \partial X_{ir} = \sum_j G_{ij} W_{rj} = (G W^T)_{ir}\)。对 \(b_j\)：每一行都加了 \(b_j\)，所以 \(\partial L / \partial b_j = \sum_i G_{ij}\)。

循环写法。 转置是从哪来的，把 \(\sum\) 写成循环最清楚。前向里 W[r][j] 被 batch 里每个样本 i 用了一次；反向就把这 \(m\) 次的影响加起来：

for i in range(m):                        # 前向
    for j in range(n):
        for r in range(k):
            Y[i][j] += X[i][r] * W[r][j]

for r in range(k):                        # 反向：dW 与 W 同形，逐个元素问"谁用过我"
    for j in range(n):
        for i in range(m):                # 对 i 求和
            dW[r][j] += X[i][r] * G[i][j]  # X 的第 r 列 · G 的第 j 列

for j in range(n):                        # b_j 加在每一行上，同样对 i 求和
    for i in range(m):
        db[j] += G[i][j]

dW[r][j] 用的是 X 的第 r 列——按列取就是转置，所以是 \(X^T G\)；db[j] 只对 \(i\) 求和，所以是按列求和。三条公式就是这三个循环。

形状记忆法。 三条公式不必背，用形状就能重建：\(\partial L / \partial W\) 必须是 \(k \times n\)，手头有 \(X\)（\(m \times k\)）和 \(G\)（\(m \times n\)），唯一能凑出 \(k \times n\) 的乘法是 \(X^T G\)；\(\partial L / \partial X\) 必须是 \(m \times k\)，手头有 \(G\)（\(m \times n\)）和 \(W\)（\(k \times n\)），唯一的凑法是 \(G W^T\)。梯度的形状与被求导的量相同这一条约束，加上”只能用手头的量做矩阵乘”，几乎总能唯一确定公式。读论文里任何一个梯度表达式时先做这个形状检查，能抓住大部分笔误。

2. 逐元素算子：ReLU

\(A = \text{ReLU}(H) = \max(H, 0)\)，逐元素。它的 Jacobian 是对角阵，对角线上是 \(\mathbb{1}[H > 0]\)；VJP 就是逐元素乘一个 0/1 掩码：

所有逐元素激活函数（GELU、SiLU、sigmoid、tanh）的反向都是”上游梯度逐元素乘导数”，区别只在导数的表达式。\(H = 0\) 处 ReLU 不可导，实践中取 0 或 1 都行——概率为零的事件。

3. softmax 与交叉熵

\(L = -\frac{1}{m}\sum_i \log p_{i, y_i}, \; p_i = \text{softmax}(z_i)\)。L0 数学系列第七篇推过对单个样本 \(\partial L_i / \partial z_i = p_i - e_{y_i}\)（\(e_{y_i}\) 是 one-hot）。对 batch 取平均后：

这里 softmax 与交叉熵合在一起求导，中间不经过 \(\partial L / \partial P\)。分开算在数学上等价，但数值上更差（\(1/p\) 在 \(p \to 0\) 时溢出），计算上也多一步。框架里 F.cross_entropy 是一个融合算子，正是为此。

四、两层网络逐层推导

把三章的规则串起来。前向五步、反向五步，每一步标出形状（\(m = 128\)）：

%% 前向五步与反向五步：上排前向，下排反向的 VJP，虚线是反向用到的前向保存量
flowchart LR
    subgraph F["前向（带 ✱ 的量要保存到反向）"]
        direction LR
        X["X ✱<br/>[128, 784]"] -->|"XW₁ + b₁"| H["H ✱（存符号）<br/>[128, 256]"] -->|"ReLU"| A["A ✱<br/>[128, 256]"] -->|"AW₂ + b₂"| Z["Z<br/>[128, 10]"] -->|"softmax · CE"| L["L<br/>标量"]
    end
    subgraph B["反向（每一步一个 VJP，从右往左）"]
        direction RL
        GZ["∂L/∂Z = (P − Y)/m<br/>[128, 10]"] -->|"· W₂ᵀ"| GA["∂L/∂A<br/>[128, 256]"] -->|"⊙ 𝟙[H > 0]"| GH["∂L/∂H<br/>[128, 256]"] -->|"· W₁ᵀ（不需要）"| GX["∂L/∂X<br/>[128, 784]"]
        GZ -->|"Aᵀ ·"| GW2["∂L/∂W₂ [256, 10]<br/>∂L/∂b₂ [10] = Σ行"]
        GH -->|"Xᵀ ·"| GW1["∂L/∂W₁ [784, 256]<br/>∂L/∂b₁ [256] = Σ行"]
    end
    A -. "保存" .-> GW2
    X -. "保存" .-> GW1
    H -. "保存符号" .-> GH
    classDef fw fill:#e3f2fd,stroke:#1565c0,color:#222
    classDef bw fill:#fff3e0,stroke:#ef6c00,color:#222
    class X,H,A,Z,L fw
    class GZ,GA,GH,GX,GW1,GW2 bw

前向五步与反向五步。上排是前向，每个箭头是一次计算，形状写在节点里；下排是反向，从 \(\partial L / \partial Z\) 出发每一步右乘一个转置（沿主链）或左乘一个转置（分叉到参数梯度）；虚线标出反向用到前向哪个保存的量——这就是第六章”激活为什么要存”。

写成公式：

\(\partial L / \partial X = G_H W_1^T\) 在数学上存在，但 \(X\) 是数据不是参数，不需要它的梯度，框架会跳过这一步。任何深度的 MLP 都是把反向 3–5 重复若干次；Transformer 的 FFN 子层就是一个两层 MLP，attention 子层多几个矩阵乘和一个 softmax，规则不变。

五、反向为什么是前向的两倍

1. 每个 Linear 反向做两个 GEMM

前向一个 Linear 做一次矩阵乘 \(XW\)：\([m, k] \times [k, n]\)，\(2mkn\) FLOPs。反向做两次：\(X^T G\)（\([k, m] \times [m, n]\)，\(2mkn\)）和 \(G W^T\)（\([m, n] \times [n, k]\)，\(2mkn\)）。反向 = 2 × 前向，对每一个 Linear 层都严格成立。逐元素算子（ReLU、norm、softmax）的计算量与 GEMM 相比可以忽略，所以整个网络也近似成立。

代入本文的网络，batch 128：

第九章的实验里用计数器验证了这两个数字——它的实现故意也算了第一层的 \(\partial L / \partial X\)（输入是图片，这一项没人用），所以比值恰好 2.00。实践中省掉它：这个两层网络里第一层占 98.7% 的权重，省掉后反向是 \(51.4 + 1.31 = 52.7\)M，比值 1.01——对浅而宽的网络”反向 = 2 × 前向”差得远；层数多时第一层占比小，比值回到接近 2，\(6ND\) 才是好的近似。

2. 训练 FLOPs 等于 6ND

对一个参数量为 \(N\) 的网络，前向每个 token 约 \(2N\) FLOPs（每个参数参与一次乘加——L0 数学系列第一篇），反向 \(4N\)，一步训练合计 \(6N\) FLOPs / token。训练 \(D\) 个 token 就是 \(6ND\)——scaling law 论文与 04 系列第十篇用的这个数字，来源就是本章的”反向做两个 GEMM”。Llama-3-8B 训 15T token：\(6 \times 8 \times 10^9 \times 15 \times 10^{12} = 7.2 \times 10^{23}\) FLOPs。

两点补充。第一，\(2N\) 忽略了 attention 里 \(QK^T\) 与 \(PV\) 这两个与参数无关、与序列长度成正比的项，短序列下可忽略，长序列下不能（04 系列第十篇算了）。第二，如果用了激活重算（下一章），反向前要再做一次前向，总量变成 \(8N\) / token——训练报告里”MFU 按 \(6ND\) 算、HFU 按 \(8ND\) 算”的区别就在这里。

六、激活为什么要存

1. 参数梯度需要本层的输入

第四章的表最后一列是关键：反向算 \(\partial L / \partial W_1 = X^T G_H\) 需要 \(X\)，算 \(\partial L / \partial W_2 = A^T G_Z\) 需要 \(A\)，算 ReLU 的反向需要 \(H\) 的符号。这些都是前向的中间结果，反向时要用，所以前向算完不能丢，必须一直留到反向走到那一层。这就是”激活显存”。

它的大小与什么成正比？每一层保存的是该层的输入，形状 \([m, \text{宽度}]\)——与 batch 大小 \(m\) 成正比，与层数成正比，与宽度成正比，与参数量无关。本文的网络在 batch 128 下：

batch 换成 4096，激活变成 17.3 MiB，权重不变。序列模型里 \(m\) 是 batch × 序列长度，所以长上下文训练的激活显存会远超权重——L1 工具箱系列给过 Llama-3-8B 在 4096 长度下仅残差流一份就是 1 GiB / 序列的锚点，精确公式在 Infra 地图 07 系列第一篇。

2. 激活重算：用计算换存储

既然激活是前向算出来的，可以不存、反向时重算。gradient checkpointing 的做法是只保存每个块（比如 Transformer 的一层）的输入，反向走到这一块时先用保存的输入重新前向一次得到块内所有中间量，再做反向。代价是多一次前向，即上一章说的 \(6N \to 8N\)，约 33% 的额外计算；收益是激活显存从”所有中间量”降到”每块一个输入”。这是训练长序列或大 batch 时的标准开关，model.gradient_checkpointing_enable() 一行。

FlashAttention 做的是同一件事的算子级版本：不保存 \([\text{seq}, \text{seq}]\) 的 attention 矩阵，反向时分块重算——04 系列第六篇讲它的 IO 复杂度。

七、梯度检查

1. 有限差分

手推的公式对不对，用数值导数验证。对参数 \(\theta\) 的某一个元素 \(\theta_j\)，中心差分

的截断误差是 \(O(\epsilon^2)\)（单侧差分是 \(O(\epsilon)\)，所以用中心）。把它与解析梯度比，用相对误差 \(\lvert g_{\text{num}} - g_{\text{ana}} \rvert / (\lvert g_{\text{num}} \rvert + \lvert g_{\text{ana}} \rvert)\)——绝对误差在梯度本身很小的元素上没有意义。

2. 为什么必须 float64

\(\epsilon\) 有两头约束：太大截断误差大，太小舍入误差大——\(L(\theta + \epsilon) - L(\theta - \epsilon)\) 是两个相近数相减，误差约为 \(L\) 的机器精度除以 \(2\epsilon\)。fp32 机器精度约 \(10^{-7}\)，取 \(\epsilon = 10^{-6}\) 时舍入误差约 \(10^{-7} / 10^{-6} = 0.1\)，检查毫无意义；fp64 机器精度约 \(10^{-16}\)，同样的 \(\epsilon\) 舍入误差 \(10^{-10}\)，截断误差 \(10^{-12}\)，可以把相对误差压到 \(10^{-7}\) 以下。所以梯度检查一定在 float64 上做、用小网络、随机抽几十个元素（全部检查要做 \(2 \times 203530\) 次前向）。第九章的结果：最大相对误差 \(1.1 \times 10^{-7}\)。

判断标准：相对误差 \(< 10^{-6}\) 通过；\(10^{-4}\) 量级要怀疑（常见于 ReLU 在 0 附近被 \(\epsilon\) 翻转符号，属于正常）；\(> 10^{-2}\) 一定有错。

3. 两个训练前的 sanity check

梯度对了，训练前还有两个几乎免费的检查。初始 loss 应接近 \(\ln C\)：10 类是 \(\ln 10 = 2.30\)，本文实验初始 2.46（Kaiming 初始化让 logits 方差略大于 1，比均匀分布稍差，正常）；远大于它说明初始化太大，远小于它说明数据泄漏或 loss 算错。LLM 上对应 \(\ln V \approx 11.8\)（L0 数学系列第五篇）。能过拟合一个小 batch：拿 16 个样本反复训，loss 应能降到接近 0；降不下去说明梯度没传到某处，或学习率不对。这两个检查在框架里同样适用。

八、Autograd 做了什么

本文手写的东西，框架用四个机制自动化了：

torch.no_grad() 关掉录带（推理与评测时省激活显存）；.detach() 把一个张量从图上摘下来（RL 里对 old logprobs 常用）；torch.utils.checkpoint 就是第六章的激活重算。第九章的实验把手写梯度与 PyTorch autograd 的结果对了一遍，差在 \(10^{-8}\) 量级——两边算的是同一组公式，差异只是浮点求和顺序。

Autograd 引擎的实现——图怎么存、多线程怎么调度、hook 在哪里——属于 Infra 地图 03 系列第三篇。算法工程师到这一层就够：知道它记了什么、存了什么、什么时候释放。

九、案例：MNIST 手写数字，NumPy 从零

问题与数据：MNIST——60,000 张训练、10,000 张测试的手写数字（28×28 灰度），1998 年 LeCun 等为比较分类器造的数据集，L2 第四、五篇已经用 KNN（2.95% 错误率）和 SVM（1.43%）做过它。这一章不用任何框架，把前八章推出来的公式写成 120 行 NumPy，训一个两层 MLP，看它能到多少、每一步的算量和内存是多少、以及公式有没有推错。它是整个系列的基座：后面五篇的每个实验都在这 120 行上加东西。

思路：Linear / ReLU / softmax_ce 三个类各自实现 forward（存下反向要用的量）和 backward（第三章的三条公式）；MLP 把它们串起来；训练循环就是 forward → loss → backward → 参数 -= lr × 梯度。训之前先做两件事：梯度检查（第七章：公式对不对）和 FLOPs / 激活计数（第五、六章：账对不对）。

1. 代码

约 120 行 NumPy，是整个系列的基座，后面每篇往上加。核心的三个类：

class Linear:
    def forward(self, X):
        self.X = X                                  # 保存激活
        return X @ self.W + self.b
    def backward(self, dY):
        self.dW[...] = self.X.T @ dY                # X^T G      2mkn
        self.db[...] = dY.sum(0)
        return dY @ self.W.T                        # G W^T      2mkn

class ReLU:
    def forward(self, X): self.mask = X > 0; return X * self.mask
    def backward(self, dY): return dY * self.mask

def softmax_ce(logits, y):
    z = logits - logits.max(1, keepdims=True)       # 数值稳定
    p = np.exp(z); p /= p.sum(1, keepdims=True)
    m = logits.shape[0]
    loss = -np.log(p[np.arange(m), y]).mean()
    d = p.copy(); d[np.arange(m), y] -= 1.0         # p - y
    return loss, d / m, p

一个数值细节：这里先算 softmax 再取 log，正确类的 logit 比最大值低 700 以上时 exp 会下溢成 0、log(0) = -inf（配套脚本加了 1e-12 兜底，代价是 loss 被截在 27.6 上下——第二篇”无归一化残差”实验里 logits 到 \(10^9\) 时打印的 loss 就是这个上限，不是真实值）。生产实现（PyTorch 的 cross_entropy）用 log-sum-exp 形式 logp = z - log(sum(exp(z))) 直接算对数概率，没有这一步下溢。

网络是 MLP([784, 256, 10])，训练循环就是 forward → softmax_ce → backward → 对每个参数 P -= lr * dP。梯度检查在一个 [20, 16, 5] 的 float64 小网络上做，每个参数随机抽 30 个元素、\(\epsilon = 10^{-6}\)、中心差分。FLOPs 用一个全局计数器在每次矩阵乘处累加。

2. 结果

grad check worst rel err: 1.11e-07
params: 203530
batch 128: fwd 52.0 MFLOPs  bwd 104.1 MFLOPs  ratio 2.00
  2*N*tokens = 52.1 MFLOPs (fwd approx ignoring bias)
  saved activations: 552 KiB  vs weights 795 KiB

与 PyTorch autograd 对齐（同一份权重、同一个 batch）：

loss numpy 2.460904 torch 2.460904
dW1 max abs diff 5.59e-09
dW2 max abs diff 7.45e-09

MNIST 上 SGD、学习率 0.1、batch 128：

epoch 1   train loss 0.4477  test acc 92.52%
epoch 5   train loss 0.1416  test acc 96.01%
epoch 10  train loss 0.0842  test acc 97.26%
epoch 15  train loss 0.0585  test acc 97.61%

每个 epoch 在笔记本 CPU 上约 0.1 秒。四个数字各验证了本文的一个结论：\(10^{-7}\) 的相对误差说明第三、四章的公式没错；2.00 的比值是第五章；552 KiB 与 795 KiB 的对比是第六章——这个小网络的激活已经与权重同量级，batch 再大 32 倍就是权重的 22 倍；97% 说明这一套确实能训。

训练过程长什么样（case_01_mnist_mlp.py 把它画出来）：

初始 loss 2.46——接近 \(\ln 10 = 2.30\)（10 类瞎猜），第七章说的 sanity check 之一：如果初始 loss 远大于它，是初始化太大；远小于它，是标签泄漏。前 100 步 loss 就从 2.3 掉到 0.5，之后越来越慢——第一个 epoch 结束 92.5%，再训 14 个 epoch 才到 97.6%。10,000 张测试图错 239 张（2.39%），最常混的是 7→2（12 次）、9→7、9→4（各 11 次）——和 L2 第四篇 KNN 错的是同一批字。

和 L2 的两个分类器比：KNN 2.95%、RBF-SVM 1.43%、这个两层 MLP 2.39%——三种完全不同的方法在同一份数据上落在同一个量级。MLP 的优势不在这个数字上，在于它学出来的特征可以继续堆：加深、加卷积（第五篇的 LeNet-5 到 0.9% 以下）、换数据（同一份代码换成文本、语音），而 KNN 和 SVM 到这里就到顶了。这也是为什么后面五篇都在这 120 行上加东西，而不是换方法。

3. 三个值得自己动手的扩展

• 把 ReLU 换成 GELU 或 SiLU，只改 backward 里的导数表达式，梯度检查应仍通过；
• 加一层，看 FLOPs 比值仍是 2、激活线性增长；
• 把梯度检查改在 float32 上做，观察相对误差跳到 \(10^{-2}\) 量级——第七章的论证。

落地还差什么：97.6% 在 MNIST 上是 1998 年的水平（LeNet-5 0.95%，今天的最好结果 0.2% 以下）；这 120 行没有的东西每一样都是后面一篇：深了训不动（第二篇）、SGD 太慢太挑学习率（第三篇）、数据少了过拟合（第四篇）、把图拉直成 784 维丢掉了”相邻像素相关”这个信息（第五篇的卷积）。生产里没有人用 NumPy 写反向传播——但读懂了这 120 行，第八章的 loss.backward() 就不再是黑盒，第五、六章的两笔账（反向两倍算量、激活占显存）就是你看训练日志时脑子里的第一反应。

十、本文小结

• 反向传播 = 按计算图反向拓扑序对每个算子算一次 VJP \(J^T v\)；从不构造 Jacobian（本文第一层的 Jacobian 有 13 GB）。
• 梯度的形状与被求导的量相同。 Linear 的三条公式 \(\partial L / \partial W = X^T G\)、\(\partial L / \partial X = G W^T\)、\(\partial L / \partial b = \sum_i G_{i,:}\) 可以从形状唯一重建；逐元素算子的反向是上游梯度乘导数；softmax 与交叉熵合并求导得 \((P - Y) / m\)。
• 反向 = 2 × 前向：每个 Linear 反向做两个 GEMM。由此训练 FLOPs \(= 6N\) / token，激活重算下 \(8N\)；本文实测比值 2.00。
• 激活必须保存到反向，因为 \(\partial L / \partial W\) 需要该层的输入。激活显存与 batch × 序列长度 × 层数 × 宽度成正比、与参数量无关；gradient checkpointing 用一次额外前向换掉它。
• 梯度检查用 float64、中心差分、相对误差，\(< 10^{-6}\) 通过；训练前看初始 loss 是否接近 \(\ln C\)、能否过拟合一个小 batch。
• Autograd 做的是录带、每个算子的 backward、saved tensors、.grad 累加；手写结果与它差 \(10^{-8}\)。
• 下一篇把这个网络加深到 64 层，看梯度在层间传播时会发生什么。

配套代码：deep-learning-foundations/01_backprop.py——本文的全部数字由它跑出（梯度检查、FLOPs、与 autograd 对齐、15 个 epoch 的训练），--quick 一分钟内跑完；第九章的训练曲线与错分样本由 case_01_mnist_mlp.py 画出；三个扩展可以直接在上面改。

十一、自测

1. 一个 Linear 层输入 \(X \in \mathbb{R}^{128 \times 784}\)、权重 \(W \in \mathbb{R}^{784 \times 256}\)，上游梯度 \(G = \partial L / \partial Y\)。\(\partial L / \partial W\)、\(\partial L / \partial X\)、\(\partial L / \partial b\) 各是什么形状、怎么算？

   答案

   \(\partial L/\partial W = X^T G\)，形状 \([784, 256]\)（与 \(W\) 同形）；\(\partial L/\partial X = G W^T\)，形状 \([128, 784]\)（与 \(X\) 同形）；\(\partial L/\partial b = \sum_i G_{i,:}\)，形状 \([256]\)。三条都能从”梯度与被求导的量同形”倒推出来。

2. softmax + 交叉熵对 logits \(Z\) 的梯度是什么？batch 大小 \(m\) 在哪里？

   答案

   \(\partial L / \partial Z = (P - Y) / m\)：\(P\) 是 softmax 输出、\(Y\) 是 one-hot 标签；除以 \(m\) 是因为 loss 对 batch 取了平均。

3. 一个 7B 模型训练 1T token，前向、反向、合计各多少 FLOPs？开了 gradient checkpointing 呢？

   答案

   前向 \(2N D = 2 \times 7 \times 10^9 \times 10^{12} = 1.4 \times 10^{22}\)；反向是前向两倍 \(2.8 \times 10^{22}\)；合计 \(6ND = 4.2 \times 10^{22}\)。开 checkpointing 多一次前向，\(8ND = 5.6 \times 10^{22}\)，多 33%。

4. 本文的两层网络在 batch 128 下激活占 552 KiB、权重 795 KiB。batch 换成 4096，两者各变成多少？

   答案

   激活与 batch 成正比：\(552 \times 32 = 17{,}664\) KiB \(\approx 17.3\) MiB；权重不变，仍是 795 KiB。激活可以远超权重，序列模型里更是如此。

5. 梯度检查为什么必须用 float64、用中心差分、看相对误差而不是绝对误差？

   答案

   float32 只有 7 位有效数字，有限差分的截断误差与舍入误差加起来在 \(10^{-3}\) 量级，分不清”对”与”错”；中心差分 \((f(x+h) - f(x-h)) / 2h\) 的误差是 \(O(h^2)\)，比单侧差分的 \(O(h)\) 准；相对误差 \(\lvert a - b \rvert / (\lvert a \rvert + \lvert b \rvert)\) 不受梯度绝对量级影响——梯度是 \(10^{-8}\) 时绝对误差 \(10^{-9}\) 已经是 10% 的错。

• 下一篇
  深度学习基础（02）：训练为什么不稳定——初始化、归一化与残差