第 4 讲：Transformer 的改进

对应原课：CMU 11-667 Fall 2025 · Lecture 4 · Thu 4-Sep-2025 · Architecture Advancement on Transformers（Jaromir Savelka、Yubin Kim） 原课 slides：Architecture Advancements on Transformers（课表上标为 transformer advancements） 阅读材料：Ramachandran et al. (2017), Searching for Activation Functions；Xiong et al. (2020), On Layer Normalization in the Transformer Architecture；Ainslie et al. (2023), GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints；Su et al. (2021), RoFormer: Enhanced Transformer with Rotary Position Embedding 配套代码：code/lecture04/（2 个模块 + 5 个演示脚本；最长的一个首次运行约 100 秒） 本讲的写法：骨架来自原课 slides；推导、“一次换一个零件”的花园世界实验、GQA 改装与长度外推实验是补充的。所有数字都来自配套代码的实际输出。

目录

• 0. 本讲导读
• 1. 前馈层：激活函数与门控
• 2. 层归一化：放在哪里、用哪一种
• 3. 注意力：为推理省下 KV cache
• 4. 位置编码：绝对、相对、旋转
• 5. 合起来：现代配方
• 6. 小结、常见误区、练习

0. 本讲导读

0.1 这一讲要回答的问题

第 2 讲搭起来的是”教科书版”的 Transformer。从 2017 年的原版到今天的 Llama 3，结构的骨架没变——嵌入、一叠”注意力 + 前馈”块、残差、归一化、输出层——但几乎每个零件都换过。slides 把可以改进的地方归成四处：

            ┌────────────────────────────────────────────────┐
            │  前馈层 FFN        ReLU → GELU → Swish → SwiGLU    │ §1
  一个块 ── │  注意力 Attention  多头 → 多查询 → 分组查询 → MLA  │ §3
            │  层归一化 LayerNorm post → pre；LayerNorm → RMSNorm │ §2
            └────────────────────────────────────────────────┘
  块之外：   位置编码            正弦 / 可学习 → 相对位置 → RoPE       §4

slides 给这一讲定的目标是：看清楚新结构的细节；理解每次升级的动机和收益；在具体场景下选对结构。这一讲对每一处改动都问三个问题：它解决了什么问题？代价是什么？在小模型上能不能看出来？——最后一个问题的答案常常是”看不出来”，而这本身就是本讲最重要的教训之一。

0.2 学完这一讲你应该能

• 写出 ReLU、GELU、Swish、SwiGLU 的公式，解释门控 FFN 怎样对齐参数量；
• 解释 post-norm 和 pre-norm 的区别，为什么 post-norm 离不开学习率预热，RMSNorm 比 LayerNorm 省了什么；
• 算出一个模型的 KV cache 大小，解释 MQA、GQA、MLA 各自怎样把它压小、代价是什么，以及怎样把已训练的 MHA 模型”改装”成 GQA；
• 证明 RoPE 让注意力分数只依赖相对位置，说出 T5 相对偏置、ALiBi 的思路，理解”长度外推”问题；
• 说出 Llama 3 的四个选择，以及”大多数改动是为了效率和稳定性”这句话的含义。

0.3 本讲的实验工具

配套代码写了一个”每个零件都能换”的模型 ModernLM（modern_transformer.py）：前馈层、KV 头数、归一化的种类和位置、位置编码、有没有偏置，全部由 ModelConfig 控制。取第 2 讲的设置时，它可以直接载入第 2 讲训练好的参数，输出逐位相同（最大差 0.0，测试困惑度同为 3.851）——这保证了下面所有的比较，都只是”换了一个零件”。

还有一件事要先记住：第 2 讲 §5.5 发现，同一个配置换一个随机种子，测试困惑度就能差 0.034。下面每当两个配置的困惑度相差几个百分点时，都要先问一句：这比噪声大吗？

1. 前馈层：激活函数与门控

1.1 三个激活函数

第 2 讲的 FFN 是 $\text{FFN}(x) = W_2\, \phi(W_1 x)$（slides 省略了偏置），$\phi$ 是激活函数。三代选择：

• ReLU（原版 Transformer）：$\text{ReLU}(x) = \max(0, x)$。
• GELU（Hendrycks & Gimpel 2016；BERT、GPT-2）：$\text{GELU}(x) = x\,\Phi(x)$，$\Phi$ 是标准正态分布的累积分布函数。直觉：用”$x$ 有多大概率比一个随机的标准正态数大”来给 $x$ 加权——大的正数几乎原样通过，大的负数几乎归零，中间平滑过渡。BERT、GPT-2 用的是它的 tanh 近似，和精确值在 $[-5, 0]$ 上最多差 $4.7 \times 10^{-4}$（代码 01 §1）。
• Swish（Ramachandran et al. 2017）：$\text{Swish}_\beta(x) = x \cdot \sigma(\beta x)$。$\beta = 1$ 时也叫 SiLU。这篇本讲的阅读材料用自动搜索在大量候选函数里找到了它：$\beta \to 0$ 时它趋于线性函数 $x/2$，$\beta \to \infty$ 时趋于 ReLU——Swish 是两者之间的一族光滑插值。

代码 01 §1 的数值：

两个值得注意的性质：

• ReLU 在负半轴上值和导数都恰好是 0。 一个神经元如果对所有输入都落在负半轴，它就再也收不到梯度——所谓”死掉的 ReLU”。GELU 和 SiLU 在负半轴上有一小段非零的值和导数，梯度不会完全断掉。
• GELU 和 Swish 是非单调的：GELU 在 $x = -0.752$ 处取最小值 $-0.170$，SiLU 在 $x = -1.279$ 处取最小值 $-0.278$——负半轴上先下降再回升。Ramachandran et al. 认为这个”小鼓包”和光滑性是 Swish 表现更好的原因之一；slides 也展示了 Swish 的梯度在零点附近是平滑变化的，而不是 ReLU 那样的阶跃。

1.2 门控 FFN：SwiGLU

slides 的第二个改进是”双线性层”（bilinear layers），即 Shazeer (2020) 的 GLU 变体：

$\odot$ 是逐元素相乘。一条支路（$W_1$）经过非线性后当作”门”，决定每个中间神经元开多大；另一条支路（$W_3$）提供”内容”。输出对输入是二次的（两个线性函数的乘积），比”线性 → 非线性 → 线性”表达能力更强。

代价是多了一个矩阵。为了公平比较，slides 强调要”缩小投影的隐藏维度，保持参数量不变”：普通 FFN 的参数是 $2 \cdot d \cdot d_{ff}$，门控 FFN 是 $3 \cdot d \cdot d_{ff}'$，取 $d_{ff}' = \frac{2}{3} d_{ff}$ 就相等。代码 01 §2（$d = 64$，带偏置）：普通 FFN（$d_{ff} = 256$）33,088 个参数；SwiGLU 若 $d_{ff}$ 不变是 49,728 个；取 $d_{ff} = 171 \approx \frac{2}{3} \times 256$ 后是 33,238 个。Llama 取 $d_{ff} = \frac{8}{3} d$（再取整到 256 的倍数），对应普通 FFN 的 $4d$。

Shazeer (2020) 在 T5 的预训练和下游任务上比较了这些变体，GEGLU 和 SwiGLU 的困惑度最低、下游效果也更好（slides 展示了 T5-base 预训练过程中的困惑度曲线）。有意思的是，论文对”为什么有效”没有给出解释，作者半开玩笑地把成功归于”神的眷顾”——很多结构改进都是实验先行、理论在后。

1.3 在花园世界上比一比

代码 01 §3：第 2 讲的标准模型（$d = 64$、4 个头、2 层、pre-norm LayerNorm、正弦位置编码），只换前馈层，参数量对齐，每种用 2 个随机种子各训练 6 遍。

测试困惑度的差别（最多 0.027）和同一种前馈层两个种子之间的差别（最多 0.036）同一量级——看不出谁更好。门控 FFN 的训练损失确实最低，方向和 Shazeer 的结果一致，但只差千分之二，在测试集上被淹没了（第 2 讲 §5.2：测试集困惑度离真实分布的差距大部分来自”没见过的搭配”，和结构无关）。

这不是实验失败，而是一个重要的提醒：这些改进的收益是在几十亿参数、几千亿 token 的规模上测出来的。 在一个 10 万参数、语法极其简单的玩具语言上，模型容量根本不是瓶颈，再好的前馈层也没有用武之地。第 12 讲的规模定律会讨论”小规模实验能多大程度上预测大规模结果”。

1.4 稀疏性，以及没有出现的”死神经元”

代码 01 §4 用验证集的前 200 篇文档跑一遍训练好的模型，看 FFN 中间层：

• ReLU 的输出有六成恰好是 0。 这是一种”免费”的稀疏性：推理时可以跳过这些神经元对应的那一列 $W_2$。Mirzadeh et al. (2023) 正是为了这个主张大模型改回 ReLU。GELU、SiLU 的输出几乎从不恰好为 0。
• 没有一个死神经元。 教科书常说 ReLU 会”死”，但在 Adam 的自适应步长、梯度裁剪和 pre-norm 之下，ReLU 很难死：作者另外试过学习率 0.01 和 0.03 各训练 2 遍，256 个神经元里最多只有 1 个从未被激活（练习 6）。死 ReLU 的问题更常见于用较大学习率的 SGD 训练、没有归一化层的网络。

2. 层归一化：放在哪里、用哪一种

2.1 LayerNorm 回顾

LayerNorm（Ba et al. 2016）把每个位置的向量归一化：

$g$ 初始化为 1、$b$ 初始化为 0。slides 的说法是：把输出对齐到标准分布，改善训练的收敛和稳定性。

2.2 post-norm vs pre-norm

第 2 讲 §3.4 提过归一化有两种放法：

post-norm（原版 2017）：  x ← LN(x + 子层(x))
pre-norm （GPT-2 起）：   x ← x + 子层(LN(x))      最后再补一个 LN

看起来只是挪了个位置，效果却大不相同。展开 $L$ 层的 pre-norm：

残差主干是一条从输入直通输出、不经过任何归一化的恒等通路，每一层只是往上”加一点”。post-norm 的主干则每过一层都被 LN 重新缩放一次，梯度要穿过 $L$ 个 LN 才能回到底层。

Xiong et al. (2020)（本讲阅读）用平均场理论分析了初始化时的梯度：post-norm 靠近输出的层，参数梯度的期望值很大，而且不随深度减小；pre-norm 的梯度随深度按 $1/\sqrt{L}$ 左右缩小。slides 的总结是：post-norm 梯度的 L2 范数大，导致大规模训练不稳定；pre-norm 梯度范数小，训练更稳定。代码 02 §2 在随机初始化的模型上测了各层 FFN 输出矩阵的梯度范数：

post-norm 最后一层的梯度在 24 层时仍约为 1，全部参数的梯度范数从 3.5 涨到 8.1；pre-norm 最后一层的梯度从 0.87 降到 0.44，总梯度范数基本不涨。

2.3 后果：post-norm 离不开预热

梯度大，意味着训练开始时（Adam 的二阶矩估计还不准、参数离合理区域很远）一大步就可能把模型推坏。原版 Transformer 用了 4000 步学习率预热，BERT 用了 1 万步——这正是 post-norm 的代价。代码 02 §3：6 层模型，训练 2 遍，不做梯度裁剪：

学习率小时两者都能训练。学习率升到 5e-3、不预热，post-norm 直接训崩：验证困惑度 42.2，和一个只看词频、完全不看上下文的 unigram 模型（44.3）差不多——模型退化成了”永远猜最常见的词”。pre-norm 毫无问题。加上 200 步预热，post-norm 也恢复正常，而且和 pre-norm 一样好。Xiong et al. 的结论正是这样：pre-norm 可以去掉预热、用更大的学习率，训练更快、更省调参。

（post-norm 并非一无是处：有研究认为训练成功的 post-norm 模型最终效果可能略好，DeepNet 等工作也提出了让深层 post-norm 稳定的初始化办法。但主流 LLM 都选择了更省心的 pre-norm。）

2.4 RMSNorm：更简单

RMSNorm（Zhang & Sennrich 2019）只做”按均方根缩放”：

不减均值（slides：“只缩放，不重新居中”），也没有偏置 $b$。代码 02 §1 验证了它们的不变性：输入变成 $3x$ 时两者输出几乎不变（只差 eps 带来的 $10^{-5}$ 级误差）；输入变成 $x + 5$ 时 LayerNorm 不变，RMSNorm 输出变化 2.7。我们手写的 RMSNorm 与 PyTorch 的 nn.RMSNorm 完全一致；输入均值为 0 时，RMSNorm 和 LayerNorm 只差 $3 \times 10^{-5}$。

LayerNorm 对”整体缩放”和”整体平移”都免疫，RMSNorm 只对缩放免疫。实践表明重新居中可有可无，而 RMSNorm 少算一次均值、少一组参数，原论文报告它在机器翻译等任务上收敛更快（slides）。代码 02 §4：2 层 pre-norm 模型，LayerNorm 的测试困惑度（2 个种子）3.852 / 3.816，RMSNorm 3.872 / 3.819——差别在随机波动之内。T5 用的”只缩放、无偏置的 LayerNorm”其实就是 RMSNorm；Llama、Qwen、DeepSeek 都用它。

3. 注意力：为推理省下 KV cache

3.1 推理的瓶颈在哪里

slides 指出标准多头注意力在推理时的问题：显存消耗巨大，每一层、每一个位置都要保留一份 K、V。第 2 讲 §4.3 说过，因果模型可以缓存前面位置的 K、V；但这个缓存本身很大，而且解码时每生成一个 token，都要把整个缓存从显存读一遍（P3 §6：解码受内存带宽限制）。缓存越小，解码越快、一张卡能同时服务的请求越多、能支持的上下文越长。

每个 token、每一层要缓存 $2 \times n_{kv} \times d_{head}$ 个数（K 和 V 各一份）。代码 03 §1（BF16，每个数 2 字节）：

Llama 2 70B 的 BF16 权重约 128 GiB，一块 80 GiB 的 GPU 本来就装不下；若用 MHA，一条 32K 的缓存还要再 80 GiB。

3.2 MHA → MQA → GQA

• 多查询注意力（MQA，Shazeer 2019）：所有查询头共用一组 K、V。KV cache 缩小到 $1/h$。slides：“进一步减小 KV cache”。代价是质量有可测的下降，训练也可能不稳定。
• 分组查询注意力（GQA，Ainslie et al. 2023）：把 $h$ 个查询头分成 $G$ 组，组内共用一组 K、V。$G = h$ 就是 MHA，$G = 1$ 就是 MQA。slides：“把 KV cache 缩小到组数那么大”。

注意 GQA 不减少注意力的计算量：查询头还是 $h$ 个，每个查询头照样对所有位置算分数；它减少的是要存、要从显存读的 K、V。代码 03 §2 用 ModernLM 带 KV cache 逐个 token 解码一条 48 个 token 的序列，和整段重算比较：

逐 token 解码时每一步只算新 token 那一行，结果和整段重算一致（浮点误差）；KV 头减半，cache 减半。

3.3 MLA：只缓存一个潜在向量

**多头潜在注意力（MLA，DeepSeek-V2）**换了一个思路：不减少头数，而是把 K、V 压缩。每个 token 先用一个下投影得到低维的潜在向量 $c_t = W_{DKV}\, x_t$（DeepSeek-V2 里是 512 维），只缓存 $c_t$；需要时再用上投影恢复每个头的 $k_{t,i} = W_{UK}^{(i)} c_t$、$v_{t,i} = W_{UV}^{(i)} c_t$。

slides 特别指出：恢复 k、v 的计算可以和 q 合并。注意力分数

可以先把 $W_{UK}$“吸收”进 q，直接和缓存的 $c_t$ 做点积，根本不需要恢复出 $k$；同理 $\sum_t a_t W_{UV} c_t = W_{UV} \sum_t a_t c_t$，$W_{UV}$ 可以并进输出矩阵。代码 03 §3 数值验证：先恢复 k 再算分数 vs 吸收进 q，最大差 $2.9 \times 10^{-6}$。

但吸收和 RoPE（§4.4）冲突：加了 RoPE 后 $k_t = R_t W_{UK} c_t$，旋转矩阵 $R_t$ 夹在中间、而且随位置变化，无法预先并进 q——代码里照样吸收，分数最多差 14.27，完全错了。DeepSeek-V2 的办法是解耦的 RoPE：另开一小段（64 维）专门承载位置信息的 q、k，这段 k 单独缓存、不参与吸收。于是每层每个 token 缓存 512 + 64 = 576 个数，约为同样头数 MHA 的 1/57（§3.1 的表）。

3.4 把训练好的 MHA 模型改装成 GQA

slides 介绍了两种使用高效注意力的方式：直接用新结构预训练，或者把一个丰富的多头注意力模型压缩、加速。后者就是 Ainslie et al. (2023) 的 uptraining：把每组内各个头的 K、V 投影矩阵取平均（mean pooling），得到一个 GQA 模型，再用原预训练预算的一小部分（他们用 5%）继续训练。他们的结果：改装后的 GQA 质量接近 MHA、生成速度接近 MQA；在几种初始化方法中，取平均好于只取第一个头，好于随机初始化。

代码 03 §4 在第 2 讲的模型（MHA，4 个头，测试困惑度 3.851，训练了 1500 步）上复现：

• 直接合并 K、V 头会让模型变差，合并得越狠（MQA）越差；随机初始化 K、V 最差——等于扔掉了学到的 K、V。
• GQA 时组内取平均明显好于只取第一个头（4.658 vs 5.259）；MQA 时直接用反而是第一个头略好，继续训练 250 步后取平均又领先。Ainslie et al. 的结论在这里大体成立，但不是每一格都成立。
• GQA 取平均后只继续训练 75 步（原预算的 5%），困惑度就回到 3.892，与原模型的差距和随机种子的波动相当；MQA 要多训练一些。这比从零训练便宜得多。
• 从零训练的 GQA、MQA 在这里一点也不比 MHA 差（3.850、3.827）：这么小的模型用不满 4 个 KV 头。MQA 的质量损失是在大模型上测出来的，GQA 才成为”质量接近 MHA、速度接近 MQA”的折中——Llama 2 70B、Llama 3 全系列、Mistral 都用 GQA。

4. 位置编码：绝对、相对、旋转

4.1 Transformer 本身对位置无感

slides 对比了三种结构的”先验”：CNN 有局部性先验（卷积核只看邻近位置），RNN 有顺序先验（一步一步往后传），而 Transformer 的注意力对位置不可知（第 2 讲 §3.5：双向注意力不加位置编码时，打乱输入只会打乱输出）。位置编码就是把对语言很有用的位置信息加回去。

4.2 绝对位置

加法式的位置编码在嵌入层把位置信息加到 token 嵌入上，有两种：

• 正弦位置编码（原版 Transformer）：不同的维度用不同的波长，用来捕捉不同尺度的相对位置（第 2 讲 §3.5）。
• 完全可学习的位置嵌入（BERT、GPT-2）：每个位置一个向量。Clark et al. (2019) 分析 BERT 的注意力发现，它学到了很强的基于位置的注意力模式（例如有的头专门看前一个词、有的专门看后一个词）。缺点是长度上限由嵌入表的大小决定，没见过的位置没有向量可查。

4.3 相对位置

slides 提出了一个语言学的先验：在语言里，只有相对位置才重要。I took my dog, Fido, to the park for his walk 里，his 和 dog 的距离是 5；在它前面加一长串 Blablabla…，距离仍然是 5，关系不变。绝对位置编码却会让这两个句子里的 his、dog 拿到完全不同的位置向量。

于是有了在注意力里编码相对位置的方法：

• 相对位置嵌入（Shaw et al. 2018）：给每个相对距离学一个向量，加到 key（和 value）上参与注意力计算。在机器翻译上比绝对位置更好（slides）。
• T5 的相对位置偏置（Raffel et al. 2020）：每个”距离桶”、每个头学一个标量，直接加到注意力分数上。代码 04 §3 打印了分桶规则：距离 0～15 每个一个桶；之后按对数分桶（20→17，32→21，64→26，100→30）；128 及以上全部落进最后一个桶（31）。近处精确、远处粗略——远处差几个词不重要。

4.4 RoPE：用旋转编码位置

旋转位置编码（RoPE，Su et al. 2021，本讲阅读）是今天几乎所有开源 LLM 的选择。slides 从点积的几何意义讲起：q、k 的长度固定时，注意力分数大致是两者夹角的余弦。希望位置相近时分数高（“转得靠近”），位置远时分数低（“转开”）。怎样让旋转只依赖相对位置？ 把 q、k 按各自位置转同样的”每步角度”：

二维的情形：把 $q$ 旋转 $m\theta$、把 $k$ 旋转 $n\theta$（$m$、$n$ 是它们的位置）：

因为旋转矩阵满足 $R_a^\top R_b = R_{b-a}$（先转回 $a$ 再转 $b$）。分数只依赖 $n - m$。用复数看更简洁：把二维向量看成复数，旋转就是乘 $e^{i m \theta}$，内积是 $\text{Re}[q\, e^{im\theta}\, \overline{k\, e^{in\theta}}] = \text{Re}[q \bar{k}\, e^{i(m-n)\theta}]$。

高维的情形：把 $d_{head}$ 个维度两两分组，每组用不同的角速度 $\theta_i = 10000^{-2i/d_{head}}$ 各转各的——这和正弦位置编码的频率完全一样。slides 引用了原版 Transformer 论文的一句话：选择正弦函数，是因为假设它能让模型容易学会按相对位置关注。RoPE 把这个想法做彻底了：用乘法（旋转）代替加法。

代码 04 §1 数值验证：同一对随机的 q、k，放在 (3, 1)、(10, 8)、(100, 98)、(1000, 998)，RoPE 的分数都是 −1.8956；而加法式正弦编码的分数是 9.26、10.76、13.38、8.26，随绝对位置变化。原因是加法展开后

只有最后一项只依赖距离，中间两项”内容 × 绝对位置”的交叉项随绝对位置变化。RoPE 没有这些交叉项。另外旋转不改变长度：$|R_m q| = |q| = 3.6594$。

远程衰减：RoFormer 论文证明了 RoPE 的分数上界随距离衰减。代码 04 §2 取 $q = k =$ 全 1 向量（64 维）：距离 0 时分数 64，距离 1、2、5、10、20、50、100、200、500、1000 时分别是 61.8、56.6、47.0、42.1、37.8、31.3、35.7、22.9、12.8、17.9。分数 $= \sum_i 2\cos(\Delta \theta_i)$：距离近时各维度”同相”、和很大；距离远了相位错开、互相抵消（带振荡）。这对应”越远的词通常越不相关”的先验。

RoPE 的其他好处：没有位置嵌入表，没有参数；只作用在 q、k 上，不影响 v；和 KV cache 天然兼容（每个 k 在写入缓存前就按自己的位置转好了）。slides 引用的 EleutherAI 实验显示，用 RoPE 的模型比用可学习位置嵌入的收敛更快。

4.5 长度外推：训练短、测试长

如果模型只在长度 $T$ 的序列上训练过，给它更长的序列会怎样？代码 04 §4：把训练文档首尾相接（文档之间放 <eos>），切成长 64 的块训练 600 步；测试时切成长 512 的序列，按位置统计平均 NLL。花园世界的文档平均只有 25 个 token，预测下一个词从来不需要 64 个以上的上下文——所以长序列上变差，只可能是因为位置编码”没见过这么远”。

（真实分布的 NLL 约为 $\ln 3.655 = 1.296$。）

• 在训练长度之内，除 NoPE 外都差不多（1.30）。
• 绝对位置编码一出训练长度就崩：可学习的位置嵌入在 64 之后的向量从来没被训练过；正弦编码虽然能算出任何位置的向量，模型却没见过这些向量和词嵌入相加后的样子。
• 相对位置编码外推得好得多。在这个实验里 RoPE 最好，T5 偏置其次，两者都随距离逐渐变差。
• NoPE 在打包数据上表现最差，连训练长度之内都差（1.388）。这和第 2 讲按文档训练时 NoPE 不输其他编码的结果相反：打包后文档的开头不再对齐序列的开头，没有位置信息的模型很难判断”本篇文档从哪里开始”。

这里有两点要和文献对照着看。第一，Press et al. (2022) 在真实语料上的比较里，T5 偏置比 RoPE 外推得好，RoPE 和正弦编码一样，超出训练长度后困惑度很快变差；我们的任务只需要很短的上下文，结论不同。第二，他们提出的 ALiBi（给注意力分数加一个随距离线性下降的偏置 $-m_h (i - j)$，每个头的斜率 $m_h$ 固定、不学习）几乎完美地外推——练习 7 在花园世界上验证了这一点。今天的长上下文模型多半仍用 RoPE，再配合位置插值（把位置按比例压缩回训练范围，Chen et al. 2023）、NTK-aware 缩放、YaRN（Peng et al. 2023）等方法，用少量长文本继续训练来扩展上下文——第 21 讲详细讨论。

5. 合起来：现代配方

5.1 Llama 3 的选择

slides 的最后一页是 Llama 3 的选择：前馈层 SwiGLU；注意力 GQA；层归一化 pre-norm + RMSNorm；位置编码 RoPE。再加上线性层不带偏置，这就是今天开源 LLM 的”标准配方”。几个有代表性的模型：

5.2 在花园世界上一次改一处

代码 05 §1 从原版（post-norm、ReLU、正弦、MHA、偏置）出发，一次改一处，改到 Llama 式（参数量都约为 10.6 万）：

两套配方各训练 3 个种子：原版 3.822 / 3.815 / 3.824（平均 3.821），Llama 式 3.848 / 3.872 / 3.799（平均 3.840）。Llama 式三个种子之间就差了 0.07，两套配方的平均值只差 0.02——在困惑度上分不出高下。

唯一稳定的信号是窗口外的物主代词：Llama 式三个种子都是 0.997～0.998，原版都是 0.986～0.989；上表里也正是加上 RoPE 的那一步从 0.984 跳到了 0.998。“找本句最近的主语”正好是一个相对位置的问题。

代码 05 §2、§3 看另外两件事：

• 稳定性：6 层、学习率 5e-3、不预热、不裁剪梯度。原版训练损失 3.809 → 3.784，最好的验证困惑度 45.320——训崩了（post-norm，§2.3）；Llama 式 1.570 → 1.325，验证困惑度 3.919。
• 推理：同样的参数量，Llama 式的 KV cache 每 token 128 个数，原版 256 个。

5.3 slides 的总结

slides 对这一讲的总结可以逐条对上我们的实验：

• 前馈层：用不同的形状改善速度—质量的权衡（SwiGLU）——收益要在大规模上才看得到（§1.3）；
• 注意力：牺牲一点质量换推理速度（GQA、MQA、MLA）——KV cache 减半、减到 1/4、减到 1/57（§3）；
• 层归一化：用简单换来训练速度和鲁棒性（pre-norm、RMSNorm）——不预热也不崩（§2.3）。

“大多数 Transformer 结构升级是为了效率和大规模训练的稳定性。简单往往获胜。” 在一个 10 万参数的玩具模型上，这些升级几乎不改变困惑度；它们的价值体现在训得稳不稳、推理快不快、能处理多长的上下文。评估一个结构改动时，只看小模型的困惑度，会把这些价值全部漏掉——这正好引出第 5 讲：怎样评估一个模型。

6. 小结、常见误区、练习

6.1 小结

1. 前馈层：ReLU → GELU（$x\Phi(x)$）→ Swish（$x\sigma(\beta x)$），后两者光滑、非单调、负半轴有梯度。门控 FFN（SwiGLU）用三个矩阵、中间层取 2/3 来对齐参数量。花园世界上四种前馈层的困惑度差别在随机波动之内；ReLU 有约六成的精确稀疏性，在 Adam + pre-norm 下没有出现死神经元。
2. 层归一化：pre-norm 的残差主干是一条恒等通路，梯度随深度变小；post-norm 靠近输出的梯度大、不随深度减小，必须预热——不预热时 6 层 post-norm 模型训崩到 unigram 的水平。RMSNorm 只缩放不居中，效果相当、计算更省。
3. 注意力：KV cache $= 2 \times n_{kv} \times d_{head}$ 个数 / token / 层。MQA 共用一组 K、V，GQA 分组共用，MLA 缓存低维潜在向量并把上投影吸收进 q（与 RoPE 冲突，需要解耦的 RoPE）。已训练的 MHA 模型可以按组平均 K、V 改装成 GQA，再训练约 5% 的预算。
4. 位置编码：绝对位置（正弦、可学习）→ 相对位置（Shaw、T5 偏置）→ RoPE（旋转 q、k，分数只依赖相对距离，上界随距离衰减）。训练长度 64、测试长度 512 时，绝对位置编码一出训练长度就崩，RoPE 和 T5 偏置逐渐变差，ALiBi 几乎不变。
5. 现代配方（Llama 3）：pre-RMSNorm + SwiGLU + RoPE + GQA + 无偏置。在小模型上困惑度分不出高下（唯一稳定的信号是 RoPE 让物主代词更准），价值在稳定性和推理效率。

6.2 常见误区

• “新结构一定在任何规模上都更好”：GLU 变体、RMSNorm 的收益是在大规模上测出来的；在小模型上它们和老配方在随机波动之内（§1.3、§5.2）。反过来，小规模实验没看出差别，也不能说明大规模上没有差别。
• “GQA 减少了注意力的计算量”：查询头数不变，计算量基本不变；它减少的是 KV cache 的显存和读取带宽（§3.2）。
• “多头改成 GQA 必须从头训练”：可以按组平均 K、V 投影、再少量继续训练（§3.4）。
• “pre-norm 和 post-norm 只是写法不同”：两者的梯度行为完全不同，决定了需不需要预热、能用多大的学习率（§2.2、§2.3）。
• “RoPE 是一种加到嵌入上的位置向量”：RoPE 不加任何东西，它旋转每一层的 q、k；v 和残差主干里没有位置信息（§4.4）。
• “RoPE 能处理任意长度”：公式上能算任意位置，但超出训练长度后效果会变差，需要位置插值、YaRN 等方法（§4.5）。
• “ReLU 一定会有大量死神经元”：在 Adam、梯度裁剪、归一化之下很少出现（§1.4）。

6.3 练习

练习 1（激活函数）：（a）求 SiLU $f(x) = x\sigma(x)$ 的导数，并说明它在负半轴上有一个极小值点。（b）证明 $\text{Swish}_\beta(x) = x\sigma(\beta x)$ 在 $\beta \to \infty$ 时逐点收敛到 ReLU，$\beta \to 0$ 时收敛到 $x/2$。

（a）$f'(x) = \sigma(x) + x\sigma(x)(1 - \sigma(x)) = \sigma(x)\,[1 + x(1 - \sigma(x))]$。$\sigma(x) > 0$，所以 $f'(x) = 0 \iff 1 + x(1 - \sigma(x)) = 0$。$x \to -\infty$ 时 $1 + x(1-\sigma(x)) \approx 1 + x < 0$，$x = 0$ 时等于 1 > 0，中间必有零点；数值求解得 $x \approx -1.278$，$f(-1.278) \approx -0.278$（代码 01 §1 用网格搜索得到 $-1.279$ 和 $-0.278$）。零点左边 $f' < 0$、右边 $f' > 0$，是极小值。

（b）$x > 0$ 时 $\sigma(\beta x) \to 1$，$x < 0$ 时 $\sigma(\beta x) \to 0$（$\beta \to \infty$），所以 $x\sigma(\beta x) \to \max(0, x)$。$\beta \to 0$ 时 $\sigma(\beta x) \to \sigma(0) = 1/2$，所以 $x\sigma(\beta x) \to x/2$。代码 01 §1 打印了 $\beta = 0.1, 1, 10$ 时的值：$\beta = 0.1$ 时 $-5 \mapsto -1.888$（接近 $-2.5$），$\beta = 10$ 时负半轴全部接近 0。

练习 2（参数对齐）：（a）证明无偏置时，$d_{ff} = \frac{8}{3}d$ 的 SwiGLU 和 $d_{ff} = 4d$ 的普通 FFN 参数量相同。（b）Llama 2 7B 的 $d = 4096$、$d_{ff} = 11008$，和 $\frac{8}{3}d$ 差多少？为什么不直接用 $\frac{8}{3}d$？

（a）普通 FFN：$2 \cdot d \cdot 4d = 8d^2$。SwiGLU：$3 \cdot d \cdot \frac{8}{3}d = 8d^2$。

（b）$\frac{8}{3} \times 4096 = 10922.7$，Llama 取 11008，多了约 0.8%。原因是硬件效率：矩阵的维度是 256 的倍数时 GPU 的矩阵乘法更快（11008 = 43 × 256）。P3 §3 的参数量公式用的就是这个 11008，与官方报告的 6,738,415,616 逐位一致。

练习 3（KV cache）：Mistral 7B 有 32 层、32 个查询头、8 个 KV 头、每头 128 维。（a）BF16 下每个 token 的 KV cache 是多少？32K 上下文、批大小 8 时一共多少？（b）如果改成 MHA 呢？

（a）每 token：$2 \times 32 \times 8 \times 128 \times 2$ 字节 $= 131,072$ 字节 $= 128$ KiB（和 Llama 3 8B 相同）。32K × 8：$128 \text{ KiB} \times 32768 \times 8 = 32$ GiB。

（b）MHA：KV 头 32 个，每 token 512 KiB，32K × 8 是 128 GiB——比模型权重（约 14 GiB）大九倍。长上下文、大批量时，KV cache 而不是权重才是显存的大头；这也是第 19 讲的 PagedAttention、KV cache 量化要解决的问题。

练习 4（RoPE 的相对性）：（a）证明二维旋转矩阵满足 $R_a^\top R_b = R_{b-a}$。（b）解释为什么 RoPE 只作用在 q、k 上，而不作用在 v 上。

（a）$R_a = \begin{pmatrix}\cos a & -\sin a \\ \sin a & \cos a\end{pmatrix}$，$R_a^\top = R_{-a}$（旋转的逆是反向旋转），而旋转的复合是角度相加：$R_{-a} R_b = R_{b-a}$。展开验证：$R_a^\top R_b$ 的左上角是 $\cos a \cos b + \sin a \sin b = \cos(b-a)$，其余同理。

（b）位置信息要影响的是”谁关注谁”，即注意力分数，而分数只由 q、k 决定。v 是”取走的内容”，给它加位置会让输出依赖绝对位置（v 不和别的东西做点积，旋转没有被抵消掉的机会），破坏相对性。

练习 5（残差主干）：写出 3 层 post-norm 的 $x_3$ 关于 $x_0$ 的展开式，和 pre-norm 的 $x_3 = x_0 + \sum_l F_l(\text{LN}(x_l))$ 对比。从 $x_3$ 到 $x_0$ 的梯度，在两种结构里各要经过什么？

post-norm（每层把两个子层合写成 $F_l$）：$x_1 = \text{LN}(x_0 + F_0(x_0))$，$x_2 = \text{LN}(x_1 + F_1(x_1))$，$x_3 = \text{LN}(x_2 + F_2(x_2))$。$\partial x_3 / \partial x_0$ 是三个”LN 的雅可比 × $(I + \partial F/\partial x)$“的乘积——每一项都经过一个 LN，没有一条不经过 LN 的路。

pre-norm：$\partial x_3 / \partial x_0 = I + \sum_l (\text{各条经过 } F_l \text{ 的路径})$。恒等项 $I$ 保证梯度至少能原样流回底层，不管中间的 $F_l$、LN 怎样缩放。这就是 pre-norm 训练更稳、各层梯度更均匀的原因（代码 02 §2）。

练习 6（动手：大学习率下的 ReLU）：把 01_ffn_activations.py 的 ReLU 模型改用学习率 0.01 和 0.03 各训练 2 遍，数一数”在验证集前 200 篇文档上从未被激活”的神经元。

学习率 0.01：两层都是 0 个（测试困惑度 3.978）；学习率 0.03：第 1 层 0 个、第 2 层 1 个（共 256 个，测试困惑度 3.976）。GELU 在这两个学习率下也都是 0 个。Adam 按每个参数的梯度大小自适应地调整步长、梯度裁剪限制了单步更新的幅度、pre-norm 让每层的输入保持在合理范围——这几样东西加起来，让”一步把神经元推进负半轴、再也回不来”变得很难发生。

练习 7（动手：ALiBi）：ALiBi（Press et al. 2022）不给嵌入加位置，而是给注意力分数加偏置 $-m_h \cdot (i - j)$，第 $h$ 个头的斜率 $m_h = 2^{-8h/H}$（$h = 1..H$），不学习。提示：ModernLM 已经支持把一个偏置张量加到注意力分数上（T5 偏置的路径），继承它并重写 _t5_bias 即可。用 04_position_encodings.py 第 4 节的设置训练并测试。

位置 0–63、64–127、128–255、256–511 的平均 NLL：1.297、1.298、1.271、1.276——几乎不随长度变化，甚至略好（更长的序列里，后面的文档有更多”前文”可以确认文档边界）。这和 Press et al. 论文的标题完全吻合：Train Short, Test Long。原因：线性偏置让模型天然”偏爱近处”，远处的 token 被压得很低，长序列里多出来的那些远距离 token 几乎不分到注意力；而且偏置是固定的函数，不存在”没训练过的位置”。代价是它强加了”越远越不重要”的先验，对真正需要远距离信息的任务可能不利。

练习 8（思考：为什么 MQA 在这里不吃亏）：§3.4 里从零训练的 MQA（3.827）甚至略好于 MHA（3.851）。这能说明大模型也应该用 MQA 吗？

不能。花园世界的语言很简单，一个 $d = 64$ 的模型里 4 个头需要的”查询方式”屈指可数（找主语、看前一个词……），一组共用的 K、V 就足以支撑它们；差别 0.024 也在随机波动之内。大模型要同时追踪成百上千种关系，KV 的表达能力会成为瓶颈——Ainslie et al. 在 T5-XXL 上测到 MQA 有明显的质量下降，GQA 才几乎追平 MHA。小规模实验不能外推出”某个省资源的改动没有代价”，正如它也不能外推出”某个改进没有收益”（§1.3）。

练习 9（思考 + 动手：两种 RoPE 写法）：RoFormer 论文把相邻两维 $(2i, 2i+1)$ 配成一对旋转；Llama 的代码把前一半和后一半配对，第 $i$ 维和第 $i + d/2$ 维一起旋转。它们训练出来的模型有区别吗？

没有本质区别。两种写法只差一个固定的维度排列 $P$：把交错排列的维度重排成”前一半是偶数维、后一半是奇数维”，交错写法就变成了分半写法。在模型里，q、k 来自可学习的投影 $W_q$、$W_k$，排列 $P$ 可以被吸收进投影矩阵——两种写法能表示的函数完全相同。数值验证：对同一对 q、k 用交错写法，和先按 $P$ 重排、再用分半写法，分数最大差 $4.8 \times 10^{-7}$。唯一要注意的是加载别人训练的权重时必须用和训练时相同的写法，否则相当于把 q、k 的维度打乱了。

6.4 延伸阅读

• 原课阅读：Ramachandran et al. (2017), Searching for Activation Functions；Xiong et al. (2020), On Layer Normalization in the Transformer Architecture；Ainslie et al. (2023), GQA；Su et al. (2021), RoFormer
• 前馈层：Hendrycks & Gimpel (2016), GELU；Shazeer (2020), GLU Variants Improve Transformer；Mirzadeh et al. (2023), ReLU Strikes Back
• 归一化：Ba et al. (2016), Layer Normalization；Zhang & Sennrich (2019), Root Mean Square Layer Normalization；Wang et al. (2022), DeepNet
• 注意力：Shazeer (2019), Fast Transformer Decoding: One Write-Head is All You Need（MQA）；DeepSeek-AI (2024), DeepSeek-V2（MLA）
• 位置编码：Shaw et al. (2018), Self-Attention with Relative Position Representations；Press et al. (2022), Train Short, Test Long（ALiBi）；Clark et al. (2019), What Does BERT Look At?；Chen et al. (2023), Extending Context Window via Positional Interpolation；Peng et al. (2023), YaRN
• 本仓库的相关内容：CMU10-423 第 4 讲（预训练与现代 Transformer）；CS336 第 3 讲（结构与超参数，有更多模型的对比）

6.5 下一讲预告

本讲反复遇到同一个问题：两个模型的困惑度差 0.02，这算不算”更好”？困惑度是不是衡量模型的唯一标准？第 5 讲讨论 LLM 的自动评估：困惑度的陷阱（为什么不同分词器的困惑度不能比）、生成任务的指标（BLEU、ROUGE、BERTScore、COMET）和它们与”真正正确”的相关性、多选题的打分方式、排序指标、多维度的基准，以及分数什么时候会骗人。