第 3 讲：架构与超参数

对应：Stanford CS336 Spring 2026 Lecture 3（Architectures, hyperparameters，Tatsunori Hashimoto 主讲）。Spring 2025 同名讲次有公开录像。 配套代码：code/lecture03/ 预计阅读时间：3–4 小时（第一遍通读 1.5 小时，第二遍对照代码 2 小时）

0. 本讲导航

0.1 学完本讲，你应该能够

1. 画出一个现代（Llama 风格）Transformer 语言模型的结构，说出它和 2017 年原始 Transformer 的五处主要不同。
2. 解释 Pre-norm vs Post-norm、LayerNorm vs RMSNorm、去掉偏置背后的理由，理解”FLOPs 不等于运行时间”。
3. 写出 SwiGLU 前馈层，推导它的隐藏维度为什么是 $\frac83 d$。
4. 推导 RoPE 为什么让注意力分数只依赖相对位置，并实现它。
5. 实现缩放点积注意力、因果掩码和多头注意力；解释为什么要除以 $\sqrt{d_k}$。
6. 说出几个关键超参数的”行业共识”（前馈倍数、头维度、宽深比、词表大小、正则化）以及例外。
7. 解释三种稳定性技巧（z-loss、QK-norm、logit 软截断）分别在防什么。
8. 解释 GQA / MQA / MLA、滑动窗口注意力为什么能降低推理成本。
9. 计算一个 Transformer 的参数量、前向 FLOPs 和训练显存（作业 1 的资源核算题）。

0.2 先修知识

0.3 本讲地图

第1节 从原始 Transformer 到"现代标准款"：改了什么？怎么学习别人的经验？
第2节 Transformer 块的直觉：残差流、注意力（交流）、前馈（思考）
   │
   ├── 架构选择 ─────────────────────────────────────────
   │   第3节 归一化：Pre/Post、LayerNorm/RMSNorm、去偏置
   │   第4节 激活与前馈层：GLU 家族与 8/3 规则
   │   第5节 位置编码：从正弦到 RoPE
   │   第6节 注意力：缩放、因果、多头
   │
   ├── 超参数与训练 ──────────────────────────────────────
   │   第7节 超参数的"共识"与例外
   │   第8节 动手训练：注意力到底带来了什么？
   │   第9节 稳定性技巧：z-loss、QK-norm、软截断
   │
   └── 第10节 注意力的变体：GQA/MQA/MLA、滑动窗口
第11节 资源核算：参数量、FLOPs、显存

本讲的主题可以用 CS336 课上的一句话概括：学习的最好方式是亲手实践，第二好的方式是学习别人的经验。 2024–2025 两年间就发布了几十个稠密模型，它们的架构大同小异。我们要做的是：看看大家都在做什么、在哪里有分歧、以及背后的理由。

1. 从原始 Transformer 到”现代标准款”

1.1 原始 Transformer（2017）的选择

1.2 你在作业 1 中实现的”现代版本”

token ids
   │
Embedding ─────────────────────────────┐ 残差流（residual stream）
   │                                   │
   ├─► RMSNorm ─► 多头注意力(含 RoPE) ─► (+)
   │                                   │
   ├─► RMSNorm ─► SwiGLU 前馈层 ──────► (+)      × L 层
   │                                   │
RMSNorm（最后）◄───────────────────────┘
   │
Linear（lm_head）→ logits

五处不同：

1. Pre-norm：归一化放在每个子层的输入处，而不是残差相加之后；
2. RMSNorm 代替 LayerNorm；
3. 去掉偏置：线性层和归一化层都没有偏置项；
4. SwiGLU 代替 ReLU 前馈层；
5. RoPE 代替正弦位置编码。

这基本就是 Llama 系列的架构，今天大多数开放模型（Qwen、Mistral、DeepSeek 的稠密部分、OLMo……）都是它的变体。下面逐一讲清楚为什么。

2. Transformer 块的直觉

在进入细节之前，先建立一个整体的心智模型：

• 残差流（residual stream）是贯穿所有层的”主干道”：每个位置有一个 $d$ 维向量，从嵌入层一路传到输出层。
• 每一层读取残差流（先归一化），做一些计算，再把结果加回残差流：$x \leftarrow x + f(\text{Norm}(x))$。
• 注意力子层负责位置之间的”交流”：每个位置从其他位置收集信息（本讲第 6 节）。
• 前馈子层负责每个位置自己的”思考”：对收集到的信息做非线性变换，各位置独立计算（第 4 节）。

残差连接的妙处在于：即使某一层什么也没学到（$f \approx 0$），信号也能原样通过；梯度也能沿着”加法”这条路直接流回前面的层。这是能训练几十上百层深度网络的关键。

3. 归一化

3.1 Pre-norm vs Post-norm：大家唯一一致同意的事

区别在于归一化是否在残差的”主干道”上。Post-norm 让主干道每一层都要经过一次归一化；Pre-norm 让主干道上只有加法，是一条畅通的高速公路。

现状：除了 BERT 这类早期模型（以及 OPT-350M 这个”奇怪的例外”），几乎所有现代大模型都用 Pre-norm。

理由：

• Xiong et al.（2020）分析发现，Post-norm 在初始化时靠近输出层的梯度很大，训练必须依赖学习率预热；Pre-norm 的梯度更均衡，可以去掉预热。
• Salazar & Nguyen（2019）等工作观察到 Pre-norm 训练中的梯度尖峰更少。
• 今天的主要理由：稳定性，以及能用更大的学习率训练大网络。

配套代码 05_train_transformer.py 演示 4 做了一个小实验（8 层 Transformer，不预热，训练 200 步）：

同样的学习率，Pre-norm 正常训练、Post-norm 训练不起来；Post-norm 降低学习率后也能训练。这正是”Pre-norm 能承受更大学习率”的具体体现。（这是小规模实验，具体”坏掉”的学习率与层数、宽度等有关：6 层、更窄的配置下，Post-norm 只是开始慢一些。）

新动向：“双重归一化”。既然把归一化放在主干道上不好，那放在子层输出、但仍在主干道之外呢？即 $x \leftarrow x + \text{Norm}(f(\text{Norm}(x)))$。Grok、Gemma 2 使用这种”前后都归一化”的方式；OLMo 2 只在子层输出处归一化（“非残差的 post-norm”）。共同点是：不在残差主干道上做归一化。

3.2 LayerNorm vs RMSNorm

RMSNorm 不减均值、没有偏置 $\beta$，只把向量的”整体大小”（均方根）归一化到 1。使用 LayerNorm 的：GPT-1/2/3、OPT、GPT-J、BLOOM；使用 RMSNorm 的：Llama 家族、PaLM、Chinchilla、T5，以及绝大多数新模型。

为什么换？ 现代的解释是：更快，而且效果一样好（有些论文甚至报告效果略好）。少算一个均值、少一个偏置参数。

但这说得通吗？ 矩阵乘法占了绝大多数 FLOPs，归一化层的 FLOPs 微不足道，省下那点计算有什么意义？答案是本课反复出现的一课：FLOPs 不等于运行时间。归一化是内存受限的运算（第 2 讲第 6 节）：它要读写整个激活张量，每个数只做几次运算。CS336 引用的数据（Ivanov et al.）：在 BERT 中，归一化这类运算只占 0.17% 的 FLOPs，却占约 25% 的运行时间。

配套代码 01_building_blocks.py 演示 3 在 CPU 上也观察到了同样的现象：RMSNorm 的 FLOPs 只有相邻矩阵乘法的 0.05%，时间却有它的约 5%；LayerNorm（我们的教学实现）又比 RMSNorm 慢约一倍。在 GPU 上矩阵乘法快得多，归一化的时间占比通常更高。

实现细节：作业 1 要求在 RMSNorm 内部先把输入转成 FP32 再计算。原因见第 2 讲：平方求和是一种累加，低精度下误差大，而且会上溢。演示 2 的一个极端例子：FP16 的激活值量级约 300 时，平方后超过 FP16 的最大值 65504 变成 inf，全程用 FP16 计算的 RMSNorm 输出全部变成了 0；先转 FP32 再算则完全正常。

3.3 去掉偏置项

大多数现代 Transformer 的线性层和归一化层都没有偏置。理由与 RMSNorm 类似：

• 省内存搬运：偏置虽然参数少，但每次都要读写；
• 优化稳定性：多个团队报告去掉偏置后训练更稳定。

而效果几乎没有损失。

4. 激活函数与前馈层

4.1 激活函数”动物园”

ReLU、GELU、Swish/SiLU、ELU、GLU、GeGLU、ReGLU、SwiGLU……它们是什么？用的人多吗？重要吗？

SiLU（也叫 Swish）$= x\cdot\text{sigmoid}(x)$，形状和 GELU 很像：大的正数几乎原样通过，大的负数趋于 0，处处光滑，在 0 附近允许少量负值通过。

4.2 门控线性单元（GLU）

GLU 的思路：在”线性 + 非线性”之外，再加一路逐元素相乘的门：

这多了一个参数矩阵 $V$。按激活函数的不同，有 ReGLU（ReLU）、GeGLU（GELU）、SwiGLU（SiLU）：

直觉上，门控让网络能做”乘法式”的交互（两路信号相乘），表达能力比”线性变换 + 逐元素非线性”更强。

4.3 8/3 规则

GLU 有三个矩阵，普通前馈层只有两个。为了参数量相同（公平比较），要缩小隐藏维度：

实际中通常再取整到 64 的倍数（GPU 上矩阵维度是 64/128 的倍数时效率最高，第 5 讲）。例如 $d=768$ 时 $d_{ff}=2048$，两种前馈层的参数量完全相同（配套代码演示 4）。

4.4 GLU 真的更好吗？

是的，相当一致地更好。 Shazeer（2020）在参数量相同的条件下比较了多种变体，GeGLU 和 SwiGLU 的困惑度稳定地优于 ReLU 和 GELU；Narang et al.（2020）等后续工作也证实了这一点。不过：

• GLU 不是必需的：GPT-3 用 GELU 也训练得很好；
• 也有少数例外，如 Nemotron-4 340B 用的是”平方 ReLU”；
• 至于为什么更好，Shazeer 在论文结尾坦言没有解释，只能归功于”神的恩典”（第 1 讲 1.4 节）。

4.5 串行层 vs 并行层

标准的 Transformer 块是串行的：先注意力，再前馈。GPT-J 首先采用了并行的写法：

两个子层读取同一个输入，可以共享一次归一化，矩阵乘法也可以融合，速度更快。PaLM、GPT-NeoX，以及较新的 Cohere Command A、Falcon 2 11B 等使用并行层，但大多数模型仍然是串行的。

5. 位置编码：从正弦到 RoPE

5.1 为什么需要位置编码

注意力本身不看位置。配套代码 02_rope.py 演示 0：把输入序列打乱，注意力的输出只是跟着换了位置，数值完全不变。“狗咬人”和”人咬狗”在它眼里是同一袋词。所以必须把位置信息注入模型。

5.2 几种位置编码

5.3 我们想要什么性质

理想的位置编码应该让注意力分数只依赖相对距离 $i - j$：

“前一个词”对理解当前词的作用，不应该因为这句话出现在文章开头还是结尾而不同。现有方法的问题：

• 正弦编码（加法）：展开后有交叉项 $\langle x + PE_m, y + PE_n\rangle = \langle x,y\rangle + \langle x, PE_n\rangle + \langle PE_m, y\rangle + \langle PE_m, PE_n\rangle$，中间两项依赖绝对位置。演示 1：同样两个词、距离都是 3，放在位置 (0,3)、(10,13)、(50,53) 时，点积分别是 25.7、20.8、17.2——随绝对位置变化。
• 可学习的绝对位置：显然不是相对的。
• 相对位置偏置：满足性质，但它不是一个内积形式，和高效的注意力实现不太兼容。

5.4 RoPE：用旋转编码位置

关键观察：内积对旋转是不变的——两个向量同时旋转同一个角度，它们的内积不变。

RoPE（Su et al. 2021）的做法：把 $d$ 维向量的分量两两配对 $(x_0,x_1), (x_2,x_3), \ldots$，每一对看成二维平面上的一个点；位置 $m$ 的向量，第 $k$ 对旋转角度 $m\theta_k$，其中

为什么这样就只依赖相对位置？ 对每一对，用 $R(\varphi)$ 表示旋转 $\varphi$ 角：

旋转矩阵的转置是反向旋转，两个旋转合并后只剩下角度差 $(n-m)\theta$。所有配对都如此，整体点积就只依赖 $m-n$。

演示 2 验证：同样的 $q, k$、距离固定为 3，放在位置 (0,3)、(10,13)、(50,53)、(1000,1003) 时，点积完全相同（都是 −10.74679）。而且旋转不改变向量长度。

复数视角（选读）：把每一对 $(x_{2k}, x_{2k+1})$ 看成复数 $z_k = x_{2k} + i x_{2k+1}$，旋转就是乘以 $e^{im\theta_k}$。复数内积中 $e^{im\theta}\cdot\overline{e^{in\theta}} = e^{i(m-n)\theta}$，只剩相对位置。演示 3 验证了复数乘法与 RoPE 实现的差异只有 $6\times10^{-8}$。

与正弦编码的区别：RoPE 是乘法（旋转），不是加法，所以没有交叉项。

5.5 RoPE 的直觉：一个有很多根指针的钟

不同配对的转速不同（$d=64$，$\Theta=10000$）：

就像一个有很多根指针的钟：转得快的”秒针”区分相邻的位置，转得慢的”时针”区分很远的位置。$\Theta$ 越大，最慢的指针越慢，能区分的距离越远——Llama 3 把 $\Theta$ 从 10000 调到了 500000 以支持长上下文。

另一个性质：同一个向量放在相距越远的两个位置，点积（总体趋势上）越小（演示 4：距离 0 时 128，距离 1024 时约 26），因为不同转速的分量逐渐”失去同步”。这给模型一个温和的偏好：近处的 token 更容易被关注。

5.6 实现要点

• 只对 query 和 key 旋转，不对 value 旋转（位置信息只需要影响”谁关注谁”）。
• 每一层的注意力都要做，而不是只在输入处加一次。
• cos/sin 表可以预先算好，存成不可训练的缓冲区。
• 边界：存在两种常见写法——相邻配对 $(x_0,x_1)$（原论文、Meta 的 Llama 实现、作业 1）和前后半配对 $(x_i, x_{i+d/2})$（GPT-NeoX、Hugging Face 的实现）。两者数学上等价（只差一个维度排列），但用一种写法训练的权重必须用同一种写法推理（演示 5）。

6. 注意力：缩放、因果、多头

6.1 缩放点积注意力

直觉：每个位置发出一个 query（“我在找什么”），与所有位置的 key（“我是什么”）比较相似度，得到的权重用来对 value（“我的内容”）加权平均。

softmax 要用数值稳定的写法（减去最大值），作业 1 要求自己实现。

6.2 为什么除以 $\sqrt{d_k}$

如果 $q, k$ 的每个分量是方差为 1 的独立随机数，点积 $q\cdot k = \sum_{i=1}^{d_k} q_i k_i$ 的方差就是 $d_k$。$d_k$ 越大，分数的差距越大，softmax 越接近”只选最大的一个”。配套代码 03_attention.py 演示 2（32 个 key 中最大注意力权重的平均值）：

不缩放时注意力退化成”硬选择”，而且 softmax 饱和区的梯度几乎为 0，难以训练。除以 $\sqrt{d_k}$ 把方差拉回 1。

6.3 因果掩码

语言模型训练时每个位置都在预测下一个 token；如果位置 $i$ 能看到 $i+1$，它直接”抄答案”就行了。所以要用因果掩码：位置 $i$ 只能关注 $0,\ldots,i$。实现上，把禁止的位置的分数设为 $-\infty$，softmax 后变成 0。演示 3 的验证：只修改最后一个 token，前面所有位置的输出变化量都是 0。

边界：如果某一行被全部掩掉，softmax 会得到 NaN（$-\infty - (-\infty)$）。因果掩码不会出现这种情况（对角线总是可见），但处理变长序列的填充掩码时要小心。

6.4 多头注意力

把 $d$ 维拆成 $h$ 个头，每个头 $d/h$ 维，各自独立做注意力，最后拼接再经过输出投影 $W_O$。一个头只能产生一种注意力模式，多个头可以同时关注不同的东西（前一个词、句子的主语、匹配的括号……），而总计算量和单个大头一样。用 einops 写就是 rearrange(x, "b s (h d) -> b h s d", h=num_heads)（第 2 讲第 4 节讲过为什么不能直接 view）。

我们的实现与 PyTorch 官方的 scaled_dot_product_attention(is_causal=True) 的最大差异为 $5\times10^{-7}$。

7. 超参数：惊人的共识与例外

7.1 前馈层倍数 $d_{ff}/d$

共识：ReLU/GELU 用 $d_{ff} = 4d$；GLU 类用约 $\frac83 d \approx 2.67d$。配套代码 06_hyperparams_and_stability.py 演示 1 汇总了真实模型：

大多数 GLU 模型在 2.673.5 之间。著名的例外是 T5 11B：$d = 1024$，$d_{ff} = 65536$，整整 64 倍！这说明大胆的选择也能工作；但 T5 的后续版本 T5 v1.1 改回了常规的 2.5 倍（GeGLU），说明 64 倍很可能不是最优的。Kaplan et al.（2020）的实验显示，这个比例在 110 之间都接近最优，是一个宽阔的”盆地”。

7.2 头数 × 头维度 = $d$

共识：几乎所有模型都满足 $h \times d_{head} = d$（比例为 1）。这不是必须的——头维度可以大于 $d/h$——但大家都这么做。例外主要来自 Google 的模型（T5 的 16 倍、LaMDA、PaLM）。这条规则”有共识但缺少验证”。

7.3 宽深比 $d/L$

模型该深还是该宽？大多数模型的 $d/L$ 在 100~200 之间（GPT-3、OPT、Mistral、Qwen、OLMo 都是 128）。Kaplan 等人的实验显示在很宽的范围内效果差别不大。实际中宽深比更多由系统因素决定：

• 太深：层与层必须串行计算，延迟高；流水线并行（第 8 讲）的流水段多。
• 太宽：单个矩阵太大，需要张量并行（第 8 讲），通信开销大。

7.4 词表大小

• 单语（英文为主）模型：3 万~5 万（原始 Transformer 37000、GPT-2/3 50257、T5 32128、Llama 32000）；
• 多语言 / 生产系统：10 万~25 万（GPT-4 约 10 万、Qwen 约 15 万、mT5 25 万、PaLM 25.6 万、Gemma 约 26 万）。

单语模型不需要很大的词表，多语言模型需要（第 1 讲第 11–12 节）。

7.5 正则化：dropout 与权重衰减

预训练需要正则化吗？反对的理由很有说服力：数据有几万亿 token，远多于参数；而且通常只过一遍数据（单个 epoch），模型很难”记住”训练集。实际情况：

较早的模型常用 dropout，新模型大多只用权重衰减。 但权重衰减的作用不是防止过拟合：Andriushchenko et al.（2023）的分析显示，它主要通过与学习率调度（如余弦退火）的相互作用来影响优化动力学，让训练后期的损失更低。

8. 动手训练：注意力到底带来了什么？

8.1 实验设计

第 1 讲的语料只有 17 KB，Transformer 在一两百步内就能把它完全背下来（训练损失接近 0），看不出注意力的作用。所以配套代码 05_train_transformer.py 设计了一个专门的实验：生成 4000 个模板化的小故事，每个故事有两个随机起名的角色，名字（从约 8000 种可能中随机组合）在故事里反复出现：

Once upon a time, there was a curious mouse named Tohalo. Tohalo loved to play in the garden. One day, Tohalo found a red ball. Then Tohalo met a dog named Fira. Fira said, “Can I dance with you, Tohalo?” …

名字第一次出现时谁也猜不到；但再次出现时，只要能回头看到前文，就能照抄。第 2 讲的窗口模型只能看最近 8 个字节，看不到几十个字节之前的名字。

8.2 结果

同样训练 600 步（单位：比特/字节，越低越好）：

在”重复名字”的字节上，窗口模型只能猜，Transformer 接近 0——它学会了从上文中把名字抄过来。在其他字节上两者差距小得多。Transformer 参数更多，但即使把窗口模型做得再大，它也看不到 8 字节之外的内容：能力的差别来自结构。

8.3 看看注意力在看哪里

演示 2 打开了训练好的模型的注意力。故事里的名字是 Nugopo，模型读到第三次出现的 ”…Then Nu”，要预测下一个字母 ‘g’：

• 第 0 层：各头的注意力很分散（最大权重 0.06~0.24），看的是一些模板性的位置；
• 最后一层：有头把 0.85 的注意力集中在前文 “Nugopo” 中的字母 g 上——恰好就是要预测的字母。

模型在做的事情是：“上次 Nu 后面跟的是 g，所以这次也是 g”。这种机制叫归纳头（induction head），被认为是 Transformer 进行上下文学习的基础之一。

8.4 生成

用 top-p 采样（下面第 8.5 节）让模型续写，提示里现编一个名字：

Once upon a time, there was a sleepy duck named Zebeta. Zebeta loved to draw in the village. One day, Zebeta found a blue kite. Then Zebeta met a frog named Lofi. Lofi said, “Can I draw with you, Zebeta?” …

训练数据里几乎不可能有 “Zebeta”，但模型在后文中一致地使用了它。（模型偶尔也会抄错一个字母，这是小模型、短训练的正常现象。）

8.5 解码：温度与 top-p

作业 1 要求实现两种解码技巧：

• 温度 $T$：logits 除以 $T$ 后再 softmax。$T<1$ 更保守，$T>1$ 更随机。
• top-p（核采样）：把概率从大到小排序，只保留累积概率刚好达到 $p$ 的最小集合，在其中重新归一化后采样。模型很确定时集合很小（接近贪心），不确定时集合变大（保留多样性），同时砍掉了概率极低的”胡言乱语”长尾。

9. 稳定性技巧：小心 softmax

训练大模型时，损失曲线有时会突然飙升（loss spike），甚至彻底发散。最常见的罪魁祸首是 softmax：它包含指数函数（容易上溢）和除法（可能除以接近 0 的数）。模型里有两处 softmax：输出层（词表上的概率）和注意力。

9.1 z-loss：输出层的 softmax

softmax 对所有 logits 加同一个常数不变，所以交叉熵对”logits 整体往上漂”毫无约束。数值漂大后，低精度的绝对精度变差：BF16 在 1000 附近的相邻间隔是 4，logits 之间的细微差别会被抹掉。配套代码演示 2：

z-loss（PaLM 使用，后来 Baichuan 2、DCLM、OLMo 2 等也使用）在损失上加一项 $10^{-4}\cdot z^2$，其中 $z = \log\sum_j e^{o_j}$ 是 softmax 归一化常数的对数。它把 logits 的整体水平”拴”在 0 附近（此时 $z\approx 0$，softmax 的分母约为 1）。

9.2 QK-norm：注意力的 softmax

注意力分数 $= q\cdot k/\sqrt d$。训练中如果 $q, k$ 的范数不断增长（都放大 $s$ 倍，分数就放大 $s^2$ 倍），softmax 会越来越尖，注意力”塌缩”成只看一个位置，梯度也随之消失。演示 3（64 个位置，均匀分布的熵为 4.16）：

QK-norm 在点积之前先对 $q$ 和 $k$ 做归一化（RMSNorm），分数的大小就与 $q, k$ 本身的范数无关了。它最早来自视觉和多模态模型，现在 DCLM、OLMo 2、Gemma 3、Qwen3 等模型都在使用。

9.3 logit 软截断

小的值几乎不变（$c=30$ 时 $10\to9.65$），大的值被平滑地压到 $\pm c$ 以内（$200\to30$）。Gemma 2 对注意力分数用 $c=50$、对最终 logits 用 $c=30$。它能防止 logits 爆炸，但也可能损失一些表达能力，影响效果（Gemma 3 就改用了 QK-norm）。

10. 注意力的变体

大多数模型对注意力头本身改动很少，主要的例外都与推理成本有关。

10.1 推理时的瓶颈：KV 缓存

训练时，整个序列并行计算，注意力的算术强度很高，GPU 能跑满。但生成时，模型一次只产生一个 token，每一步都要读取之前所有 token 的 key 和 value。为了不重复计算，把它们存起来，叫 KV 缓存。于是每生成一个 token，都要把整个 KV 缓存从显存里读一遍，而只做很少的计算——内存受限。KV 缓存的大小直接决定了能支持多长的上下文、多大的批次，以及生成的速度（第 10 讲详细分析）。

每个 token 的 KV 缓存 $= 2 \times \text{层数} \times \text{KV 头数} \times \text{头维度} \times \text{字节数}$。

10.2 MQA 与 GQA

• MQA（多查询注意力，Shazeer 2019）：所有查询头共享一组 K、V。
• GQA（分组查询注意力，Ainslie et al. 2023）：查询头分组，每组共享一组 K、V。MHA 和 MQA 是它的两个极端。

配套代码 03_attention.py 演示 4 计算了真实模型的 KV 缓存（BF16）：

效果方面：MQA 有小幅的困惑度损失，GQA 几乎没有损失。GQA 提供了一个在”表达能力”和”推理效率”之间调节的旋钮，已成为新模型的标配。

10.3 MLA：多头潜在注意力（DeepSeek-V2/V3）

MLA 更进一步：把 K、V 表示成一个低维”潜在向量” $c^{KV}$ 的函数，KV 缓存里只存这个潜在向量，需要时再投影回 K、V。把 K 的上投影矩阵合并进 Q 的投影，计算时甚至不需要显式地恢复 K。

一个麻烦是 RoPE 与这种缓存方式冲突：旋转矩阵夹在两个投影矩阵中间，无法合并。DeepSeek 的解决办法是保留少量不压缩、专门用来做旋转的 key 维度。

10.4 稀疏与滑动窗口注意力

完整注意力的计算量随序列长度平方增长。滑动窗口注意力让每个位置只看最近 $w$ 个位置，计算量变成 $n\times w$；堆叠 $L$ 层后，信息仍然可以”接力”传出约 $L\times w$ 的距离。$n = 32768$、$w = 4096$ 时，需要计算的分数从约 5.4 亿降到约 1.3 亿。GPT-3（部分层使用稀疏注意力）、Mistral 等使用过这类设计。

当前的标准做法：交替使用局部和全局注意力。例如 Cohere Command A 每 4 层用一次完整注意力，其余层用滑动窗口；有意思的是，它的全局层不用任何位置编码（NoPE），局部层用 RoPE——长程信息靠全局层，短程信息靠局部层。Llama 4、Gemma 3（每 5 个局部层配 1 个全局层）、OLMo 3 等也采用了局部 + 全局的交替结构。更激进的做法（把一部分注意力层换成线性注意力或状态空间模型）是下一讲的内容。

11. 资源核算：作业 1 的核心题

11.1 参数量

对作业 1 的架构（Pre-norm、RMSNorm、SwiGLU、不共享嵌入、无偏置）：

配套代码 04_transformer_lm.py 用真实模型核对了这个公式。

11.2 前向 FLOPs

一条长度为 $T$ 的序列（按 $2mkn$ 计）：

这个公式与 FLOP 计数器的结果完全一致。把它用在 GPT-2 各个尺寸上（$T = 1024$，$d_{ff} = 4d$，用作业 1 的架构）：

观察：

• 模型越大，前馈层占比越高，输出层占比越低（$Vd$ 相对于 $Ld^2$ 越来越不重要）。
• 上下文从 1K 增加到 16K 时，与 $T^2$ 成正比的注意力核心从 7% 变成了最大的一项——长上下文昂贵的根源。
• 除注意力核心外，每一项都约等于 $2T\times(\text{对应参数量})$，这就是”前向约 $2N$ FLOPs 每 token”的来源。

11.3 训练显存

用 AdamW、全部 FP32 训练 GPT-2 XL（作业 1 的题目）：

• 参数 8.5 GB + 梯度 8.5 GB + AdamW 状态 17.0 GB = 34.0 GB；
• 激活值：每个样本（1024 个 token）约 18.0 GB，其中注意力分数和 softmax 结果（$T^2$ 项）约占 56%。

激活值的估算公式（每层保存的主要张量）：

总显存 $\approx 34.0 + 18.0\times B$ GB，80 GB 最多放下批次 2。在一个小配置上用保存张量钩子实测，实际保存的激活与公式估算相差 7%（我们教学风格的 softmax、RoPE 会多保存一些中间结果）。

怎么办？混合精度、激活检查点（第 2 讲），FlashAttention 不保存 $T\times T$ 矩阵（第 5 讲），多卡切分状态和参数（第 7 讲）——后面的系统单元基本都在解决这个问题。

12. 本讲小结

1. 现代 Transformer = 原始 Transformer + Pre-norm + RMSNorm + 无偏置 + SwiGLU + RoPE，这是大多数开放模型的共同骨架。
2. Pre-norm 让残差主干道畅通，训练更稳定、能用更大学习率；新动向是在主干道之外加归一化。
3. RMSNorm 与去偏置的主要收益是减少内存搬运：FLOPs 不等于运行时间。
4. GLU 类前馈层稳定地优于 ReLU/GELU；为保持参数量相同，$d_{ff}\approx\frac83 d$。
5. RoPE 通过旋转 q、k 让注意力分数只依赖相对位置；多种转速像钟的指针；只作用于 q、k，每层都做。
6. 注意力除以 $\sqrt{d_k}$ 防止 softmax 饱和；因果掩码防止”偷看答案”；多头提供多种注意力模式。
7. 超参数高度一致：$d_{ff}=4d$ 或 $\frac83 d$，$h\cdot d_{head}=d$，$d/L\approx100\sim200$，单语词表 3–5 万、多语言 10–25 万；新模型少用 dropout、仍用权重衰减（主要影响优化动力学）。
8. 训练不稳定多来自 softmax：z-loss（输出层）、QK-norm（注意力）、软截断（两者）。
9. 推理成本推动了注意力变体：GQA/MQA/MLA 缩小 KV 缓存；滑动窗口降低计算；局部 + 全局交替是当前主流。
10. 实验表明，注意力让模型能从上文中”抄写”信息（归纳头）——这是固定窗口模型做不到的。

13. 与作业 1 的关系

作业 1 在 TinyStories 上的推荐配置是：词表 10000、上下文 256、$d=512$、$d_{ff}=1344$、4 层、16 头、$\Theta=10000$。

14. 练习

练习 1（参数量）：计算作业 1 TinyStories 配置（$V=10000$、$d=512$、$L=4$、$d_{ff}=1344$）的参数量。其中嵌入和输出层占多少？

$N = 2\times10000\times512 + 4\times(4\times512^2 + 3\times512\times1344 + 2\times512) + 512 = 10{,}240{,}000 + 12{,}455{,}936 + 512 = 22{,}696{,}448$。嵌入和输出层共 1024 万，约占 45%——小模型的参数有将近一半在词表上（第 1 讲第 11 节）。可以用 04_transformer_lm.py 中的 param_count_formula 验证。

练习 2（8/3 规则）：$d=512$ 时，$\frac83 d$ 是多少？向上、向下取整到 64 的倍数分别是多少？作业 1 用的 1344 属于哪一种？

$\frac83\times512 \approx 1365.3$；向上取整到 1408（$=22\times64$），向下取整到 1344（$=21\times64$）。作业 1 用的是向下取整。两者都满足”约为 8/3 倍且是 64 的倍数”。

练习 3（RoPE 推导）★：对二维情形证明 $\langle R(m\theta)q, R(n\theta)k\rangle = \langle R((m-n)\theta)q, k\rangle$。提示：旋转矩阵满足 $R(\alpha)^\top = R(-\alpha)$ 和 $R(\alpha)R(\beta) = R(\alpha+\beta)$。

$\langle R(m\theta)q, R(n\theta)k\rangle = q^\top R(m\theta)^\top R(n\theta)k = q^\top R(-m\theta)R(n\theta)k = q^\top R((n-m)\theta)k$。而 $\langle R((m-n)\theta)q, k\rangle = q^\top R((m-n)\theta)^\top k = q^\top R((n-m)\theta)k$，两者相等，只依赖 $m-n$。

练习 4（为什么除以 $\sqrt{d_k}$）：设 $q_i, k_i$ 独立、均值 0、方差 1。证明 $\text{Var}(q\cdot k) = d_k$。

$q_ik_i$ 的均值为 $E[q_i]E[k_i]=0$，方差为 $E[q_i^2k_i^2] = E[q_i^2]E[k_i^2] = 1$。$d_k$ 个独立项之和的方差为 $d_k$。除以 $\sqrt{d_k}$ 后方差为 1。

练习 5（KV 缓存）：一个 40 层、40 个注意力头、头维度 128 的模型，批次 8、上下文 8192，BF16 存储。用 MHA 时 KV 缓存多大？改用 8 个 KV 头的 GQA 呢？

MHA：每 token $2\times40\times40\times128\times2 = 819{,}200$ 字节，总计 $819200\times8192\times8 \approx 53.7$ GB。GQA（8 个 KV 头）：缩小到 1/5，约 10.7 GB。

练习 6（FLOPs 构成）：对练习 1 的配置、$T = 256$，一次前向传播中注意力投影、注意力核心、前馈层、输出层各占多少 FLOPs？哪一项最大？为什么和 GPT-2 XL 的情况不同？

注意力投影 $4\times8\times256\times512^2 \approx 2.15$ GFLOPs；注意力核心 $4\times4\times256^2\times512 \approx 0.54$ GFLOPs；前馈层 $4\times6\times256\times512\times1344 \approx 4.23$ GFLOPs；输出层 $2\times256\times512\times10000 \approx 2.62$ GFLOPs。前馈层最大，但输出层占了约 27%，远高于 GPT-2 XL 的 3.6%——模型小时 $Vd$ 相对于 $Ld^2$ 很大。

练习 7（Pre vs Post）：用”残差主干道”的直觉解释，为什么 Post-norm 的深层网络更难训练。

Pre-norm 中 $x_L = x_0 + \sum_l f_l(\cdot)$，从输出到任意一层都有一条”恒等”的梯度通路，梯度可以不受阻碍地流回去。Post-norm 中每一层都要经过归一化 $x_{l+1} = \text{Norm}(x_l + f_l(x_l))$，梯度要连续穿过 $L$ 个归一化层，每次都会被重新缩放，层数多时梯度的大小变得难以控制，需要小学习率和预热来弥补。

练习 8（z-loss）：为什么 softmax 对 logits 的整体平移不变？既然不变，为什么还要用 z-loss 约束它？

$\frac{e^{o_i+c}}{\sum_j e^{o_j+c}} = \frac{e^c e^{o_i}}{e^c\sum_j e^{o_j}}$，$e^c$ 约掉了。数学上不变，但数值上不是：浮点数的绝对精度随数值增大而变差，低精度（BF16）下 logits 整体很大时，它们之间的细微差别会被舍入掉（本讲 9.1 节的表格）；而且 exp 更容易上溢。z-loss 让 logits 保持在 0 附近，远离这些数值问题。

练习 9（消融实验，作业 1 的内容）★★：在 05_train_transformer.py 的实验上做以下修改，比较”重复名字”上的损失： (a) 去掉 RoPE（theta=None，即 NoPE）； (b) 把 SwiGLU 换成参数量相同的 ReLU 前馈层（01_building_blocks.py 中的 ReLUFFN）； (c) 把所有 RMSNorm 换成恒等映射（去掉归一化），并尝试不同的学习率； (d) 给 MultiHeadSelfAttention 加上 QK-norm。

练习 10（思考）：本讲的实验中，Transformer 在”重复名字”上几乎做到了 0 损失。你能设计一个它也做不好的任务吗？（提示：想想上下文长度、以及需要”多步推理”才能得到答案的情况。）

15. 延伸阅读

• CS336 Spring 2026 Lecture 3 课件（GitHub stanford-cs336/lectures 中的 lecture_03.pdf）：包含大量模型配置的汇总表。
• Vaswani et al.，Attention Is All You Need（2017）。
• Xiong et al.，On Layer Normalization in the Transformer Architecture（2020）：Pre-norm vs Post-norm。
• Zhang & Sennrich，Root Mean Square Layer Normalization（2019）。
• Shazeer，GLU Variants Improve Transformer（2020）；Narang et al.，Do Transformer Modifications Transfer Across Implementations and Applications?（2021）。
• Su et al.，RoFormer: Enhanced Transformer with Rotary Position Embedding（2021）。
• Ainslie et al.，GQA（2023）；Shazeer，Fast Transformer Decoding: One Write-Head is All You Need（2019，MQA）。
• Olsson et al.，In-context Learning and Induction Heads（Anthropic，2022）：归纳头。
• Touvron et al.，Llama / Llama 2 论文；Llama 3 Herd of Models（2024）：现代架构的”标准答案”。

16. 下一讲预告

本讲的 Transformer 有两个”天生”的成本问题：注意力随上下文长度平方增长；每个 token 都要经过全部参数。下一讲（注意力替代方案与混合专家）介绍两类更激进的改动：

• 线性注意力：把 $(QK^\top)V$ 改写成 $Q(K^\top V)$，训练时像 Transformer 一样并行，推理时像 RNN 一样只需常数大小的状态；以及它的发展：RetNet、Mamba-2、Gated DeltaNet，和”线性层 + 注意力层”的混合架构。
• 混合专家（MoE）：把一个大前馈层换成许多个”专家”，每个 token 只用其中几个——参数多、计算少。路由怎么做？专家怎么保持负载均衡？DeepSeek 的 MoE 为什么这样设计？

下一讲之前，建议你：跑通 code/lecture03/ 下的 6 个脚本，完成练习 1、3、6，并尝试练习 9 中的至少一个消融实验。