Transformer 架构及其变体

第 7 章 最终得到了一套稳定的文本到词元 ID 映射。本章继续追踪这些 ID 如何经过仅解码器 Transformer：先进入嵌入层，再穿过一叠残差块，经过最终归一化，最后投影到词表。2017 年提出的 Transformer 在编码器-解码器系统中引入了注意力与前馈网络组成的基本块 (Vaswani et al. 2017)。现代因果语言模型沿用这套结构，但在归一化、位置处理、前馈网络、注意力头和缓存表示上各有不同。

这些差异并非表面风格。它们决定哪些位置可以交换信息、梯度怎样穿过深层网络、每个块包含多少参数，以及服务进程要为每条活跃序列保留多少状态。因此，一份有用的架构说明必须写明张量形状、掩码语义和内存单位，不能只画一个写着“Transformer”的方框。

范围：因果解码器

这里设一个批次包含 $B$ 条长度为 $S$ 的序列，模型宽度为 $d$。分词器从大小为 $V$ 的词表中输出整数 ID： $T\in\{0,\ldots,V-1\}^{B\times S}$。嵌入表把它映射为初始残差流 $X_0\in\mathbb{R}^{B\times S\times d}$。经过 $L$ 个块和最终归一化后，输出矩阵生成覆盖整个词表的 logits：

$$\begin{aligned} \widehat{X}_L &= \operatorname{Norm}_{f}(X_L), \\ Z &= \widehat{X}_LW_U, \\ p(t_{i+1}=v\mid t_{\le i}) &= \operatorname{softmax}(Z_i)_v. \end{aligned}$$

其中，$X_L$ 是最后一个块的输出，$\operatorname{Norm}_{f}$ 是最终归一化，$W_U\in\mathbb{R}^{d\times V}$ 是输出矩阵，也称反嵌入矩阵，$Z_i\in\mathbb{R}^{V}$ 是位置 $i$ 的 logits，$v$ 是候选的下一词元 ID。因果条件 $t_{\le i}$ 表示位置 $i$ 只能依赖位置 $i$ 及其之前的词元。编码器模型使用双向掩码，编码器-解码器模型还会加入交叉注意力。它们与这里的模型共享若干部件，但信息流契约并不相同。

输入嵌入与 $W_U^\top$ 可以共享参数，这种设计叫作权重绑定 (Press and Wolf 2017)。绑定会省去一张 $Vd$ 矩阵，但也要求输入与输出两种用途共用同一组参数行。这是一项架构选择，不是由模型大小决定的规则。

一个块会更新同一条残差流两次

常见的前置归一化解码器块包含两次残差更新：

$$\begin{aligned} A_\ell &= \operatorname{Attn}_\ell(\operatorname{Norm}_{\ell,1}(X_\ell)), \\ U_\ell &= X_\ell+A_\ell, \\ F_\ell &= \operatorname{FFN}_\ell(\operatorname{Norm}_{\ell,2}(U_\ell)), \\ X_{\ell+1} &= U_\ell+F_\ell. \end{aligned}$$

其中，$\ell\in\{0,\ldots,L-1\}$ 是块索引，$X_\ell$ 是块的输入，$A_\ell$ 和 $F_\ell$ 分别是注意力与前馈网络产生的更新，$U_\ell$ 是加入注意力更新后的状态，$X_{\ell+1}$ 是块的输出。注意力在掩码允许的位置之间混合信息，前馈网络（FFN）则独立变换每个位置。两者都返回与残差流形状相同的张量，残差相加才有定义。

残差流（residual stream）是穿过整个堆叠的宽度为 $d$ 的状态。称它为“流”有助于理解信息怎样逐层更新，但它不是独立的记忆模块，只是每个子层都会读取和修改的一组隐藏向量。

图中采用的是前置归一化。原始 Transformer 则在每次残差相加后做归一化 (Vaswani et al. 2017)：

$$\begin{aligned} A_\ell &= \operatorname{Attn}_\ell(X_\ell), \\ U_\ell &= \operatorname{Norm}_{\ell,1}(X_\ell+A_\ell), \\ F_\ell &= \operatorname{FFN}_\ell(U_\ell), \\ X_{\ell+1} &= \operatorname{Norm}_{\ell,2}(U_\ell+F_\ell). \end{aligned}$$

这里各符号的形状与前置归一化块相同，只是归一化器移到了残差相加之后。Xiong 等人解释了为什么前置归一化在初始化时具有更平稳的梯度，并可能减少对学习率预热的依赖 (Xiong et al. 2020)。这不代表后置归一化无效，初始化、残差缩放、深度和优化器设置都会改变比较结果。应记录准确的块方程，而不能只写“前置归一化”或“后置归一化”。

归一化控制尺度

LayerNorm 会把隐藏向量居中，并除以其标准差 (Ba et al. 2016)。均方根层归一化（RMSNorm）由 Zhang 和 Sennrich 于 2019 年提出，它省去均值减除，按均方根做归一化 (Zhang and Sennrich 2019)：

$$\begin{gathered} r(x)=\frac{x}{\sqrt{\frac{1}{d}\sum_{j=1}^{d}x_j^2+\varepsilon}}, \\ \operatorname{RMSNorm}(x)=\gamma\odot r(x). \end{gathered}$$

其中，$x\in\mathbb{R}^{d}$ 是一个词元的隐藏向量，$r(x)$ 是按均方根缩放后的向量，$x_j$ 是第 $j$ 个分量，$\gamma\in\mathbb{R}^{d}$ 是学到的逐元素缩放，$\varepsilon>0$ 用来防止除零，$\odot$ 表示逐元素乘法。实现必须固定 $\varepsilon$、累加精度，以及是否包含可学习偏置。这些细节会影响检查点兼容性和低精度稳定性。

有些架构还会在查询与键做点积前分别归一化。QK 归一化解决的问题不同于残差流归一化，它限制的是注意力分数的尺度。模型清单应分别写明这两处归一化。

前馈网络提供逐位置容量

原始 Transformer 使用两张矩阵，并在中间放置 ReLU。常见的门控替代方案是 Swish 门控线性单元（SwiGLU） (Shazeer 2020)：

$$\begin{gathered} G(x)=\operatorname{SiLU}(xW_g)\odot(xW_u), \\ \operatorname{SwiGLU}(x)=G(x)W_d. \end{gathered}$$

其中，$x\in\mathbb{R}^{d}$ 是一个残差向量，$G(x)$ 是逐元素门控后的中间表示，$W_g,W_u\in\mathbb{R}^{d\times d_f}$ 分别是门控投影和值投影，$W_d\in\mathbb{R}^{d_f\times d}$ 把结果投影回模型宽度，$d_f$ 是 FFN 隐藏宽度，$\operatorname{SiLU}(z)=z/(1+e^{-z})$ 按元素计算。忽略偏置时，这种门控 FFN 有 $3dd_f$ 个参数，两矩阵前馈网络有 $2dd_f$ 个参数。

这个计数解释了常说的“三分之二”调整。隐藏宽度为 $4d$ 的普通 FFN 有 $8d^2$ 个参数。要让三矩阵门控 FFN 与之匹配，需要在按硬件友好倍数取整前令 $d_f=8d/3$。这个比例是参数预算规则，不是 SwiGLU 的固有属性。实际模型会选择不同的扩展比例；矩阵参数量相同，也不保证不同内核或硬件上的运行时间相同。

注意力与 FFN 的作用不同，但“一个负责存知识、另一个只负责路由”属于解释，不是张量不变量。消融实验可以定位行为，前向方程本身不能指定语义归属。

带显式形状的因果自注意力

对一条序列，设 $X\in\mathbb{R}^{S\times d}$，并假定查询头、键头和值头的宽度相同。令 $H_q$ 为查询头数，$H_{kv}$ 为不同的键值头数，$d_h$ 为单个头的宽度。学到的投影产生

$$\begin{aligned} Q &\in \mathbb{R}^{S\times H_q\times d_h}, \\ K,V &\in \mathbb{R}^{S\times H_{kv}\times d_h}. \end{aligned}$$

每个查询头 $a$ 都会映射到一个键值头 $g(a)$。对查询位置 $i$，其输出为

$$\begin{aligned} s_{ij}^{(a)} &=\frac{q_i^{(a)\top}k_j^{(g(a))}}{\sqrt{d_h}} \\ &\quad+M_{ij}, \\ \alpha_{ij}^{(a)} &=\operatorname{softmax}_{j}\bigl(s_{ij}^{(a)}\bigr), \\ z_i^{(a)} &=\sum_{j=1}^{S}\alpha_{ij}^{(a)}v_j^{(g(a))}. \end{aligned}$$

其中，$q_i^{(a)}$ 是位置 $i$ 上查询头 $a$ 的查询，$k_j^{(g(a))}$ 和 $v_j^{(g(a))}$ 是位置 $j$ 上选中的键和值，$s_{ij}^{(a)}$ 是加入掩码后的分数，$\alpha_{ij}^{(a)}$ 是归一化后的注意力权重，$M$ 是掩码。普通因果注意力在 $j\le i$ 时取 $M_{ij}=0$，在 $j>i$ 时取 $M_{ij}=-\infty$。负无穷项经过 softmax 后概率为零。填充、打包文档、滑动窗口和前缀语言模型需要不同的掩码。

在常见初始化假设下，除以 $\sqrt{d_h}$ 可以防止分数方差随头宽度增长 (Vaswani et al. 2017)。把 $H_q$ 个查询头的输出拼接起来，再经过输出投影，就得到宽度为 $d$ 的残差更新。

当 $H_qd_h=d$ 且省略偏置时，查询投影和输出投影合计有 $2d^2$ 个参数，键投影和值投影合计有 $2dH_{kv}d_h$ 个参数。标准 MHA 因此使用 $4d^2$ 个投影参数。GQA 与 MQA 会减少 K/V 投影项和缓存，但稠密查询-键交互仍按 $S^2H_qd_h$ 扩展。

完整序列上的稠密注意力层会形成 $S^2$ 个查询-键关系，因此算术量随 $S$ 的平方增长。朴素实现还会显式生成 $S\times S$ 分数矩阵，但数学定义并不要求保留这块临时存储。

位置信息通过查询和键进入注意力

不带掩码、只看内容的自注意力具有置换等变性，打乱输入位置只会同样打乱输出。因果掩码虽然规定了方向，却没有为相同内容向量显式表示距离。旋转位置嵌入（RoPE）会按位置旋转查询和键中成对的坐标 (Su et al. 2024)。设参与旋转的宽度 $d_R$ 为偶数，定义

$$\begin{aligned} \omega_r &= \theta^{-2r/d_R}, \\ R_p^{(r)} &= \begin{bmatrix} \cos(p\omega_r)&-\sin(p\omega_r)\\ \sin(p\omega_r)&\cos(p\omega_r) \end{bmatrix}, \\ &\hspace{3em}0\le r<d_R/2. \end{aligned}$$

其中，$\theta>0$ 是频率基数，$r$ 表示一对旋转坐标，$\omega_r$ 是该坐标对的角频率，$p$ 是词元位置。把所有 $R_p^{(r)}$ 组成分块对角矩阵 $R_p$，则

$$\begin{aligned} q_i' &= R_iq_i, \\ k_j' &= R_jk_j, \\ q_i'{}^\top k_j' &= q_i^\top R_{j-i}k_j. \end{aligned}$$

这里 $q_i$ 和 $k_j$ 是旋转前的查询与键，$R_i$ 和 $R_j$ 使用配置好的频率方案。最后一个恒等式表明，点积中的位置部分取决于相对偏移 $j-i$。这并不意味着完整注意力分数只由距离决定，内容向量仍然参与计算。

频率基数、旋转维度占比，以及上下文扩展采用的任何缩放规则，都属于检查点契约。训练后更改它们属于适配实验，而不是元数据修改。RoPE 很常见，但不是唯一有效的位置方案。ALiBi 会对每个因果上有效的分数加入按头设置的线性距离惩罚 (Press et al. 2022)：

$$\begin{aligned} s_{ij}^{(a)} &=\frac{q_i^{(a)\top}k_j^{(a)}}{\sqrt{d_h}} \\ &\quad-m_a(i-j)+M_{ij}. \end{aligned}$$

其中，$m_a>0$ 是头 $a$ 的固定斜率，对允许访问的历史位置有 $i-j\ge0$，$M_{ij}$ 是前文定义的因果掩码。ALiBi 不引入可学习的位置向量。无论采用哪一种方法，长上下文表现都取决于训练长度、频率或斜率配置、后续适配和评估分布。

预填充创建缓存，解码复用缓存

自回归服务会显露出整序列训练图没有展示的两个阶段。第一遍就是 预填充（prefill） 阶段：它一次读完整个提示词，在因果掩码下生成各位置的 logits，并把提示词的键和值写入缓存。最后一个提示词位置的 logits 用于预测第一个生成词元。不同位置可以并行计算，尽管服务系统可能把很长的预填充分块执行。第二阶段是 解码（decode）：每条序列一次处理一个新位置，让新查询读取已经保存的历史键和值，选择或采样一个词元，把该位置的键和值追加到缓存，再重复这个过程。

在前缀固定的标准因果自注意力中，某一层过去位置的键和值不会因后续解码而改变。复用它们可以避免反复投影整个前缀。不过，稠密解码的每一步仍要读取保留的缓存，因此缓存既占容量，也消耗带宽。对长度不齐的批次，持久键值（KV）张量的有效载荷记作

$$\begin{aligned} C_\ell &= H_{kv,\ell}(d_{k,\ell}+d_{v,\ell}), \\ M_{\mathrm{KV}} &= b\sum_{r=1}^{B}S_r \sum_{\ell=1}^{L}C_\ell\quad\text{字节}. \end{aligned}$$

其中，$C_\ell$ 是第 $\ell$ 层每个位置保存的键值元素数，$B$ 是驻留序列数，$S_r$ 是第 $r$ 条序列的缓存长度，$L$ 是注意力层数，$H_{kv,\ell}$ 是第 $\ell$ 层不同的 KV 头数，$d_{k,\ell}$ 和 $d_{v,\ell}$ 是每个头的键和值元素数，$b$ 是每个缓存元素占用的字节数。若所有序列长度都为 $S$，并且 $d_k=d_v=d_h$，公式可化为 $2LBSH_{kv}d_hb$。真实服务还会产生分配块、元数据、填充或碎片、量化尺度、推测解码或束搜索状态、前缀共享和卸载等开销。上式计算的是稠密张量的有效载荷，不是进程总内存。

复用缓存后，每层为一个新词元执行稠密注意力的成本与历史长度 $S$ 线性相关，这里省略头数和维度因子。若提示词长度为 $P$，之后生成 $T$ 个词元，整个解码过程的注意力工作量与 $TP+T^2$ 成正比。缓存复用避免了重复投影前缀，却不会让长解码变成常数时间。

缓存随驻留词元数线性增长。它是否超过模型权重，取决于上述所有变量，以及权重和缓存精度。因此，不能把某一个交叉点说成 Transformer 的普遍属性。

图 8.4 采用明确假设：32 层、头宽度 128、每个 KV 元素两字节、批大小为一，以及 70 亿个两字节权重参数。MHA、GQA 与 MQA 三条线只改变 $H_{kv}$，图中每个数值都直接来自缓存公式。

下面的可执行单元会输出每词元缓存有效载荷，以及缓存等于所述权重字节数时的上下文长度。可以修改批大小、精度或 KV 头数，测试不同的部署配置。

layers = 32
head_width = 128
bytes_per_kv_element = 2
batch_size = 1
parameter_count = 7_000_000_000
bytes_per_weight = 2
weight_bytes = parameter_count * bytes_per_weight

def kv_payload_bytes(sequence_length, kv_heads):
    return (
        2
        * layers
        * batch_size
        * sequence_length
        * kv_heads
        * head_width
        * bytes_per_kv_element
    )

for name, kv_heads in [("MHA", 32), ("GQA", 8), ("MQA", 1)]:
    bytes_per_token = kv_payload_bytes(1, kv_heads)
    crossover_tokens = weight_bytes / bytes_per_token
    cache_at_128k_gib = kv_payload_bytes(128_000, kv_heads) / 2**30
    print(
        f"{name}: {bytes_per_token / 1024:.0f} KiB/词元；"
        f"128k 缓存={cache_at_128k_gib:.2f} GiB；"
        f"在 {crossover_tokens:,.0f} 个词元时等于权重字节数"
    )

MHA、GQA 与 MQA 改变 KV 头的共享方式

其中，多头注意力（MHA）使用 $H_{kv}=H_q$，每个查询头都有自己的键头和值头 (Vaswani et al. 2017)。多查询注意力（MQA）使用 $H_{kv}=1$，所有查询头共享一组键和值 (Shazeer 2019)。分组查询注意力（GQA）取中间值 $1<H_{kv}<H_q$，把多个查询头映射到同一个 KV 头 (Ainslie et al. 2023)。

在稠密有效载荷公式下，其他维度不变时，从 MHA 改为 GQA 会把缓存缩小 $H_q/H_{kv}$ 倍；改为 MQA 则缩小 $H_q$ 倍。这些是精确的存储比例，不是质量结论。Ainslie 等人报告称，可以把 MHA 检查点使用原始预训练算力的 5% 续训为 GQA，并在其实验中取得接近 MHA 的质量和接近 MQA 的速度 (Ainslie et al. 2023)。这说明 GQA 是有用的设计点，但不能证明某个 KV 头数适合所有模型和上下文长度。

MLA 改变缓存表示

DeepSeek-V2 引入的多头潜在注意力（MLA）主要不是通过共享传统 KV 头来节省缓存 (DeepSeek-AI 2024)。它把每个残差向量投影成低维潜在向量 $c_i^{KV}\in\mathbb{R}^{d_c}$，再从中生成内容键和值。系统还会保留一小段独立的 RoPE 键分量，因为位置相关的旋转使完整键投影无法全部吸收到后续矩阵乘法中。

在 DeepSeek-V2 的构造里，每个位置的有效载荷为

$$M_{\mathrm{MLA}}=LBS(d_c+d_h^R)b\quad\text{字节}.$$

其中，$d_c$ 是联合 KV 潜在宽度，$d_h^R$ 是解耦的 RoPE 键宽度，$L$、$B$、$S$ 和 $b$ 与稠密缓存公式中的含义相同。DeepSeek-V2 取 $d_c=4d_h$、$d_h^R=d_h/2$，即每个位置、每层缓存 $4.5d_h$ 个元素。传统等宽 MQA 缓存的是 $2d_h$ 个元素，所以 MLA 并非普遍比 MQA 更小。优化后的解码可以把上投影矩阵吸收到查询投影和输出投影中，无需显式重建完整的逐头 K 和 V。

DeepSeek 报告称，相比其 MHA 基线，MLA 所需缓存更少，并在受控消融列出的大多数指标上表现更好。这些结果支持论文中的架构和训练配方，但不是潜在注意力支配 GQA 的普遍定理。

区分算术量、临时内存与持久状态

多种技术都被称为“高效注意力”，但它们改变的是不同资源：

FlashAttention 注意力内核（FlashAttention）会分块计算精确注意力，因此无需把完整分数矩阵写入高带宽内存 (Dao et al. 2022)。它减少训练和预填充阶段的 IO 与临时激活内存，却不会自行改变解码时保留的历史键值元素数量。

稀疏注意力会改变数学连接方式。Native Sparse Attention 同时训练三条带门控的分支：压缩的粗粒度词元、选中的细粒度块，以及局部滑动窗口 (Yuan et al. 2025)。DeepSeek Sparse Attention（DSA）是另一种较晚的设计。它先用低维 lightning indexer 为历史词元排序，再让主 MLA 路径读取选出的 top-$k$ 子集 (DeepSeek-AI 2025)。主要注意力从全部 $S$ 个历史位置转向 $k\ll S$ 个位置，但 DeepSeek-V3.2 报告的索引器仍以较小维度计算所有查询与键的分数，因此仍是二次复杂度。选择器错误、继续训练成本和保留的缓存状态也必须计入。稀疏选择不会自动删除完整缓存。

把架构记录为一份契约

一份可复现的解码器规格至少应写明：

• 词表规模、嵌入宽度、输出投影形状和权重绑定方式；
• 块数、模型宽度、FFN 宽度、激活函数，以及所有偏置设置；
• 残差顺序、归一化位置、类型、epsilon 与计算精度；
• 查询头数、KV 头数、头宽度，以及查询到 KV 的映射；
• 注意力掩码、文档边界行为，以及任何局部或稀疏模式；
• 位置机制、RoPE 基数与旋转维度，以及缩放规则；
• 缓存表示、元素类型、量化尺度粒度，以及每词元有效载荷字节数；
• 最终归一化，以及任何 logit 缩放或裁剪。

“类似 LLaMA”或“使用 GQA”都不足以确定这些字段。序列化模型配置、检查点张量、训练代码和服务引擎必须遵守同一份契约。

在完整训练前完成验证

1. 核对形状与参数量。 推导每个投影的形状，并与序列化张量中的参数总量核对，其中包括共享参数。
2. 测试因果性。 改动未来词元，确认更早位置的 logits 不变；若模型使用填充、打包文档或前缀掩码，也要覆盖这些情况。
3. 比较参考注意力实现。 在小张量上，把优化内核与直观的高精度实现比较，并声明数值容差。
4. 测试缓存等价性。 在每个位置确认逐词元解码 logits 与完整前缀前向传播在容差内一致。
5. 测量缓存字节数。 在不同批大小和序列长度下，对照声明公式比较已分配内存与有效载荷。
6. 覆盖位置边界。 测试训练长度、部署上限，以及前缀缓存或序列打包使用的位置偏移。
7. 消融架构选择。 在参数量、词元数和算力预算匹配的条件下，比较归一化、FFN、头共享与位置方案。
8. 做端到端基准测试。 报告训练吞吐量、预填充延迟、解码延迟、最大驻留批大小，以及目标工作负载上的质量。

Transformer 块可以压缩成一页图，但它的契约从词元 ID 一直延伸到服务内存。只有把数学运算、张量布局和运行时表示分开说明，才能得到真正精确的架构定义。

延伸阅读