必须恰好是 0 的那个梯度：八步手推 softmax 的反向传播

上一篇结尾写的是：下一段开始手推那二十多步。工具齐了。

这一篇用八次。从 loss 出发，一路倒推回 logits —— 正好是 F.cross_entropy 内部那一段。

但先说清楚为什么要手推。

backward() 从不出错，也从不告诉你什么

loss.backward() 的承诺是「你只管写 forward，梯度我来算」。这个承诺在数值上永远成立 —— PyTorch 算出来的梯度就是对的。它漏的是另一头：

backward() 帮你算梯度，但不帮你判断梯度合不合理。 要判断，就得知道梯度是怎么来的。

这就是 Karpathy 说的漏抽象（leaky abstraction）—— 抽象层在正常情况下让你不必关心底下， 但一旦出事，不懂底下就修不了。第十四到十七篇教的是看仪表盘，这一篇教的是拆发动机。

第四篇写过一句「自动求导只保证数学正确，不保证数值稳定」，说的是同一件事的另一半。

把 forward 摊成 18 行

要手推，先得能定位。原来三四行写完的前向，现在一行一个原子操作，每行右边标着形状 —— 第十八篇说过，形状就是推梯度的主要线索。

这套代码有两处反直觉的设计，都是为了判卷：

① 退回 2 层，不是第十七篇那个 6 层。 n_embd=10、n_hidden=64、batch=32， 七个参数张量共 4137 个。要手推二十多个中间量，层数必须压到最少 —— 多一层 BN 就多十步反向。

② b1 故意初始化成非零。

b1 = torch.randn(n_hidden, generator=g) * 0.1     # 不是 zeros

正常训练时 bias 设 0 是标准做法（前面接了 BN 时甚至可以不要）。但这一段在做自动判卷，规则不同：

如果某个量恰好是 0，你推错的公式算出来可能碰巧也等于 0，判卷就给你一个假的通过。 全用非零随机数，任何错误都会立刻暴露。bngain / bnbias 加扰动是同样的道理。

为什么要拆这么细

第四篇讲过一次这个取舍 —— 当时是在 micrograd 里拆 tanh：算子打包得越大，写起来越省、跑起来越快， 但出错时定位不到行。那时是可选项，这里是必需品：

如果前向写成一句 F.cross_entropy(logits, Yb)，你只能验证最终的 dlogits 对不对； 推错了不知道错在哪一步。摊开之后每个中间量各有各的 .grad：

logprobs        | exact: True
probs           | exact: True
counts_sum_inv  | exact: False    ← 定位到这里

摊开是为了让判卷能定位到行。

为什么这八步里一个参数都没有

第十四到十七篇一直在讲「哪些层有参数」，但从没横着切过一刀。按每种层干什么重排，只有三类：

参数几乎全长在混合层上。 这个网络 4137 个参数里，混合层占 4009 个（96.9%）， 归一化只占 128 个 γβ（3.1%），非线性一个都没有。

而这一篇要推的八步，三类一个都不是 —— 从 logits 往后全是纯算术：减最大值、exp、求和、 取倒数、相乘、取对数、挑一列。从这里往后不再有神经元、不再有参数，所以八步里不会出现 任何 dW。参数的反向要等下一段的 dW2 / db2 才开始。

cmp：怎么给自己判卷

标准答案来自开跑前那一次 loss.backward()。中间张量得先 retain_grad() —— 和第十七篇画图 ② 踩的是同一个机制：PyTorch 默认不保存中间梯度，用完就扔。

def cmp(s, dt, t):
    ex = torch.all(dt == t.grad).item()        # 逐位完全相等
    app = torch.allclose(dt, t.grad)           # 浮点容差内相等
    maxdiff = (dt - t.grad).abs().max().item()
    print(f'{s:15s} | exact: {ex} | approximate: {app} | maxdiff: {maxdiff}')

❗ 两列不一样，不代表你错了

关键是浮点加法不满足结合律：(a+b)+c ≠ a+(b+c)，差在最后几位。 你的求和顺序和 PyTorch 内部的不一定一样，结果就会差一点点。

💡 每推一步就立刻 cmp，不要攒五步再验。 梯度是往回传的，第 3 步错了会让第 4、5 步全错 —— 攒着验只会让你找不到源头。

八步走一遍

四种形态上一篇讲过了，这里只标不解释。 每一步真正要看的，是它新带来的那件事。

① dlogprobs：前向挑出来的，反向放回去

反向传播的起点，不需要任何上游梯度。

前向那一行是 loss = -logprobs[range(n), Yb].mean() —— 从 32×27 的表里每行只挑一个数， 就是正确答案那一列。没被挑到的格子改成任何值 loss 都不变，梯度是 0。

dlogprobs = torch.zeros_like(logprobs)   # 先铺一张同形的全 0 表
dlogprobs[range(n), Yb] = -1.0/n         # 只往挑过的那些位置填

🔑 这两行里最值得记的是第二行的索引表达式 —— 它和前向那行一字不差。

前向「从 A 挑出一部分」，反向就是「把梯度放回原位，其余补 0」。 索引表达式原样复用，不用另想办法。

先铺 0 再填，而不是反过来，是因为绝大多数位置本来就是 0。 zeros_like 而不是 torch.zeros(32, 27)：不用手写形状，换个 batch size 也不会错。

三个容易搞错的地方：

一个好用的自检：dlogprobs.sum() 应该正好等于 −1.0 —— .mean() 让每个是 −1/32，32 个加起来是 −1。

② dprobs：穿过 log

形态 A，局部导数 1/x，一句话的事：dprobs = (1.0 / probs) * dlogprobs。

值得多看一眼的是这一步把三个一模一样的梯度拉开了差距。第 ① 步填进去的全是 −1/32， 而 1/p 在小值区会爆炸：

差 99 倍。猜得越差，推得越狠。

这不是谁设计的规则，就是 1/x 在小值区爆炸的自然结果 —— 交叉熵天然会「重点管教」 错得离谱的样本，已经答对的几乎不再推它。第七篇从最大似然推出交叉熵时，这条性质还看不出来； 要到反向这一步才显形。

③ dcounts_sum_inv：第一次遇到广播

前向是 probs = counts * counts_sum_inv，形状 (32,27) = (32,27) * (32,1) —— counts_sum_inv 被广播了 27 份。按上一篇的判据：变量比输出少一个下标 c，形态 B，要 sum 掉 c。

dcounts_sum_inv = (counts * dprobs).sum(1, keepdim=True)     # (32,27) → (32,1)

轴为什么是 1 不是 0，上一篇整节都在讲；这里只记一个更快的判据： 变量的哪一维是 1，就对哪一维求和。 counts_sum_inv 是 (32, 1)，第 1 维是 1，所以 sum(1)。 搞反了 cmp 立刻报形状不匹配 —— 错得响亮，不会藏着。

④⑤⑥ dcounts：分叉点

counts 在前向里出现在两个等号右边（counts_sum 那行和 probs 那行），所以有两条路。 这件事上一篇用同一组数字算透了（路一是 0、路二不是，所以漏掉第二条很难查），这里不重复。

要补的是操作顺序上的三个坑：

顺带解释了一件事：为什么 dcounts 之后梯度就不稀疏了。 第二条路给每一个格子都加了一个非零值 —— 前面几步一直是「一行只有一个非零」，到这里被填满。

第 ⑤ 步穿过 counts_sum**-1，局部导数是 −S⁻²，符号会翻正。 上一篇专门有一小节讲它（连同「别写成 S⁻¹ = S⁻¹ 那句废话」），这里不重推。

另外，前向本可以直接写 counts / counts_sum，拆成「取倒数 + 相乘」是专为手推做的安排 —— 除法要用商法则，两个分支都麻烦；拆开后每步只剩一个简单导数。 第十七篇那个网络里的 bnvar_inv = (bnvar + 1e-5)**-0.5 是同一个设计。

⑦ dnorm_logits：穿过 exp

形态 A，局部导数 exp(x) —— 而它就是前向已经算好的 counts 本身，不用重算：

dnorm_logits = counts * dcounts

⚠️ 前提是 dcounts 必须是完整的。如果在第 ④ 步之后、第 ⑥ 步之前就算这一步， 用的是只有一半的梯度，后面全错。

❗ 算出来之后，每行加起来恰好是 0

print(dnorm_logits.sum(1).abs().max())    # 1e-9 量级

笔记是从 softmax 的性质反推的：softmax 的输出每行永远加起来等于 1，既然这个和是个常数， 那么「整行同时加上一个量」就不会改变任何概率 —— 所以整行梯度的和必须是 0。

这条其实能直接算出来，不用靠观察。 把前面几步串起来：

Σ_c e·(T·dp + dS)                    ← dnorm_logits 一行的和
  = T·Σ_c e·dp  +  dS·Σ_c e          ← 拆开
  = T·dT        +  (−S⁻²·dT)·S       ← 前一项就是第 ③ 步，后一项是第 ⑤ 步
  = dT·(1/S − 1/S)
  = 0

恒等于 0，和具体数值无关。 两个角度互相印证：一边是 softmax 的结构性质， 一边是把三步的公式代进去消掉。那个 1e-9 不是「约等于 0」，是浮点把精确的 0 算成了这样。

⑧ dlogits 与 dlogit_maxes：本段的压轴

这一步涉及两个前向操作：

logit_maxes = logits.max(1, keepdim=True).values   # 操作 1
norm_logits = logits - logit_maxes                 # 操作 2

依赖关系画出来是这样：

logits ──①──→ logit_maxes ──┐
   │                        ③
   └────────②───────────────┴──→ norm_logits

🔑 三条边，所以是三个求导 —— 但只输出两个梯度，因为有两条边都指向 dlogits。

顺序只能是 ② → ③ → ①：①要用 dlogit_maxes 当上游，而它是③算出来的。

② 减法对 logits：.clone() 不能省

norm_logits = logits - logit_maxes 对 logits 求导，系数 +1，形态 A，梯度原样传过去：

dlogits = dnorm_logits.clone()

⚠️ 写成 dlogits = dnorm_logits 就错了 —— 那只是给同一块内存起了个别名。 下面给 dlogits 加第二条路时，dnorm_logits 会被连带改掉。

这个错的症状很有迷惑性：cmp('norm_logits', ...) 在这一行之前是绿的， 跑到后面再回头看就变红了 —— 而你什么都没动过它。 这是 Python 的引用语义，跟梯度没关系，但在一长串 d 变量里特别难查。

第十篇讲 view 共享 storage 时提过一句「调试时手动改了某个值，要知道它可能连带改了别处」。 同一件事，这次出现在反向里。

③ 减法对 maxes：形态 B

同一行，对 logit_maxes 求导。它缺下标 c，所以要 sum 掉 c，系数是 −1：

dlogit_maxes = (-dnorm_logits).sum(1, keepdim=True)    # (32,27) → (32,1)

① max 的梯度是路由

logit_maxes = logits.max(1).values。最大值那一格稍微变大一点，m 就跟着变大同样多； 其他位置动一动（只要没超过最大值），m 纹丝不动。所以：

∂m[b]/∂logits[b,c] = 1（c 是这一行的最大值所在）或 0（其他）

只有赢家拿到梯度，其余全是 0。 这叫路由 —— 梯度不是被分配，是被转发给一个人。

dlogits += F.one_hot(logits.max(1).indices, num_classes=logits.shape[1]) * dlogit_maxes

one_hot 在这里的作用就是把「谁是最大值」变成一张 0/1 表。 第七篇讲 one-hot 时它还是输入的编码方式，这次它是一个开关。

❗ 而 dlogit_maxes 必须恰好是 0

logit_maxes | exact: False | approximate: True | maxdiff ~1e-9

原因第十篇已经推过：softmax 有平移不变性，softmax(z) = softmax(z − c) 对任意常数成立， 分子分母同时乘的 exp(−c) 会约掉。既然 logit_maxes 根本不影响最终输出，它的梯度当然是 0。

从第 ⑦ 步的结果也能直接看出来 —— 那一步刚验证过「每行加起来是 0」， 而这一步做的就是把每行加起来再取负。两条路，同一个结论。

不是精确的 0，还是因为浮点加法不满足结合律：27 个数抵消之后剩最后几位。 这正是 cmp 要有两列的理由。

⚠️ 这里还不能 cmp('logits', ...) —— logits 上游还有 logits = h @ W2 + b2 没接上。 数值上 dlogits 其实已经完整了，想提前验也可以。

三条可以带走的操作规则

八个式子会忘，这三条不会。

① 要求谁的梯度，就搜它在前向里出现在哪些「等号右边」

counts_sum_inv = counts_sum**-1      # 它被「生产」出来的地方 —— 不算
probs = counts * counts_sum_inv      # 它被「使用」的地方 —— 这才是战场

因为梯度是从下游流回来的，所以只看「被使用」的地方。

这条一下子回答了两个问题：出现在几个等号右边，就有几条路（要不要 +=）； 而那一行的下标关系，决定了是哪种形态。上一篇的判别流程第一步问的正是这个。

顺带也解决了命名混淆：d + 变量名 = 那个变量的梯度。 第 ③ 步叫 dcounts_sum_inv，求的就是 counts_sum_inv 的梯度 —— 尽管那一行代码里 同时出现了 counts、dprobs、counts_sum_inv 三个东西，角色完全不同： 一个是求导对象，一个是上游梯度，一个只是局部导数的值。

② 「有几个求导」= 依赖图上有几条边，不是有几个变量

第 ⑧ 步就是这条的教材：三个变量、三条边、三个求导，但只输出两个梯度。 数变量会算成两个，数边才对。

③ 「梯度应该是 0」本身是一个可以验证的性质

这是全篇最值钱的一条，而且可以直接迁移。

减最大值是一个纯数值技巧 —— 它的全部意义是防止 exp 溢出，数学上不改变任何结果。 而「不改变结果」这句话，是有可执行的检验的：它的梯度必须是 0。

反过来用：

你在自己的代码里加了个「应该不影响结果」的东西（截断、缩放、稳定化、重参数化）， 跑一遍反向看它的梯度 —— 是 0 → 这确实是个纯技巧，模型没变； 不是 0 → 你偷偷改变了模型，那就不是技巧而是 bug。

backward() 不会告诉你哪个梯度「应该」是多少。但你知道的时候，它就成了一把尺子。

收官

五种反向模式在这八步里全部出现过：

后面不会再有第六种。 剩下的全是这五种的重复组合 —— 只是下一段的 BatchNorm 那里， 同样的模式要叠十层，而且那才是唯一让样本互相影响的地方（第十六篇算过那笔账）。