Transformer 架构及其变体

分词之后，模型终于拿到一串整数 id。前沿稠密模型用一组几乎完全相同的 Transformer 块，把这些 id 变成预测，而这个块自 2023 年以来几乎没变过 (Touvron et al. 2023; Bai et al. 2023)。它内部的大多数问题已经有稳定答案：怎么归一化、怎么激活、怎么编码位置。唯有一个问题仍在变化，而且它不在前向传播的数学里。这是一笔生成期间驻留在内存里的成本，也就是 KV 缓存的大小。它也是今天仍能推动架构师改动这个块的主要力量。因此，这一章分两层看：一层看已经定型的前向传播，一层看那项仍然开放的成本，正是它驱动了从 MHA 经 MQA、GQA 到 MLA 的注意力变体。

Transformer 块的两项职责

一个 Transformer 块要做两件不同的事，而且做的方式得经得起被堆叠上百层。它必须让词元之间交换信息，因为如果位置 $i$ 看不到位置 $j$，一个语言模型就毫无用处；它还必须变换每个词元的特征，模型的大部分知识主要存放在这里。把这两个操作朴素地堆起来行不通：深层残差网络很难保持可训练，激活与归一化的选择决定了梯度（也就是告诉每个参数该如何调整以降低损失的那个信号）能否存活到最底层，而注意力，这个唯一跨位置混合的操作，计算上花费 $O(n^2)$，还在生成期间留下一个无界增长的缓存。

读懂一个块的正确方式，不是「先注意力，再 MLP」，而是两个子层，它们各自读取一条共享的残差流，算出一个增量，再把它写回去：

$$x = x + \text{sublayer}(\text{norm}(x))$$

其中 $x$ 是进入子层前的残差流，$\text{norm}(x)$ 是先归一化后的副本，$\text{sublayer}(\cdot)$ 算出的不是新的状态，而是要加回原流的增量。这个式子把前置归一化块的设计直觉压缩成一行：保留一条不被破坏的恒等通路，同时允许每个子层向共享状态写入一小步修改。残差流（residual stream）是模型的工作记忆，一个宽度为 $d_{\text{model}}$ 的向量，从嵌入一路携带到最后一层。每个块都会读取它，并向它累加一次修改。深度就是这样的读取、修改和写回有多少趟；宽度则是这条流一次能容纳多少。这一框架让其余的设计自然推导出来。图 8.1 就以这种方式画出一个前置归一化块：残差流作为一条恒等路径笔直向下，每个子层从它分叉而出，经过一次归一化，算出一个增量，再把增量加回去。

第一个子层是注意力，唯一让词元之间交换信息的地方，也是本章后半部分的主题。第二个是逐位置的 MLP（也称前馈网络或 FFN），对每个词元独立施加，它持有大部分参数：一个前馈层扩展到约 $d_{\text{model}}$ 的四倍再投影回来，每个块花费约 $8 \cdot d_{\text{model}}^2$，让注意力的投影相形见绌。知识主要存放在 MLP，混合则全部发生在注意力。

已经定型的三项选择

下面三项选择让这样一摞深层块保持可训练，而且都已经定型。它们不是活跃的争论，而是一份记录：共识堆叠如何由一系列代价很小的改进，一步步装配而成。

归一化有两条正交的轴线。一条是形式。LayerNorm (Ba et al. 2016) 用可学习的增益与偏置做重新居中与重新缩放；均方根层归一化（RMSNorm） (Zhang and Sennrich 2019) 只做重新缩放，弃掉了均值减除与偏置。均值居中带来的收益很小，所以 RMSNorm 更便宜，如今成了默认。另一条是位置。后置归一化是原始设计，把归一化放在残差相加之后 (Vaswani et al. 2017)；前置归一化把归一化放进分支内部，让一条恒等路径笔直沿残差流而下。前置归一化保留了从堆叠顶端到底端的稳定梯度路径，正是它让深层模型无需精心设计的预热（在训练最初若干步里把学习率慢慢抬升上来）就能训练 (Xiong et al. 2020)。

激活函数有一条清晰的演进路线：ReLU，简单但容易出死单元，然后是 GeLU (Hendrycks and Gimpel 2016)，平滑的 BERT 与 GPT-2 时代默认，再然后是门控线性单元，其中 Swish 门控线性单元（SwiGLU） 是现代的赢家 (Shazeer 2020)。SwiGLU 计算 $(\text{Swish}(xW_1) \odot xV)\, W_2$，是三个矩阵而非两个。容易漏掉的细节是：隐藏维度被按约 $2/3$ 缩放，让门控前馈网络（FFN）保持与普通 $4\times$ 前馈网络相同的参数量，这正是让同等条件比较公平的原因。

位置是第三项。注意力本身对置换不变，所以块必须被告知词元位于何处。旋转位置嵌入（RoPE） 把查询与键向量按与绝对位置成正比的角度旋转，在二维子空间里成对地进行，使得点积 $q \cdot k$ 只依赖于相对偏移 $i - j$ (Su et al. 2021)。一个基础频率参数 theta，设定旋转转得多快，并悄然支配着上下文随后能扩展多远。ALiBi 走的是另一条路：在 softmax 之前的分数上加一个静态的、每个头各自的线性惩罚 $-m \cdot (i - j)$，完全不带可学习的位置参数 (Press et al. 2022)。两者的对比，是向量上的旋转还是分数上的偏置。RoPE 成了稠密模型的默认。

还有一个从外部把整个堆叠框住的选择。嵌入与反嵌入位于两端：输入嵌入把词元 id 映射为向量，LM 头把最终的残差状态映射为词表 logits（模型在 softmax 把它们变成概率之前，为每个词元产生的原始分数）。权重绑定让两者共享同一个矩阵 (Press and Wolf 2017)。这是一处真实的取舍，而非默认。绑定省下 $\text{vocab} \cdot d$ 个参数，并对小模型起正则化作用，但一旦词表很大，这点节省就微不足道，于是大多数现代大模型都让两者不绑定。

这些选择装配起来，就是整个已定型的块。Vaswani 等交付的原始 Transformer 带后置归一化和普通前馈网络 (Vaswani et al. 2017)。后置归一化在深度上若无大量预热则会发散，于是前置归一化取而代之 (Xiong et al. 2020)。RMSNorm 弃掉了 LayerNorm 的均值减除，质量上没有可测代价 (Zhang and Sennrich 2019)。SwiGLU 在 $2/3$ 重新缩放下略胜 GeLU (Shazeer 2020)。RoPE 取代 ALiBi 成了位置默认 (Su et al. 2021)。到 LLaMA 与 Qwen 之时，共识已经定型：前置归一化、RMSNorm、SwiGLU 前馈网络、RoPE，以及任何地方都不带偏置项 (Touvron et al. 2023; Bai et al. 2023)。稠密块自 2023 年以来变化很少，主要变化转移到了数据、规模与后训练。

那唯一不肯定型的成本

现在轮到注意力本身，那个跨位置混合的子层，也是本章活跃约束的来源。缩放点积注意力把每个词元投影成一个查询、键与值，然后计算 (Vaswani et al. 2017)

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

其中 $Q$ 是查询，$K$ 是键，$V$ 是值，$QK^\top$ 给出每个位置对其他位置的相似度分数。softmax 把每行分数变成一组和为一的正权重，再对 $V$ 做加权求和，于是一个词元得到的是它从其他词元读回来的信息。$\sqrt{d}$ 这个除数防止 logit 随维度增长，免得 softmax 饱和。一个因果掩码把上三角清零，使一个词元只关注过去。成本在计算与分数矩阵两方面都是 $O(n^2)$，这正是下游每一个效率技巧背后的压力。

多头注意力把 $d$ 切成 $h$ 个头，让注意力能同时关注多个子空间，拼接它们的输出再投影一次 (Vaswani et al. 2017)。头在参数上成本很低，在推断时却有明确代价，因为每个头都携带自己的键和值，而这些必须保留下来。

保留它们不是细节，而是本章其余成本的来源。要看清原因，得弄清这个块到底如何把提示词变成文本：语言模型一次只生成一个词元，而且分两遍来做，两遍的形状大不相同。第一遍就是 预填充（prefill） 阶段：它一次读完整个提示词，每个位置都关注所有更早的位置，并一起算出全部提示词词元的键和值。第二遍就是 解码（decode） 阶段：它一次发出一个词元，每一步都把刚生成的词元作为下一步输入，以产生后续词元。解码是自回归的：一步的输出就是下一步的输入，所以序列每步增长一个位置。

解码正是保留键和值得到回报的地方。一个词元的键和值一旦算出便不再改变，因此与其在每一步对整个序列重算注意力（那是为每个词元重复一遍的二次方计算量），模型把它们缓存下来，每一步只为那个新词元对所有过去词元的缓存做注意力计算。缓存就是不重算所付的代价。这正是这些向量必须保留的原因，也是为什么它们的成本以内存而非计算来衡量。这一预填充与解码拆分在服务期的后果，它设定的延迟与它强加的调度，正是 第 31 章 的主题；这里只需明白：是解码给了缓存存在的理由。

这笔成本是本章的核心。缓存大小是

$$2 \cdot \text{layers} \cdot \text{heads} \cdot d_{\text{head}} \cdot \text{seq} \cdot \text{batch} \cdot \text{dtype}$$

这里前面的 $2$ 分别对应键和值，$\text{layers}$ 和 $\text{heads}$ 是每层每头都要保存一份缓存，$d_{\text{head}}$ 是每个头的宽度，$\text{seq}$ 与 $\text{batch}$ 分别是上下文长度和批大小，$\text{dtype}$ 是每个数占用的字节数。这个乘积随上下文长度与批大小线性增长，在长上下文下会超过模型权重本身。图 8.4 画出这一线性增长，以及它与一条水平权重线的交叉点，每个变体的斜率由它保留多少 KV 头决定。这个事实，正是解释 MQA、GQA 与 MLA 的关键。它们各自回答的问题是：在质量明显下降之前，能删掉多少 KV 缓存？

拿一个 7B 级模型试一试：把每个变体的 KV 缓存对上下文作图，读出它在哪里超过权重。改动 n_kv，看共享头如何把交叉点推远。

import numpy as np
import matplotlib.pyplot as plt

layers, d_head, dtype, params = 32, 128, 2, 7e9  # fp16 字节数，7B 权重
weights_gb = params * dtype / 1e9
seq = np.linspace(0, 200_000, 400)

for name, n_kv in [("MHA", 32), ("GQA", 8), ("MQA", 1)]:
    per_tok = 2 * layers * n_kv * d_head * dtype  # 每词元字节数
    cache_gb = per_tok * seq / 1e9
    crossover = weights_gb * 1e9 / per_tok
    print(f"{name:3s}: {per_tok/1024:6.0f} KB/词元，在 {crossover:>7.0f} 个词元处越过权重大小")
    plt.plot(seq / 1000, cache_gb, label=name)

plt.axhline(weights_gb, ls="--", color="gray", label=f"权重 ({weights_gb:.0f} GB)")
plt.xlabel("上下文长度（千词元）"); plt.ylabel("KV 缓存（GB）")
plt.legend(); plt.title("KV 缓存与上下文长度，批量=1"); plt.show()

第一种做法最激进。多查询注意力让所有查询头共享单个键/值头，把缓存削减 $h$ 倍 (Shazeer 2019)。削减是巨大的，但压到单个 KV 头会破坏训练稳定性，并让质量非平滑地下降。第二种做法介于两端之间。分组查询注意力在 $g$ 个查询头组内共享 K 与 V，其中 $g$ 是 KV 头的组数：一端恢复成 MHA（$g = h$），另一端恢复成 MQA（$g = 1$），并能从一个现有的 MHA 检查点廉价地续训 (Ainslie et al. 2023)。GQA 以多数节省换回多数质量，成了 LLaMA 与 Qwen 家族的默认 (Touvron et al. 2023; Bai et al. 2023)。图 8.5 展示了其机制：缓存为每个不同的 KV 头持有一份 K 与 V，所以节省完全来自让多少个查询头去共享同一个 KV 头。

多头潜在注意力在 DeepSeek-V2 中提出，从另一个角度处理同一笔成本 (DeepSeek-AI 2024)。它不共享头，而是把 K 与 V 向下投影到一个低秩潜在向量，用它来缓存，替代完整的逐头 K 与 V，并在使用时即时解压，旁边还并行携带一个解耦的 RoPE 分量。它的卖点是以 MHA 级质量，实现 GQA 级、甚至更好的缓存节省。它是 图 8.6 所勾勒的缓存压缩前沿上的第二条分支。

第三条分支更年轻，针对的是另一项成本。共享与压缩缩小的都是每个缓存词元的开销，可每个过去词元仍然被关注，$O(n^2)$ 的分数计算原封不动。稀疏注意力放弃了这个前提：每个查询只关注过去词元中被选出的一个子集，而不是全部，于是注意力计算量与长上下文成本一并被削减。给一个用全注意力训练好的模型事后套上稀疏模式，历来要折损质量，所以当前的设计把选择本身放进训练里学。DeepSeek 的原生稀疏注意力（NSA）让层级化的词元选择与模型其余部分联合训练 (Yuan et al. 2025)；它的后继者 DeepSeek 稀疏注意力（DSA）于 2025 年 9 月随生产模型 DeepSeek-V3.2-Exp 上线，用一个轻量索引器为每个查询挑出要读的词元 (DeepSeek-AI 2025)。稀疏对服务与长上下文行为意味着什么，留到 第 35 章 再谈；在这里，它是前沿的第三条分支，也是三者中第一条直接削减二次方计算本身的分支。

之所以在这里提到 FlashAttention 注意力内核（FlashAttention），是为了说明这些变体没有解决的是什么。它让注意力保持精确，又靠从不把 $O(n^2)$ 的分数矩阵写进高带宽内存，让它变得可负担，但它并不缩小 KV 缓存，因为那是一种内存占用，不是计算核 (Dao et al. 2022)。该核以及缓存分页的服务期机制，留到 第 33 章 与 第 32 章 再谈。

承载成本的架构参数

上面每个选择都是一个带拐点的平衡。并排摆开，它们就是架构师真正需要调整的架构参数。

• 归一化位置，前置还是后置。 后置归一化只要能训练起来，可以达到略好的最终损失，但在深度上若无大量预热与谨慎初始化就不稳定。前置归一化用一丝峰值质量，换来一条稳健可训练的梯度路径，在实践中胜出 (Xiong et al. 2020)。
• 归一化形式，LayerNorm 还是 RMSNorm。 均值居中带来的收益很小。RMSNorm 弃掉它，换来更低成本和少一组参数，几乎没有质量代价，如今为默认 (Zhang and Sennrich 2019)。
• 激活函数选择。 SwiGLU 的增益虽小但稳定，并且在 $2/3$ 隐藏维度重新缩放下几乎不增加参数量 (Shazeer 2020)。代价是多了第三个权重矩阵，以及对核函数与并行而言略不规整的形状。
• 位置方案是一道单向门。 基础方案在运行起始时选定，事后更改代价高昂。它约束着上下文随后能扩展多远、模型外推得多好。RoPE 的 theta 基数是一个小参数，长上下文后果却很大，其机制见 第 35 章。
• 权重绑定。 一处真实的取舍：参数节省与小模型正则化，对上规模时的边际成本和那份让大模型选择不绑定的自由 (Press and Wolf 2017)。
• 深度还是宽度。 在固定参数预算下，更深意味着更多的顺序组合，但训练与并行更难，伴随流水线气泡与串行延迟；更宽意味着每步更多算力，对张量并行更友好，也更易保持稳定。它在 第 10 章 中与系统层协同设计，而非孤立求解，其中每 FLOP 质量的一面由 第 5 章 设定。
• KV 缓存与质量，主导性的轴线。 MHA 以完整缓存给出完整质量，GQA 以缓存的一小部分给出多数质量，MQA 以质量与稳定性风险给出最小缓存，而 MLA 声称通过潜在压缩打破这一权衡。正确的取舍点取决于目标上下文长度与批量，因为缓存随两者一同扩展。

装配这个块，以及它如何失效

一旦上述选择定下来，装配共识稠密块大体上只是核对账目：在每个子层的输入上用 RMSNorm 做前置归一化、一个隐藏维度按 $2/3$ 重新缩放的 SwiGLU 前馈网络、在点积之前对查询与键施加 RoPE、为 KV 头做 GQA 分组，以及任何地方都不带偏置项。自 LLaMA 以来，几乎每一个开源稠密模型都能看出这一形状 (Touvron et al. 2023)。

失效模式值得逐一点名，因为它们各在不同时刻出现。深度上的后置归一化会在训练早期发散或停滞，这正是前置归一化被发明出来要修复的经典症状 (Xiong et al. 2020)。为短上下文选定的位置方案，或过小的 RoPE theta，会限制日后的长上下文行为，代价只在第一次提升上下文窗口时浮现。切换到门控前馈网络时忘了 $2/3$ 重新缩放，会悄悄放大或缩小参数预算，并破坏同等条件的比较。不绑定、未正则化的输出 logits 可能无界漂移，这正是 z-loss 防护它们的原因，机制见 第 5 章。而过度压缩到 MQA 会破坏训练稳定性、并让质量非平滑地下降，这正是 GQA 之所以存在的原因 (Ainslie et al. 2023)：最便宜的缓存，并不一定等于最好的缓存。

延伸阅读