前置知识：把公式接到张量和梯度

读语言模型代码时，最常见的障碍并不是复杂定理，而是丢失了“这一维表示什么”。本章用一个具体批次补齐概率、矩阵和求导。目标是你能检查代码，而不只是看懂讲义上的公式。后面的 CPU 实验采用 NumPy 显式反传；迁移到 PyTorch 时，数学没有改变，框架只接管梯度累积。

一批文本如何表示

设批大小 $B=2$、序列长度 $T=4$、词表 $V=7$、隐藏维度 $d=8$。输入 ids 形状是 [2,4]，元素是整数编号，不是连续特征。嵌入表 $E\in {\mathbb{R}}^{7\times 8}$ 经索引后得到 $X=E[ids]\in {\mathbb{R}}^{2\times 4\times 8}$。线性层 $W\in {\mathbb{R}}^{8\times 7}$ 把最后一维映射成词表 logits，前两个维度保留，结果为 [2,4,7]。

线性层不需要把批次和时间维写进参数。它对每个位置应用相同映射，参数共享使任意位置都可以学习同一组语言规律。需要矩阵梯度时，才把 [B,T,d] 临时展成 [BT,d]；这不应改变 token 的相邻关系。

# 形状审计，不依赖具体框架
x = embedding[ids]             # [B,T,d]
logits = x @ output_weight     # [B,T,V]
flat_x = x.reshape(B*T, d)

广播也必须有语义。例如增益 gain[d] 可以作用于所有位置；因果 mask [T,T] 可以广播到 [B,H,T,T]；目标 targets[B,T] 却不能直接与 logits 相减。目标先对应一个词表位置，概念上相当于 one-hot [B,T,V]。

从概率到稳定交叉熵

logits $z_i$ 是未归一化分数。softmax 给出 $p_i=e^{z_i}/\sum _j{e}^{z_j}$。若目标是 $y$，负对数似然为

$$\ell =-\log p_y=\log \sum _j{e}^{z_j}-z_y.$$

直接计算 $e^{1000}$ 会溢出。令 $m=\underset{j}{max}{z}_j$，利用 softmax 的平移不变性，写成

$$\ell =\log \sum _j{e}^{z_j-m}-(z_y-m).$$

本地 cross_entropy 正是这个式子。它没有先算概率再取 log，因为很小的概率可能先下溢到零。对所有 logits 加同一个常数，概率与 loss 应保持不变，这是很有用的数值测试。当 logits 全零时，$p_y=1/V$，loss 为 $\log V$；这也是训练前的量级参照，但随机初始化不保证恰好等于均匀分布。

最值得手推的一个梯度

因为 $\partial \log \sum _j{e}^{z_j}/\partial z_i=p_i$，所以

$$\frac{\partial \ell }{\partial z_i}=p_i-\mathbf{1}[i=y].$$

平均 $BT$ 个 token 后，还需除以 $BT$。漏除会让梯度随批大小成比例放大，等价于偷偷改学习率。设 $Z=XW$，链式法则给出 ${\mathrm{\nabla }}_{W}L=X^T{\mathrm{\nabla }}_{Z}L$，${\mathrm{\nabla }}_{X}L={\mathrm{\nabla }}_{Z}L{W}^T$。参数在所有位置共享，因此权重梯度要汇总各位置的贡献。

同理，同一个 token 在输入出现多次，其嵌入梯度必须累加。NumPy 的高级索引 g[ids] += dx 对重复索引并不提供所需的累加语义。本地实现使用 np.add.at(g, ids, dx)；PyTorch 的 embedding backward 会处理这种累积。这个细节体现了“数学式正确”与“数组代码正确”的区别。

梯度检查的意义和边界

对某个参数标量 ${\theta }_i$，中心差分

$$g_i^{\mathrm{num}}=\frac{L({\theta }_i+ϵ)-L({\theta }_i-ϵ)}{2ϵ}$$

可以独立检查手写反传。本地测试用 float64、$ϵ={10}^{-5}$，在每个参数张量抽取若干元素，覆盖两层模型的残差、旋转、归一化与重复嵌入。它不会证明全部参数永远正确，但比“loss 恰好下降”强：错误梯度也可能暂时下降。

flowchart LR
  A[相同输入和参数] --> B[解析反向传播]
  A --> C[单个参数加减 epsilon]
  C --> D[两次前向求差分]
  B --> E[比较绝对与相对误差]
  D --> E

查看流程图文本

flowchart LR
  A[相同输入和参数] --> B[解析反向传播]
  A --> C[单个参数加减 epsilon]
  C --> D[两次前向求差分]
  B --> E[比较绝对与相对误差]
  D --> E

检查时冻结随机性，关闭 dropout，不在两次前向之间更新参数。float32 下太小的 $ϵ$ 容易被舍入吞没，太大又偏离局部导数；检查失败要先分清数值精度与公式错误。

对照 PyTorch 的最小训练语义

# 说明框架语义的片段；本地可执行入口仍是 train.py
optimizer.zero_grad(set_to_none=True)
logits = model(input_ids)
loss = torch.nn.functional.cross_entropy(
    logits.reshape(-1, vocab_size), targets.reshape(-1)
)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()

backward() 默认累积到参数梯度，zero_grad() 因而是一次更新周期的起点；梯度累积训练则有意在多个 microbatch 之间不清零，并按等效总批次归一化。eval() 改变 dropout 等层的行为，no_grad() 关闭梯度记录，两者不是同一件事。

自测

排查张量程序时，先写一份极小输入通常比打印大模型更有效。例如让每个位置带上唯一的整数标签，执行 reshape、transpose、再逆变换，确认标签回到原位置；给线性层设置单位矩阵，检查它是否保留对应特征。归一化则可以先输入全零向量、所有分量相同的向量和只有一个非零分量的向量，观察 epsilon 与缩放的作用。测试不仅要覆盖“随机正常输入”，还要让轴错误、广播错误和重复索引错误产生容易辨认的输出。

1. [2,4,8] @ [8,7] 的输出是什么？若只计算最后一个位置的 loss，仍然是完整 teacher-forcing 训练吗？
2. 为什么每一行 logits 的交叉熵梯度之和应为零？
3. 一个 embedding ID 出现三次，其参数是否只更新一次？

展开答案

1. 输出 [2,4,7]。只计算最后一位会丢掉其余七个位置的监督；它可以定义另一种目标，但不能声称与八个位置平均的目标相同。2. 概率之和为一，one-hot 之和也为一；二者相减的和为零，也对应整体平移不改变 loss。3. 优化器每步更新该参数一次，但更新所用梯度必须包含三次出现的贡献之和。

来源：PyTorch 自动微分说明、官方 A1 2025 快照。下一章：分词。