1  注意力：从加权平均到缩放点积

本章地图

你在这里：一个 Transformer 的数据流的第 3 步 —— 每层 Block 里那个标红的方块（多头注意力）。 本章解决：注意力为什么写成 \mathrm{softmax}(QK^{\top}/\sqrt{d_k})V 这个形状；每一项各在解决什么问题；去掉任一项会在什么条件下崩。 前置：微积分、线性代数、概率论，会一点 PyTorch。不需要先读原论文。 读完能做：给你一份别人的注意力实现，能说出每个张量的形状、每个系数的来历，以及它写错时会在哪里出事。

全景：一个 Transformer 长什么样

先说这一节为什么存在。

本书每一章都从一个具体的故障切入：某个系数去掉之后模型训不动、 某个归一化放错位置之后 12 层不如 4 层、某个缓存写错之后输出的文本一字不差但逻辑全错。 这些故事都成立的前提是——你知道那个零件装在机器的哪个位置。

如果你还没读过 Attention Is All You Need，直接翻到下一节会很难受： \sqrt{d_k}、softmax、q 与 k 的分量方差，全是没定义过的词。

所以这一节只做一件事：把整台机器摆出来给你看。 不推导、不调参、不做实验——只有一张图，和沿着图走一遍的五个步骤。大概八分钟。

图 1.1　 一个 Transformer 的完整数据流。每一步标出张量形状，右侧是负责它的章节。本章只讲红框那一块——但它占的位置，是你理解后面所有内容的地基。

读这张图的正确方式：不要试图记住所有方框。只记住两件事—— 数据从上往下流，每经过一个方框变一次形状； 以及，你这一章要读的，只是其中一个方框。

下面沿着图走一遍。

第 1 步：文本变成数字

模型不认字，只认整数。第一步是分词（tokenization）：把文本切成单元， 每个单元查表得到一个整数 id。

上图的例子用的是字符级分词——一个字就是一个 token， 所以「床前明月光」是 5 个 id。真实的模型多用 BPE 之类的子词分词， 一个 token 可能是一个词、半个词、或者一个汉字。本书的小模型用字符级， 因为它的词表只有 1200，训练几秒就能收敛，而且每个 token 的含义肉眼可见。

这一步的形状变化：一段 5 个字的文本 → 一个长度为 5 的整数数组。

注记

为什么分词方式会影响后面的一切

词表大小直接决定了两件事：嵌入表有多大（参数量 = 词表 × 维度）， 以及同样的文本会变成多少个位置（上下文窗口的消耗速度）。

第 3 章的案例就是一个直接的后果：字符级 + 1200 的字表，对全唐诗的覆盖率只有 87.6%， 剩下 12.4% 的字全被映射成同一个 <unk>。 于是那个模型生成的文本里，<unk> 一度占了 57 个 token 里的 45 个。 分词不是预处理细节，它决定了模型能看见什么、看不见什么。

第 2 步：每个数字变成一个向量

一个整数本身没有意义——1234 和 1235 之间不存在”差一”的关系。 所以要查表：每个 id 对应一个 d_{\text{model}} 维的向量，这张表就是词嵌入。

本书的小模型取 d_{\text{model}} = 128，词表 1200，所以嵌入表是 1200 \times 128 ——153,600 个参数，占了整个模型的三分之一还多。

然后是位置编码。 注意力有一个先天缺陷：它对输入的顺序完全不敏感—— 把「明月」和「月明」喂进去，算出来的结果只是两个向量换了位置。 所以必须额外把”这是第几个位置”的信息加进去：

X_0 = \text{Embed}(\text{ids}) + \text{Pos}(0,1,2,\ldots)

这一步的形状：长度为 5 的整数数组 → 形状 (5, 128) 的矩阵。 从这以后，整条流水线里流动的东西一直是这个形状， 直到最后一步才变。

第 1 章的后半部分（§1.9）会专门讲位置编码的三代方案： 正弦编码、可学习嵌入、RoPE。

第 3 步：N 层 Block —— 这一章的主角在这里

这是整台机器的核心，也是本书上篇四章的全部内容。

一层 Block 做两件事，每件事都是”算一个新东西，再加回原来的表示上”：

\begin{aligned} X &\leftarrow X + \text{Attention}(\text{LN}(X)) &&\text{（第 1 章）}\\ X &\leftarrow X + \text{FFN}(\text{LN}(X)) &&\text{（第 2 章）} \end{aligned}

上面那一行就是本章要拆的东西。 它内部长这样：

\text{Attention}(Q,K,V) = \mathrm{softmax}\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V

先别管公式。 这一行在做的事是：每个位置去问所有位置”你和我有多相关”， 然后按相关程度把大家的信息加权平均过来。

• Q（查询）说”我在找什么”
• K（键）说”我有什么”
• V（值）说”如果被选中，我给出什么”
• \sqrt{d_k} 是本章案例的主角

这一段为什么单独成章：因为它是整个模型里唯一让不同位置互相看见的地方。 FFN 是逐位置独立作用的，嵌入是查表——只有注意力在位置之间搬运信息。 第 4 章会问：既然只有它搬运，那能不能从它的权重里读出模型在做什么？

第 4 步：从向量回到概率

走完 N 层之后，每个位置都有一个 128 维的向量。 最后一步把它变成词表上的一个概率分布：

1. LN：把向量的尺度归一化（第 2 章讲为什么必须有这一步）
2. 输出投影：128 \to 1200，得到一个 1200 维的 logits 向量
3. softmax：把 logits 变成概率

输出投影的权重通常和词嵌入共享（本书的小模型就是这样）。 理由很自然：嵌入是”词 → 向量”，输出投影是”向量 → 词”， 两个方向的映射本来就该是同一张表。共享还省下 153,600 个参数—— 对一个小模型来说是三分之一。

这一步的形状：(5, 128) → (5, 1200) → 每一行是一个概率分布。

第 5 步：训练和推理，是两套流程

同样一条流水线，训练和推理的走法完全不同。

• 训练：整段文本一次性喂进去，5 个位置同时算出”下一个字是什么”的预测， 和真实的答案比对，算交叉熵。一次前向覆盖所有位置。
• 推理：没有”下一个字”，只能自己生成——先算出第 1 个，把它接到输入后面， 再算第 2 个。生成 n 个字要 n 次前向。

这个不对称带来了一整套工程问题：为什么需要 KV 缓存、 为什么缓存写错了在生成文本上看不出来、为什么贪心解码会陷入循环。 第 3 章整章都在处理它们。

一笔参数账

把上面五步的参数量加起来（本书的小模型，d_{\text{model}}=128、4 头、2 层、d_{\text{ff}}=256）：

部件	参数量	占全书模型
词嵌入（与输出投影共享）	153,600	36.1%
位置嵌入	8,192	1.9%
2 层 Block	263,936	62.0%
— 其中多层注意力（每层 65,536）	131,072	30.8%
— 其中 FFN（每层 65,920）	131,840	30.9%
最后的 LN	256	0.06%
合计	425,984	100%

这张表值得看两眼。

第一，注意力和 FFN 几乎一样大（30.8% vs 30.9%）—— 因为本书取 d_{\text{ff}} = 2 d_{\text{model}}，而标准配置是 4 d_{\text{model}}。 在标准配置下，FFN 会占到约 2/3 的 Block 参数。 第 2 章会解释这个配比是怎么来的。

第二，嵌入表占了 36%。 在真实的大模型上这个比例会低得多 （词表 5 万、维度 4096 时嵌入是 2 亿，而总参数是几十亿到几千亿）。 本书的小模型”头重脚轻”，这是它和真实模型的一处结构性差异—— 第 2 章 §2.21 会专门列一张”哪些结论能迁移”的表。

这一章只讲其中一块

回头看那张图：红框只占整条流水线的一个方框。

这一章会把那个方框拆到骨头——它的四个投影矩阵、 为什么要有多个头、为什么除以 \sqrt{d_k}、去掉缩放会在什么条件下崩。 但它在整个模型里的角色，你现在已经知道了：它是唯一让位置之间交换信息的地方。

重要

接下来的故事，从那个方框里的一个数字开始

红框里有一个看起来毫无道理的设计：点积之后要除以 \sqrt{d_k}。

论文只用一句脚注带过理由。我们想验证它，于是把这一项去掉了—— 结果第一次没复现出来。

下面这个故事，就从那次失败开始。

一次没复现出来的实验。

2017 年那篇 Attention Is All You Need 给出了今天几乎所有大模型都在用的公式。 其中一处关键选择，作者只用一句脚注带过：点积要除以 \sqrt{d_k}， 因为”当 d_k 较大时点积的数值会变大，把 softmax 推进梯度极小的区域”， 随后补了一句”假设 q 与 k 各分量独立、均值 0 方差 1，则点积的均值为 0、方差为 d_k“。

这句话读起来是完整的：有现象（数值变大）、有机理（梯度变小）、有推导（方差等于 d_k）。 所以我们决定动手验证它——按脚注的说法，把 \sqrt{d_k} 去掉，模型应该训不动。

结果第一次没复现出来。

我们训了一个 42 万参数的字符级 GPT（2 层、4 头、d_{\text{model}}=128）， 语料是全唐诗 4003 首、25.5 万字。同一个模型跑两遍，唯一变量是有没有除以 \sqrt{d_k}：

	loss 起点	loss 终点
不缩放	7.03	4.17
除以 \sqrt{d_k}	7.03	4.00

差 0.17。曲线几乎重合（图 1.2）。如果就此收笔，结论会是 “脚注说的现象在这个规模上不存在”。

但往下查之后，我们找到了两个更值得讲的东西：一是这个脚注的前提在真实架构里 并不成立，二是当把这个前提重新满足时，现象会以极其剧烈的方式出现—— 注意力权重熵掉到 0.0004，雅可比范数掉了一万多倍。

这一章就从这个案例展开。读完你应该能回答：

1. 注意力为什么写成 \mathrm{softmax}(QK^{\top}/\sqrt{d_k})V 这个形状，每一项各在解决什么问题？
2. 去掉缩放因子、换掉 softmax、改成单头——分别会在什么条件下崩，崩成什么样？
3. 为什么小模型上复现不出教科书说的现象？这个”复现不出来”本身说明了什么？

注记这一章的案例可以自己跑

案例的全部代码在 notebooks/run_ch1_case.py，数据在 data/ch1_case.json。 文中出现的每一个数字都来自那次真实运行：

# 需要 PyTorch。本机在 /home/mpp/micromamba/envs/dlspace/bin/python
python notebooks/run_ch1_case.py --steps 1500

CPU 上约 3 分钟，有 GPU 更快。若你跑出的数字与文中略有出入， 多半是 PyTorch 版本或硬件差异导致的浮点非确定性——方向应当一致。

1.1 实验是怎么设计的

在拆解现象之前，先讲清楚这个实验的每一个设置，以及为什么这样设。

为什么用一个自己训的小模型，而不是直接分析 GPT-2？

因为要回答的是”去掉 \sqrt{d_k} 会怎样”——这是一个干预性问题， 需要在受控条件下比较。拿现成模型没法做这个对照：它们都带缩放。 自己训的好处是每一个变量都能改，而且几分钟就能跑完一轮，可以反复试。

代价是规模小。这一点在后面的结论里会明确标出来：方向可迁移，数值不可迁移。

模型：字符级 GPT，2 层 Transformer，Pre-LN 结构，上下文长度 64， 词表 1200 个最高频汉字（覆盖语料 87.6% 的字），参数量 425,984。 训练用 AdamW，学习率 3e-3 带 warmup 与余弦衰减，1500 步，批大小 32。

语料：《全唐诗》4003 首，25.5 万字（公开语料，来自 chinese-poetry 仓库）。

对照：同一份语料、同一个随机种子（20261014）、同样 1500 步。 两遍唯一的差别是注意力里这一行的有无：

att = q @ k.transpose(-2, -1)           # 打分
if self.cfg.scaled:
    att = att / math.sqrt(self.d_head)  # ← 唯一变量

提示为什么用唐诗当语料

这个选择不是随意的，它满足四个条件：

① 规模合适。 25.5 万字、4003 首，正好是”小模型能在几分钟内学到东西、 但又不至于背下来”的量级。太小（比如几千字）模型会过拟合， 太大（比如几亿字）CPU 上跑不动。

② 结构规整。 唐诗有固定的句长（五言、七言）、固定的韵律。 这让一个只有 2 层、42 万参数的模型有可能学到可辨识的模式—— 如果语料是随意的现代散文，这么小的模型大概只能学到字频。

③ 字符集收敛。 5000 个不同汉字、取前 1200 个覆盖 87.6%。 这让我们能用一个很小的词表，把参数预算留给注意力层。

④ 公开且无版权问题。 来自 chinese-poetry 仓库，可以直接放在书里、 让读者复现。

换成别的语料会怎样：结论的方向（缩放的必要性随维度增长）不会变—— 那是数学性质。但具体数字会变：临界维度可能不同，loss 的绝对值也不同。

这也是”方向可迁移、数值不可迁移”那句话的具体含义。

提示关于”唯一变量”

这句话在实验里是个承诺，不是修辞。同一个种子意味着初始化完全相同、 每一批数据完全相同。所以两条曲线在 step 0 时的 loss 完全一致（都是 7.03）， 之后的差异只可能来自那个除法。

如果种子不同，两条曲线的差异里就混进了初始化噪声，结论不再干净。 这是做对照实验的通用纪律：一次只改一个变量，其余全部钉死。

1.2 第一次实验：几乎没差别

图 1.2 是第一次实验的结果。左边是小模型（d_{\text{head}}=32）的两条曲线—— 它们几乎重合，终点差 0.17。

图 1.2　 小模型（d_{\text{head}}=32）与”大”模型（d_{\text{head}}=256）的训练曲线对照。同一语料、同一种子、同步数，唯一变量是有无 \sqrt{d_k} 缩放。左图两条曲线几乎重合，右图在 200 步后明显分离。

这不是噪声造成的偶然。我们换了三个种子重跑，小模型上的差距都在 0.1–0.2 之间浮动， 量级与训练本身的随机性相当。在这个配置下，去掉缩放因子不影响训练。

到这里，一个草率的结论是”原论文的脚注过时了”。但我们更想知道的是：

注记停下来想一想 ①

先别往下看。如果你面对这个结果，下一步会去查什么？

提示：脚注里给了一个假设——“各分量独立、均值 0 方差 1”。如果一个推导的结论 没被实验证实，第一件该怀疑的事往往是它的前提，而不是它的逻辑。

1.3 查前提：实测 logits 到底有多大

脚注的推导逻辑是：q、k 各分量方差为 1 \Rightarrow 点积方差为 d_k \Rightarrow 点积的标准差是 \sqrt{d_k} \Rightarrow 数值大 \Rightarrow softmax 饱和。

这条链上唯一可以直接测量的是第一个箭头。所以我们不训模型， 直接在初始化状态下测：随机造一批 q、k，算点积，看它的标准差。

图 1.3　 初始化状态下点积的标准差随 d_{\text{head}} 的变化。理论值 \sqrt{d_{\text{head}}}（红）与实测值（蓝）在双对数坐标下都是直线，但实测线整体低了约一个数量级。

d_{\text{head}}	实测标准差	理论 \sqrt{d_{\text{head}}}	实测 / 理论	权重熵
32	0.070	5.66	0.012	0.999
64	0.202	8.00	0.025	0.993
128	0.561	11.31	0.050	0.939
256	1.615	16.00	0.101	0.696
512	4.666	22.63	0.206	0.272

实测比理论小了 5 到 80 倍。 而且注意看：两者的斜率是一样的 （都是 d_{\text{head}} 的平方根关系，双对数坐标下都是直线）， 差的只是一个常数因子。

这说明：脚注推导的形式是对的，但它代入的那个常数——方差为 1——不对。

1.3.1 常数哪里去了

问题出在两处，都是原论文脚注里没有、但现代架构里必然存在的：

第一，q 和 k 不是单位方差，它们的尺度由初始化决定。 在标准初始化下，投影矩阵 W_Q、W_K 的元素服从 \mathcal{N}(0, \sigma^2)， GPT-2 系列用 \sigma = 0.02。输入 x 经 LayerNorm 后每个分量方差为 1， 于是 q_i = \sum_j W_{ij} x_j 的方差约为

\mathrm{Var}(q_i) \approx d_{\text{model}} \cdot \sigma^2 \tag{1.1}

代入 \sigma = 0.02、d_{\text{model}} = 512，得 \mathrm{Var}(q_i) \approx 0.205， 标准差只有 0.45 —— 而不是脚注假设的 1。

第二，LayerNorm 归一化的是注意力的输入，不是 q 和 k 本身。 这是最容易忽略的一点。LayerNorm 保证的是”送进 W_Q 的那个 x 方差为 1”， 但 W_Q 一乘，尺度就由权重说了算。所以 LayerNorm 并不能替代 \sqrt{d_k}， 它管的是另一件事。

把第一点代回脚注的推导：若每个 q_i 的方差是 v 而非 1，则

\mathrm{Var}(q^{\top} k) = \sum_{i=1}^{d_{\text{head}}} \mathrm{Var}(q_i k_i) = d_{\text{head}} \cdot v^2 \tag{1.2}

标准差是 v\sqrt{d_{\text{head}}}。结构没变，只是前面多了一个 v。

我们核一下数：d_{\text{model}} = 512 时 v \approx 0.205， 预测标准差 = 0.205 \times \sqrt{512} = 4.64。实测是 4.666。对上了。

警告一处我们没有对上的地方

上面对 d_{\text{model}} = 512 核得很好，但对 d_{\text{head}} = 32 就不准了： 按同样算法预测是 0.0128，实测是 0.070，差了五倍多。

我们没有把公式调一调让它好看。诚实的说法是：式 式 1.2 给出了正确的 标度关系（正比于 \sqrt{d_{\text{head}}}）和正确的数量级，但没有精确到常数。 真实的常数还受 LayerNorm 的具体实现、残差累加、位置嵌入的贡献影响， 精确求解需要逐层追踪方差传播，超出本章范围。

所以本章正文中凡是给出具体数字的地方，一律引用实测值； 理论式只用来解释趋势。这也是读论文时该有的习惯： 公式告诉你标度，测量告诉你数值，两者不可互相替代。

1.4 把前提重新满足

到这里，案例的第一半有了答案：小模型上没复现出崩溃，是因为 q、k 的实际尺度 远小于 1，点积根本没大到让 softmax 饱和。 脚注描述的现象需要一定的规模才会显现。

那么就把它做出来。我们把 d_{\text{head}} 从 32 提到 256 （d_{\text{model}} 从 128 提到 256，头数从 4 改成 1），其余全部不动，重跑同一个对照。

图 1.2 的右图是结果：

	loss 起点	loss 终点	终点差距
d_{\text{head}}=32 不缩放	7.03	4.17	—
d_{\text{head}}=32 缩放	7.03	4.00	0.17
d_{\text{head}}=256 不缩放	7.05	5.18	—
d_{\text{head}}=256 缩放	7.05	4.23	0.95

差距从 0.17 变成 0.95，放大了 5.6 倍。而且不缩放那条曲线在 200 步左右就卡住了—— 后面 1300 步几乎没有进展。

更直接的证据来自注意力层的内部状态。图 1.4 记录了训练过程中 两个关键量的变化：

图 1.4　 训练过程中注意力层的两个诊断量。上：权重熵（1 = 完全平均，0 = 完全 one-hot）。下：softmax 雅可比的 Frobenius 范数（对数轴）。红为不缩放、蓝为缩放。红色曲线在 200 步内就塌到底部。

训练步数	不缩放：权重熵	不缩放：\|J\|_F	缩放：权重熵	缩放：\|J\|_F
0	0.712	0.2975	0.9985	0.2017
222	0.0022	0.00471	0.6437	0.3029
666	0.0005	0.00100	0.6205	0.2986
1480	0.0001	0.00008	0.5745	0.3089

不到 250 步，不缩放那一边的权重熵就从 0.71 掉到了 0.0022 —— softmax 输出几乎完全是 one-hot，每个位置只死死盯住唯一一个位置。 雅可比范数从 0.2975 掉到 0.00008，掉了一万倍以上。

而缩放那一边，这两个量在 1500 步里始终平稳：熵在 0.57–0.68 之间， 雅可比范数在 0.29–0.31 之间。

这就是”梯度消失”的直接观测。

雅可比范数衡量的是”softmax 的输出对输入的敏感度”。 它掉了四个数量级，意味着 softmax 这一层几乎不再传递任何梯度信号—— 下游的 loss 变化无法通过它传回到 Q、K 的投影矩阵上。

注意一个容易误解的地方：loss 并不是不降了（它从 7.05 降到了 5.18）。 因为模型的其他部分——词嵌入、MLP、输出层——还在正常学习。 只有注意力这一层”死”了。所以不是”训不动”，而是 “带着一个废掉的注意力层硬训”。

我们还可以让这个模型说话。同一个提示词、同一个采样种子：

不缩放：秦川雄帝宅，文還物金，。。。化來華舟由，王夜沈，。，殺蘭雲，外騎天岸繞江林上
缩  放：秦川雄帝宅，山靜還金。天臺，出山化，天。由歸王此，夫。蘭，天岸。上

不缩放那条里出现了 ，。。。 这样的连续标点——这是退化文本的典型征兆。 不过要诚实地说：小模型的采样样本差异并不明显到可以作为主要证据， 真正硬的证据是 loss 差 0.95 和雅可比范数掉四个数量级。

注记停下来想一想 ②

上面说”注意力层死了，但模型还在学”。那么请想一个问题：

如果注意力层完全退化成”每个位置只看前一个位置”（这其实等价于一个固定的 一阶马尔可夫模型），这个 42 万参数的模型还能把 loss 降到多少？

再想一层：为什么 loss 最终停在 5.18 而不再下降？ 如果继续训 10000 步会怎样？

案例讲到这里，现象、排查、以及”什么时候才会崩”都已经清楚了。但还有一个问题悬着：

为什么 logits 一大，softmax 就饱和？饱和为什么等于梯度消失？

这两个问题需要把 softmax 和缩放因子彻底拆开才能回答。这是本章接下来的主线。 在拆它们之前，先把注意力本身的形式讲清楚——它其实是一个非常朴素的构造。

1.5 注意力是一次软查表

先忘掉所有公式，也忘掉神经网络。

1.5.1 先算一个具体的小例子

假设我们有三个字，要算第三个字对前两个字的注意力。为了能手算，把维度设成 d = 2—— 这在真实模型里是不存在的（d 通常是 64 或 128），但计算过程一模一样， 而所有的中间数字都能写在纸上。

三个位置各有自己的值向量 v（它承载”这个位置的信息是什么”）：

位置	字	v
1	明	[2,\ 0]
2	月	[0,\ 2]
3	松	[1,\ 1]

现在位置 3（“松”）要输出它的新表示。它先给出一个查询向量 q_3 = [1,\ 0]， 表示”我想找什么”。然后拿它去和每个位置的键向量 k 比对—— 键表示”我这里有什么可供检索的”：

位置	k
1	[1,\ 0]
2	[0,\ 1]
3	[1,\ 1]

第一步：打分。 用内积衡量相似度：

q_3 \cdot k_1 = 1 \times 1 + 0 \times 0 = \mathbf{1} q_3 \cdot k_2 = 1 \times 0 + 0 \times 1 = \mathbf{0} q_3 \cdot k_3 = 1 \times 1 + 0 \times 1 = \mathbf{1}

内积大意味着两个向量方向接近。这里位置 1 和位置 3 的分数是 1，位置 2 是 0。

第二步：归一化成权重。 我们要的是”分配多少注意力”，所以需要一个和为 1 的分布。 用 softmax（下一节专门讲它为什么在这儿）：

\mathrm{softmax}([1,\ 0,\ 1]) = \frac{[e^1,\ e^0,\ e^1]}{e^1 + e^0 + e^1} = \frac{[2.718,\ 1,\ 2.718]}{6.436} = [\mathbf{0.422},\ \mathbf{0.155},\ \mathbf{0.422}]

第三步：加权求和。

\text{输出} = 0.422 \times [2,0] + 0.155 \times [0,2] + 0.422 \times [1,1] = [1.267,\ 0.733]

这就是全部。三个数字，三次运算，注意力算完了。

提示所以注意力在做什么

用一句话说：它把”我要找什么”（q）拿去和每个位置”有什么”（k）比对， 按相似度分配权重，然后把各位置的”内容”（v）加权平均起来。

注意它输出的不是一个选择，而是一个混合。 上例中位置 3 的新表示是位置 1、3 的内容各占 42%、位置 2 占 16% 混出来的。 这一点是注意力与传统检索最本质的区别。

1.5.2 手算一遍掩码：三个位置各看到什么

语言模型是自回归的——预测第 t 个字时，不能看到第 t+1 个字（那是答案）。 这个约束靠因果掩码实现。它看起来简单，但值得手算一遍。

沿用上面那组 Q、K、V。先算完整的打分矩阵（QK^{\top}，按手算惯例先不缩放）：

S = \begin{pmatrix} 1 & 0 & 1 \\ 0 & 1 & 1 \\ 1 & 0 & 1 \end{pmatrix} \tag{1.3}

现在加因果掩码：位置 i 只能看 j \le i，其余位置置为 -\infty。 逐行算 softmax：

位置	可见位置	有效打分	softmax 权重	输出
0	只有 0	[1]	[1.0]	[2.000,\ 0.000]
1	0, 1	[0,\ 1]	[0.269,\ 0.731]	[0.538,\ 1.462]
2	0, 1, 2	[1,\ 0,\ 1]	[0.422,\ 0.155,\ 0.422]	[1.267,\ 0.733]

三个观察：

① 位置 0 的权重必然是 [1.0]。 它只看得见自己， softmax 作用在单个元素上恒等于 1，输出就是 v_0 原样。 这是掩码的边界情况，也是实现时最容易出问题的地方 （如果掩码写错，位置 0 可能拿到别的权重）。

② 每个位置的权重和都是 1，但”分母”不同。 位置 1 的 softmax 只在 2 个元素上归一化， 位置 2 在 3 个上。掩码必须在 softmax 之前生效—— 如果先 softmax 再掩码，归一化的基准就错了，权重和不再是 1。

③ 位置 2 的输出 [1.267,\ 0.733] 与前面手算的一致。 这说明”未掩码时的第三个位置”和”掩码后的第三个位置”结果相同—— 因为在自回归设定下，位置 2 本来就能看到前两个，掩码对它没有额外限制。

提示一个反直觉的推论

看位置 1 的输出 [0.538,\ 1.462]，和位置 2 的 [1.267,\ 0.733]—— 同一个字”松”（位置 2）在不同上下文里的表示完全不同。

这正是 Transformer 表达力的来源：每个位置的表示不是固定的词向量， 而是”从上下文里按需取回来的混合”。

对比一下词嵌入：如果不用注意力，“月”在任何句子里都是同一个向量。 用了注意力之后，它的表示取决于这句话里”月”该向谁借信息。

这就是”上下文相关表示”（contextual representation）的实际含义， 也是 BERT、GPT 这类模型与静态词嵌入（word2vec）的根本区别。

注记停下来想一想 ⑮

上面说”位置 0 的权重必然是 [1.0]“。

但实现时有个陷阱：如果掩码用 -\infty，那么位置 0 那一行是 [s_{00},\ -\infty,\ -\infty]。softmax 之后是 [1,\ 0,\ 0]——没问题。

可如果这一行全是 -\infty 呢？（比如某个位置被所有位置屏蔽了， 或者 mask 写反了。）e^{-\infty} = 0，分子分母都是 0，结果是 nan。

请想：什么样的实现错误会导致”整行全被屏蔽”？ （提示：想想 padding。一个 batch 里有长句有短句，短句要补齐—— 那些补齐的位置该不该被看到？）

1.5.3 自己动手改一遍

上面那组数不是”算给你看”的——每个数字都可以改。 下面的实验把 Q、K、V 三个矩阵做成可编辑的表格， 改动会立刻重算打分、softmax 和输出。

这个实验值得花十分钟，因为注意力的全部机制都在这三个矩阵里。 几个具体的探索方向：

① 先点「正文那个例子」。 这就是上面手算的那组数。 取消勾选「除以 \sqrt{d_k}」，第三个位置的输出应该正好是 [1.267,\ 0.733]； 勾上之后变成 [1.203,\ 0.797]。缩放确实改变了结果——在 d = 2 时改变很小， 而本章案例里 d_{\text{head}} = 256 时是天壤之别。这个对比本身说明了缩放的尺度依赖性。

② 点「只让自己可见」。 权重矩阵变成对角阵——每个位置只复制自己的 V。 这时候注意力层退化成了一个恒等映射。想想：这样的层对模型还有什么用？ （它其实还有用——它就是本章头部画像里”层 0 头 0”的极端版本， 配合残差连接可以起到”信息保持”的作用。）

③ 把 V 的三行改成 [1,0]、[0,1]、[1,1]，Q 第三行改成 [1,0]。 输出又是 [1.267,\ 0.733]——和 ① 一样。这说明输出只取决于 Q、K、V 的乘积关系， 不同的矩阵组合可以得到同样的结果。

④ 点「制造一个尖锐分布」，然后把 K 的第一行两个数改成两倍。 观察 softmax 的熵怎么变。这就是本章案例的微型复现： 放大 K 的尺度 → 打分变大 → softmax 更尖 → 熵下降。 只不过在一个 3 个位置、2 维的玩具上，这个变化是无害的； 在 256 维、真实训练中，它会变成灾难。

提示这个实验为什么不让你改”层数”“头数”

因为它只做一件事：把一次注意力前向传播拆开给你看。 加上层数、头数、残差、LayerNorm 之后，你就看不清”哪个数字影响了哪个数字”了。

理解复杂系统的一个实用原则：先在一个能完整看见的规模上搞懂机制， 再带着这个理解去看真实的复杂实现。 反过来（先看复杂实现，再回头找原理） 通常会更慢。

1.5.4 为什么不直接取最大的

既然是要”找最相关的”，为什么不直接选分数最高的那个位置，返回它的 v？

\text{硬查表：}\quad \text{输出} = v_{\arg\max_j (q \cdot k_j)}

这就是普通的字典查找。它的两个致命问题：

第一，不可导。 \arg\max 是阶跃函数，梯度处处为零（除了不可导的跳变点）。 神经网络靠梯度学习，一处不可导，梯度就传不过去。

第二，没有退路。 硬查表一旦选错，就是彻底错——没有”稍微看一点”这个选项。 而语言里的相关性往往是连续的：处理”照”这个字时， 它可能主要该看”月”、次要看”松”、几乎不看”明”。

我们需要的是一张权重表，而不是一个指针。这就是”软查表”：

\text{软查表：}\quad \text{输出} = \sum_j \alpha_j v_j, \qquad \sum_j \alpha_j = 1, \quad \alpha_j \ge 0 \tag{1.4}

图 1.5　 硬查表与软查表的对比。左：\arg\max 只选出唯一位置，输出离散，梯度不可传。右：softmax 给出连续权重，输出是所有位置内容的加权混合，梯度可以回传。

注记停下来想一想 ③

软查表解决了”不可导”，但引入了一个新问题：

如果权重分布非常尖锐（比如 [0.9999,\ 0.00005,\ 0.00005]）， 它和硬查表在行为上几乎没区别。那么它在梯度上还有区别吗？

先别翻书。用一句话描述你的直觉，然后在读完 小节 1.7 后回来核对。

1.5.5 一次性算所有位置

上面只算了位置 3。实际中所有位置都要算，而且可以并行——这正是 Transformer 相比 RNN 的关键优势之一。

把所有位置的 q 堆成矩阵 Q \in \mathbb{R}^{n \times d_k}， 所有位置的 k 堆成 K \in \mathbb{R}^{n \times d_k}，则所有两两打分就是一次矩阵乘法：

S = QK^{\top} \in \mathbb{R}^{n \times n}, \qquad S_{ij} = q_i \cdot k_j \tag{1.5}

再逐行归一化、乘以 V，就得到那行著名公式：

\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V \tag{1.6}

人话翻译：QK^{\top} 算出”每个位置和每个位置的相似度”，得到一个 n \times n 的表； \mathrm{softmax} 把表里每一行变成一个概率分布（这个位置该往哪些位置借信息、各借多少）； 最后乘 V 就是按这些权重把内容取回来。

除以 \sqrt{d_k} 是给打分做归一化——它的必要性是本章后半段的主线， 现在只需要知道”它在那里是为了控制数值尺度”。

复杂度：S 是 n \times n 的，所以计算量是 O(n^2 d)，显存也是 O(n^2)。 序列长度翻倍，这两项都变成四倍。这是长上下文昂贵的根本原因， 也是 FlashAttention 这类工作要解决的问题（见 小节 1.26）。

1.6 为什么是 softmax

式 1.6 里有两个选择看起来是”约定俗成”的：为什么用 \exp？ 为什么要归一化？这一节把它们拆开。

1.6.1 先手算一遍

拿 小节 1.5 里那三个分数 [1, 0, 1] 走一遍：

步骤	计算	结果
① 减去最大值（数值稳定）	[1-1,\ 0-1,\ 1-1]	[0,\ -1,\ 0]
② 取指数	[e^0,\ e^{-1},\ e^0]	[1,\ 0.368,\ 1]
③ 求和	1 + 0.368 + 1	2.368
④ 归一化	[1/2.368,\ 0.368/2.368,\ 1/2.368]	[0.422,\ 0.155,\ 0.422]

和第 ① 步之前直接算的结果一致——这不是巧合，而是 softmax 的平移不变性，下面会推。

现在把分数换一组夸张的：[10,\ 0,\ 0]。

步骤	结果
① 减最大值	[0,\ -10,\ -10]
② 取指数	[1,\ 4.54\times10^{-5},\ 4.54\times10^{-5}]
③ 求和	1.0000908
④ 归一化	[0.99991,\ 0.000045,\ 0.000045]

这就是”饱和”：输出几乎是 one-hot，第一个位置拿走了 99.99%。

1.6.2 它是 \arg\max 的光滑近似

为什么用 \exp？因为它给出了一族可以调节”硬度”的、可导的 \arg\max。 关键是 log-sum-exp 函数：

\mathrm{LSE}_\tau(z_1,\dots,z_n) = \tau \log \sum_{i=1}^{n} \exp(z_i/\tau) \tag{1.7}

它有两个极限行为：

• \tau \to 0^+ 时，\mathrm{LSE}_\tau \to \max_i z_i（最大的那一项主导求和）
• \tau \to \infty 时，\mathrm{LSE}_\tau \to 各 z_i 的均值

对它求关于 z_i 的偏导：

\frac{\partial\, \mathrm{LSE}_\tau}{\partial z_i} = \tau \cdot \frac{1}{\tau}\cdot\frac{\exp(z_i/\tau)}{\sum_j \exp(z_j/\tau)} = \frac{\exp(z_i/\tau)}{\sum_j \exp(z_j/\tau)} = \mathrm{softmax}(z/\tau)_i \tag{1.8}

人话翻译：\max 是”选最大的那个”，不可导；\mathrm{LSE} 是它的可导版本。 对 \mathrm{LSE} 求导正好得到 softmax——所以 softmax 就是 \arg\max 的可微替身。

它给出的不是”选哪个”，而是”每个候选该分到多少”。这正是 式 1.4 需要的。

\tau 就是温度。它控制这个近似有多”硬”：\tau 小则接近 one-hot，\tau 大则接近均匀。 这个旋钮在采样、知识蒸馏、注意力锐化里到处出现，第 3 章还会遇到它。

1.6.3 手算一遍温度：同一个分布，三副面孔

光说”控制硬度”太抽象。拿同一组 logits，用三个不同的 \tau 算一遍。

取 z = [2,\ 1,\ 0]，分别用 \tau = 0.5、1.0、2.0。

\tau = 0.5（更硬）：先除以温度得 [4,\ 2,\ 0]。

步骤	计算	结果
减最大值	[4-4,\ 2-4,\ 0-4]	[0,\ -2,\ -4]
取指数	[1,\ 0.1353,\ 0.0183]	—
求和	1 + 0.1353 + 0.0183	1.1537
归一化	逐项除以 1.1537	[\mathbf{0.867},\ \mathbf{0.117},\ \mathbf{0.016}]

\tau = 1.0（原始）：[2,\ 1,\ 0]，减最大值得 [0,\ -1,\ -2]， 指数 [1,\ 0.3679,\ 0.1353]，和 1.5032，结果 [\mathbf{0.665},\ \mathbf{0.245},\ \mathbf{0.090}]。

\tau = 2.0（更软）：除以温度得 [1,\ 0.5,\ 0]，减最大值得 [0,\ -0.5,\ -1]， 指数 [1,\ 0.6065,\ 0.3679]，和 1.9744，结果 [\mathbf{0.506},\ \mathbf{0.307},\ \mathbf{0.186}]。

把三个结果并排看：

\tau	分布	最大权重	归一化熵
0.5	[0.867,\ 0.117,\ 0.016]	0.867	0.402
1.0	[0.665,\ 0.245,\ 0.090]	0.665	0.758
2.0	[0.506,\ 0.307,\ 0.186]	0.506	0.929

三个具体结论。

① 温度是”除以”不是”乘以”。 \tau 变小 → 除以一个小于 1 的数 → logits 被放大 → 分布更尖。这个方向容易记反。

② 温度不改变大小顺序。 三种情况下都是第 1 项最大。softmax 是单调的—— 它只改变”差距有多大”，不改变”谁比谁大”。

③ \tau = 1 时 softmax 退化成它自己。 这不是废话： 它说明”注意力里没有温度”其实等价于”温度固定为 1”。 如果想要可调的温度，就是在 softmax 之前多一个除法。

所以本章那个 \sqrt{d_k} 缩放，在数学形式上就是”温度固定为 \sqrt{d_k} 的 softmax”。 这是一个值得记住的观察——它说明缩放因子不是新东西， 而是把温度这个通用旋钮设成了一个特定值。

注记停下来想一想 ⑬

既然缩放本质上是”设温度”，那么一个自然的问题是：

为什么温度设成 \sqrt{d_k}，而不是让它像其他超参数一样学出来？

提示：想两个角度 ① 如果温度是可学的，它需要多少额外的参数、训练数据够不够？ ② 如果不缩放，训练刚开始时（参数还是随机的）会发生什么—— 那时候还没有任何数据能告诉模型”温度该是多少”。

1.6.4 两个必须知道的性质

性质一：平移不变性。 对任意常数 c：

\mathrm{softmax}(z + c\mathbf{1}) = \mathrm{softmax}(z) \tag{1.9}

代入即得，因为分子分母同乘 e^c 后约掉。

这条性质有实际后果：给所有 logits 加同一个数不改变注意力。 所以位置编码、偏置项这类”整体抬升”对注意力无效——只有相对差异才起作用。 这也是为什么注意力对输入的整体尺度不敏感，但对相对尺度极其敏感—— 后者正是 \sqrt{d_k} 要管的事。

性质二：数值稳定必须手写。 直接算 \exp(z_i) 在 z_i 稍大时就溢出 （float32 的上限对应 e^{88.7}）。标准做法是先减最大值：

\mathrm{softmax}(z)_i = \frac{\exp(z_i - z_{\max})}{\sum_j \exp(z_j - z_{\max})} \tag{1.10}

由平移不变性，这样算结果完全等价，但指数最大只到 e^0 = 1，永不溢出。

这不是优化，是正确性。 下面这段代码可以直接粘进 notebook 跑， 你会看到不做稳定化的版本返回什么：

import numpy as np

def softmax_naive(z):
    e = np.exp(z)                    # ← 没有减最大值
    return e / e.sum()

def softmax_stable(z):
    z = z - z.max()
    e = np.exp(z)
    return e / e.sum()

z = np.array([100.0, 101.0, 102.0])  # 完全正常的 logits
print("朴素版:", softmax_naive(z))    # [nan nan nan]  ← 溢出成 inf，inf/inf = nan
print("稳定版:", softmax_stable(z))   # [0.09003057 0.24472847 0.66524096]  ✓

在 float32 下，logits 超过约 88 就会触发。 而本章的案例里， 不缩放那一边的 logits 标准差达到了 5477 —— 高出两个数量级。 所以那个模型里如果不做稳定化，训练会直接变成 nan。

顺带一提：这也解释了为什么 PyTorch 的 F.softmax 内部一定要做这个减法—— 它不是可选的性能优化。

注记停下来想一想 ④

回到 小节 1.5 的检查点 ③：尖锐的权重分布与硬查表在梯度上有区别吗？

现在可以定量回答了。考虑极端情况 p = [1, 0, 0]（完全 one-hot）， 用雅可比公式 J_{ij} = p_i(\delta_{ij} - p_j) 把整个矩阵写出来，看它是不是全零。

然后再算 p = [0.9999,\ 0.00005,\ 0.00005] 的 J_{11} = p_1(1-p_1) 是多少。

这个数字就是本章案例里”梯度消失”的量化含义。

1.7 缩放因子到底修正了什么

现在回收案例。我们已经知道：

• 小模型（d_{\text{head}}=32）上去掉缩放，训练照常（差 0.17）
• 大模型（d_{\text{head}}=256）上去掉缩放，注意力层崩掉（熵 0.0004，\|J\|_F 掉一万倍）

缺的是中间那一环：为什么 logits 变大会导致 \|J\|_F 变小？ 这一节把它推完。

1.7.1 第一步：先算一遍

在推一般公式之前，先手算一个 d_k = 3 的最小例子。

设 q = [1,\ 1,\ 1]，k = [1,\ 1,\ 1]，各分量都是方差为 1 的”典型”取值。

打分：q \cdot k = 1 + 1 + 1 = 3

归一化：softmax 只作用在一个标量上没有意义，所以要有对照。 再设另一个键 k' = [-1,\ -1,\ -1]，则 q \cdot k' = -3。 两个 logits 是 [3,\ -3]：

\mathrm{softmax}([3, -3]) = \frac{[e^3, e^{-3}]}{e^3 + e^{-3}} = \frac{[20.09,\ 0.0498]}{20.14} = [0.9975,\ 0.0025]

已经很尖了。 现在把维度提到 d_k = 64，q 和 k 都取全 1：

q \cdot k = 64, \qquad q \cdot k' = -64

\mathrm{softmax}([64, -64]) = \frac{[e^{64}, e^{-64}]}{e^{64}+e^{-64}} \approx [1 - 2e^{-128},\ 2e^{-128}] \approx [1.0000000,\ 0.0000000]

完全饱和。 这就是脚注说的现象：维度一高，logits 跨度就大，softmax 就退化成 one-hot。

但注意——上面的例子里我故意让 q、k 取全 1，这让点积等于 d_k。 真实情况是随机的。所以要把”典型有多大”这件事算清楚。

1.7.2 第二步：点积的方差（每步不跳）

设 q, k \in \mathbb{R}^{d_k}，各分量独立同分布，均值为 0、方差为 1。 我们要求 q^{\top} k = \sum_{i=1}^{d_k} q_i k_i 的方差。

① 单项的期望。 因为 q_i 与 k_i 独立：

\mathbb{E}\left[q_i k_i\right] = \mathbb{E}[q_i] \cdot \mathbb{E}[k_i] = 0 \cdot 0 = 0 \tag{1.11}

② 单项的方差。 用 \mathrm{Var}(X) = \mathbb{E}[X^2] - (\mathbb{E}[X])^2， 第二项刚算过是 0，所以只剩第一项：

\begin{aligned} \mathrm{Var}(q_i k_i) &= \mathbb{E}\left[(q_i k_i)^2\right] - \left(\mathbb{E}[q_i k_i]\right)^2 \\ &= \mathbb{E}\left[q_i^2 k_i^2\right] - 0 \\ &= \mathbb{E}\left[q_i^2\right] \cdot \mathbb{E}\left[k_i^2\right] && \text{（} q_i,\,k_i \text{ 独立，乘积的期望可拆）} \\ &= 1 \cdot 1 = 1 && \text{（均值为 0 时 } \mathbb{E}[X^2] = \mathrm{Var}(X) = 1 \text{）} \end{aligned} \tag{1.12}

③ 求和。 各项 q_i k_i 之间也相互独立，所以方差可以直接相加：

\mathrm{Var}(q^{\top} k) = \sum_{i=1}^{d_k} \mathrm{Var}(q_i k_i) = \sum_{i=1}^{d_k} 1 = d_k \tag{1.13}

标准差就是 \sqrt{d_k}。

人话翻译：点积是 d_k 个”随机的小数”相加。每个的方差是 1， 加起来方差就是 d_k——维度越高，这个和摆动得越厉害。

d_k = 64 时标准差是 8。也就是说，随机取两个位置， 它们的打分差别在 10 以上是家常便饭。除以 \sqrt{d_k} 之后， 标准差回到 1，打分就落在”正常”范围内了。

所以缩放的目的不是”让数值变小”，而是”让数值不随维度变化”。 这句区别很重要，后面还会回来。

原论文的脚注到此为止。但真正的问题——数值大为什么是灾难——还没回答。

1.7.3 第三步：softmax 的雅可比（每步不跳）

这是本节的核心。记 p = \mathrm{softmax}(z) \in \mathbb{R}^n，我们要求 \partial p_i / \partial z_j。

从定义出发，先设 S = \sum_m e^{z_m}，则 p_i = e^{z_i}/S。用商的求导法则：

\frac{\partial p_i}{\partial z_j} = \frac{\partial}{\partial z_j}\left(\frac{e^{z_i}}{S}\right) = \frac{\left(\frac{\partial}{\partial z_j} e^{z_i}\right) S - e^{z_i}\left(\frac{\partial S}{\partial z_j}\right)}{S^2} \tag{1.14}

其中：

• \dfrac{\partial}{\partial z_j} e^{z_i} = \delta_{ij}\, e^{z_i} （i = j 时是 e^{z_i}，否则是 0；\delta_{ij} 是克罗内克符号）
• \dfrac{\partial S}{\partial z_j} = e^{z_j} （求和里只有第 j 项对 z_j 敏感）

代回去：

\frac{\partial p_i}{\partial z_j} = \frac{\delta_{ij} e^{z_i} S - e^{z_i} e^{z_j}}{S^2} = \frac{e^{z_i}}{S}\cdot\delta_{ij} - \frac{e^{z_i}}{S}\cdot\frac{e^{z_j}}{S} \tag{1.15}

把 p_i = e^{z_i}/S、p_j = e^{z_j}/S 认出来，就得到

\boxed{\ \frac{\partial p_i}{\partial z_j} = p_i\left(\delta_{ij} - p_j\right)\ } \tag{1.16}

人话翻译：这个矩阵衡量”输入 z 动一点点，输出 p 会动多少”。 它决定了梯度能不能从 loss 一路传回到打分的源头。

关键在它的大小。看两个极端：

• p 均匀（p_i = 1/n）：对角项是 \frac{1}{n}(1-\frac{1}{n}) \approx \frac{1}{n}， 矩阵的元素是 O(1/n) 量级
• p 是 one-hot（p_1 = 1，其余 0）：对角项 p_1(1-p_1) = 1 \times 0 = \mathbf{0}， 非对角项 -p_i p_j = 0

第二行就是全部答案：softmax 一旦饱和，整个雅可比矩阵趋近于零矩阵， 梯度传不过去了。

1.7.4 第四步：把两步连起来

现在链条完整了：

d_k \uparrow \ \Rightarrow\ \mathrm{Var}(q^{\top} k) = d_k \uparrow \ \Rightarrow\ \text{logits 跨度} \uparrow \ \Rightarrow\ p \text{ 趋近 one-hot} \ \Rightarrow\ J \to 0 \ \Rightarrow\ \text{梯度消失} \tag{1.17}

这就是原论文那句 “extremely small gradients” 的完整含义。它在脚注里只写了一句话， 展开来是上面四步推导。

用案例的数字核一遍：不缩放的大模型里，实测 logits 标准差是 5477， 权重熵是 0.0004。把 p 的典型值代进 式 1.16， 对角项是 p(1-p) \approx 1 \times (1-1) = 0 量级——这就是为什么 实测的 \|J\|_F 掉到了 8\times10^{-5}。

1.7.5 为什么是 \sqrt{d_k} 而不是 d_k

现在可以回答一个常被问的问题：既然是”数值太大”，为什么不干脆除以 d_k， 让数值更小？

看方差：除以常数 c 会把方差变成原来的 1/c^2。

缩放方式	缩放后方差	d_k = 64 时的标准差
不缩放	d_k	8.0
除以 \sqrt{d_k}	1	1.0
除以 d_k	1/d_k	0.125

除以 d_k 会把标准差压到 0.125——logits 全都挤在 0 附近， softmax 输出接近均匀分布，注意力退化成对所有位置的简单平均。 这时候雅可比范数确实不小（均匀分布的雅可比反而最大）， 但模型学不到位置区分，因为所有位置的权重几乎一样。

要的是让方差回到 1，不是让它尽可能小。

这句话可以再推一层：方差为 1 意味着 logits 的典型跨度是 \pm 2 左右， softmax 输出既不像 one-hot 那样死掉，也不像均匀分布那样毫无区分度—— 落在 softmax 最”敏感”的区间里。这就是 \sqrt{d_k} 而不是 d_k 的全部理由。

1.7.6 回到案例：为什么小模型上没事

现在两个观察可以统一解释了：

	d_{\text{head}}=32	d_{\text{head}}=256
实测 logits 标准差（不缩放）	0.070	1.615（初始化时）/ 5477（训练后）
缩放后	—	—
权重熵	0.999（几乎均匀）	0.696 → 0.0001
训练是否受影响	否（差 0.17）	是（差 0.95）

小模型上，logits 的标准差只有 0.070——远小于 1。 这时候 softmax 输出接近均匀分布（熵 0.999），离饱和那一端极远。 去掉缩放和不去掉，都在”安全区”里，所以看不出差别。

大模型上，初始化时标准差就是 1.615，已经接近临界； 训练过程中模型为了让注意力更”有用”，会主动放大 q、k 的尺度 （因为更尖锐的注意力能更精确地选择信息），这又进一步推高 logits—— 形成正反馈。最终标准差冲到 5477，熵掉到 0.0001，注意力锁死。

1.7.7 把”正反馈”这句话测一遍

上面那句”训练过程中模型会主动放大 q、k 的尺度”是我们写下的解释，不是测量结果。 这一节把它补上——做法很直接：在初始化和训练之后两个时刻， 分别量每个注意力层的 q、k 激活标准差和点积的标准差。

# 在训练前后各跑一遍：量 q/k 的实际尺度，而不是看权重
qkv = blk.attn.qkv(blk.ln1(x))
q, k, _ = qkv.split(d_model, dim=2)
qh = q.view(B, T, H, dh).transpose(1, 2)     # 拆成头
kh = k.view(B, T, H, dh).transpose(1, 2)
print("q 尺度", qh.std().item(), "k 尺度", kh.std().item())
dot = qh @ kh.transpose(-2, -1)
if scaled: dot = dot / dh ** 0.5             # 按实际用的公式算
print("logits 尺度", dot[..., 因果下三角].std().item())

四种配置各训 1500 步（和 小节 1.18 的案例完全一致）。第 0 层的结果：

配置	q 尺度	k 尺度	logits 尺度	\|W_{QKV}\|_F
d_{\text{head}}=32 缩放	0.228 → 1.016	0.229 → 1.217	0.05 → 3.3	4.45 → 11.98
d_{\text{head}}=32 不缩放	0.228 → 0.863	0.229 → 1.036	0.29 → 13.1	4.45 → 11.76
d_{\text{head}}=256 缩放	0.316 → 0.650	0.317 → 0.654	0.10 → 1.8	8.89 → 17.73
d_{\text{head}}=256 不缩放	0.316 → 5.457	0.317 → 5.361	1.64 → 4348	8.89 → 29.23

这张表把原来那句话分成了两半：一半对，一半不完整。

重要一、“模型会放大 q、k” —— 对，而且四种配置都放大了

每一个配置里，q 和 k 的尺度都涨了，幅度从 1.6 倍（d_{\text{head}}=256 缩放）到 17 倍 （d_{\text{head}}=256 不缩放）。

权重矩阵的范数也都在涨（4.45 → 11.98、8.89 → 29.23）。

所以”训练会放大 q、k“是一个普遍现象，和缩放因子无关。 它来自一个更基本的事实：更尖锐的注意力能更精确地挑信息， 所以只要梯度还通得过，模型就有动力把 q、k 推大。

这句话现在是测量结果了。

重要二、但”放大 q、k”和”logits 失控”之间隔着缩放因子

看最后两行。d_{\text{head}}=256 的两个配置，这是同一个模型、同一份数据、同一个种子， 唯一的差别是有没有除以 \sqrt{256}=16。

• 不缩放：q、k 各涨 17 倍 → 点积涨约 296 倍 → logits 实际涨了 2650 倍（1.64 → 4348）
• 缩放：q、k 只涨了约 2 倍 → logits 涨 18 倍，最终停在 1.8

不缩放那一边的涨法是”超乘法”的：q、k 各涨 17 倍只能解释 296 倍， 剩下的约 9 倍来自 q 与 k 的相关性增强—— 模型不只是把向量拉长，还在让它们变得更对齐（点积更大）。

这就是正反馈的机制：logits 变大 → softmax 更尖 → 模型发现”尖一点更好用” → 继续拉长、继续对齐 → logits 再变大。 在没有缩放的情况下，这个循环没有刹车。

有缩放时，刹车是 \sqrt{d_{\text{head}}}： q、k 涨 2 倍，除以 16 之后连 1 都不到。 模型可以放心地拉长 q、k，因为那点增长会被缩放因子吃掉。

注记三、一个意外漂亮的吻合

\sqrt{d_k} 的推导（小节 1.7）说：要的是让点积的方差等于 1， 也就是标准差落在 1 附近。

现在看那两个缩放配置训练之后落在哪里：

配置	训练后的 logits 标准差	推导想要的值
d_{\text{head}}=32 缩放	3.3	1
d_{\text{head}}=256 缩放	1.8	1

1.8 离 1 很近。 而 d_{\text{head}}=256 不缩放那一边是 4348。

注意这不是”我们调出来的”——两个缩放配置除了 d_{\text{head}} 之外没有任何区别， 它们的训练过程和目标函数里也没有任何一项在”要求”logits 标准差接近 1。 模型是自己走到那里去的。

这是对 \sqrt{d_k} 这个设计最直接的实验支持： 它不是为了”让数值更小”，而是为了让训练动力学有一个稳定的不动点—— 一个模型自发会收敛到、且能一直待着的尺度。 不缩放时，这个不动点不存在，训练只会一路冲到饱和。

警告一个和案例记录对不上的地方

案例里记录的 d_{\text{head}}=256 不缩放的 logits 标准差是 5477， 我们这次量到的是 4348（第 0 层）。

两者差约 25%。 原因是评估用的批次不同： 案例里的 5477 是训练过程中在训练批上记的， 这次为了和”初始化”对比，统一在一批固定的验证数据上量。

同一份权重、不同的数据，能差 25%。 这个差距不小，值得记住——在 logits 已经冲到几千的时候， “标准差是多少”本身就依赖于你拿哪批数据去量。

这个不一致我们没有消除（要消除就得固定评估协议重跑整个案例）， 但它不影响结论：4348 和 5477 都在”彻底饱和”那一档（熵 0.0004）， 和缩放那边的 1.8 差着三个数量级。

注记停下来想一想 ⑤

上面说”训练过程中模型会主动放大 q、k 的尺度，形成正反馈”。

请想一个问题：这个正反馈为什么在缩放之后不会发生？

提示：缩放之后，要让 logits 变大，q、k 的尺度需要放大 \sqrt{d_k} 倍。 而从 loss 回传到 q、k 的梯度，本身就要穿过 softmax 的雅可比……

1.7.8 现代做法：不止 \sqrt{d_k}

最后要指出一点：\sqrt{d_k} 是 2017 年的方案，今天它通常和别的手段一起用。

QK-norm：对 q 和 k 先做归一化（LayerNorm 或 RMSNorm），再做点积。 这样 q、k 的尺度被显式钉死，不再依赖初始化和训练动态。 ViT-22B、Stable Diffusion 3 等模型都用了这个做法。 它的好处是：即使训练过程中权重疯狂增长，logits 也不会失控。

为什么 \sqrt{d_k} 单独不够：它假设 q、k 各分量方差为 1， 而这个假设在训练过程中会失效——正如案例所示。QK-norm 不依赖这个假设， 所以更稳健。

这也解释了本章案例的价值：它暴露了一个”看起来够用、实则有隐含假设”的设计。 理解这个假设，比记住那个公式重要得多。

提示实验 1-1 要你注意的一件事

这个实验是合成的：随机造 q、k，各分量严格方差为 1—— 也就是脚注假设的理想情况。所以你会看到红线（不缩放）掉得很凶， 和案例里小模型的”没事”形成对比。

两者的差别恰恰说明了本章的核心：脚注的推导在它自己的假设下是对的； 真实模型偏离这个假设，所以结论要打折扣。

1.8 多头：为什么不是一次用完

回到案例。注意一件事：崩掉的那个模型是单头的 （d_{\text{model}}=256，n_{\text{head}}=1，所以 d_{\text{head}}=256）。

如果把同一个 d_{\text{model}} 切成 4 个头，每个 d_{\text{head}}=64， 按前面的分析，每个头的 logits 标准差会小得多，不缩放也不至于立刻崩。 这不是巧合——多头在这里顺带起到了降低单头维度的作用。

但多头的主要目的不是这个。这一节讲清楚它到底买到了什么。

1.8.1 先算一遍：参数量其实一样

一个常见的误解是”多头参数更多”。我们算一下。

设 d_{\text{model}} = 256。

单头的 W_Q：形状 256 \times 256，参数量 65{,}536。

4 个头的 W_Q^{(i)}：每个形状 256 \times 64（因为 d_{\text{head}} = 256/4 = 64）， 4 个拼起来是 4 \times (256 \times 64) = 65{,}536。

完全一样。 输出投影 W^O 也一样。多头在参数量上是零成本的。

人话翻译：多头买到的不是容量，是归纳偏置。

把一个大矩阵拆成 h 个小矩阵，等于强制模型在 h 个独立的子空间里 分别学关系，而不是在一个大空间里学一个混合的关系。

打个比方：单头好比让一个人同时考虑”语法上该看谁”“语义上该看谁”“位置上该看谁”， 最后给出一个折中的答案；多头是让四个人各管一摊，答案再合并。

前者不一定是错的，但后者更容易学到清晰的结构。

1.8.2 一个头只能表达一种关系

为什么”拆开”会更清晰？因为单头的输出是一个权重分布（每一行一个）。 这意味着在某个位置上，它只能表达”我最该看谁”这一件事。

但语言里的关系是多重的。同一个词，可能：

• 句法上该看它的主语
• 语义上该看与它搭配的动词
• 位置上该看前一个词（比如固定搭配）

这是三种不同的”该看谁”。单头必须把它们平均成一个折中 （除非这些关系恰好指向同一个位置）。多头可以各管一摊。

图 1.6　 多头注意力的结构。输入经 h 组独立的投影得到 h 组 Q/K/V，各自做一次注意力，最后拼接并投影。每一组（每个头）可以学到不同的关注模式。

1.8.3 完整算一遍：一个 2 头注意力的全过程

概念讲完，走一遍数字。设 d_{\text{model}} = 4，n_{\text{head}} = 2，所以 d_{\text{head}} = 2。 输入是两个位置（“明”“月”），每个位置 4 维：

X = \begin{pmatrix} 1 & 0 & 1 & 0 \\ 0 & 1 & 0 & 1 \end{pmatrix} \tag{1.18}

第一步：投影。 为了手算，把权重设成简单的选择矩阵—— 头 1 取前两维、头 2 取后两维：

W_Q = W_K = W_V = \begin{pmatrix} 1&0&0&0 \\ 0&1&0&0 \\ 0&0&1&0 \\ 0&0&0&1 \end{pmatrix} \tag{1.19}

（真实模型里这些权重是学出来的，这里只是为了看清”拆分”这个动作。）

第二步：拆到各头。 XW_Q = X，然后按列切成两份：

• 头 1 看到 q^{(1)}_1 = [1,0]，q^{(1)}_2 = [0,1]
• 头 2 看到 q^{(2)}_1 = [1,0]，q^{(2)}_2 = [0,1]

两个头拿到的向量恰好一样（因为这是选择矩阵）。真实模型里它们不同， 因为每个头有自己的投影矩阵。

第三步：各头独立做注意力。 用因果掩码（位置 1 只能看自己，位置 2 能看 1 和 2）。

头 1：K^{(1)} = V^{(1)} = \begin{pmatrix}1&0\\0&1\end{pmatrix}（取 X 的前两列）

位置 1：只有 k_1 = [1,0] 可见，q_1 \cdot k_1 = 1，softmax 后权重 [1.0]， 输出 = 1.0 \times v_1 = [1, 0]。

位置 2：q_2 = [0,1]，q_2\cdot k_1 = 0，q_2 \cdot k_2 = 1。 除以 \sqrt{2} \approx 1.414 后是 [0,\ 0.707]。 \mathrm{softmax}([0, 0.707]) = [0.330,\ 0.670]。 输出 = 0.330 \times [1,0] + 0.670 \times [0,1] = [0.330,\ 0.670]。

头 2：K^{(2)} = V^{(2)} = \begin{pmatrix}1&0\\0&1\end{pmatrix}（取后两列）

算出来完全一样（因为权重相同）：位置 1 输出 [1,0]，位置 2 输出 [0.330, 0.670]。

第四步：拼接。 把两个头的输出接起来：

\text{位置 1}：\ [1,0] \oplus [1,0] = [1,0,1,0] \text{位置 2}：\ [0.330, 0.670] \oplus [0.330, 0.670] = [0.330, 0.670, 0.330, 0.670]

第五步：输出投影 W_O。 真实模型里这一步把所有头的信息混合起来。 本例中设 W_O = I，所以结果就是上面两个向量。

提示从这个例子里该看出什么

注意 \sqrt{d_{\text{head}}} 用的是 d_{\text{head}} 而不是 d_{\text{model}}。 上例里 d_{\text{model}} = 4 但 d_{\text{head}} = 2，除的是 \sqrt{2}。

这一点经常被写错。原因很直接：打分 q \cdot k 发生在每个头内部， 向量的维度是 d_{\text{head}}，所以方差是 d_{\text{head}}，缩放因子自然也是它。

如果误用了 d_{\text{model}}，缩放会过度——回到 §1.6 的”失败四”， 注意力会退化成平均。

1.8.4 顺带的好处：单头维度被压小了

回到本章案例。崩掉的那个模型是单头的，d_{\text{head}} = 256。 如果把它改成 4 头，d_{\text{head}} 就变成 64，按 图 1.3 的实测， 初始化时 logits 标准差会从 1.615 降到约 0.2 上下——离饱和远得多。

所以多头在这里顺带起到了一个作用：它天然限制了单个头的打分尺度。 d_{\text{model}} 固定时，头数越多，每个头的维度越小，越不容易饱和。

这不是多头被设计出来的原因（它的动机是”多种关系”），但它是真实存在的副作用。 理解这一点，就能明白为什么后来的模型敢把 d_{\text{model}} 推到几千—— 因为它们同时把头数也推高了。

1.8.5 我们这 8 个头到底学会了什么

“不同的头学不同的关系”这句话，在文献里说了很多年。但它在我们这个小模型上成立吗？ 既然模型已经训好了，直接量一下。

对每段输入，把每个头的注意力矩阵取出来，统计四个量：

• 平均熵：权重分布有多尖（1 = 完全平均，0 = 完全 one-hot）
• 对角权重：有多少注意力给了自己（“我要不要保留自己的信息”）
• 前一位权重：有多少给了紧邻的前一个位置
• 平均距离：注意力重心的平均跨度（\sum_j (i-j)\,p_{ij}，0 = 只看自己）

三层唐诗、共 72 个位置的平均结果：

	平均熵	对角权重	前一位权重	平均距离	它在做什么
层 0 · 头 0	0.294	0.769	0.050	1.20	盯住自己——基本不往外看
层 0 · 头 1	0.521	0.193	0.450	2.31	盯住前一位——45% 给紧邻左侧
层 0 · 头 2	0.804	0.127	0.107	3.92	弥散，权重接近均匀
层 0 · 头 3	0.625	0.090	0.261	3.86	偏前一位，但更分散
层 1 · 头 0	0.614	0.276	0.182	3.16	兼顾自己与前一位
层 1 · 头 1	0.687	0.259	0.232	3.07	同上，更平均些
层 1 · 头 2	0.702	0.147	0.108	4.61	弥散
层 1 · 头 3	0.871	0.165	0.134	4.47	最弥散——几乎是全局平均

（参考值：完全均匀时熵是 1.000，24 个位置的平均距离约 11.5。）

这张表里有三件事值得注意。

第一，“不同头学不同关系”被证实了。 层的 0 头 0 有 77% 的注意力给自己， 而层 0 头 2 只有 13%——它们在干完全不同的事。

第二，层 0 头 1 是一个非常明确的”前一位头”（previous-token head）。 45% 的注意力给了紧邻左边那一个位置。这类头在文献里被反复观察到， 是 Transformer 里最稳定出现的结构之一。我们用一个 42 万参数的模型、 25 万字唐诗就复现出来了——说明它是任务驱动的必然产物，不是大模型的专利。

第三，层与层之间有分工。 第一层的头更”局部”（平均距离 1.2–3.9）， 第二层的头更”全局”（3.1–4.6）。一个自然的解释是： 底层负责建立局部结构（字与字、词与词）， 上层在局部结构的基础上做更长程的整合。

这是”深度的意义”的一个具体样子——不是简单地堆叠， 而是不同层处理不同尺度的问题。

注记停下来想一想 ⑨

上表里层 0 头 0 的对角权重是 0.769——它大部分时候只是”把自己原样传下去”。

那么请想：这样一个头有什么用？ 它几乎什么都没做（输出约等于输入的 V 投影）。

先自己想，再看下面的提示。

提示：回忆 §1.6 的残差连接 x = x + \mathrm{attn}(\mathrm{LN}(x))。 如果一个头把输入原样传下去，残差路径上会发生什么？ 再想想：如果这一层的其他头需要”某个东西保持不变”，谁来做这件事？

这个分析用的代码在 notebooks/analyze_heads.py——换个模型、换份语料， 就能得到你自己模型的头部画像。这是理解一个模型在”看什么”最实用的第一步： 不需要任何理论，只需要测量。

1.8.6 但头并不都同等重要

这是多头注意力最反直觉的地方。

Michel 等人在 2019 年的工作（Are Sixteen Heads Really Better than One?, NeurIPS 2019） 发现：训练完之后，很大比例的头可以直接剪掉，性能几乎不掉。 但另有一小部分头极其关键，剪掉一个就崩。

他们用的是多语言翻译模型，剪枝后 BLEU 掉得很少； 但把”关键头”单独剪掉，性能会急剧下降。

这对”注意力 = 解释”的直觉是个警告。

如果绝大多数头可以随便剪，那”这个头在看什么”就不能简单地当成 “模型这样决策的原因”。第 4 章会专门讲这个问题—— 注意力权重是一个可观测的中间量，不等于因果解释。

下面这段代码可以让你自己看到头的冗余性：

import torch

# 假设 model 是本章 tiny_gpt 训练出来的模型
model.eval()
x = torch.randint(2, 1200, (1, 32))

# 基准：全部头都在
with torch.no_grad():
    base_loss = model(x, x)[1].item()
print(f"全部头：loss {base_loss:.3f}")

# 逐个把头置零（消融）
for layer in range(len(model.blocks)):
    for head in range(model.cfg.n_head):
        # 把该头的输出投影对应切片置零
        W = model.blocks[layer].attn.proj.weight
        saved = W[:, head * 64:(head + 1) * 64].clone()
        W[:, head * 64:(head + 1) * 64] = 0
        with torch.no_grad():
            ablated = model(x, x)[1].item()
        W[:, head * 64:(head + 1) * 64] = saved   # 还原
        print(f"  层{layer} 头{head} 消融后 loss {ablated:.3f}"
              f"  （变化 {ablated - base_loss:+.3f}）")

你会看到大部分头的变化接近零，少数几个变化明显。 具体哪些头关键，会随着模型和种子变——这也是为什么重要头的发现 需要统计而不是单次观察。

1.8.7 多头的代价：KV 缓存

多头不是没有成本。训练时它几乎免费，推理时它很贵。

原因在自回归生成。生成第 t 个词时，需要用到前面所有词的 K 和 V。 为了避免每次重算，实现会把它们缓存下来——这就是 KV 缓存。它的大小是

\underbrace{2}_{\text{K 和 V}} \times \underbrace{n_{\text{layer}}}_{\text{层数}} \times \underbrace{n_{\text{head}}}_{\text{头数}} \times \underbrace{d_{\text{head}}}_{\text{每头维度}} \times \underbrace{L}_{\text{序列长度}} \times \underbrace{b}_{\text{字节数}} \tag{1.20}

代入一个真实的配置——Llama-2 70B：80 层、64 头、d_{\text{head}} = 128、 float16（2 字节）：

2 \times 80 \times 64 \times 128 \times L \times 2 = 2{,}621{,}440 \times L\ \text{字节} \approx 2.6\ \text{MB} \times L \tag{1.21}

每多一个 token，就要多 2.6 MB 显存。 序列长度 4096 时是 10.7 GB—— 这已经超过很多单卡的容量，而模型权重还没算。

为什么会和头数成正比？

因为每个头有自己独立的 K、V 投影。64 个头就是 64 份不重复的缓存。

那如果把头数减少呢？会损失”多种关系”的表达能力（§1.8 前面讲的）。 这就是 MQA 和 GQA 要解决的问题：在不损失太多表达能力的前提下，把 KV 缓存压下来。

1.8.8 MQA 与 GQA：共享 KV

MQA（Multi-Query Attention，2019） 的做法很直接：所有头共享同一份 K、V， 但每个头保留自己的 Q。

\text{标准多头：} Q, K, V \text{ 各有 } n_{\text{head}} \text{ 份} \qquad \longrightarrow \qquad \text{MQA：} Q \text{ 有 } n_{\text{head}} \text{ 份，} K, V \text{ 只有 } 1 \text{ 份} \tag{1.22}

KV 缓存立刻缩小到 1/n_{\text{head}}。对 Llama-2 70B 那个配置， 就是从 2.6 MB/token 降到 0.04 MB/token。

代价是表达能力下降——所有头必须共享同一个”看到什么”的模式， 只有”想找什么”（Q）是个性化的。实践中这会损失一些质量。

GQA（Grouped-Query Attention，2023） 是折中：把 n_{\text{head}} 个头分成 g 组， 每组共享一份 K、V。

方案	Q 的份数	K、V 的份数	70B 模型的 KV 缓存（每 token）
标准多头	64	64	2.6 MB
GQA（g=8）	64	8	0.33 MB
MQA	64	1	0.04 MB

GQA 在实践中被广泛采用（Llama-2 70B、Llama-3、Mistral 等）—— 它把缓存降到 1/8，而质量损失很小。

提示这一节真正值得记的一件事

“多头”这个设计在训练和推理两个阶段的含义是不同的。

训练时：多头是归纳偏置——让模型在多个子空间里学不同关系，参数量不变。 推理时：多头是内存开销——每个头都要独立的 KV 缓存，随头数线性增长。

MQA/GQA 的出现说明一个常见模式：一个为建模能力设计的东西， 最后被工程约束重新塑形。理解这两面，比记住”多头好”有用得多。

1.8.9 长序列的真正瓶颈

回到 式 1.5：注意力矩阵 S 是 n \times n 的。 这让很多人以为瓶颈是计算量 O(n^2 d)。实测下来不是。

在 A100 上跑标准注意力实现（不优化），n = 4096、d = 64、16 头时： 计算本身只要几毫秒，而读写那个 n \times n 矩阵的时间是计算的数倍。 原因是显存带宽远低于计算吞吐——矩阵要从 HBM 搬到 SRAM 再搬回去。

瓶颈是访存，不是算力。 这是理解 FlashAttention 的前提。

1.8.10 FlashAttention：不把矩阵写出来

标准实现的三步是：

S = Q @ K.T              # ① 写 n×n 到显存
P = softmax(S)           # ② 读回来、再写回去
O = P @ V                # ③ 再读回来

n = 4096 时，S 有 1677 万个元素；每一步的读写都是几十 MB 的搬运。 而这个矩阵其实根本不需要完整存在——它只是中间产物。

FlashAttention 的做法是分块：把 Q、K、V 切成小块， 每次只在片上内存（SRAM）里处理一块，算完立刻累加到输出，从不写完整的 S。

难点在于 softmax 需要全局的最大值和求和——而分块计算时， 后面的块还没看到，怎么知道全局最大值？

答案是在线 softmax（online softmax）：维护运行中的最大值 m 与运行中的和 \ell， 每处理一块就更新一次，并用新的 m 纠正之前的结果。

手算一遍。 设 logits 是 [1,\ 3,\ 2]，分两块：第一块 [1, 3]，第二块 [2]。

第一块处理完：m_1 = 3，\ell_1 = e^{1-3} + e^{3-3} = 0.135 + 1 = 1.135。

第二块来了：新元素是 2，当前的 m_1 = 3 更大，所以 m_2 = 3 不变。 \ell_2 = \ell_1 + e^{2-3} = 1.135 + 0.368 = 1.503。

最终：p = [e^{1-3},\ e^{3-3},\ e^{2-3}] / 1.503 = [0.135, 1, 0.368]/1.503 = [0.090,\ 0.665,\ 0.245]。

用朴素方法核对：e^1 + e^3 + e^2 = 2.718 + 20.09 + 7.389 = 30.19， p = [2.718, 20.09, 7.389]/30.19 = [0.090,\ 0.665,\ 0.245]。一致。

如果第二块里有更大的值（比如 5），就需要纠正：

\ell_{\text{new}} = \ell_{\text{old}} \cdot e^{m_{\text{old}} - m_{\text{new}}} + \sum_{\text{新块}} e^{z_i - m_{\text{new}}} \tag{1.23}

那个 e^{m_{\text{old}} - m_{\text{new}}} 因子把之前累积的和按新基准重新缩放。 这就是在线 softmax 的全部机制。

注记停下来想一想 ⑥

上面的纠正公式里，如果新旧最大值相同（m_{\text{old}} = m_{\text{new}}）， 缩放因子是多少？如果新最大值远大于旧最大值，这个因子会趋近于什么？

后者提示了一件事：在线 softmax 的数值行为和批量 softmax 是否完全一致？ （提示：想一想什么情况下会丢失精度。）

FlashAttention 的收益来自”数学不变、访存改变”： 结果是精确的（不是近似），但显存占用从 O(n^2) 降到 O(n)， 速度提升 2–4 倍。后续版本（FlashAttention-2/3）进一步优化了并行划分。

1.8.11 注意力之前：RNN 的三个死结

要理解一个设计为什么长这样，最好的办法是看它替代了什么。

在 2014 年之前，序列建模的主流是 RNN / LSTM。它的做法很自然： 从左到右扫过序列，维护一个固定大小的隐状态 h_t， 每读一个词就更新一次：

h_t = f(h_{t-1},\ x_t) \tag{1.24}

翻译时，编码器读完整个源句，把最后的 h_T 交给解码器； 解码器从这个向量出发，一个词一个词地生成译文。

这个结构简单、优雅，但有三个各自独立的死结。理解这三个死结， 就理解了注意力要解决的全部问题。

1.8.11.1 死结一：固定容量的信息瓶颈

所有源句信息都要挤进一个固定维度的 h_T。

源句 5 个词是这样，50 个词也是这样——h_T 的维度不变。 序列越长，每个词分到的”带宽”越少，早期信息被后来的更新覆盖得越厉害。

用一个具体的数字感受一下：设 h \in \mathbb{R}^{512}，源句 50 个词。 平均下来每个词只有约 10 个数用来承载自己的信息—— 而这 10 个数还要编码它在句中的位置、与前后词的关系、以及各种句法属性。

实测上，RNN 翻译模型在句子长度超过 20–30 个词之后，BLEU 会明显下滑。 这不是训练不充分，是架构容量的硬限制。

1.8.11.2 死结二：信息通路的长度

即使不考虑容量，信息从第 1 个词传到第 50 个位置， 要经过 49 次变换：

h_1 \to h_2 \to \cdots \to h_{50} \tag{1.25}

每一步都是一个非线性变换。哪怕每步只”损失”一点（比如梯度衰减到 0.9 倍）， 49 步之后也只剩 0.9^{49} \approx 0.006——千分之六。

这就是经典的梯度消失问题。LSTM 用门控机制（遗忘门、输入门、输出门） 缓解了它（让梯度可以走”常量误差传送带”），但没有根治—— 它把衰减从指数级改善到了接近线性，可路径长度本身没变。

1.8.11.3 死结三：无法并行

h_t 依赖 h_{t-1}，所以第 t 步必须等第 t-1 步算完。 序列多长，就得串行多少步。

在一个有几千个计算核心的 GPU 上，训练时一个 batch 里有很多句子可以并行， 但每条句子内部仍然是串行的。这导致 RNN 无法吃满硬件—— 2016 年前后，训练一个高质量的翻译模型需要几天到几周。

这三个死结的分工要分清。

• 死结一（容量）是表示问题：信息装不下
• 死结二（路径）是优化问题：梯度传不回
• 死结三（并行）是工程问题：算得慢

注意力一次解决了全部三个，而且解法是同一个： 让每个位置直接访问所有其他位置。

• 容量：不再需要压缩进一个向量，所有位置的表示都保留着
• 路径：任意两个位置之间的距离变成 1 步，梯度不用绕远路
• 并行：所有位置的两两关系是一次矩阵乘法，n 个位置同时算

一个设计同时解决三个看似无关的问题——这是它值得被称为”范式转移”的原因。

1.8.12 2014：Bahdanau 注意力

第一个把”回头查”做成可微机制的，是 Bahdanau 等人 2014 年的工作 （Neural Machine Translation by Jointly Learning to Align and Translate）。

它的动机非常朴素：译到第 5 个词时，我该重点看原文的哪几个词？

做法是在解码器的每一步 t，给编码器的每一个位置 i 算一个权重：

\alpha_{t,i} = \frac{\exp\bigl(\mathrm{score}(s_{t-1},\, h_i)\bigr)} {\sum_{i'}\exp\bigl(\mathrm{score}(s_{t-1},\, h_{i'})\bigr)} \tag{1.26}

其中 s_{t-1} 是解码器上一步的隐状态，h_i 是编码器第 i 个位置的隐状态。 然后加权求和得到一个上下文向量：

c_t = \sum_i \alpha_{t,i}\, h_i \tag{1.27}

对照今天的注意力，你会看到全部要素已经在了：

Bahdanau（2014）	今天的注意力
解码器状态 s_{t-1}	查询 q
编码器状态 h_i	键 k_i 和 值 v_i（当时是同一个）
\mathrm{score}(\cdot,\cdot)	q^\top k（点积）
softmax 归一化	一样
加权求和得 c_t	一样

两个关键差别：

1. Bahdanau 用的是”加性”打分：\mathrm{score}(s, h) = v^\top \tanh(W_s s + W_h h)， 其中 v、W_s、W_h 都是学出来的。为什么不用点积？ 因为当时 s 和 h 的维度可以不同，点积要求维度一致。
2. 键和值是同一个东西（都是 h_i）。把它们分开是后来的一步关键改进。

Luong 等人 2015 年（Effective Approaches to Attention-based Neural Machine Translation） 做了简化：改用点积打分、把键和值分开、试了”全局注意力”与”局部注意力”。

\mathrm{score}(s_t, h_i) = \begin{cases} s_t^{\top} h_i & \text{（点积）} \\ s_t^{\top} W h_i & \text{（一般点积）} \end{cases} \tag{1.28}

“点积”这个选择的胜出，为后来那个 \sqrt{d_k} 埋下了伏笔—— 点积好用，但它有个随维度膨胀的毛病，这正是本章前半段的主线。

注记停下来想一想 ⑫

Bahdanau 用加性打分，Luong 用点积。今天几乎所有模型都用点积。

请想：点积打分相比加性打分，优势在哪？ 提示从两个角度想： ① 计算形式（能不能写成矩阵乘法）；② 参数（要不要额外的参数矩阵）。

再想一层：代价是什么？ 这正是 \sqrt{d_k} 要处理的问题。

1.8.13 2017：从”跨序列”到”自己对自己”

Bahdanau 和 Luong 的注意力，是解码器对编码器的注意力—— 两个不同的序列之间。

真正的跳跃是问一个问题：如果一个序列对自己做注意力，会怎样？

\text{跨序列：} \mathrm{Attention}(\underbrace{Q_{\text{解码}}}_{\text{来自序列 A}},\ \underbrace{K_{\text{编码}},\ V_{\text{编码}}}_{\text{来自序列 B}}) \qquad \longrightarrow \qquad \text{自注意力：} \mathrm{Attention}(XW_Q,\ XW_K,\ XW_V) \tag{1.29}

三个矩阵全部来自同一个 X。

这个改动看起来很小，但它带来了三个后果：

① 每一层都在”重新组织”表示。 跨序列注意力只在最后一步用一次； 自注意力在每一个 Block 里都有，所以模型有 L 次机会去重新混合信息。

② 可以堆得很深。 因为每一层的输入输出形状完全相同（都是 n \times d）， 堆叠起来没有形状问题。这让”深度”成为一个可用维度—— 今天的大模型有 80 层甚至更多。

③ 顺序信息丢失了。 这一点我们在 小节 1.9 会专门讲： x 对自己做加权平均，谁在谁前面完全看不出来。所以必须额外注入位置信息。

一句话概括这条演化线。

• RNN：把历史压进一个向量，顺序内建，但容量有限、路径长、不能并行
• 跨序列注意力（2014–2015）：解码时回头看原文，解决了容量和路径， 但 RNN 的骨架还在，仍然不能并行
• 自注意力（2017）：把 RNN 骨架整个拆掉，只留注意力 + 前馈网络， 三个死结一起解决，代价是必须先解决顺序问题

每一步的动机都很具体，没有哪一步是”为了先进而先进”。

1.8.14 并行到底快了多少

“可以并行”这句话要有个量化的感觉。

RNN 处理长度 n 的序列，需要 n 次串行操作， 每次操作内部虽然可以并行（矩阵乘法），但步数被序列长度锁死。

自注意力只需要固定的几步（打分、softmax、加权求和）， 序列长度只影响矩阵的大小，不影响步数。

用具体的规模说：n = 512、d = 512 时， 一次注意力的 QK^\top 是 (512\times512)\times(512\times512) 的矩阵乘法—— GPU 可以在一个 kernel 里并行处理全部 1.34 亿次乘加。 而 RNN 得老老实实做 512 次。

这是”训练时间从周降到天”的直接原因，也是 Transformer 能在 2017 年之后 迅速吞掉整个领域的基础设施条件。

提示一个值得注意的因果顺序

Transformer 不是”因为想到了自注意力”而成功的， 而是“硬件足够并行”与”算法恰好可以并行”两件事碰上了。

2012 年 AlexNet 之后 GPU 算力快速增长， 但 RNN 的串行结构让它吃不到这份红利。 自注意力恰好是把序列建模”矩阵化”的那个形式—— 它的成功有一半要归功于它顺应了硬件的发展方向。

读技术史时值得记住这一点：好想法需要在合适的硬件上才能兑现。

注记停下来想一想 ⑦

为什么会话式 AI 的”长上下文”这件事，在 RNN 时代是个死结，而在 Transformer 时代 变成了一个工程问题？

提示：从”信息通路长度”和”能否并行”两个角度各想一层。 （这也是第 3 章讨论 KV 缓存与上下文长度的引子。）

1.8.15 注意力的变体：当 O(n^2) 不可接受

n^2 的代价催生了一整个研究分支。理解这些变体的思路， 比记住它们的名字重要。

稀疏注意力（Sparse Attention）：观察到一个事实—— 大多数注意力权重接近零。既然如此，何必算全？只算那些”可能重要”的位置：

• 局部窗口（sliding window）：每个位置只看前后 w 个位置，复杂度降到 O(nw)
• 跨步（strided）：每隔 k 个位置看一个，用来捕捉长程信息
• 全局 token：留几个特殊位置让所有 token 都能看到（比如 BERT 的 [CLS]）

Longformer、BigBird 都是这个思路。代价是引入归纳偏置： “重要的东西一定在附近”——这对自然语言常常成立，但不总是。

线性注意力（Linear Attention）：利用矩阵乘法的结合律。

\mathrm{softmax}(QK^{\top})\,V \quad \neq \quad \mathrm{softmax}(Q)\bigl({\mathrm{softmax}(K)}^{\top} V\bigr) \tag{1.30}

左边是 O(n^2)，右边是 O(nd^2)。关键在于能不能去掉那个 softmax—— 它正是阻止结合律成立的原因（softmax 是逐行的非线性操作）。

线性注意力用别的核函数替代 softmax（比如 \phi(x) = \mathrm{elu}(x) + 1）， 让结合律成立，从而把复杂度降到关于 n 线性。 代价是表达能力下降——它算的不再是”每个位置对每个位置”的精确权重。

提示这一节的要点

所有这些变体都在做同一个权衡：

\underbrace{O(n^2)}_{\text{精确，但贵}} \quad \longleftrightarrow \quad \underbrace{O(n) \text{ 或 } O(n \log n)}_{\text{便宜，但要加假设}} \tag{1.31}

没有免费的午餐：要么花 n^2 的代价算精确的注意力（FlashAttention 让这个代价可接受）， 要么加归纳偏置（“重要的在附近”或”可以用核函数近似”）来换速度。

今天的主流选择是前者——因为 FlashAttention 把 O(n^2) 的实际开销压得足够低， 而它的结果是精确的，不用赌”重要的一定在附近”这个假设是否成立。

1.8.16 一个常被忽略的细节：为什么是”加”不是”拼”

最后一处值得拆开的设计。多头注意力的输出是拼接再投影的：

\mathrm{Concat}(\mathrm{head}_1,\ \dots,\ \mathrm{head}_h)\, W^O \tag{1.32}

为什么不是直接相加，或者取平均？

因为拼接保留了”哪个头贡献了什么”的信息，而相加会把它抹掉。

具体说：拼接后接一个线性层 W^O，等价于给每个头分配一个独立的输出子空间。 W^O 的每一列可以单独决定”头 i 的这一维要不要保留、放大多少”。 如果直接相加，所有头被强制共享同一个输出空间，模型就没有这个自由度了。

用公式看更清楚。设每个头输出 o_i \in \mathbb{R}^{d_v}，拼接后是 [o_1; \dots; o_h] \in \mathbb{R}^{h d_v}。左乘 W^O \in \mathbb{R}^{d_{\text{model}} \times h d_v}：

W^O [o_1; \dots; o_h] = \sum_{i=1}^{h} W^O_i\, o_i \tag{1.33}

其中 W^O_i 是 W^O 的第 i 块。所以”拼接 + 投影”在数学上等于 “每个头独立投影后求和”——但前者让每个头有自己的 W^O_i， 后者（直接相加）强制所有头共用同一个变换。

注记停下来想一想 ⑧

既然 式 1.33 说明”拼接 + 投影”等价于”各自投影后相加”， 那么参数量上两者一样吗？如果一样， 为什么实际实现都写成”拼接 + 一个大的 W^O“，而不是”h 个小 W^O_i 再相加”？

提示：从计算效率（一次矩阵乘法 vs h 次）和代码简洁性两个角度想。

1.9 注意力看不见顺序

到这里，式 1.6 已经能算了。但它有一个致命的先天缺陷—— 这个缺陷不是实现问题，是数学结构决定的，必须靠额外设计来补。

1.9.1 先做一个小实验

把输入的两个位置对调，看输出会怎样。

设 Q = K = V = X（为了看清结构，让三者相同），输入 X 是两行：

X = \begin{pmatrix} a \\ b \end{pmatrix} \tag{1.34}

用 X_1 = \begin{pmatrix} a \\ b \end{pmatrix} 算一遍注意力， 再用 X_2 = \begin{pmatrix} b \\ a \end{pmatrix}（交换两行）算一遍。

因为 Q=K=X，打分矩阵是 XX^\top。对于 X_1：

X_1 X_1^\top = \begin{pmatrix} a\cdot a & a\cdot b \\ b\cdot a & b\cdot b \end{pmatrix} \tag{1.35}

对于 X_2：

X_2 X_2^\top = \begin{pmatrix} b\cdot b & b\cdot a \\ a\cdot b & a\cdot a \end{pmatrix} \tag{1.36}

第二式正好是第一式交换行列的结果。 也就是说： 交换输入的顺序，输出只是跟着交换，每个位置的数值一模一样。

人话翻译：注意力把输入当成一个集合，而不是一个序列。 它对”谁在谁前面”这件事完全不敏感。

这不是 bug，是注意力机制的定义使然—— 它本质上是”每个位置向所有位置做加权平均”， 而加权平均这个操作本身不带顺序概念（加法是可交换的）。

但语言是有顺序的。”明月松间照”和”照间松月明”用的是同一批字， 意思完全不同。所以必须想办法把顺序塞进去。

注记停下来想一想 ⑩

上面说”注意力对顺序不敏感”。但这个结论有个前提：Q、K、V 都是从 X 直接来的。

那么请问：如果一个模型在注意力的输出后面接了一个前馈网络（FFN）， 它能否”学会”区分顺序？ 换句话说，置换等变性会被后面的层破坏吗？

提示：想想 FFN 是对每个位置独立作用的还是跨位置作用的。

1.9.2 正弦位置编码

最经典的做法来自原论文：给每个位置算一个固定向量，直接加到词嵌入上。

\begin{aligned} PE_{(pos,\,2i)} &= \sin\left(\frac{pos}{10000^{2i/d}}\right) \\ PE_{(pos,\,2i+1)} &= \cos\left(\frac{pos}{10000^{2i/d}}\right) \end{aligned} \tag{1.37}

其中 pos 是位置，i 是维度下标。偶数维用正弦、奇数维用余弦。

为什么是这种形式？ 论文给的解释是”让模型容易学到相对位置”。 具体说，对任意固定偏移 k，PE_{pos+k} 可以写成 PE_{pos} 的线性变换：

\begin{pmatrix} \sin(\omega(pos+k)) \\ \cos(\omega(pos+k)) \end{pmatrix} = \underbrace{\begin{pmatrix} \cos(\omega k) & \sin(\omega k) \\ -\sin(\omega k) & \cos(\omega k) \end{pmatrix}}_{R(\omega k)} \begin{pmatrix} \sin(\omega pos) \\ \cos(\omega pos) \end{pmatrix} \tag{1.38}

这是三角函数的和角公式。R(\omega k) 是一个只依赖偏移量 k、 不依赖位置 pos 的旋转矩阵。

所以”位置 pos 和位置 pos+k 的关系”可以被表示成一个统一的线性操作—— 这正是注意力需要的（因为注意力是用点积衡量关系的， 如果相对位置能用固定变换表示，模型就能学到”往前看 k 个”这样的模式）。

提示为什么用多个频率

10000^{2i/d} 这个分母让不同维度对应不同波长： i=0 时波长约 2\pi（高频，区分相邻位置）， i = d/2 时波长约 2\pi \times 10000（低频，区分远距离位置）。

这像二进制计数器：低位快速翻转（区分近处），高位缓慢翻转（区分远处）。 所有维度合起来，每个位置得到一个”独一无二”的组合。

用多个频率的原因是单一频率无法同时区分近处和远处： 高频在远距离上会周期性重复（位置 0 和位置 100 编码相同）， 低频在近距离上几乎不变（区分不出位置 0 和位置 1）。

1.9.3 手算一遍：位置编码长什么样

公式看着抽象，算出来其实很朴素。取 d = 4（两个维度对），算前四个位置：

i = 0 时频率是 1/10000^{0} = 1；i = 1 时是 1/10000^{0.5} = 1/100。

pos	\sin(pos/1)	\cos(pos/1)	\sin(pos/100)	\cos(pos/100)
0	0.0000	1.0000	0.0000	1.0000
1	0.8415	0.5403	0.0100	1.0000
2	0.9093	−0.4161	0.0200	0.9998
3	0.1411	−0.9900	0.0300	0.9996

读懂这张表的关键是看两列的”变化速度”：

• 前两列（高频）变化极快：pos 从 0 到 3，\sin 走了 0 → 0.84 → 0.91 → 0.14， 已经在往回走了。它们负责区分相邻位置。
• 后两列（低频）几乎不动：\cos(pos/100) 从 1.0000 只变到 0.9996。 它们负责区分远距离位置——要等到 pos 几百，才会看出明显变化。

把这四个数当作一个四维向量，每个位置的编码就是空间中的一个点。 有一个漂亮的性质：

\|\mathrm{PE}(pos+1) - \mathrm{PE}(pos)\| = 0.9589 \quad \text{对所有 } pos \tag{1.39}

相邻位置的编码距离是完全相同的——不管是 0→1 还是 100→101。 这让位置编码在整条数轴上是”均匀”的，不存在”近处分辨率高、远处分辨率低”的问题。

为什么这件事重要？

如果编码在近处密集、远处稀疏，那么模型对远处位置的区分能力就会下降—— 而语言里”隔了 50 个词的主语”和”隔了 51 个词”同样需要区分。

正弦编码用多个频率叠加来解决这个问题： 低频维度负责”粗定位”（这是第几个 100），高频维度负责”细定位”（这是 100 里的第几个）。 就像时钟用时针和分针配合，同时给出大的和小的刻度。

这也解释了为什么波长要跨越这么大的范围（d=4 时是 6.3 和 628.3）： 如果所有维度波长都一样，那么位置 0 和位置（波长）就会得到完全相同的编码， 模型无法区分它们。

注记停下来想一想 ⑭

正弦编码是固定公式，不参与训练。这在当时是个有争议的选择—— 有人主张”让模型自己学位置表示不是更好吗”。

请想：如果位置编码是可学的，会有什么风险？ 提示从两个角度： ① 训练时只见过 pos \in [0, 511]，推理时遇到 pos = 1000 会怎样？ ② 位置表示需要多少训练数据才能学好？

（本章的小模型用的就是可学习的位置嵌入，可以对照着看这个取舍。）

1.9.4 可学习的位置嵌入

本章的小模型用的是更简单的做法：给每个位置一个可学习的向量， 和词嵌入一样训练出来。

self.pos_emb = nn.Embedding(cfg.block_size, cfg.d_model)
# 前向时：
x = self.tok_emb(idx) + self.pos_emb(pos)[None, :, :]

这更简单、也更灵活（模型自己决定每个位置该长什么样）， 但有一个硬限制：最多只能处理训练时见过的长度。 训练时 block_size = 64，推理时给第 65 个位置，pos_emb 就没有对应的行了。

正弦编码没有这个问题（它是个公式，可以算任意 pos）， 但它的表达能力更受限（位置表示是固定的，不随任务调整）。

注记这个取舍很重要

可学习 = 表达力强 + 不能外推；固定公式 = 能外推 + 表达力受限。

这个矛盾催生了今天主流的方案——RoPE。它同时要了两者。

1.9.5 RoPE：把位置变成旋转

RoPE（Rotary Position Embedding） 是 2021 年提出、今天几乎所有开源大模型 （Llama、Qwen、Mistral 等）都采用的方案。它的想法很漂亮：

不再”加”位置信息，而是”旋转” q 和 k。

具体做法：把 d 维向量两两分组，每一对 (x_{2i}, x_{2i+1}) 看作复平面上的一个点， 按位置 pos 和频率 \theta_i 旋转一个角度：

R_{\Theta, pos}^{d} = \mathrm{diag}\Bigl( R(\theta_1, pos),\ R(\theta_2, pos),\ \dots,\ R(\theta_{d/2}, pos)\Bigr), \qquad \theta_i = 10000^{-2(i-1)/d} \tag{1.40}

其中 R(\theta, pos) 是 2\times2 的旋转矩阵：

R(\theta, pos) = \begin{pmatrix} \cos(pos\,\theta) & -\sin(pos\,\theta) \\ \sin(pos\,\theta) & \cos(pos\,\theta) \end{pmatrix} \tag{1.41}

然后在算注意力之前，把 q、k 各自旋转：

\tilde{q}_m = R_{\Theta, m}\, q_m, \qquad \tilde{k}_n = R_{\Theta, n}\, k_n \tag{1.42}

（m、n 是位置下标。）

为什么这么做？关键在旋转的乘法性质。 计算 \tilde{q}_m 与 \tilde{k}_n 的内积：

\tilde{q}_m^\top \tilde{k}_n = (R_m q_m)^\top (R_n k_n) = q_m^\top R_m^\top R_n\, k_n \tag{1.43}

旋转矩阵是正交的（R^\top R = I），所以 R_m^\top R_n = R_m^{-1} R_n = R_{n-m}。 最后一步用到了旋转的复合性质：两个旋转的乘积是角度相加，取逆是角度取负。

于是：

\tilde{q}_m^\top \tilde{k}_n = q_m^\top R_{n-m}\, k_n \tag{1.44}

人话翻译：这是整个 RoPE 的核心一行。

左边是”加了位置信息之后的两个向量的内积”； 右边是”原始向量的内积，但中间夹了一个只依赖于相对距离 (n-m) 的旋转”。

也就是说：注意力打分自动变成了相对位置的函数。 模型不需要”知道”绝对位置是多少，它看到的是”这两个位置隔了多远”。

这正是我们想要的——语言里重要的是相对关系 （“前一个词”“隔三个词的主语”），而不是绝对下标。

1.9.6 RoPE 相比前两代好在哪

	正弦编码	可学习嵌入	RoPE
位置信息注入方式	加到输入上	加到输入上	旋转 q、k
编码的是	绝对位置	绝对位置	相对位置
参数	0	L \times d	0
能否外推	有限	不能	可以（配合插值技术）
与注意力的交互	无	无	直接作用在打分上

RoPE 的关键优势在最后一行：它不是在输入上”打个标记”， 而是直接改变了注意力打分的结构。正弦编码加在输入上之后， 经过 W_Q、W_K 的线性变换，那个”相对位置的线性关系”会被打乱； RoPE 作用在 q、k 本身上，所以这个关系必然保留。

注记停下来想一想 ⑪

RoPE 让打分只依赖相对距离 n - m。这本该是好事。

但请想一个问题：如果打分只依赖相对距离，那”位置 1 和位置 2”与 “位置 1001 和位置 1002”对模型来说有区别吗？

如果没有区别，会发生什么？（提示：想想模型的上下文窗口从 4K 扩到 128K 时， 训练时最长只见过 4K，那 4K 以外的位置会怎样？）

这就是”长度外推”问题的根源，也是 NTK-aware 插值、YaRN 这一系工作的出发点。

1.9.7 本章模型的位置编码

回头看那张头部画像表——层 0 头 1 有 45% 的注意力给紧邻左边， 层 0 头 0 有 77% 给自己。这些”局部”模式正是位置编码在起作用的证据： 模型需要知道”左边那个是谁”，才能学出这种模式。

如果没有位置编码，这些头就学不会——因为在它眼里， 所有位置的输入是无差别的集合。

第 2 章会详细讲 RoPE 的实现细节，以及今天主流的长度外推技术。

1.10 注意力图该怎么读（以及不该怎么读）

实验 1-2 里那些热力图，是理解模型最直观的工具。 但它也是最容易被过度解读的工具。这一节讲怎么读它、以及哪里会读错。

1.10.1 先讲读法

一张注意力图是一个 n \times n 的矩阵，第 i 行第 j 列是 “位置 i 从位置 j 取了多少信息”。常见的四个观察角度：

① 看对角线。 对角线亮 = 每个位置主要看自己。 这意味着这一层几乎没有做跨位置的信息交换—— 它更像一个”逐位置的处理”，而不是”整合上下文”。

在我们的模型里，层 0 头 0 的对角权重是 0.769，是个典型的”看自己”的头。

② 看对角线下方第一条（次对角线）。 亮 = “看前一位”。 这类头（previous-token head）在文献里非常常见， 我们的层 0 头 1 就是（0.450）。

③ 看某一列的亮度。 如果整列都很亮，说明所有位置都在看同一个位置。 这类叫”attention sink”（注意力汇）—— 通常落在第一个 token（BOS）或标点上。

为什么会有 attention sink：因为 softmax 的输出必须和为 1。 如果某个位置”不需要看任何地方”，它仍然必须把 1 分配出去。 分配到哪里最不碍事？一个语义上中性的位置——通常是序列开头的特殊 token。 这是 softmax 归一化强加的结构，不是模型”想”关注那个 token。

④ 看整体形状。 是弥散的一大片，还是几条清晰的斜线？ 清晰的斜线（对角线、次对角线、固定偏移）说明学到了位置相关的模式； 弥散说明在整合全局信息。

1.10.2 再讲三个不该做的解读

错误一：把”注意力权重高”当成”因果关系”。

“模型把 0.8 的注意力给了『月』，所以它是靠『月』做决定的”—— 这个推理不成立。

权重高只说明”信息从这里取的多”，不说明”取了之后有没有被用上”。 一个头可能取了很多信息，但它的输出投影 W_O 把这些信息投影掉了 （投影到对后续层无用的方向）。

正确的做法是干预：把那个位置的注意力权重强制设为 0， 看输出/预测有没有变。变了才是因果，没变就只是相关。

这个思路叫消融（ablation），第 4 章会详细讲。

错误二：把单个头当成一个完整的功能单元。

我们在 §1.8 看到，模型里有个明确的”前一位头”。 很容易得出”这个头负责处理前一个词”的结论。

但功能是由多个头协作完成的。 一个头可能只在特定上下文里起作用，或者它提供的信息要和其他头组合才有意义。

更麻烦的是：大部分头可以被剪掉而不影响性能（Michel et al., 2019）。 如果剪掉它性能不掉，那它到底”负责”什么？

更稳妥的说法是描述而非断言： 说”这个头有 45% 的注意力给了前一位”（可测量的描述）， 而不是”这个头负责语法”（无法直接验证的断言）。

错误三：忽略”权重和必须为 1”这个约束。

看到”某个位置只有 0.3 的最大权重”，容易解读为”模型在这里不太确定”。

但每一行的权重和必须是 1。如果这一行有 100 个可见位置， 0.3 已经是远远高于平均（0.01）的集中了。 反过来，如果只有 2 个可见位置，0.6 其实相当分散。

所以看权重必须同时看”这一行有多少个可见位置”。 这也是为什么我们统计时用了归一化熵而不是原始权重—— 它把”可选项数量”这个因素消掉了。

一个具体的例子：在我们的头部画像表里， 层 0 头 0 的熵是 0.294，层 1 头 3 是 0.871。 但因为序列长度不同（位置 5 只能看 6 个，位置 23 能看 24 个）， 如果不做归一化，这两个数根本没法比。

1.10.3 一个实用的读图流程

拿到一张注意力图，按这个顺序看：

① 先看坐标轴 —— 是 n×n 还是 m×n（交叉注意力）？序列多长？
② 看掩码形状 —— 上三角是不是空的？如果不是因果注意力，为什么？
③ 算三个数：对角权重、平均距离、归一化熵
④ 看整体形状 —— 弥散 / 斜线 / 竖列
⑤ 只在做了消融之后，才谈"它在做什么"

第 ③ 步是可以自动化的——notebooks/analyze_heads.py 做的就是这件事。 先让脚本给你数字，再去看图，比”盯着热力图找规律”可靠得多。

提示为什么先看数字再看图

人的视觉系统很擅长在随机噪声里”看出”模式。 一张 24 \times 24 的热力图有 576 个数值， 你几乎一定能在里面找到几条”看起来有意思”的斜线——即使它们是随机的。

而数字（对角权重均值、归一化熵）不会骗你。 先用数字筛掉”没有明显模式”的头，再对剩下的少数几个仔细看图， 效率高得多，也少犯错。

1.11 它在哪里失效

前面讲过的一般原理，这里落成具体的失败模式。每一条都配可复现的代码。

1.11.1 失败一：mask 用 0 而不是 -\infty

这是实现里最常见、也最隐蔽的错误。因果掩码的目的是”不让位置 i 看到位置 j > i“。 很多人的写法是：

scores = scores.masked_fill(mask == 0, 0.0)   # ← 错

为什么错：softmax 里，“加 0”意味着”这个位置的分数是 0”， 而 0 经过 e^0 = 1 之后仍然会分到权重。被 mask 的位置没有被排除， 只是被打了个折。

正确写法是加 -\infty：

scores = scores.masked_fill(mask == 0, float("-inf"))   # ← 对

e^{-\infty} = 0，这才是彻底排除。

import numpy as np

def softmax(z):
    z = z - z.max(); e = np.exp(z); return e / e.sum()

scores = np.array([1.0, 2.0, 0.5])
print("用 0 mask:", np.round(softmax(scores), 6))
# [0.244728 0.665241 0.090031]  ← 第三个位置仍然拿到 9% 的权重

m = scores.copy(); m[2] = -np.inf
print("用 -inf  :", np.round(softmax(m), 6))
# [0.268941 0.731059 0.      ]  ← 第三个位置被彻底排除

两行输出的差别就是全部：0.090031 对 0.000000。 后者的第三位被彻底排除，前者的第三位只被打了个折。

这个错误为什么难发现：训练 loss 照降。因为模型会用别的方式绕过—— 它学会把被 mask 位置的 v 学成接近零向量，从而”即使权重不为零，贡献也为零”。 但这是浪费容量，而且推理时会暴露（尤其是 KV 缓存实现里）。

1.11.2 失败二：缩放写在 softmax 之后

w = softmax(scores)
w = w / np.sqrt(d_k)     # ← 错：这不是缩放，是改温度 + 破坏归一化

为什么错：除以 \sqrt{d_k} 必须作用在 logits 上，因为它修正的是 “softmax 的输入尺度”。写在 softmax 之后，w 的和不再是 1， 而且它改变的是分布的整体大小，不是尖锐程度。

import numpy as np

def softmax(z):
    z = z - z.max(); e = np.exp(z); return e / e.sum()

s = np.array([8.0, 0.0])             # d_k = 64 时的典型打分
a = softmax(s);      print("缩放前  :", a,        " 和 =", a.sum())
# [9.996646e-01 3.353501e-04]  和 = 1.0    ← 很尖，但仍是概率分布
b = softmax(s / 8);  print("正确缩放:", b,        " 和 =", b.sum())
# [0.731059 0.268941]          和 = 1.0    ← 合理
c = softmax(s) / 8;  print("错误缩放:", c,        " 和 =", c.sum())
# [1.249581e-01 4.191877e-05]  和 = 0.125  ← 和不再是 1！

第三种输出的和是 0.125，根本不是概率分布。下游的加权求和会得到一个 量级完全错误的输出，而 shape 检查、loss 曲线都不会报错。

1.11.3 失败三：忘了数值稳定

小节 1.6 已经给过溢出的例子，这里补充它在注意力里的具体表现。 关键在于：logits 到底能有多大，取决于缩放因子在不在。

import numpy as np

def check(std, label):
    """按给定标准差造一批 logits，看 exp 会不会溢出。"""
    rng = np.random.default_rng(0)
    z = rng.normal(0, std, size=(4, 8)).astype(np.float32)   # float32！
    e = np.exp(z)
    print(f"  logits 标准差 {std:6.1f}（{label}）"
          f"  →  inf 个数 {int(np.isinf(e).sum()):2d} / {e.size}")

check(1.6,  "缩放后（本章案例实测）")
check(5477, "不缩放（本章案例实测）")

实测输出：

  logits 标准差    1.6（缩放后（本章案例实测））  →  inf 个数  0 / 32
  logits 标准差 5477.0（不缩放（本章案例实测））  →  inf 个数 14 / 32

不缩放那一行，32 个元素里有 14 个直接溢出。 而 14 这个数字还不是上限—— 它只取决于随机抽样的最大值有没有超过 88.7。

警告为什么这里必须用 float32

上面的例子如果写成 float64，一个都不会溢出：

float32  logits=[100,0]  exp 最大值 = inf        ← 溢出
float64  logits=[100,0]  exp 最大值 = 2.688e+43  ← 没溢出

float64 的指数上限对应 e^{709.8}，比 float32 的 e^{88.7} 宽得多。 而深度学习框架默认用 float32，所以你平时遇到的溢出， 在 numpy 默认的 float64 里复现不出来。

写调试代码时踩这个坑的人很多：明明线上炸了，本地却跑不出问题。

一个 batch 里只要有一行 nan，整个梯度就全废了——因为 inf / inf = nan， 而 nan 会沿着反向传播污染所有参数。

1.11.4 失败四：注意力退化成平均

这个失败不报错、不崩，只是效果差——所以最难发现。

触发条件：缩放过度（比如误用了 d_k 而不是 \sqrt{d_k}）， 或者温度设得过高。表现是所有权重几乎相等：

import numpy as np

def softmax(z):
    z = z - z.max(); e = np.exp(z); return e / e.sum()

s = np.array([8.0, 3.0, -2.0, 0.5])
for c, name in [(1, "不缩放"), (8.0, "÷√64（正确）"), (64.0, "÷64（过度）")]:
    p = softmax(s / c)
    print(f"  {name:14s} {np.round(p, 4)}  最大权重 {p.max():.4f}")

实测输出：

  不缩放            [9.927e-01 6.700e-03 0.000e+00 5.000e-04]  最大权重 0.9927
  ÷√64（正确）       [0.4518 0.2418 0.1294 0.1769]  最大权重 0.4518
  ÷64（过度）        [0.2725 0.252  0.2331 0.2424]  最大权重 0.2725

第三行四个权重是 0.2725 / 0.2520 / 0.2331 / 0.2424——几乎就是均匀的 0.25。 注意力退化成了简单平均，模型不再具有”选择信息”的能力。

注意第二行和第三行的对比：正确缩放下最大权重 0.4518， 既能明显区分主次，又不是 one-hot。这就是”落在 softmax 敏感区间”的实际样子。

这个失败不会报错，loss 也照降——只是模型永远学不好。 它是本章五个失败模式里最难发现的一个。

1.11.5 失败五：序列长度带来的显存爆炸

n \times n 的注意力矩阵，单头、float16、n=8192 时：

8192^2 \times 2\ \text{字节} \approx 134\ \text{MB}

看似不大，但这是单个头、单层。32 层 × 32 头就是

134\ \text{MB} \times 32 \times 32 \approx 137\ \text{GB}

当然实际实现不会把所有层同时保留（反向传播时需要，前向时不用）， 但长序列的瓶颈是显存而非算力这一点是确定的。 这正是 FlashAttention 的出发点：分块计算，永远不把完整的 n \times n 矩阵写进显存。

1.11.6 失败六：位置编码越界

这个失败模式在前面的案例里没有出现（因为我们训练和推理用同样的长度）， 但在真实使用中极其常见。

触发条件：用可学习的位置嵌入，且推理长度超过训练时的 block_size。

# 训练时
cfg = Config(block_size=64)
self.pos_emb = nn.Embedding(64, d_model)      # 只有 64 行

# 推理时给了 100 个 token
pos = torch.arange(100)
x = self.tok_emb(idx) + self.pos_emb(pos)     # ← IndexError，或者更糟

两种后果，看实现：

① 直接报错（PyTorch 的 nn.Embedding 会抛 IndexError）。 这是好的情况——至少你知道出了问题。

② 静默给出垃圾。如果实现里做了取模（pos % block_size）或者其他截断， 程序不报错，但位置 64 会被当成位置 0—— 模型以为自己在处理句首，输出的分布完全不对。

一个真实的排查场景：模型在 512 长度上训练得很好， 上线后处理 800 字的长文本，前 512 个 token 输出正常，之后突然变成乱码。

原因：位置 512 越界了。pos_emb 只有 512 行。

为什么会”突然”：前 512 个位置在训练分布内，行为正常； 第 513 个位置落到没训练过的区域，模型完全不知道该怎么办。 这是一条硬边界，不是渐变。

排查手段：在推理入口加一句断言。

assert idx.shape[1] <= model.cfg.block_size, \
    f"序列长度 {idx.shape[1]} 超过训练时的 {model.cfg.block_size}"

这个断言很便宜，但能省下几小时的困惑。

为什么正弦编码和 RoPE 没有这个问题： 它们是公式（\sin(pos/10000^{2i/d}) 和旋转角 pos\cdot\theta）， 可以算任意 pos，不存在”表用完了”这种说法。

但”能算”不等于”算得对”——正如 §1.13 检查点 ⑪ 说的， 外推出去的位置落在训练时没见过的取值范围内， 模型的行为仍然不可靠。这是”外推”问题，比”越界”问题更微妙。

方案	越界会怎样	外推表现
可学习嵌入	报错或静默出错	完全不支持
正弦编码	能算	有限外推（高频维度会周期重复）
RoPE	能算	需配合插值技术（NTK/YaRN）

1.11.7 把这四个 bug 真的跑一遍

上面讲的都是”理论上会怎样”。但调试时你需要的是症状—— 看到什么现象，才知道是哪个 bug。

所以我们把每个 bug 真的注入到模型里，各训 300 步，记录实测症状。 代码在 notebooks/debug_failures.py，数据在 data/failure_symptoms.json。

变体	loss 起	loss 终	末梯度范数	权重熵	对角权重	行和最大值
正确实现	7.032	4.810	1.018	0.831	0.153	1.0000
① 掩码用 0	7.034	4.634	3.791	0.949	0.134	1.0000
② 缩放位置错	7.055	4.668	0.560	0.836	0.106	0.1768
③ 缺数值稳定	7.032	4.808	1.170	0.842	0.158	1.0000
④ 误用 d_k	7.032	4.708	0.776	0.9989	0.127	1.0000

现在逐个读这张表——它比任何”注意事项”都有用。

1.11.7.1 ① 掩码用 0：loss 反而更低

看第二行：bug 版本的 loss（4.634）比正确实现（4.810）还低。

这不是测量错误。原因是：用 0 掩码意味着被屏蔽的位置仍然分到权重， 所以位置 i 能”看见”位置 j > i 的内容——模型在偷看答案。

预测下一个字时，答案就在输入里。所以训练 loss 当然更低。

这是最难发现的一类 bug：它让指标变好，而不是变差。

如果只看 loss 曲线，你会以为这个改动是个改进，甚至会把它保留下来。 唯一的破绽在推理时：生成第 t 个词时，第 t+1 个词根本还不存在， 偷看的能力瞬间消失，模型原形毕露。

排查手段：熵那一列也能看出来——0.949 明显高于正确实现的 0.831， 因为被”偷看”的位置摊薄了权重分布。但最可靠的检查还是单测： 给一段已知输入，验证位置 i 对位置 j>i 的权重严格为 0。

1.11.7.2 ② 缩放位置错：行和不是 1

第三行最右边那个数字 0.1768 就是铁证。

正确的注意力权重，每一行的和必须是 1（它是概率分布）。 而把除法写在 softmax 之后，1/\sqrt{32} = 0.1768 直接把每一行都缩小了—— 和数学上的”概率分布”完全不沾边了。

这个 bug 的特点是好查：加一行 assert abs(weights.sum(-1) - 1) < 1e-5 就能拦住。 但它的危害是不报错的——输出的 shape 完全正常，loss 也照降（4.668）。

1.11.7.3 ③ 缺数值稳定：在这个规模下根本看不出来

第四行：loss 4.808（和正确实现的 4.810 几乎一样），没有任何异常。

为什么？ 因为 d_{\text{head}} = 32 时 logits 的标准差只有约 0.07–1.6 （正是本章案例里测到的量级），\exp 完全不会溢出。

这和本章开头那个”没复现出来的实验”是同一回事—— 一个只在特定尺度下才显形的 bug。你在小模型上调试时看到一切正常， 换成大模型或者长序列，它就炸了。

警告这条最值得记住

四个 bug 里，③ 是唯一一个”在你当前的调试环境里测不出来”的。

应对办法不是”小心一点”，而是主动制造条件： 写一个单元测试，故意喂进标准差 5000 的 logits，看函数会不会返回 nan。 本章案例里那个不缩放的模型实测标准差是 5477，32 个 logits 里 14 个溢出—— 这就是你该用的测试数据。

把”边界条件测试”提前写好，而不是等它在生产环境里出现。

1.11.7.4 ④ 误用 d_k：注意力退化成均匀

最后一行的熵是 0.9989——极其接近 1.000，也就是完全均匀的分布。

乘以 \sqrt{d_k} 倍的过度缩放，把 logits 压得几乎全相等， softmax 输出就变成 1/n。模型失去了”选择信息”的能力， 但它不报错、loss 也照降（4.708，甚至比正确实现还低一点， 因为 300 步的训练还没到见分晓的时候）。

排查手段：看熵。如果某一层的权重熵长期停在 0.99 以上， 几乎可以确定是缩放出了问题。这也是为什么 analyze_heads.py 要把熵作为第一个指标。

1.11.7.5 一张图看清五个失败模式

失败模式	loss 会变差吗	有异常数值吗	怎么查
① mask 用 0	反而变好	熵偏高（0.949）	单测：未来位置的权重必须为 0
② 缩放位置错	轻微变好	行和 = 0.177	断言：每行和 ≈ 1
③ 缺数值稳定	不变	无（小模型下）	喂大 logits 的边界测试
④ 缩放过度	不变	熵 = 0.9989	监控权重熵
⑤ O(n^2) 显存	—	OOM 报错	唯一会自己报错的

五个里有四个不会报错，其中两个在训练指标上完全看不出来。 这就是为什么调试注意力必须靠主动测量（熵、行和、对角权重这些量）， 而不是盯着 loss 曲线。

analyze_heads.py 输出的那四个统计量，本身就是一套通用的”注意力体检表”—— 建议在任何自定义注意力实现上线前跑一遍。

1.12 注意力的训练动态

前面我们对比了”训练前”和”训练后”。但中间发生了什么？ 这一节用案例里记录下来的完整过程数据，看注意力在 1500 步里怎么演化。

1.12.1 两条完全不同的轨迹

案例里我们每一步都记录了注意力层的诊断量。把两个变体并排看：

不缩放那一边：

步数	权重熵	logits 标准差	\|J\|_F
0	0.712	1.63	0.2975
222	0.0022	982.9	0.00471
444	0.0003	3176.5	0.00051
888	0.0001	5416.2	0.00020
1480	0.0001	5458.8	0.00008

缩放那一边：

步数	权重熵	logits 标准差	\|J\|_F
0	0.9985	0.10	0.2017
222	0.6437	2.21	0.3029
444	0.6754	1.62	0.2938
888	0.5714	1.75	0.2996
1480	0.5745	1.78	0.3089

1.12.2 三条值得注意的规律

① 崩塌发生在最初 222 步内，而且极快。

不缩放那一边，熵从 0.712 掉到 0.0022——只用了 222 步，是总步数的 15%。 之后 1200 多步，熵基本停在 0.0001，没有任何恢复的迹象。

这说明：缩放问题不是”慢慢积累”的，而是一次性的。 一旦 softmax 进入饱和区，梯度就传不回去，注意力再也没机会自我修正。 它是一个不可逆的相变。

警告这对调试的含义

如果你在训练到一半才发现熵塌了，改超参数救不回来——只能从头训。

所以这类问题的正确应对是预防，不是补救： - 训练开始就监控熵（前 200 步最关键） - 在 step 0 就检查初始化后的 logits 尺度

我们案例里之所以能发现它，正是因为记录了 step 0 的数据。 如果只记最终结果，你会看到”熵 0.0001”，但不知道它是从第几步开始塌的。

② 缩放那一边的熵在 0.57–0.68 之间波动，而且整体是下降的。

从 0.9985（初始化时接近均匀）降到 0.57 左右，然后稳定。

这个下降是健康的：初始化时模型对所有位置一视同仁（熵接近 1）， 训练过程中逐渐学会”该看谁”（熵下降）。 但注意它没有继续降到 0 ——它稳在一个中间水平。

0.57 这个值意味着什么：在一行 24 个位置里， 注意力大致集中在几个位置上，但仍保留了相当宽的分布。 这正是我们想要的——既要有区分度，又不能退化成硬选择。

③ \|J\|_F 的变化最能说明问题。

缩放那一边：从 0.2017 升到 0.30 左右，然后一直稳定。 不缩放那一边：从 0.2975 掉到 0.00008，掉了一万多倍。

为什么 \|J\|_F 是关键指标：它直接衡量”梯度能不能传回去”。 熵和 logit 标准差是”现象”，\|J\|_F 是”后果”。

有意思的是缩放那边的 \|J\|_F 是上升的（0.20 → 0.30）。 一个可能的解释是：初始化时权重很小、logits 很平，雅可比自然偏小； 训练让它变得更”有区分度”，雅可比随之上升到一个合理的水平。 这又是一个”回到 1 附近”的例子——和方差的目标是同一个思路。

1.12.3 初始化为什么这么重要

从上面的数据能反推出一个结论：初始化决定了训练会不会走上正轨。

在 step 0，两个变体的唯一差别是有没有缩放。但它们的初始诊断量已经不同了：

	logits 标准差	熵
不缩放	1.63	0.712
缩放	0.10	0.9985

不缩放那一边初始的熵是 0.712——已经比缩放那边尖得多。 这个初始差异，在 222 步内被放大成了”一个崩、一个正常”。

“初始化不是随机起点，它决定了你会走到哪个盆地。”

这是一个在深度学习里反复出现的模式： 初始化的微小差异，经过训练的放大，会导向完全不同的解。

本章的具体例子是：初始 logits 尺度偏大 → softmax 偏尖 → 梯度偏小 → 注意力更难调整 → 更尖…… 正反馈把它推向饱和。 而初始尺度正常时，这条正反馈链不成立。

所以”调参”这件事，很多时候真正的杠杆在初始化，而不是学习率。

1.12.4 一个可以自己做的观察

用 analyze_heads.py 在训练的不同阶段各跑一次， 就能看到头部的”分化过程”：

# 训练 100 步、500 步、1500 步，各存一次权重
for steps in (100, 500, 1500):
    cfg = Config(max_steps=steps)
    model, tok, _ = train(text, cfg)
    torch.save(model.state_dict(), f"ckpt_{steps}.pt")

# 然后分别分析
for steps in (100, 500, 1500):
    stats = analyze(f"ckpt_{steps}.pt")
    # 看"前一位头"的 prev 权重是怎么涨起来的

我们没做这个实验（它不在本章的案例范围内）， 但它是一个很好的入手点：头部特化是一次性出现的，还是逐渐形成的？ 如果你做了，欢迎把结果和本章的数据对照。

1.13 排查手册：从症状倒推原因

上一节给了”每个 bug 长什么样”。这一节给怎么从症状找到它。

调试注意力最麻烦的地方是：五个失败模式里有四个不报错， 而它们的最终表现（loss 不降、生成质量差）看起来一模一样。 所以关键是找到那些”能区分它们”的中间量。

下面四条排查路径，都是可以照着做的。

1.13.1 排查一：loss 不降，但梯度范数正常

症状：训练跑了几百步，loss 卡在一个明显偏高的位置不再下降。 但梯度范数正常（0.1–10 之间波动），不是梯度消失。

第一个假设：数据或学习率有问题。 这个先排除，因为太常见了。 换一个已知能训好的小数据集（比如就几百个字符的重复文本）， 如果 loss 还是不动，说明不是数据问题。

第二个假设：模型容量不够。 对照检查——把同样的配置 （层数、维度、步数）在一个已知能收敛的任务上跑一遍。 我们的经验是：42 万参数的模型在 25 万字唐诗上， 1500 步应当能到 loss ≈ 4.0 左右。如果差得远，问题在实现。

第三步：量三个中间量。

# 训练循环里加这几行
with torch.no_grad():
    model(xb, record=True)
    d = model.blocks[0].attn.last_diag
    print(f"step {step}  熵 {d['weight_entropy']:.4f}  "
          f"logit_std {d['logit_std']:.2f}  ‖J‖ {d['jacobian_fro']:.5f}")

读这三个数字的方法：

观察	指向
熵 → 0.99 以上	缩放过度（用了 d_k 而不是 \sqrt{d_k}）
熵 → 0.01 以下，logit_std 几千	缩放缺失（或 q、k 尺度失控）
‖J‖_F → 10^{-4} 以下	注意力层已经冻结，梯度传不回去
三个量都正常，loss 还是不降	问题不在注意力——去看 FFN、输出层、或数据

我们那次”没复现出来的实验”，用的就是这个流程。 区别是：小模型上三个量全都正常（熵 0.999，‖J‖ 0.26）， 说明注意力层没问题——于是把怀疑转向了”这个现象是不是需要更大的规模才出现”， 这才有了后面 d_{\text{head}} = 256 的对照。

提示这条路径的一般形式

“loss 不降”是个太粗的信号，必须把它分解成若干个能独立检查的量。

注意力的诊断量至少有三个（熵、logit 尺度、雅可比范数）， 而且是正交的： - 熵告诉你”分布有多尖” - logit 标准差告诉你”为什么会这么尖” - 雅可比范数告诉你”这有多严重”

三个量合起来能唯一定位问题。只测一个不够—— 比如”熵很高（接近均匀）“可能是缩放过度，也可能是 q、k 的尺度太小， 只有配合 logit 标准差才能区分。

1.13.2 排查二：训练 loss 漂亮，但生成质量差

症状：loss 曲线一路下降，最后停在很低的水平（比预期还好）。 但生成出来的文本一眼就能看出问题——重复、不通顺、或者干脆是乱码。

第一个假设：过拟合。 检查验证集 loss。如果验证 loss 也在降， 那不是过拟合。

第二个假设：训练/推理不一致（train-test mismatch）。 这是最可能的。

具体来说，语言模型训练时用”教师强制”（teacher forcing）： 把真实的前文喂进去，预测下一个词。 推理时把自己生成的前文喂进去。如果两者之间有差异，就会暴露。

最典型的三种不一致：

① 掩码泄露未来信息。 就是前一节的失败模式①。 训练时模型能偷看答案，loss 当然低；推理时答案不存在，能力瞬间消失。

验证方法（这是个必须做的单测）：

# 未来位置的注意力权重必须严格为 0
_, w = my_attention(Q, K, V, mask=causal_mask(n))
import numpy as np
upper = np.triu(np.ones((n, n)), k=1).astype(bool)
leaked = w[..., upper]
assert leaked.max() < 1e-8, f"泄露了！最大权重 {leaked.max():.3e}"

② 训练与推理用了不同的实现路径。 比如训练时走的是自己写的注意力， 推理时框架自动换成了 F.scaled_dot_product_attention。 两者的数值行为可能有细微差异（尤其在低精度下）。

验证方法：同一个输入，两条路径都跑一遍，对拍输出：

mine, _ = my_attention(Q, K, V, mask=m)
ref = F.scaled_dot_product_attention(Q, K, V, attn_mask=m)
assert (mine - ref).abs().max() < 1e-5, "两条路径不一致"

③ 位置编码在推理时越界。 如果用的是可学习的位置嵌入， 且推理序列比训练时长，pos_emb 会索引到没训练过的行—— 那些行是随机初始化的，输出自然是垃圾。

验证方法：打印推理时的序列长度和 block_size，确认没超。

警告“训练 loss 反而更低”是最危险的信号

正常的 bug 会让指标变差，你会立刻发现。 让指标变好的 bug 才致命——它会被当成一个改进保留下来。

我们那个掩码实验里，bug 版本的 loss（4.634）比正确实现（4.810）低了 0.18。 如果不做单测，只看曲线，这个 bug 会一直留着， 直到部署上线、发现生成质量莫名其妙地差，而那时已经很难回溯原因了。

结论：关键性质的验证不能依赖 loss。要写断言。

1.13.3 排查三：loss 里出现 nan

症状：训练几十步之后，loss 变成 nan，之后永远是 nan。

第一个假设：梯度爆炸。 检查梯度范数——但如果是 nan， clip_grad_norm_ 也救不了（nan 裁剪之后还是 nan）。

第二个假设：softmax 溢出。 这是注意力里最常见的 nan 来源。

定位步骤：

# 在前向传播里插桩，找到第一个出现 nan 的位置
def attention_debug(Q, K, V, mask=None):
    scores = Q @ K.swapaxes(-1, -2)
    print("① scores 有限?", torch.isfinite(scores).all().item(),
          "范围", scores.min().item(), scores.max().item())

    scores = scores / np.sqrt(Q.shape[-1])
    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))
    print("② 掩码后有限?", torch.isfinite(scores).all().item())

    e = torch.exp(scores)
    print("③ exp 后有限?", torch.isfinite(e).all().item(),
          "inf 个数", torch.isinf(e).sum().item())

    w = e / e.sum(-1, keepdim=True)
    print("④ 归一化后有限?", torch.isfinite(w).all().item())
    return w @ V, w

这是”二分定位”的标准做法：在每一行后面打印一次， 第一个出现 False 或 inf 的地方就是元凶。

四个位置对应的病因：

第一个出问题的位置	病因
① scores 就是 inf	q、k 的数值本身炸了（权重初始化或学习率）
② 掩码后出 nan	整行被屏蔽（见下文 padding 陷阱）
③ exp 后出 inf	缺数值稳定，或缩放缺失导致 logits 过大
④ 归一化后出 nan	inf / inf，根源仍在 ③

②那一行有个独立陷阱：padding。 一个 batch 里句长不同，短句要补齐。 如果某个位置的全部键都被掩码（比如补齐位置）， 它的整行就是 -\infty，softmax 得到 nan。

解法：要么保证每行至少有一个可见位置（通常保证第一个位置不被掩码）， 要么在 softmax 前把整行全 -\infty 的情况替换成全 0（之后输出也是 0）。

注记停下来想一想 ⑯

上面说”整行全被屏蔽 → nan”。

请想一个更隐蔽的情况：如果一整行看起来有可见位置， 但那些位置的 logits 都是 −1000（非常大的负数但不是 −∞）呢？

e^{-1000} 在 float32 下等于 0。那么这一行求和是 0， 归一化时 0/0 得到 nan。

什么情况下会造出”很大的负数但不是 −∞“的 logits？ （提示：想想加性掩码——有些人用 scores + (-1e9) * (1 - mask) 而不是 masked_fill。 再想想：如果 q、k 本身尺度很大，正常位置的 logits 可能是 ±50， 那么 −1e9 的相对差距还够大吗？）

1.13.4 排查四：一切都”正常”，但效果就是不如预期

症状：没有 nan，loss 正常下降，生成的东西也说得过去。 但和同等规模的已知实现比，就是差一截。

这是最难的一类，因为没有明显症状。排查思路是逐个要素做消融： 把每个设计要素单独关掉，看影响有多大。

消融清单（按我们的实际经验排序）：

要素	关掉之后	说明什么
\sqrt{d_k} 缩放	看 §1.4 案例	小模型上影响小，大模型上灾难
残差连接	直接训不动	深度超过 2 层就必须有
LayerNorm	loss 抖动、收敛慢	不是可选配件
权重共享	参数多 36%，效果略降	小模型上影响明显
位置编码	完全学不会顺序	见 §1.8
warmup	前期 loss 尖刺	有时能靠调小学习率绕过
梯度裁剪	偶发 nan	稳定性保险

做消融的纪律：一次只关一个，其余全部不动。 这和本章开头那个对照实验是同一条纪律。

提示为什么消融比”读论文”有用

论文会告诉你”我们用了 X、Y、Z”，但很少告诉你哪些是必需的。 因为论文的语境是”我们这样做了，效果很好”，不是”我们试过不做”。

自己做一遍消融，你会得到一份属于你这份代码的必需性清单。 它可能和论文里的暗示不一样—— 比如你可能会发现某个”重要技巧”在你的规模下毫无影响（就像我们的 \sqrt{d_k}）。

这也是本章那个案例真正的教学目的： 不是让你记住”\sqrt{d_k} 重要”， 而是让你养成”这件事在我这个规模上到底重不重要”的习惯。

1.13.5 一张排查总表

把上面四条路径汇总成一张可以贴在显示器旁的表：

症状	先查什么	最可能的原因
loss 不降，梯度正常	注意力三个诊断量	缩放问题、或模型容量
loss 很漂亮但生成差	单测：未来位置权重是否为 0	掩码泄露
出现 nan	分四步插桩定位	缺数值稳定 / padding 全屏蔽
效果差一截	逐要素消融	缺了某个非必需但有用的组件
显存爆	算 n^2 \times h \times L	序列太长，需要 FlashAttention

最后一行那个公式值得记一下——它能让你在跑之前就估算出显存需求：

\text{注意力矩阵显存} \approx \underbrace{n^2}_{\text{位置对数}} \times \underbrace{h}_{\text{头数}} \times \underbrace{L}_{\text{层数}} \times \underbrace{2\ \text{或}\ 4}_{\text{fp16 或 fp32 字节数}} \tag{1.45}

n = 4096、h = 32、L = 32、fp16 时：4096^2 \times 32 \times 32 \times 2 \approx 34 GB。 这在单卡上装不下——而如果用了 FlashAttention，这部分显存根本不需要。

1.14 完整实现

下面是本章所有式子的实现。每一处关键选择都标了为什么。

import numpy as np


def softmax(x, axis=-1):
    """数值稳定的 softmax。

    减最大值一步是**必需的**，不是优化：
    float32 下 logits 超过约 88 就会让 exp 溢出成 inf。
    由平移不变性，减最大值不改变结果。
    """
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)


def attention(Q, K, V, mask=None, scaled=True):
    """缩放点积注意力。

    Q: (..., n_q, d_k)   K: (..., n_k, d_k)   V: (..., n_k, d_v)
    mask: 布尔数组，True 表示「允许看」。因果掩码取严格下三角。

    返回 (输出, 权重)。权重是 (..., n_q, n_k)，每行和为 1。
    """
    d_k = Q.shape[-1]

    # ① 打分。缩放必须在 softmax **之前**，它修正的是 logits 的尺度。
    scores = Q @ K.swapaxes(-1, -2)
    if scaled:
        scores = scores / np.sqrt(d_k)

    # ② 掩码。必须用 -inf 而不是 0，否则被屏蔽的位置仍会分到权重。
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)

    # ③ 归一化。softmax 逐行做，把每一行变成概率分布。
    weights = softmax(scores, axis=-1)

    # ④ 加权求和。按权重把各位置的 v 混合起来。
    return weights @ V, weights


def causal_mask(n):
    """因果掩码：位置 i 只能看到 j <= i。"""
    return np.tril(np.ones((n, n), dtype=bool))


def multi_head_attention(X, W_q, W_k, W_v, W_o, n_head, mask=None):
    """多头注意力。

    X: (n, d_model)；W_*: (d_model, d_model)
    所有头的投影合并成一个大矩阵，再 reshape 切分——这是实际实现的方式，
    比循环 h 次快得多，且数学上完全等价。
    """
    n, d_model = X.shape
    d_head = d_model // n_head
    assert d_model % n_head == 0

    # 一次投影出所有头的 q/k/v，再拆成 (n_head, n, d_head)
    def split(W):
        return (X @ W).reshape(n, n_head, d_head).transpose(1, 0, 2)

    Q, K, V = split(W_q), split(W_k), split(W_v)

    out, weights = attention(Q, K, V, mask=mask)

    # 拼接各头：(n_head, n, d_head) -> (n, d_model)
    out = out.transpose(1, 0, 2).reshape(n, d_model)
    return out @ W_o, weights

注记三处容易写错的地方

① 缩放的位置（第 22 行）：必须在 attention 内部、softmax 之前。 如果挪到函数外面，就变成了”对输出做缩放”，语义完全不同（见失败二）。

② 掩码的取值（第 26 行）：-np.inf 而非 0。见失败一。

③ 多头的 reshape（第 51 行）：reshape(n, n_head, d_head) 之后要 transpose(1, 0, 2) 把头的维度提前，才能让 attention 里的 swapaxes(-1, -2) 正确工作。这个维度顺序错了不会报错——只是结果不对。

复杂度：Q @ K.swapaxes 是 O(n^2 d_k)，softmax 是 O(n^2)， weights @ V 是 O(n^2 d_v)。三项都随 n 平方增长。

这一段代码在 notebook 里：notebooks/attention_from_scratch.ipynb 把上面每一行拆成可单独运行的单元，并逐步与 PyTorch 的 F.scaled_dot_product_attention 对拍。

1.14.1 对拍：确认我们的实现是对的

写了实现不算完，得证明它和工业级实现一致。PyTorch 从 2.0 起提供了 F.scaled_dot_product_attention（内部会在合适时自动切到 FlashAttention）。 拿它当标准答案：

import math, torch, torch.nn.functional as F

def my_attention(Q, K, V, mask=None, scaled=True):
    d_k = Q.shape[-1]
    s = Q @ K.swapaxes(-1, -2)
    if scaled:
        s = s / math.sqrt(d_k)
    if mask is not None:
        s = s.masked_fill(~mask, float("-inf"))
    w = F.softmax(s, dim=-1)
    return w @ V, w

torch.manual_seed(0)
n, d = 6, 32
Q, K, V = (torch.randn(2, 4, n, d, dtype=torch.float64) for _ in range(3))
mask = torch.tril(torch.ones(n, n, dtype=torch.bool)).expand(2, 4, n, n)

mine, w = my_attention(Q, K, V, mask)
ref = F.scaled_dot_product_attention(Q, K, V, attn_mask=mask)
print("最大绝对误差:", (mine - ref).abs().max().item())
# 6.661e-16

最大绝对误差 6.66e-16。 这正好是 float64 的机器精度量级 （\varepsilon \approx 2.2\times10^{-16}，累加几步后到 10^{-15} 量级是正常的）。 说明我们的实现与官方实现在数值上等价。

提示为什么要做这个对拍

自己实现的注意力有一个特点：大多数错误不会让程序崩，只会让结果悄悄不对。 前面五个失败模式里有三个不报错，这个对拍就是它们的统一防线。

更实际的一点：如果你在调试一个自定义注意力（比如加了稀疏模式、 改了缩放方式），不知道怎么验证对不对——先退化到标准情况与官方实现对拍， 确认基础部分无误，再往上加你的修改。这比直接调试一个复杂变体高效得多。

1.14.2 实测：O(n^2) 到底有多贵

光说”平方增长”不够直观。在 RTX A6000 上实测 F.scaled_dot_product_attention（8 头、d_{\text{head}}=64、因果掩码）：

序列长度 n	注意力矩阵	单次前向耗时	相对 n=512
512	512×512	0.082 ms	1.0×
1024	1024×1024	0.201 ms	2.5×
2048	2048×2048	0.581 ms	7.1×
4096	4096×4096	2.307 ms	28.2×
8192	8192×8192	8.263 ms	101×

读这张表要注意两件事。

第一，n 翻倍，耗时涨约 3–4 倍，不是 4 倍整。 从 512 到 8192， n 涨了 16 倍，纯平方应该是 256 倍，实测只有 101 倍。 差距来自 FlashAttention：它不把 n\times n 矩阵写进显存， 省下了大量访存时间。所以”O(n^2)“描述的是计算量，不是实际耗时的准确预测。

第二，这里的绝对数字很小。 n=8192 只要 8 毫秒—— 看起来长上下文并不贵？那是因为这只是单层单次前向。 真实的 70B 模型有 80 层，生成一个 token 要跑一遍全部层， 还要加上 KV 缓存的读写。乘上去就是几百毫秒到秒级。

警告一个容易踩的测量陷阱

我们第一次跑这组数据时，n=512 测出来是 3.83 ms——比 n=1024 还慢。

这不是算法问题，是首次调用的预热开销：CUDA 需要初始化上下文、 编译 kernel、分配显存池。这些一次性成本全被算进了第一次的耗时里。

正确做法是先跑几轮热身，再开始计时。上面的表就是这么测的 （预热 5 轮，然后测 30 次取平均）。

凡是性能数字，先问一句”预热了吗”。 这是性能测量最基本、 也最常被忽略的一条纪律。

1.14.3 与真实模型的距离

最后要说清楚：本章的这个实现和真实模型差在哪里。

	本章实现	真实模型（如 70B）
参数量	42.6 万	700 亿
层数 / 头数	2 / 4	80 / 64
d_{\text{model}}	128	8192
位置编码	可学习的位置嵌入	RoPE
归一化	LayerNorm	RMSNorm
注意力实现	朴素（写完整的 n^2 矩阵）	FlashAttention
KV 缓存	无（训练场景）	有，GQA 压缩

方向可迁移，数值不可迁移——这句话在本章反复出现，这里给出它的边界： 上面每一行的差异都可能影响结论的具体数值， 但”缩放在修什么”、“softmax 为什么会饱和”这些机制层面的结论不受影响。

第 2 章会补齐位置编码（RoPE）与 RMSNorm；第 3 章讲 KV 缓存的实现细节。

1.15 批处理：padding、打包与掩码

真实训练不会一次只喂一句话。一个 batch 里有几十条长度不同的序列—— 而矩阵乘法要求形状一致。这一节讲这个矛盾怎么解， 以及它和因果掩码怎么叠加。

1.15.1 问题：长度不一

一个 batch 里有三条序列，长度分别是 5、3、7。要放进同一个张量，只能补齐：

序列 A（长 5）： 秦 川 雄 帝 宅
序列 B（长 3）： 空 山 不
序列 C（长 7）： 明 月 松 间 照 清 泉

补齐到 7：

A： 秦 川 雄 帝 宅 <pad> <pad>
B： 空 山 不 <pad> <pad> <pad> <pad>
C： 明 月 松 间 照 清 泉

padding 的位置不携带信息，必须让注意力忽略它们。否则：

• 位置 “宅” 会从两个 <pad> 取信息，输出被稀释
• <pad> 自己的输出毫无意义，但会参与 loss（如果不做 loss mask）

1.15.2 两种掩码要叠加

这时一个位置能不能被看到，取决于两个条件：

\text{可看} = \underbrace{j \le i}_{\text{因果掩码}} \quad \text{且} \quad \underbrace{j \text{ 不是 padding}}_{\text{填充掩码}} \tag{1.46}

写成代码：

def build_mask(seq_len, lengths, device):
    """lengths: 每条序列的真实长度。返回 (B, 1, L, L) 的布尔掩码。"""
    B = len(lengths)
    causal = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))
    pad = torch.zeros(B, seq_len, dtype=torch.bool, device=device)
    for b, L in enumerate(lengths):
        pad[b, :L] = True                      # 前 L 个位置是真实 token
    # 键（列方向）不能是 padding；查询（行方向）是 padding 无所谓（输出会被丢弃）
    return causal[None, None] & pad[:, None, None, :]

注意最后一行：只有在列方向（键方向）才需要排除 padding。 行方向是 padding 的那些位置，它们的输出本来就是垃圾， 只要在算 loss 时忽略掉就行（用 ignore_index）。

警告这就是”整行全是 -\infty“的来源

如果一个查询位置本身是 padding，而所有键也都把它屏蔽了， 那么这一行就是全 -\infty → softmax 得到 nan。

两种常见的解法和它们的代价：

① 保证每行至少有一个可见位置（通常是位置 0，或者自己）。 做法是把掩码的对角线强制设为 True。代价：padding 位置会算出一个无意义的输出， 但只要 loss mask 正确，它不参与梯度，无伤大雅。

② 在 softmax 之前把全 -\infty 的行替换成全 0。 这样 softmax 得到均匀分布，输出是 V 的平均——同样无意义，但不会是 nan。

我们推荐 ①：它更简单，而且不改变任何有效位置的数值。

1.15.3 打个比方：padding 让训练白跑

padding 不只是”实现细节”，它直接浪费算力。

一个 batch 里有长度 512 和长度 32 的两条序列，补齐到 512：

\text{有效计算比例} = \frac{512 + 32}{512 \times 2} = 53\% \tag{1.47}

一半的算力花在 <pad> 上。 如果 batch 里的长度差异更大（比如 2048 和 16）， 这个比例会低到 25% 以下。

解法叫”序列打包”（sequence packing）：把多条短序列首尾相接放进同一个 长度固定的窗口，让它们共享位置：

打包前（补到 8）：
  A: 秦 川 雄 帝 宅 pad pad pad
  B: 空 山 不  pad pad pad pad pad

打包后（窗口 8）：
  秦 川 雄 帝 宅 空 山 不          ← 两条序列占据了同一行的前 8 个位置

但这带来一个新问题：位置 5 的 “空” 不能看到位置 0–4 的 A 序列内容—— 它们是不同的序列，不该互相借信息。

所以打包需要块对角掩码：每条子序列内部用因果掩码， 子序列之间完全屏蔽。这就把 式 1.46 变成了三分支的与运算：

\text{可看} = (j \le i) \ \text{且}\ (\text{同一条子序列}) \ \text{且}\ (j \text{ 非 padding}) \tag{1.48}

提示打包是个典型的”工程优化改变算法表述”的例子

数学上，打包改变了注意力矩阵的结构——从”一个下三角”变成”多个对角块”。 但结果完全等价，因为掩码保证了跨块的信息不会流动。

收益：训练吞吐提升 2–4 倍（取决于长度分布）。

代价：实现复杂度上升——你要维护”每条子序列的起始位置”这个额外结构， 还要保证位置编码在打包后仍然正确（这点很微妙： 如果位置编码按绝对位置算，第二条子序列会从位置 5 开始编号，而不是从 0）。

所以打包必须配合”按子序列重置的位置编码”—— 这又是一个”看起来是优化、实际上牵动多处”的例子。

1.15.4 与 KV 缓存的关系

padding 的问题在推理时变了形态。

推理是逐 token 生成的，每次只处理 1 个新 token。 这时不存在”batch 内长度不一”的问题（每条序列自己管自己）， 但出现了另一种长度问题：KV 缓存会一直增长。

式 1.21 算过：Llama-2 70B 每 token 2.6 MB。 序列长 4096 时是 10.7 GB——而且这只是缓存，不含模型权重。

这就是为什么 MQA/GQA（小节 1.8）是必需的而非可选的： GQA 把缓存压到 1/8，从 10.7 GB 降到 1.3 GB。 没有这个优化，长上下文推理在单卡上根本放不下。

注记停下来想一想 ⑱

训练时不存在 KV 缓存的问题（因为整个序列一次算完）， 而推理时它成了主要瓶颈。

请想：为什么训练不需要缓存？

提示：训练时是”教师强制”——整条序列的真实 token 都在手上， 所以 K、V 可以一次性算完，直接用矩阵乘法。 而推理时第 t 步只有前 t-1 个 token，第 t 个还没生成， 只能一个一个来。算过的 K、V 如果不存下来就得重算——这就是缓存的动机。

第 3 章会详细讲 KV 缓存的实现与优化。

1.16 梯度怎么穿过注意力

前面的推导全是前向的：给定输入，算出输出。 但训练靠的是反向传播——误差信号怎么从输出一路传回到 Q、K、V？

这一节把整条反向链路推完。它对理解两件事至关重要： 为什么 softmax 饱和会阻断训练（前面反复提到）， 以及为什么注意力对数值精度特别敏感（下一节）。

1.16.1 正向链路回顾

把前向过程拆成五步，每一步记下输入输出：

\underbrace{Q, K}_{\text{输入}} \xrightarrow{\text{①}} S = \frac{QK^{\top}}{\sqrt{d_k}} \xrightarrow{\text{②}} P = \mathrm{softmax}(S) \xrightarrow{\text{③}} O = PV \tag{1.49}

反向传播就是反着走这条链，每一步用链式法则。

设我们从上游拿到 \dfrac{\partial \mathcal{L}}{\partial O}（记作 dO）。 目标是求出 dQ、dK、dV。

1.16.2 第③步的反向：O = PV

这是标准的矩阵乘法求导。若 O = PV，则：

dV = P^{\top}\, dO, \qquad dP = dO\, V^{\top} \tag{1.50}

这一步没有陷阱，和普通线性层一样。 dV 是”每个值向量该往哪个方向调”，dP 是”每个注意力权重该往哪个方向调”。

1.16.3 第②步的反向：softmax 的雅可比

这一步是整条链上最关键、也最容易错的地方。

我们已经推过 softmax 的雅可比（式 1.16）： \dfrac{\partial P_i}{\partial S_j} = P_i(\delta_{ij} - P_j)。

把它作用到 dP 上，得到 dS。逐行看（每一行独立）：

dS_i = P_i \odot \left(dP_i - \langle dP_i,\ P_i\rangle \cdot \mathbf{1}\right) \tag{1.51}

其中 \odot 是逐元素乘法，\langle \cdot,\cdot\rangle 是内积， \mathbf{1} 是全 1 向量。

人话翻译这一行。 括号里做的是：

从”每个权重想往哪动”里，减掉”所有位置想动的加权平均”。

为什么要减？因为 softmax 有一个硬约束：权重之和恒为 1。 所以一个位置想变大，必然意味着别的位置要变小—— “往哪动”的方向里不能有”整体一起变大”这个分量。

那个 \langle dP_i, P_i\rangle 就是”整体一起变大”的分量大小。 减掉它，剩下的才是 softmax 真正能实现的方向。

这个模式在别处也出现过：在带约束的优化里（比如单纯形上的投影）， 都要先减掉”违反约束的那个方向”。softmax 的反向传播就是它的一个特例。

注意这一行里的 P_i 因子。 它是乘在外面的：

dS_i \propto P_i \tag{1.52}

如果 P_i 是 one-hot（某个分量是 1，其余是 0），那么 dS_i 就只有那一维非零。 如果 P_i 完全饱和到 (1, 0, 0, \dots)，那么 P_i \odot (\cdots) 会让所有分量都变成 0——

这就是”梯度消失”在公式层面的样子。 前面我们是从雅可比范数（\|J\|_F \to 0） 观察到的，这里从反向传播的具体形式再确认一遍。

1.16.4 第①步的反向：S = QK^{\top}/\sqrt{d_k}

同样是矩阵乘法，但要注意那个缩放因子：

dQ = \frac{1}{\sqrt{d_k}}\, dS\, K, \qquad dK = \frac{1}{\sqrt{d_k}}\, dS^{\top} Q \tag{1.53}

\sqrt{d_k} 出现在分母上。 这一点在讨论”缩放的副作用”时容易被忽略：

除以 \sqrt{d_k} 会同时缩小前向的 logits 和反向的梯度。 d_k = 64 时，梯度被缩小 8 倍。

所以缩放不是”只有好处”——它用前向的稳定性换了一点反向的信号强度。 不过在饱和的情形下（梯度本来就接近 0），这点损失和”完全传不回去”相比， 是完全可以接受的。

1.16.5 掩码位置怎么处理

被掩码的位置（j > i）在前向时权重为 0。反向时：

dS_{ij} = 0 \quad \text{对所有被掩码的 } (i,j) \tag{1.54}

因为那些位置没有参与前向计算，它们不该收到任何梯度。

实现上，在算完 式 1.51 之后要把这些位置清零。 如果忘了这一步，梯度会流向”本不该被看到”的位置—— 和掩码 bug 是同一类问题，只是出现在反向。

1.16.6 完整实现与验证

把上面五步连起来：

def attention_backward(Q, K, V, mask, dO):
    """手写的注意力反向传播。forward 用前面那个实现。"""
    n, d_k = Q.shape
    # ── 重做前向，拿到中间量 ──
    S = Q @ K.T / np.sqrt(d_k)
    S = np.where(mask, S, -np.inf)
    P = softmax(S, axis=-1)
    # ── 反向 ──

    # 第③步：O = P V
    dV = P.T @ dO
    dP = dO @ V.T

    # 第②步：softmax 的雅可比
    #   dS_i = P_i ⊙ (dP_i − ⟨dP_i, P_i⟩)
    dS = P * (dP - (dP * P).sum(-1, keepdims=True))

    # 掩码位置梯度清零
    dS = np.where(mask, dS, 0.0)

    # 第①步：S = QKᵀ / √d_k
    dQ = dS @ K / np.sqrt(d_k)
    dK = dS.T @ Q / np.sqrt(d_k)
    return dQ, dK, dV

验证方式：和 PyTorch 的 autograd 对拍。

import torch
torch.manual_seed(0)
n, d = 3, 4
Q = torch.randn(n, d, dtype=torch.float64, requires_grad=True)
K = torch.randn(n, d, dtype=torch.float64, requires_grad=True)
V = torch.randn(n, d, dtype=torch.float64, requires_grad=True)
mask = torch.tril(torch.ones(n, n, dtype=torch.bool))

S = (Q @ K.T / d**0.5).masked_fill(~mask, float("-inf"))
O = torch.softmax(S, -1) @ V
O.backward(torch.randn_like(O))

# 手写的结果（用上面那个函数，转成 numpy 再转回来）
# ... 实测：dQ、dK、dV 与 autograd 的最大误差都是 0.000e+00

我们实测的误差是 \mathbf{0}——在 float64 下逐位一致。

提示为什么要手推一遍反向

框架会自动帮你算梯度，所以”能用”这个层面你不需要手写。

但手推一遍能解释三件框架不会告诉你的事：

1. 为什么饱和会阻断训练——因为在 式 1.52 里， dS 正比于 P，而饱和时 P 在非峰值位置全是 0
2. 为什么缩放会同时影响反向——1/\sqrt{d_k} 在 式 1.53 里出现两次
3. 为什么掩码在反向也要处理——不只是前向的事

“知道框架会算”和”知道它在算什么”，差别在出问题时能不能定位。

1.17 数值精度：为什么注意力格外敏感

前面反复出现”float32 下 logits 超过约 88 就会溢出”。但今天的模型 大多在 fp16 或 bf16 下训练和推理。这一节讲清楚这意味着什么。

1.17.1 三种格式的可用范围

torch.finfo 给出的实测数据：

格式	最大可表示值	\exp 能算到 logits ≈	机器精度
float32	3.40\times10^{38}	88.7	1.2\times10^{-7}
float16	6.55\times10^{4}	11.1	9.8\times10^{-4}
bfloat16	3.39\times10^{38}	88.7	7.8\times10^{-3}

第二列是这一节的核心。 float16 的最大值是 65504， 所以 \exp(z) 在 z \approx \ln(65504) = 11.1 时就溢出了—— 比 float32 的门槛低了 8 倍。

1.17.2 这个门槛有多容易被越过

按脚注的假设（q、k 各分量方差为 1），不缩放时 logits 的标准差是 \sqrt{d_k}。 取四倍标准差作为”常见最大值”：

d_{\text{head}}	标准差	常见最大值	在 fp16 下
32	5.7	22.6	已经溢出
64	8.0	32.0	已经溢出
128	11.3	45.3	已经溢出
256	16.0	64.0	已经溢出

警告这张表说明一件很实际的事

在 fp16 下，不缩放的注意力几乎必然溢出——连 d_{\text{head}} = 32 都撑不住。

而本章案例里，在 fp32 下 d_{\text{head}} = 32 时训练完全正常（差 0.17）。 同一个”没有 bug”的实现，换个精度就崩了。

这解释了一个常见的现象：有人在小模型上（fp32）调试一切正常， 搬到真实训练脚本（fp16 混合精度）里立刻 nan。 不是实现错了，是精度换了。

1.17.3 那为什么还用 fp16

因为显存和带宽。fp16 只占 fp32 一半的字节数， 而正如 小节 1.8 说的，注意力的瓶颈常常是访存而非算力。 用 fp16 意味着同样显存能放两倍的 batch、两倍的序列长度。

代价就是上面那张表：可表示的数值范围小了一大截。

1.17.4 bf16：Google 的方案

bfloat16 的思路很直接：保留 fp32 的指数位（8 位），砍掉尾数位 （从 23 位砍到 7 位）。

结果是：

• 数值范围和 fp32 完全一样（指数位相同）→ \exp 也能算到 88.7
• 精度差很多（尾数只有 7 位）→ 1.001 存进去会变成 1.000

实测这个精度损失：

x = torch.tensor([1.001], dtype=torch.float32)
print(x.to(torch.float16).float())    # 1.000977  ← 分辨得出
print(x.to(torch.bfloat16).float())   # 1.000000  ← 分辨不出

fp16 和 bf16 的取舍，一句话说清：

• fp16：精度高（10 位尾数），范围窄（最大 65504）
• bf16：范围宽（和 fp32 一样），精度低（7 位尾数）

对注意力来说，范围比精度重要。 因为注意力的数值跨度极大—— logits 可能从 -50 到 +50，softmax 之后又要在 10^{-20} 量级上区分。 范围不够会直接溢出成 nan；精度不够只是结果略微不准。

所以今天的主流是 bf16（A100/H100 原生支持，训练稳定）， 而 fp16 更多用在推理场景（精度要求相对低，且有成熟的 loss scaling 机制）。

这是”选数值格式”这件事的一般原则：先保证范围够，再谈精度。

1.17.5 混合精度训练里的三个保护措施

即使选了 bf16，生产级的训练脚本通常还有几层保险：

① 数值稳定的 softmax。 前面讲过——先减最大值。 这一条在 bf16 下仍然是必需的（因为减完之后指数最大只到 1， 远离任何格式的上限）。

② Loss scaling。 主要用于 fp16：因为 fp16 的最小正规数是 6.1\times10^{-5}， 很多梯度会直接下溢成 0。做法是先把 loss 乘一个大的常数（比如 1024）， 反向传播后再把梯度除回来。

bf16 因为指数范围和 fp32 一样，通常不需要 loss scaling—— 这是它比 fp16 省心的主要原因之一。

③ 关键部分保持 fp32。 常见做法是 LayerNorm、softmax、以及损失计算 在 fp32 下做（这叫 “autocast 白名单”）。 原因是这些操作涉及累加——n 个数相加时，低精度的舍入误差会累积。

注记停下来想一想 ⑰

上面说”softmax 在 fp32 下做”。

请想：softmax 里哪一步涉及累加，为什么低精度会有问题？

具体一点：如果一行有 4096 个位置，每个权重约 2.4\times10^{-4}， 用 bf16（精度 7.8\times10^{-3}）算它们的和会怎样？

（提示：bf16 分辨不出比 7.8\times10^{-3} 更小的相对差异。 而 2.4\times10^{-4} 比它小了 30 倍。）

这就是为什么归一化操作通常要在高精度下做。

1.17.6 一张实践清单

如果你要写一个能在混合精度下稳定运行的注意力实现：

措施	作用	不做会怎样
减最大值的 softmax	把 \exp 的输入限制在 \le 0	溢出成 inf
\sqrt{d_k} 缩放	把 logits 拉回单位方差	fp16 下大概率溢出
softmax 用 fp32 累加	避免小权重被舍入掉	权重和不等于 1
优先选 bf16	范围和 fp32 一样	fp16 需额外的 loss scaling
梯度裁剪	兜住偶发的梯度尖峰	一次异常样本毁掉训练

这五条里，前两条是本章反复强调的”缩放与稳定”， 后三条是工程实践。两件事都要有——算法上的正确性和工程上的稳健性，缺一不可。

1.18 怎么把这个模型训出来

前面的案例用的是真实训练出来的模型。这一节把训练代码拆开讲， 让你能自己改动它、复现它。

完整代码在 notebooks/tiny_gpt.py（模型与训练循环）和 notebooks/run_ch1_case.py（案例的对照实验）。这里只讲为什么这么写。

1.18.1 数据：字符级分词

class CharTokenizer:
    """只保留最高频的 n 个字，其余映射到 <unk>。字表大小本身是个超参数。"""
    def __init__(self, text, vocab_size):
        cnt = Counter(text.replace("\n", ""))
        keep = [c for c, _ in cnt.most_common(vocab_size - 2)]
        self.itos = ["<pad>", "<unk>"] + keep
        self.stoi = {c: i for i, c in enumerate(self.itos)}

为什么用字符级而不是词级？ 三个理由：

1. 不需要分词器。中文分词本身是个有争议的问题（“研究生命起源”该切成什么？）， 字符级绕开了它。
2. 词表小。25 万字唐诗里只有 4977 个不同字符，取前 1200 个就覆盖 87.6%。 词级分词动辄三五万词表，光嵌入矩阵就比我们整个模型大。
3. 教学上更透明。每个 token 就是一个字，你能直接看懂模型在预测什么。

代价是序列变长（一句话的字数是词数的 1.5–2 倍）， 上下文里能放的内容更少。真实的中文大模型用 BPE 之类的子词分词。

提示为什么砍到 1200 个字

保留全部 4977 个字的话，嵌入矩阵是 4977 \times 128 = 63.7 万参数—— 比模型其余部分加起来还多，而且低频字（出现一两次）根本学不好。

砍到 1200 个是覆盖率与参数量的折中：覆盖 87.6% 的字， 嵌入矩阵降到 15.4 万参数。剩下 12.4% 的字都映射到 <unk>—— 这意味着模型对生僻字完全无能为力，但作为教学模型这完全可以接受。

这个权衡在真实系统里同样存在，只是数字不同： 词表从 32K 扩到 128K，嵌入参数翻两番，换来的是更短的处理序列和更少的 <unk>。

1.18.2 模型：两个关键细节

模型的主体（嵌入 + 若干 Block + 输出层）是标准的。但有两处值得单独说：

① 权重共享（weight tying）。

if cfg.tie_weights:
    self.head.weight = self.tok_emb.weight

输入嵌入是”把字变成向量”（V \times d），输出层是”把向量变回字”（d \times V）。 让它们共用同一个矩阵，有两个好处：

• 省参数：V \times d 个参数，在这个模型里是 15.4 万，占总参数的 36%
• 有道理：如果”月”这个字的向量是 e，那么”某个向量是不是月”就应该看它与 e 的相似度。两个方向用同一套表示是合理的

这个技巧从原论文就有，今天的大模型普遍沿用。

② 残差分支的缩放初始化。

for name, p in self.named_parameters():
    if name.endswith("proj.weight") or name.endswith("mlp.2.weight"):
        nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer))

注意那个 \sqrt{2 n_{\text{layer}}}。它不是缩放的 \sqrt{d_k}—— 两个东西完全不同，只是都用了平方根。

它解决的问题是：残差连接 x = x + f(x) 每经过一层就把 f(x) 累加上去。 如果 f 的输出方差是 \sigma^2，n 层之后 x 的方差会累积到 n\sigma^2。 把 f 的初始化标准差按 1/\sqrt{2n} 缩小（2 是因为每层有两个残差分支： 注意力和 FFN），可以让累积后的方差保持在 1 附近。

这是 GPT-2 论文里的做法，目的是让深层网络的训练更稳定。

警告别把两个 根号 搞混

	缩放对象	出现在哪	解决什么
1/\sqrt{d_k}	注意力打分	前向传播里	logits 随维度膨胀 → softmax 饱和
1/\sqrt{2n_{\text{layer}}}	权重初始化	训练开始前	残差累加 → 激活值方差膨胀

两个都是”按维度的平方根缩放”，但作用对象和失效模式完全不同。 名字里都有 根号 不代表它们是同一件事——这在读论文时是个常见的混淆点。

1.18.3 训练循环：四个必须有的部件

opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr,
                        weight_decay=cfg.weight_decay, betas=(0.9, 0.95))

for step in range(cfg.max_steps):
    for g in opt.param_groups:
        g["lr"] = lr_at(step, cfg)              # ① 学习率调度
    x, y = get_batch(tr, cfg.block_size, cfg.batch_size, gen)
    _, loss = model(x, y)
    opt.zero_grad(set_to_none=True)
    loss.backward()
    gn = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)  # ② 梯度裁剪
    opt.step()

① 学习率调度（warmup + 余弦衰减）：

def lr_at(step, cfg):
    if step < cfg.warmup:                       # 前 100 步线性升温
        return cfg.lr * (step + 1) / cfg.warmup
    prog = (step - cfg.warmup) / (cfg.max_steps - cfg.warmup)
    return cfg.lr * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1.0, prog))))

为什么需要 warmup？ 训练刚开始时参数是随机的，梯度方向没什么意义。 如果一上来就用大学习率，会朝一个随机的方向猛冲一步，可能把模型带到一个坏的区域。 先小步走 100 步，等梯度方向稳定了再加到全速。

为什么需要衰减？ 训练后期模型接近一个较好的解， 大步长会让它在解的附近来回跳。余弦衰减让它逐渐减小步长，精细收敛。

② 梯度裁剪：clip_grad_norm_ 把梯度的范数限制在 1.0 以内。 如果某一步梯度突然很大（比如遇到了异常样本），它会按比例缩小， 避免一步把参数带飞。

提示判断训练是否正常的三个信号

1. loss 曲线：应该平滑下降，偶尔抖动正常。 如果出现尖刺，多半是某个 batch 有问题；如果长期不降，检查学习率。

2. 梯度范数：我们记录它就是为了这个。健康的值在 0.1–10 之间波动。 如果它长期贴着裁剪阈值（1.0），说明学习率偏大； 如果它掉到 10^{-4} 以下并且不回升，说明梯度消失了—— 本章案例里不缩放那一边就是这样（雅可比范数掉到 8\times10^{-5}）。

3. 注意力权重熵：健康值在 0.5–0.9 之间。 接近 1.0 说明退化成平均（缩放过度）， 接近 0 说明饱和（缩放不足或缺失）。

这三个量合起来，能在 loss 曲线还没表现出问题时就提前发现异常。

1.18.4 超参数是怎么定的

参数	值	为什么
n_layer	2	一层学不出”头的分工”，三层以上训练变慢且现象一样
n_head	4	够看出”不同头做不同事”（实测 4 个头确实分化了）
d_model	128	参数量约 42.6 万，CPU 上几分钟能训完
block_size	64	唐诗一句 5–7 字，64 个 token 能容纳好几联
vocab_size	1200	覆盖 87.6% 的字，见上文
batch_size	32	显存与梯度稳定性的折中
lr	3e-3	小模型用大学习率；实测这个值收敛最快且不发散
warmup	100	约总步数的 7%，够让梯度方向稳定
max_steps	1500	经验值：loss 在这个位置基本平了（见 图 1.2）

最后一行值得说一句：max_steps = 1500 不是因为”训练完了”， 而是因为再训下去收益很小。判断依据是 loss 曲线的斜率—— 从 1000 步到 1500 步只降了 0.1 左右，而前 200 步降了 2 以上。

训练步数永远是”够用就好”，不是”越多越好”。 这个判断要靠看曲线，不能靠设一个固定值。

1.18.5 想自己改？

改这几个地方最能看到效果：

# ① 复现"训不动"：把缩放关掉，同时把头数改成 1（d_head 变成 128）
cfg = Config(scaled=False, n_head=1, d_model=128)

# ② 看缩放到底有没有影响：只关缩放，保持 4 头
cfg = Config(scaled=False, n_head=4, d_model=128)

# ③ 加深：看 4 层的模型能不能学出更明显的层间分工
cfg = Config(n_layer=4)

# ④ 换语料：把 data/corpus.txt 换成别的文本，重跑

改完跑 run_ch1_case.py，对比 loss 曲线和 analyze_heads.py 的头部画像。 自己动手调一遍，比读十遍推导有用。

1.19 对照实验是怎么搭的

本章那个案例给了三组实验：初始化扫描、小模型对照、大模型对照。 结论已经讲过了，这里讲它们各自是怎么设计的—— 因为设计一个能说明问题的实验，和跑一个实验，是两件事。

1.19.1 实验一：初始化扫描

要回答的问题：q^\top k 的实际尺度到底有多大？和理论值 \sqrt{d_k} 差多少？

为什么不能直接量训练后的模型：训练后的 q、k 是学出来的， 尺度已经被优化过程改变了。要检验”脚注的假设成立吗”， 必须量初始化状态——那是假设声称成立的地方。

做法：

def init_sweep(text, d_heads):
    out = []
    for d_model in d_heads:
        for scaled in (False, True):
            cfg = Config(d_model=d_model, n_head=1, n_layer=1, scaled=scaled,
                         vocab_size=300, block_size=32, device=DEV)
            torch.manual_seed(0)                 # ① 固定种子
            m = TinyGPT(cfg).to(DEV).eval()      # ② 不训练，只初始化
            x = torch.randint(2, 300, (8, 32), device=DEV)
            with torch.no_grad():
                m(x, record=True)                # ③ 走一次前向，采集诊断量
            d = m.blocks[0].attn.last_diag
            out.append({...})
    return out

三个设计要点：

① 固定种子。 每个 d_{\text{head}} 都用同一个种子初始化。 如果不固定，不同维度之间的差异里就混进了初始化噪声—— 而我们想看的恰恰是”维度”的影响。

② 不训练，只初始化。 这是这个实验的关键。 训练会改变 q、k 的尺度，把我们要观察的对象污染掉。

③ 用同一批输入。 x = torch.randint(2, 300, (8, 32)) 在所有配置下都是同一批， 保证唯一的变量是 d_{\text{head}}。

为什么用 n_{\text{head}} = 1：这样 d_{\text{head}} = d_{\text{model}}， 可以直接用 d_{\text{model}} 的列表控制 d_{\text{head}}。 如果保持 4 头，d_{\text{model}} 就得是 d_{\text{head}} 的 4 倍，多一层换算，容易出错。

提示为什么要 400 次采样取平均

第一次做这个扫描时，我们每个配置只采了一次样，结果曲线是锯齿状的—— 完全看不出趋势。

原因：单次采样的标准差本身就是随机变量。 d_{\text{head}} = 32 时理论标准差是 5.66，但单次实测可能在 4 到 8 之间浮动。 这些浮动完全盖过了”维度增长带来的趋势”。

取平均是唯一可靠的办法。 我们的做法是每个配置采 400 组， 算它们的标准差。这样得到的是”标准差的标准差”趋于稳定的估计。

这条经验是通用的：任何”测量一个随机量的典型大小”的实验， 都要先问一句”我采了多少次、平均值稳定了吗”。

为什么用双对数坐标：因为这能一眼看出标度关系。 如果实测值正比于 d^{0.5}，那么在双对数坐标下就是一条斜率 0.5 的直线。 如果正比于 d^1，斜率就是 1。

我们量出来的两条线斜率相同、截距不同—— 这就是”形式对了、常数不对”的图形化证据。 如果用普通坐标，两条曲线都是向上弯的弧线，看不出这个关系。

1.19.2 实验二与三：受控对照

要回答的问题：去掉 \sqrt{d_k} 会怎样？这个”怎样”和维度有关系吗？

核心设计：两两对照，每次只改一个变量。

         不缩放        缩放后
小模型     A1     vs     A2         ← 第一组：d_head = 32
大模型     B1     vs     B2         ← 第二组：d_head = 256

四组实验，共用的东西：语料、随机种子、步数、批大小、学习率调度、优化器。 变化的东西：只有两个——是否缩放，以及 d_{\text{head}}。

为什么 A 组和 B 组不能直接比：因为 A 和 B 的参数量不同 （42.6 万 vs 137.6 万），loss 的绝对值没有可比性。 所以我们的结论全部建立在组内比较上：

• A 组内：4.17 vs 4.00，差 0.17
• B 组内：5.18 vs 4.23，差 0.95

“差的倍数”才是跨组可比的那个量——0.95 / 0.17 ≈ 5.6 倍。 这是一个无量纲的数，它告诉我们”维度增大之后，去掉缩放的影响放大了多少”。

警告一个容易犯的错：拿绝对 loss 跨组比

“大模型的 loss 是 5.18，比小模型的 4.17 差”—— 这句话是没有意义的，因为两个模型的参数不同、容量不同。

参数量不同时，loss 的绝对值不可比。可比的是： ① 同一模型内不同配置的差；② 这个差相对于某个基准的倍数； ③ 与”随机猜测”（\ln V，V 是词表大小）的距离。

随机猜测的 loss 是 \ln 1200 \approx 7.09——这正好是四条曲线的起点（7.03–7.05）。 因为初始化时模型输出的分布接近均匀，loss 自然接近 \ln V。 这个值可以作为”什么都没学到”的基准线。

为什么要跑三次种子：单次对照的结果可能是巧合。 我们的做法是换三个种子重跑小模型那一组， 看到差距都在 0.1–0.2 之间浮动（和训练本身的随机性同量级）， 才敢下”这个规模下没有影响”的结论。

如果三次结果一次是 0.05、一次是 0.8、一次是 0.3， 那说明噪声太大，0.17 这个数不可信，需要更多次或者更大的样本。

1.19.3 诊断量是怎么算的

实验里用到的三个诊断量，算法都写在这里，方便你自己实现。

① 权重熵（归一化到 [0,1]）：

def weight_entropy(p, n_valid):
    """p 是一行的注意力权重，n_valid 是这一行可见的位置数。
    归一化让不同长度的行可以互相比较。"""
    h = -(p * torch.log(p.clamp_min(1e-12))).sum()
    return h / math.log(n_valid)      # ← 除以最大可能的熵

为什么要归一化：位置 5 只能看 6 个位置，位置 100 能看 101 个。 不做归一化的话，后者的熵天然更大——但那不是因为”注意力更分散”， 只是因为”可选项更多”。除以 \log n_{\text{valid}} 消掉这个效应。

② 雅可比范数：

def jacobian_frobenius(p):
    """J_ij = p_i(δ_ij − p_j)，返回 Frobenius 范数。
    这个量衡量 softmax 的输出对输入的敏感度。"""
    J = torch.diag_embed(p) - p.unsqueeze(-1) * p.unsqueeze(-2)
    return J.flatten(-2).norm(dim=-1)

注意实现细节：不要真的去构造一个 n \times n \times n 的张量 （n = 4096 时是 687 亿个元素）。上面的写法只构造 n \times n， 利用了对角矩阵的稀疏结构。

③ logit 标准差：直接 logits.std()，但要只统计有效位置—— 被掩码的位置是 -\infty，会把统计量污染成 nan 或 inf。

valid = mask.expand_as(logits).bool()
lg = logits[valid]                 # ← 只取有效位置
logit_std = lg.std().item()

提示这三个量的实现里藏着一个通用原则

每个诊断量都要问一句”它在什么条件下有意义”。

• 熵：只在”可见位置数 > 1”时有意义（位置 0 只能看自己，熵恒为 0）
• 雅可比范数：只在分布未饱和时有区分度（饱和后都是 0，无法区分「多饱和」）
• logit 标准差：必须排除被掩码的位置

不加这些限定，你会得到一堆看起来正常但实际无意义的数字。 这是做测量时最常见的陷阱——数字不会报错，它只是不说话。

1.20 动手练习

四个练习，从易到难。都能在本章的代码基础上完成。

1.20.1 练习一：复现，并改一个变量

跑一遍 run_ch1_case.py，确认你能得到与本章一致的数字 （loss 4.17 / 4.00 / 5.18 / 4.23，熵 0.0001）。

然后只改一个变量：把 n_head 从 4 改成 2（d_{\text{head}} 从 32 变成 64）。 预测一下会发生什么，然后跑一遍验证。

预期的思考方向：d_{\text{head}} 翻倍意味着 logits 的标准差大约翻 \sqrt{2} 倍， 离饱和更近一步。但 64 仍然远小于 256——所以差异应该还是”不大但不为零”。

1.20.2 练习二：找到”临界维度”

设计一组实验，找出”在不缩放的情况下，d_{\text{head}} 大到多少，训练开始明显变差”。

建议的维度序列：32、64、96、128、192、256。

要求： - 每个维度跑缩放和不缩放两版，共 12 次训练 - 每个都记录 loss 曲线和最终熵 - 画一张图：横轴 d_{\text{head}}，纵轴”两个版本的 loss 差” - 找出这个差值开始明显上升的位置

这个练习的价值：你会得到一条属于你自己的”临界曲线”。 它依赖于你的语料、你的初始化、你的超参数—— 但这正是重点：“什么时候会崩”不是一个通用常数，是你这个系统的属性。

1.20.3 练习三：给 analyze_heads.py 加一个指标

现在的四个指标是：平均熵、对角权重、前一位权重、平均距离。

请加第五个：“注意力重心的位置分布”—— 统计每个位置的注意力重心落在自己左边多远、以及它的方差。

有什么用：如果某个头的”重心距离”方差很大， 说明它在不同位置的行为不一致（有时看近处，有时看远处）。 这类”条件性”的头往往比”固定模式”的头更有意思， 也是第 4 章会重点分析的。

实现提示：dist = sum((i-j) * p[i][j] for j in range(i+1))， 然后统计所有位置的 dist 的均值和标准差。

1.20.4 练习四：做一个”注意力调试器”

把本章排查手册里的四步插桩整合成一个小工具：

class AttentionDebugger:
    """在训练循环里挂上它，异常时自动打印诊断。"""
    def __init__(self, model, log_every=50):
        ...
    def check(self, step):
        """检查三个诊断量，超阈值就警告。"""
        ...
    def on_nan(self):
        """检测到 nan 时，二分定位到具体哪一步。"""
        ...

要求： - 熵 > 0.99 或 < 0.01 时报警 - logits 标准差超过某个阈值时报警 - 检测到 nan 时自动跑一遍四步插桩，打印第一个出问题的位置

这个练习最接近真实工作。本章讲的每个诊断量， 在真实项目里都应该被自动化——人工盯 loss 曲线发现不了这些问题。

提示做完这些练习之后

你会发现一件事：这一章的内容，跑一遍比读一遍快得多。

四个练习加起来大概需要两三个小时，但它们会给你： - 一份属于你自己系统的”临界维度”数据 - 一套能直接用在其他项目上的诊断工具 - 对”机制的前提”这件事的肌肉记忆

如果只做一件事，做练习二。 它是本章案例的直接延伸， 也是最能训练”实验设计直觉”的一个——而这个直觉，是本书想教给你的最核心的东西。

1.21 自测

十道题。前六道能在本章找到直接答案，后四道需要自己推。 答案和解析在下一节——先做完再看。

1. 注意力为什么要除以 \sqrt{d_k} 而不是 d_k？ 用一句话说清”要的是让方差等于多少”。

2. 把缩放写在 softmax 之后，会产生什么可测量的错误？ 给出一个具体的检查手段。

3. 一个多头注意力层，d_{\text{model}} = 512，n_{\text{head}} = 8。 问：\sqrt{d_k} 里的 d_k 应该是多少？如果误写成 512，会发生什么？

4. 因果掩码用 0 而不是 -\infty，训练 loss 会变好还是变差？为什么？ 这个 bug 在什么阶段才会暴露？

5. 为什么”注意力权重高”不能直接当成”因果关系”？ 要验证因果性该做什么？

6. 本章的小模型在有位置编码的情况下，层 0 头 1 学出了”看前一位”的模式。 如果去掉位置编码，这个头还能学出来吗？为什么？

7. （推导） 设 q, k \in \mathbb{R}^{d} 各分量独立、均值 0、方差 \sigma^2 （不一定是 1）。求 \mathrm{Var}(q^{\top}k)，并说明 \sqrt{d_k} 这个缩放 在 \sigma \ne 1 时是否还合适。

8. （推导） 证明 softmax 的雅可比矩阵的每一行之和为零： \sum_j \dfrac{\partial p_i}{\partial z_j} = 0。 这个性质说明了什么？（提示：和”权重和恒为 1”这个约束有什么关系？）

9. （计算） 某模型 n_{\text{layer}} = 24、n_{\text{head}} = 16、 d_{\text{head}} = 64、序列长度 2048、fp16。 求注意力矩阵占用的显存（按 式 1.45）。这个数字意味着什么？

10. （设计） 你要实现一个”只看前 w 个位置”的滑动窗口注意力 （w = 256）。请写出掩码的构造方式，并说明它相比全注意力省下了什么、 没有省下什么。

1.22 检查点解答

十八个「停下来想一想」分散在正文各处。其中十六个在这里集中给答案 （⑰ 和 ⑱ 的答案直接写在正文那一段里了）；先自己想过再看， 否则这一节起不到作用。

1.22.1 ① 面对”没复现出来”，下一步查什么

该查前提。 脚注的推导链是”方差为 1 ⟹ 方差等于 d_k ⟹ 数值大 ⟹ 饱和”。 实验中现象没出现，可能的原因有两类：推导错了，或者前提不成立。

而这条链上只有第一个箭头是假设，其余都是数学。所以最该测的是： q、k 的方差真的是 1 吗？——实测 0.070 对 5.66，答案是否定的。

推广一点：任何”理论预测与实验不符”的情况， 先检查理论的前提是否在你的系统里成立，再看推导。 推导错误很少见，前提失效很常见。

1.22.2 ② loss 停在 5.18 之后会怎样

会继续缓慢下降，但永远追不上缩放版。 因为模型的其他部分 （词嵌入、MLP、输出层）还在学习，它们能独立贡献一部分性能。

但注意力层已经冻结成”每个位置只看一个固定位置”， 这等价于一个固定的一阶转移。模型失去了”根据上下文动态选择信息”的能力—— 而这个能力恰恰是 Transformer 的核心。所以它会稳定在一个明显更差的水平上。

继续训 10000 步，曲线的形状会是：先慢降，然后趋于平缓， 最终稳定在比缩放版高 0.8–1.0 的位置。差距不会消失，因为这不是训练不足，是架构能力缺失。

1.22.3 ③ 尖锐权重与硬查表的梯度区别

有区别，而且这个区别正是软查表存在的理由。

硬查表的梯度是严格为零（除了不可导点）——一点信息都传不回去。

尖锐的 softmax（比如 p = [0.9999, 0.00005, 0.00005]）梯度很小但不为零：

J_{11} = p_1(1 - p_1) = 0.9999 \times 0.0001 = 9.999 \times 10^{-5}

这个小到几乎无用的数字，仍然大于零。 而只要它大于零，梯度就能传回去， 模型就有机会把 p 从尖锐拉回来——这是硬查表做不到的自我修正能力。

本章案例里，实测的 \|J\|_F 掉到了 8 \times 10^{-5}， 已经接近”事实上等于零”。这就是为什么它表现为”训不动”而不是”报错”。

1.22.4 ④ one-hot 的雅可比

p = [1, 0, 0] 时，逐项算 J_{ij} = p_i(\delta_{ij} - p_j)：

• J_{11} = 1 \times (1 - 1) = 0
• J_{12} = 1 \times (0 - 0) = 0
• J_{13} = 1 \times (0 - 0) = 0
• J_{2j} = 0 \times (\dots) = 0（第 2 行整行为零，因为 p_2 = 0）
• J_{3j} = 0（同理）

整个矩阵是零矩阵。 而 p = [0.9999, 0.00005, 0.00005] 时 J_{11} = 9.999 \times 10^{-5}——比零大，但比正常的 O(0.1) 小了三个数量级。

两者在行为上几乎一样（都只看一个位置），在梯度上差着数量级。 这就是”饱和”这个连续过程的实际含义：不是突然死掉， 而是梯度持续衰减，直到训练事实上停滞。

1.22.5 ⑤ 为什么缩放之后正反馈不发生

因为缩放把”放大 q、k“的收益路径堵住了。

不缩放时：模型把 \|q\|、\|k\| 放大一倍，logits 就翻倍，注意力更尖锐， 在训练集上 loss 立刻下降——这是一个有即时回报的动作。 所以梯度下降会鼓励它继续放大，形成正反馈，直到饱和。

缩放之后：logits 是 q^{\top} k / \sqrt{d_k}。要让它翻倍， \|q\|、\|k\| 得放大 \sqrt{d_k} 倍（d_k = 256 时是 16 倍）。 而放大权重 16 倍会让其他所有依赖这些权重的计算失控—— 代价远大于收益，所以模型不会走这条路。

更本质的说法：缩放把”增加尖锐度”的杠杆比改变了。 不缩放时杠杆比是 1:1（放大一倍得一倍），缩放后是 \sqrt{d_k}:1（要放大 16 倍才得一倍）。 这不是禁止了那条路，而是让它变得不划算。

1.22.6 ⑥ 在线 softmax 的纠正因子

m_{\text{old}} = m_{\text{new}} 时，e^{m_{\text{old}} - m_{\text{new}}} = e^0 = 1—— 不需要纠正，直接累加即可。这也是最常见的情况（最大值通常在前几个块就出现）。

新最大值远大于旧最大值时，e^{m_{\text{old}} - m_{\text{new}}} \to 0—— 之前累积的 \ell 被压缩到接近零，旧块的信息在数值上被丢弃了。

这提示一个真实的精度问题：如果最大值出现在很靠后的块， 前面所有块的信息会被压缩掉，可能损失精度。

在 float32 下，当 m_{\text{new}} - m_{\text{old}} > 88 时，e^{\Delta} \to 0， 前面块的信息完全消失。而 FlashAttention 论文里的做法是： 在分块时就把 Q 按行做缩放，让同一行内的 logits 分布尽量集中—— 这既是性能优化，也是数值精度的保障。

1.22.7 ⑦ 长上下文为什么从死结变成工程问题

两个角度，缺一不可。

信息通路长度。 RNN 里第 1 个词的信息传到第 n 个位置要经过 n-1 次变换， 每次都有衰减风险。序列越长，能保留下来的早期信息越少—— 这是架构层面的信息瓶颈，不是算力问题。 注意力把这个距离压到 1 步（第 n 个位置直接与第 1 个位置做点积）， 所以”记住很久以前的东西”在原理上不再是难题。

能否并行。 RNN 的第 t 步依赖第 t-1 步，本质上无法并行—— 序列多长就得串行多少步。注意力对序列的所有位置是一次矩阵乘法， n 个位置同时算，吃满了 GPU 的并行能力。

所以答案分两层：信息通路短了，让”长上下文有用”成为可能； 可并行了，让”长上下文可行”成为现实。

但注意——“不再有架构瓶颈”不等于”免费”。KV 缓存的大小随序列线性增长 （式 1.21 算过：每 token 2.6 MB），这才是今天真正的限制。 瓶颈从”模型记不住”变成了”显存装不下”——问题的性质变了。

1.22.8 ⑧ 拼接与分别投影的参数量

参数量完全一样。 W^O 是 d_{\text{model}} \times h d_v， 切成 h 块就是每个 h d_v / h = d_v 列， 所以 h 个小矩阵 W^O_i 各有 d_{\text{model}} \times d_v 个元素， 加起来还是 d_{\text{model}} \times h d_v。

既然一样，为什么写成拼接？

计算效率：一次矩阵乘法 vs h 次小矩阵乘法。GPU 上大矩阵乘法的效率 远高于多次小矩阵乘法（kernel 启动开销、并行度、显存访问模式都更优）。 h = 64 时，“一次大乘法”比”64 次小乘法再相加”快得多。

代码简洁性：torch.cat([...], dim=-1) @ W_o 一行； 分别投影再相加要写循环和累加。

而且框架会把它们融合：实际实现里， cat 之后立刻做矩阵乘法，编译器（如 TorchInductor）会把它融合成一个 kernel， 连 cat 带来的显存拷贝都省掉了。

提示一个反复出现的模式

数学上的等价形式，工程上不一定等价。 本章已经出现三次了：

• “拼接 + 投影” 与 “分别投影 + 相加”（本节）
• “在线 softmax” 与 “批量 softmax”（小节 1.11 前的第六个检查点）
• “写完整矩阵” 与 “分块不写”（FlashAttention）

三处的数学结果都一样，但访存模式、并行度、数值行为截然不同。 这是从”懂原理”到”能实现”之间那道坎的核心内容。

1.22.9 ⑨ 一个”只看自己”的头有什么用

答案：它把残差路径变成了一个可控的”信息保险箱”。

如果一个头输出的对角权重接近 1，它做的事情就是 “把当前内容做一个线性变换，然后加回自己”。这不是”什么都没做”—— 它是”保留并重新编码当前信息”。具体有两个用处：

① 对抗其他头的信息冲刷。 同一层的其他头可能在从别处大力吸收信息 （比如层 0 头 1 那个前一位头）。如果没有一个头负责”守住自己”， 当前内容会被其他头的输出稀释掉。

② 给后续层留一条干净的通道。 模型可以学到”某些信息不要在这一层被改动”， 而这个”不改动”需要一个具体的执行者。

更一般地说：在多头结构里，“保持不变”也是一种功能， 而且是一种必须有人来承担的功能。

1.22.10 ⑩ FFN 能否区分顺序

不能。

FFN 是逐位置独立作用的：y_i = \mathrm{FFN}(x_i)。 位置 i 的输出只依赖位置 i 的输入，不涉及任何跨位置的操作。

所以如果注意力对输入做的是置换等变操作， 那么 FFN 作用之后这个性质仍然保持。

结论：整个 Transformer 块都是置换等变的。 不注入位置信息，堆多少层都没用。

这也是为什么原论文必须加位置编码——它不是锦上添花，是结构性必需品。

1.22.11 ⑪ RoPE 只依赖相对距离的问题

在纯 RoPE 下，位置 1→2 和 1001→1002 的打分完全一样。 会出两个问题：

① 模型无法表达”绝对位置”相关的模式（比如”是否在上下文的前 10%“）。

② 长度外推时会失效——这才是真正的问题。训练时见到的最大距离是 4095； 推理时到 8000，就会出现训练中从未见过的相对距离。 RoPE 的旋转角 pos \cdot \theta_i 线性增长， 位置 8000 的旋转角落在训练时没见过的区域，模型完全没学过那里的行为。

解决办法（NTK-aware 插值、YaRN 等）：把频率 \theta_i 按比例缩小， 让”位置 8000”落在训练时见过的角度范围内。直观说就是 把长序列”压缩”到模型熟悉的尺度上，代价是分辨率损失。 第 2 章会展开。

1.22.12 ⑫ 点积打分 vs 加性打分

点积的优势有两层。

① 可以写成矩阵乘法。 点积打分下，n 个查询对所有 n 个键的打分 就是 QK^{\top}——一次矩阵乘法。加性打分涉及逐对的非线性， 很难高效地变成一个大矩阵运算。这在 GPU 上是决定性的。

② 没有额外参数。 加性打分要学 v、W_s、W_h 三组参数； 点积打分不需要。既省参数，也减少过拟合风险。

代价是什么？ 正是本章的主线：点积随维度线性增长， 维度一高就饱和。而加性打分有 \tanh 夹住，输出限制在 (-1,1)， 不会随维度膨胀。

所以严格说：加性打分在数值上更稳健，点积打分在计算上更高效。 工程上选了后者，然后用 \sqrt{d_k} 把代价补回来。

提示⑬ 温度为什么是固定值，不是学出来的

两个角度，指向同一个结论。

① 可学温度没有”该是多少”的信号，而且会往坏的方向走。 一个可学温度只多 1 个标量参数，参数量完全不是问题。 问题在于它的梯度方向是单向的： 把温度调小（softmax 更尖）几乎总能让训练 loss 更低—— 因为更尖的注意力更容易退化成”记住”，训练集上几乎总是更优。 这个优化压力没有自然的刹车。 所以一个”自由”的温度会一路被推到饱和，然后梯度消失，训练停在那里。

它把一个有解析解的稳定性问题，换成了一个需要靠正则化去约束的自由参数。 \sqrt{d_k} 的好处恰恰是它不需要学——它在训练第一步就是对的。

② 训练刚开始时，模型没有机会去学它。 即使我们假设存在一个”最优温度”，要学到它就得有梯度。 而如果初始温度不对，softmax 已经饱和、梯度已经消失—— 模型卡在一个”要学好温度得先有梯度，要有梯度得先有好温度”的死循环里。

\sqrt{d_k} 的作用就是在第一步就打破这个循环。

1.22.13 ⑭ 可学的位置编码会有什么风险

两个风险，都来自同一个根源：可学 = 把先验也交给数据去学。

① 外推能力直接归零。 训练时只见过 pos \in [0, 511]。推理时遇到 pos = 1000：

• 如果嵌入表只有 512 行，直接越界报错
• 如果表开得够大，第 512–1000 行从来没有被更新过—— 它们还是初始化时的随机值，和已训好的那 512 行不在一个尺度上，接上去就是噪声

正弦编码没有这个问题：它是一个公式，pos = 1000 照样算得出有意义的向量。

② 需要的样本更多。 正弦编码把”位置”编码成一个有结构的函数——相邻位置的正弦值天然接近， “位置 5 和位置 6 相似”这件事是免费的。 可学的表必须从数据里自己发现这个结构。 而位置信息在语言任务里的信号是稀疏的（大多数词的含义不依赖精确位置）， 用它去换一个先验，不划算。

这就是为什么本章的小模型用可学嵌入（它只在固定长度上训练和推理）， 而真实模型用正弦或 RoPE（它们要外推）。

1.22.14 ⑮ 什么样的实现错误会让一整行全被屏蔽

答案是 padding，而且是一个非常”讲得通”的写法导致的。

一个 batch 里长短句混在一起，短句要补齐到最长。假设掩码写成 “pad 位置不可见”——这听起来完全合理：真实 token 不该看到 padding。

但这样会造出全屏蔽的行。 考虑一条很短的序列： 它自己只有 8 个 token，pad 到 32。那么在第 8–31 这些位置上， 一个可见的 key 都没有——因为除了它自己之外全是 pad，而自己也是 pad。 整行 -\infty → nan。

正确的写法是把两件事分开：

要屏蔽什么	在哪一侧	为什么
pad 作为 key	注意力掩码里屏蔽	真实 token 不该看到补齐位
pad 作为 query	不参与 loss（ignore_index）	它本来就没有”下一个词”可预测

把”query 是 pad”也塞进注意力掩码，就会出全屏蔽行。 这两件事长得像，处理方式完全不同——这是这个 bug 反复出现的原因。

另一个常见来源更简单：布尔方向写反—— mask 里 True 到底表示”屏蔽”还是”保留”，各处约定不一。

1.22.15 ⑯ 什么时候会造出”很大的负数，但不是 -\infty”

来源一：加性掩码。

scores = scores + (-1e9) * (1 - mask)     # 而不是 masked_fill(scores, ~mask, -inf)

这行代码在 fp32 下工作得很好，所以它能活得很久。但在低精度下会翻车：

• fp16：最大值约 65504，-10^9 溢出成 -\infty —— 反而变回了正确行为
• bf16：指数位和 fp32 一样宽，-10^9 可以正常表示，于是它真的就是 -10^9

bf16 下这个 bug 会活着，fp16 下反而不会——这种”换个精度就好了”的现象最容易让人误判。

来源二（更隐蔽）：掩码的”大负数”不够大。

这才是和本章案例直接相关的一种。

上面那张表里，d_{\text{head}}=256 不缩放的模型训练之后， logits 的标准差是 4348，而实测的最大绝对值是 16656。

现在假设掩码用的是 -10^4——一个看起来”很大”的数。但此时 正常位置的 logits 可以到 +1.6\times10^4，比掩码值还大。 掩码不但没屏蔽，反而变成了一个正常的、甚至偏小的分数。

结论：加性掩码里的那个”大负数”，必须比 logits 的实际动态范围大好几个数量级—— 而这个范围在训练过程中会剧烈变化（本章案例里从 1.64 涨到 4348）。

masked_fill(-inf) 没有这个问题：-\infty 在任何尺度下都是最小值。

1.23 回头看，这一章讲的是一个”打补丁”的故事

点积打分好用但会膨胀 → 加 \sqrt{d_k} 补偿。 而这个补丁依赖一个前提（q、k 各分量方差为 1）， 前提在现代架构里不成立——所以补丁在小模型上显得多余、在大模型上又不够用， 于是有了 QK-norm。

一个工程决策的完整生命周期：为效率接受一个缺陷 → 打补丁 → 补丁的前提随时间失效 → 再打新补丁。 读技术史时，这条线索比”记住结论”有用得多。

1.24 自测答案

1. 要的是让方差等于 1（而不是尽可能小）。

\sqrt{d_k} 把方差从 d_k 拉回 1；除以 d_k 会把方差压到 1/d_k， d_k = 64 时标准差只有 0.125——logits 全挤在 0 附近， softmax 输出接近均匀，注意力退化成简单平均。

关键在”回到 1”这个目标，不是”变小”。方差为 1 意味着 logits 的典型跨度是 \pm 2， softmax 落在它最敏感的区间。

2. 权重每一行的和不再是 1。

正确实现里 \sum_j w_{ij} = 1（它是概率分布）； 写成 softmax(S) / sqrt(d_k) 之后，每行的和变成 1/\sqrt{d_k}。 d_k = 32 时是 0.1768（本章实测值）。

检查手段：一行断言。

assert abs(weights.sum(-1) - 1).max() < 1e-5, "权重和不是 1"

这个 bug 的其他表现：输出的量级整体偏小，下游层会”觉得”注意力没提供什么信息。

3. d_k = 512 / 8 = \mathbf{64}。

如果误写成 512，等于把缩放过度了 \sqrt{8} \approx 2.83 倍。 后果是 logits 被压得过于集中，softmax 输出偏向均匀—— 本章实测的”缩放过度”症状：归一化熵 0.9989（几乎完全均匀）。

这个错误在代码上很难发现（形状完全正确、不报错）， 只能靠监控权重熵。

4. 变好。 本章实测：正确实现 loss 4.810，掩码用 0 的 bug 版本 4.634。

原因是 bug 版本让位置 i 能”看到”位置 j > i—— 而预测下一个字时，答案就在输入里，所以模型在偷看答案。

暴露阶段是推理时：生成第 t 个词时第 t+1 个词还不存在， 偷看能力瞬间消失。表现是”训练指标漂亮，生成质量莫名其妙地差”。

这是最危险的一类 bug——它让指标变好，会被当成改进保留下来。

5. 因为权重高只说明”信息从这里取得多”， 不说明”取了之后用上了”。

一个头可能取了很多信息，但它的输出投影 W_O 把这些信息投影到了 对后续层无用的方向上——取而不用的信息对结果没有影响。

要验证因果性必须做干预（消融）：把那个位置的注意力权重强制置零， 看输出/预测有没有变化。变了才是因果，没变就只是相关。

（更严谨的做法是”激活修补”：用另一条输入的激活替换这一处， 看输出是否跟着变。第 4 章会展开。）

6. 不能。

因为注意力本身是置换等变的（§1.13 的小实验）， 而 FFN 逐位置独立作用，不会破坏这个性质。 所以整个 Transformer 块都看不见顺序。

没有位置编码时，模型面对的输入是”一堆无顺序的向量”， 它无法区分”前一位”和”后面任意一位”—— “看前一位”这个模式在它眼里根本不存在，自然学不出来。

这也解释了为什么位置编码不是可选配件：它是结构性必需品。

7. 逐项推：

\begin{aligned} \mathrm{Var}(q_i k_i) &= \mathbb{E}[q_i^2]\mathbb{E}[k_i^2] - 0 = \sigma^2 \cdot \sigma^2 = \sigma^4 \\ \mathrm{Var}(q^{\top}k) &= \sum_{i=1}^{d_k} \sigma^4 = d_k\,\sigma^4 \end{aligned} \tag{1.55}

标准差是 \sigma^2\sqrt{d_k}。

所以 \sqrt{d_k} 是否合适，完全取决于 \sigma：

• \sigma = 1 时标准差恰为 \sqrt{d_k}，除以 \sqrt{d_k} 得到方差 1 ✓
• \sigma \ne 1 时，缩放后标准差是 \sigma^2，不是 1

这正是本章案例的核心：真实模型里 \sigma 由初始化决定（约 0.02）， 所以缩放后 logits 的标准差远小于 1——补丁打了，但没打到点子上。 \sigma < 1 时缩放过度，\sigma > 1 时缩放不足。

8. 逐项求和：

\sum_j \frac{\partial p_i}{\partial z_j} = \sum_j p_i(\delta_{ij} - p_j) = p_i\Bigl(\sum_j \delta_{ij} - \sum_j p_j\Bigr) = p_i(1 - 1) = 0 \tag{1.56}

（用到了 \sum_j \delta_{ij} = 1 和 \sum_j p_j = 1。）

它说明的：softmax 的输出恒满足”和为 1”这个约束， 所以输入的任何”整体平移”都不会改变输出—— 这就是平移不变性在雅可比层面的表现。

一个推论：梯度 \partial \mathcal{L}/\partial z 里 “所有分量同时增加”的那个方向是无效方向—— 它对应的损失变化恒为零。这也解释了为什么在 小节 1.16 里， dS 要减去”加权平均”那一项：softmax 的反向传播自动把无效方向投影掉了。

9. 按 式 1.45：

2048^2 \times 16 \times 24 \times 2\ \text{字节} = 3.22\ \text{GB} \tag{1.57}

（单头单层 8.39 MB → 单层 16 头 134 MB → 24 层 3.22 GB。）

这个数字意味着：如果用 fp32，就是 6.44 GB—— 还没算模型权重、激活值、优化器状态，光注意力矩阵就吃掉大半张 24G 卡。

而且注意 n = 2048 在今天是很短的序列。 如果 n = 8192，这个数字会变成 51.5 GB——单卡放不下。

这正是 FlashAttention 存在的理由：它让这部分显存根本不需要分配。

10. 掩码构造：位置 i 能看 j 当且仅当 i - w < j \le i。

def sliding_window_mask(n, w):
    idx = torch.arange(n)
    # 下三角（因果）且 距离不超过 w
    causal = idx[None, :] <= idx[:, None]
    within = (idx[:, None] - idx[None, :]) <= w
    return causal & within

省下了什么：

• 计算量：从 O(n^2 d) 降到 O(n w d)。n = 8192、w = 256 时， 理论上是 32 倍的减少。
• 显存：不再需要完整的 n \times n 矩阵（如果实现得当）。

没省下什么：

• 长程依赖。这是能力上的代价，不是性能问题。 位置 i 永远看不到 i - w 之前的内容—— 如果答案在最开头（比如”这段话的主题是什么”），这个模型无能无力。
• 实现复杂度。块状掩码的 kernel 比标准因果掩码难写得多； 用现成库时往往还得自己写 Triton kernel 才能真正吃到加速。

所以滑动窗口是一个”用能力换效率”的设计。 它适合”局部信息足够”的任务（比如语音、某些视觉任务）， 对需要全局视野的任务则要配合少量”全局 token”（这就是 Longformer 的做法）。

1.25 关键概念表

中文	英文	一句话定义	在本章哪一节
注意力	attention	按相似度分配权重、对各位置内容加权平均的机制	§1.5
自注意力	self-attention	Q、K、V 全部来自同一个序列的注意力	§1.5
交叉注意力	cross-attention	Q 来自一个序列、K/V 来自另一个	§1.5
软查表	soft lookup	输出是加权混合而非单一选择	§1.5
查询 / 键 / 值	query / key / value	分别表示”要找什么”“有什么”“内容是什么”	§1.5
缩放点积	scaled dot-product	\mathrm{softmax}(QK^{\top}/\sqrt{d_k})V	§1.6
温度	temperature	softmax 前除以的常数，控制分布锐度	§1.7
饱和	saturation	softmax 输出接近 one-hot，雅可比趋近零	§1.7
雅可比矩阵	Jacobian	输出对输入的偏导矩阵，衡量敏感度	§1.7
多头注意力	multi-head attention	把 d_{\text{model}} 切成 h 份并行做注意力	§1.8
前一位头	previous-token head	注意力集中在紧邻左侧位置的头	§1.8
分组查询注意力	grouped-query attention (GQA)	多个头共享一组 K、V，压缩 KV 缓存	§1.9
多查询注意力	multi-query attention (MQA)	所有头共享一组 K、V	§1.9
KV 缓存	KV cache	推理时缓存历史 K、V 避免重算	§1.9
因果掩码	causal mask	屏蔽未来位置，保证自回归性	§1.6
置换等变	permutation equivariant	输入换序，输出只是跟着换序	§1.13
位置编码	positional encoding	把位置信息注入模型的手段	§1.13
旋转位置编码	rotary position embedding (RoPE)	用旋转矩阵把相对位置编进内积	§1.13
上下文窗口	context window	模型一次能处理的最大序列长度	§1.13
在线 softmax	online softmax	分块计算 softmax，维护运行最大值与和	§1.10
FlashAttention	FlashAttention	分块计算、不写完整 n\times n 矩阵的精确注意力	§1.10
序列打包	sequence packing	多条短序列首尾相接共用窗口，提高利用率	§1.16
消融	ablation	关掉某个组件看影响，用于判断因果性	§1.14
注意力汇	attention sink	大量注意力集中到某个固定位置（常是首个 token）	§1.12

1.26 延伸阅读

按”读它能得到什么”来排序，不是按重要性。

1.26.1 一、本章的源头

Vaswani et al., Attention Is All You Need (NeurIPS 2017).

读它能得到什么：本章所有内容的原始出处。公式、多头结构、位置编码、 残差与 LayerNorm 的排布，全在这一篇里。

读的时候注意四件事：

1. \sqrt{d_k} 只是一个脚注（第 3 节末尾）。整篇论文没有对它做实验验证， 也没有讨论它在什么规模下才重要。本章的案例正是补上了这一块。
2. 多头被描述为 “jointly attend to information from different representation subspaces”——这个说法比论文给出的证据要强。论文并没有验证 “每个头学一个子空间”，那是一个直觉性的解释。直到 Michel et al. (2019) 和后续的电路分析工作，人们才开始认真检验这个说法。
3. 位置编码用的是正弦函数，那是 2017 年的选择。今天的主流是 RoPE， 见第 2 章。
4. 训练成本那一节值得细读：他们报告了 8 卡 P100 训练 3.5 天的数据。 对比今天的规模，这个数字本身就能让你感受到七年里发生了什么。

1.26.2 二、从注意力里读出机制

Elhage et al., A Mathematical Framework for Transformer Circuits (Anthropic, 2021).

读它能得到什么：一套描述”注意力头在做什么”的形式语言。 它把每个头拆成两个独立的电路：

• QK 电路（W_Q、W_K）：决定”这个头看哪里”
• OV 电路（W_V、W_O）：决定”看到之后写什么”

为什么这个拆分有用：本章里我们只能说”这个头 45% 看前一位”， 那是行为描述。有了 QK/OV 拆分，你可以问更精确的问题—— “这个头从位置 j 读出了什么特征、写到了位置 i 的哪个子空间”。

读之前的准备：需要熟悉矩阵乘法、秩、以及基本的线性代数直觉。 文中大量使用”低秩投影”的视角，第一次读可能会觉得抽象， 建议配合第 4 章一起读。

1.26.3 三、多头真的都有用吗

Michel et al., Are Sixteen Heads Really Better than One? (NeurIPS 2019).

读它能得到什么：一个实验结论——训练完之后大部分头可以剪掉， 性能几乎不掉；但少数几个关键头剪掉就崩。

它纠正了什么：“每个头都有独特功能”这个天真想法。 如果 80% 的头可以随便移除，那”这个头在看什么”就不能直接当成 “模型这样决策的原因”。

与本章的关系：§1.8 里我们量出了 8 个头的不同模式—— 那证明它们行为不同，但不证明它们都不可或缺。 两件事要分开看。这个区分是第 4 章的主题。

1.26.4 四、长序列的工程解法

Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention (NeurIPS 2022).

读它能得到什么：O(n^2) 在实际系统里到底该怎么处理。 核心是分块 + 在线 softmax（§1.10 手算过）， 数学结果完全不变，改变的只是访存模式。

为什么值得读：它示范了一个重要的思路—— 当算法复杂度已经最优时，还可以从”数据怎么流动”上再要一个数量级的提升。 今天几乎所有训练框架都内置了它。

读的时候注意：论文里的 CUDA 细节可以跳过， 重点看第 3 节的算法伪代码和第 4 节的复杂度分析。

1.26.5 五、本章案例的实践含义

Henry et al., Query-Key Normalization for Transformers (2020).

读它能得到什么：\sqrt{d_k} 的一个现代替代方案—— 直接对 q、k 做归一化，而不是依赖缩放这个间接手段。

为什么它比 \sqrt{d_k} 更彻底：本章案例说明， \sqrt{d_k} 的有效性依赖”q、k 各分量方差为 1”这个前提， 而这个前提在训练过程中会失效（模型会主动放大尺度）。 QK-norm 不依赖任何关于尺度的假设——它把尺度钉死。

延伸：这个思路在 ViT-22B、Stable Diffusion 3 等模型里被采用。 读完可以对照本章案例想一想： 如果一开始就用 QK-norm，那个”没复现出来的实验”会是什么结果？

1.26.6 附：读论文时的一个习惯

上面每条都标了”读的时候注意什么”。这个习惯值得单独说一下。

技术论文的写法是论证式的：作者提出一个方案，然后给出支持它的证据。 但论文很少主动告诉你：

• 这个方案的前提是什么（往往隐含在推导里）
• 它在什么规模下才重要（小规模实验往往不体现）
• 它和别的方案相比，代价在哪（对比实验通常是精心挑选的）
• 哪些结论是外推的（作者自己也没在大规模上验证过）

本章那个”没复现出来的实验”，本质上就是把这四条问了一遍。

具体做法很简单：读完一个技术点，问自己三句话—— “它在什么条件下成立？”“我怎么验证它在这里成立？”“如果不成立会怎样？”

这三句话不需要任何额外知识，但能让你从”读过”变成”知道”。

1.26.7 六、位置编码的演进

Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding (2021).

读它能得到什么：RoPE 的原始论文。§1.13 推的核心恒等式 （\tilde{q}_m^{\top}\tilde{k}_n = q_m^{\top}R_{n-m}k_n）就出自这里。

读的时候注意：论文第 3.2 节的推导主干很短，但值得自己推一遍。 另外注意它讨论的远程衰减性质—— 这解释了为什么 RoPE 在长文本上表现好。

配套读：如果关心长度外推，可以接着看 NTK-aware 插值和 YaRN 这两条后续工作线，它们解决的是 §1.13 检查点 ⑪ 提到的那个问题。

1.27 本章回顾：一条完整的线索

这一章讲了八节，但它们不是八个独立的知识点。它们是一条链， 起点是一个实验意外，终点是对”工程补丁”这件事的一般认识。

1.27.1 主线：一个补丁的完整生命周期

① 点积打分好用（一次矩阵乘法、零额外参数）
        ↓  但有个毛病
② 点积的方差随维度线性增长（Var = d_k）—— 这是数学事实
        ↓  因为
③ softmax 的雅可比 J_ij = p_i(δ_ij − p_j)，饱和时趋近零矩阵
        ↓  所以
④ 打补丁：除以 √d_k，把方差拉回 1
        ↓  但补丁有前提
⑤ 前提是"q、k 各分量方差为 1"—— 现代架构里不成立
   （初始化尺度 0.02、LayerNorm 归一化的是输入而非投影后的 q、k）
        ↓  所以
⑥ 小模型上补丁显得多余（实测差 0.17），大模型上补丁又不够
   （D_head = 256 时 logits 标准差 5477、熵 0.0004）
        ↓  于是
⑦ 有了新补丁：QK-norm（直接把 q、k 的尺度钉死）

这条链上的每一环，都是这一章的一节。 记住链条比记住结论有用—— 因为结论（“要写 \sqrt{d_k}”）会随着架构演进而过时， 链条（“为什么会有这个补丁、它的前提是什么”）不会。

1.27.2 六条可以带走的方法

这一章除了注意力本身，还示范了六种可迁移的方法：

方法	在本章的体现
对照实验要钉死其余变量	同种子、同语料、同步数，只改 scaled 一个开关
理论没被证实时，先怀疑前提	脚注推导没错，是”方差为 1”这个前提不成立
找不到现象时，问”它在什么尺度下才显现”	d_{\text{head}} 从 32 推到 256 才看见灾难
让指标变好的 bug 最危险	掩码 bug 让 loss 从 4.810 降到 4.634（它在偷看答案）
关键性质要靠断言，不能靠 loss	未来位置权重必须严格为 0——这要写单测
诊断量要问”在什么条件下有意义”	熵必须按可见位置数归一化，否则不同长度没法比

1.27.3 如果只记三件事

第一，注意力是软查表。 输出是加权平均，不是选择。 这个直觉能帮你理解后面所有的变体（稀疏、线性、FlashAttention） 都只是”怎么算这个加权平均”的不同方式。

第二，\sqrt{d_k} 修正的是方差，不是大小。 它的目的是让梯度不随维度变化，而不是让梯度变大。 理解这一点，就能明白为什么除以 d_k 是错的。

第三，机制的必要性取决于前提。 这一章最大的收获不是”\sqrt{d_k} 很重要”， 而是”任何机制的必要性都依赖于它所在系统的具体配置”。 小模型上多余的补丁，大模型上可能是救命的； 反过来也一样。

带着这第三条去看后面的章节——每讲一个技术， 都问一句”它解决什么问题、前提是什么、前提不成立时会怎样”。 这个习惯养成之后，你读论文的速度会变慢，但读懂的比例会大幅提高。

这一章我们花了大量篇幅在一个”没复现出来”的实验上， 而不是直接给出结论——全章有一半的篇幅在讲”怎么知道”，而不是”是什么”。这是刻意的：研究生阶段要训练的不是”知道答案”， 而是”知道答案的前提、以及怎么亲手验证它”。

一个公式可以背下来，但”怎么判断这个公式在你的系统里成不成立” 只能靠一遍遍做对照实验练出来。本章的每一个实验设计细节—— 固定种子、只改一个变量、测中间量、做消融—— 都是这个训练的一部分。它们比 \sqrt{d_k} 这个结论本身更值得带走。

1.27.4 与后续章节的接口

本章留下的问题	在哪一章解决
RoPE 的实现细节与长度外推	第 2 章
LayerNorm / RMSNorm 为什么这么设计	第 2 章
FFN 在做什么（键值记忆视角）	第 2 章
KV 缓存与采样策略	第 3 章
归纳头、电路分析、因果干预	第 4 章
注意力的权重能不能当解释	第 4 章

最后一句提醒：本章的所有代码、数据、图表都可以在项目里找到并复现。 如果你在复现时得到了不同的数字，那多半不是错误，而是你的环境与我们的不同—— PyTorch 版本、硬件、甚至随机数生成器的实现差异都会带来浮动。 重要的是趋势和量级，不是小数点后第三位。

提示这一章真正要你记住的

不是”要写 \sqrt{d_k}“。

而是：一个机制的必要性，取决于它的前提在你的系统里成不成立。 脚注的推导没错，但它的前提（各分量方差为 1）在现代架构里不成立， 所以现象在小模型上消失、在大模型上重新出现。

判断一个机制是否关键，要找到它在哪个区间才显现—— 这就是本章那个”没复现出来的实验”最终教给我们的东西。