第 5 章　逐层推导因果注意力、MLP 和残差

5.1 LayerNorm 沿哪个维度计算

一个位置的特征为 $x\in {\mathbb{R}}^C$。LayerNorm 计算该位置内的均值 $\mu =C^{-1}\sum _c{x}_c$ 和方差 ${\sigma }^2=C^{-1}\sum _c(x_c-\mu {)}^2$，再输出

$$y_c={\gamma }_c\frac{x_c-\mu }{\sqrt{{\sigma }^2+{10}^{-5}}}+{\beta }_c.$$

$\gamma$ 是长度 $C$ 的可学习缩放，$\beta$ 在 bias=False 时不存在。每个批次、每个时间位置独立归一化；它不把不同句子或不同时间位置混在一起。因此训练和生成时不需要像 BatchNorm 那样维护运行均值。源码自定义 LayerNorm 的主要原因是提供可选偏置接口，归一化运算仍调用 F.layer_norm。LayerNorm

分母中的 ${10}^{-5}$ 防止方差过小时除零。若 $x=[1,1]$，中心化后都为 0；输出是偏置，而不是 NaN。训练时 $\gamma$ 和 $\beta$ 获得梯度，输入也通过均值、方差参与求导；不能把统计量当常数手写一个不完整反向过程。

5.2 Q、K、V 是三次线性投影，合成一次执行

输入 $X$ 为 $(B,T,C)$，c_attn 输出 $(B,T,3C)$，沿最后一维拆成 Q、K、V。令每头宽度 $d=C/H$，每个张量先 view 为 $(B,T,H,d)$，再把时间和头轴交换，得到 $(B,H,T,d)$。断言 $CmodH=0$ 就是在保证这个分组有整数大小。

对某个批次与头，查询位置 $i$ 对键位置 $j$ 的分数为

$$s_{ij}=\frac{q_i\cdot k_j}{\sqrt{d}}.$$

点积把 $d$ 个乘积相加；若这些分量方差近似稳定，未经缩放的点积方差随 $d$ 增大。除以 $\sqrt{d}$ 有助于避免 softmax 一开始就极端尖锐。随后对每个查询行做 softmax，得到注意力权重 $a_{ij}$，输出 $o_i=\sum _j{a}_{ij}{v}_j$。Q 和 K 决定“从哪里读”，V 提供被读出的内容；这三个角色是计算上的区别，不意味着某些神经元永久代表固定语义。

操作	输入形状	输出形状
q @ k.transpose(-2,-1)	$(B,H,T,d)$ 与 $(B,H,d,T)$	$(B,H,T,T)$
行 softmax	$(B,H,T,T)$	相同，每行和为 1
att @ v	$(B,H,T,T)$ 与 $(B,H,T,d)$	$(B,H,T,d)$
合并各头	转轴、连续化、view	$(B,T,C)$
输出投影	$(B,T,C)$	$(B,T,C)$

transpose 通常只改变步幅，不把数据重新排成连续存储。源码在合并头之前显式 .contiguous()，然后 .view(B,T,C)；如果省掉连续化，view 可能因步幅不兼容报错，或者不得不用能复制的 reshape。

5.3 掩码必须在 softmax 之前作用

自回归训练在位置 $i$ 只能访问 $j\le i$。手工路径创建下三角 0/1 掩码，把未来位置的分数改成负无穷，再做 softmax：

$$a_{ij}=\{\begin{array}{ll}{\displaystyle \frac{e^{s_{ij}}}{\sum _{k=0}^i{e}^{s_{ik}}}},& j\le i,\\ 0,& j>i.\end{array}$$

例如第二个位置可以看到第一个和第二个位置。若未归一化分数为 $[0,\log 2,100]$，屏蔽未来后权重为 $[1/3,2/3,0]$。未来位置即使有极大分数也不应改变结果。如果先对三项 softmax 再乘掩码，保留项的和会远小于 1，已经不是同一个注意力运算。

可切换原始尺寸，滚动查看细节

查看大图

尺寸读取中

适合窗口：完整显示图片，不放大超过原始尺寸。 按 Esc 关闭。

此图未提供文字说明

源码的手工实现把 mask 注册为 buffer，形状 $(1,1,T_{max},T_{max})$，可在批次和头维广播；它不是可学习参数，但会随模型移动设备。masked_fill 是对分数设置禁区，不是把未来的输入 embedding 清零。层与层都保持同一因果关系，才能保证深层也不泄漏。

本书修改输入后半段、比较前半段 logits，最大误差实际为 0。这检验的是“在固定权重下，前缀输出不依赖后缀”，比只看掩码图更接近真正需要的性质。测试时关闭 dropout，否则两次前向可能因随机掩码不同而出现差异。

5.4 名为 flash，不保证用了 FlashAttention 内核

self.flash 只是检查 PyTorch 是否有 scaled_dot_product_attention 函数。存在时传入 is_causal=True，训练态使用配置 dropout，评估态明确传 0。这个标记不探测硬件，也不证明 CPU 上运行 CUDA FlashAttention。具体后端由 PyTorch、设备、dtype 和形状决定；本书 CPU 测试只能验证 SDPA 接口结果。固定实现

为了核对两条路径，本书复制同一个 double 模型，给副本显式注册下三角 buffer，再设 flash=False。仅翻转标记会失败，因为初始化时有 SDPA 的模型没有建立手工 mask。两条路径 logits 最大差为 $2.78\times {10}^{-17}$，属于浮点舍入量级。带 dropout 时随机数消费顺序可能不同，不能要求逐位相等。

手工路径显式保留 $(B,H,T,T)$ 分数和权重，长度增大时内存增长快。优化后端可以减少中间张量读写，但不会改变模型所表达的因果注意力定义。性能结论必须另测，不能从误差测试推出加速倍数。

5.5 MLP 改每个位置，注意力连接位置

MLP 的第一层把 $C$ 扩为 $4C$，经 GELU，再投影回 $C$，最后 dropout。GELU 可写为 $\mathrm{GELU}(x)=x\mathrm{\Phi }(x)$，$\mathrm{\Phi }$ 是标准正态分布累积分布函数。它是平滑非线性，不是简单截断负数的 ReLU。若去掉非线性，两层线性映射可以合成为一个线性映射，扩宽再压回的表达作用会大幅改变。

MLP 对每个位置使用同一组参数，但位置之间不直接通信。位置间混合由注意力承担，特征间非线性组合由 MLP 承担。Block.forward 依次执行

$$U=X+\mathrm{Attention}({\mathrm{LN}}_1(X)),Y=U+\mathrm{MLP}({\mathrm{LN}}_2(U)).$$

这是 pre-LN 结构：先归一化，再计算分支，最后加回原始残差流。每次相加双方都是 $(B,T,C)$。残差让梯度包含一条恒等路径，减少必须连续穿过所有非线性变换的压力；它不保证任何深度、学习率都稳定。

补充练习。 将 attention 输出投影改成 $(C/2)$ 维，再保留 x + attn(...)，会怎样？若把因果 mask 去掉而训练 loss 更低，能说模型更好吗？

答案。 残差相加的形状不兼容，需要额外投影才能定义运算，而那已经改变结构。去掉 mask 后模型可读到目标对应的未来输入，训练任务变得泄漏；更低损失不能作为自回归生成改善的证据。用前缀不变性测试可直接揭示这种反例。

本章来源

• model.py：LayerNorm、CausalSelfAttention、MLP、Block
• PyTorch SDPA 官方接口
• 本书核对程序：examples/verify_mechanics.py，结果 qa/author/mechanics.json。