第 26 章高效、稀疏与替代架构地图

26.1 究竟是哪项资源不够用了

红蓝钥匙的例子要求模型在很长的前缀末尾，用到开头出现的颜色。第 25 章调整了位置、归一化和前馈运算，却没有消除跨越长前缀传递信息的需要。前缀越长，生成时需要保存的上下文可能越多，位置之间的比较也可能越贵。若把前馈网络加宽，还会多出另一笔开销：更多参数需要存储和参与计算。

这是三个不同的问题。少存一些键和值，不等于减少查询位置；少比较一些位置对，不等于减少前馈参数；一个 token 只调用部分专家，也不等于其他专家已经从内存消失。我们先看改动针对哪项资源，再判断架构名称意味着什么。

回想第 17 章的几个对象。查询 Q 用来寻找应当参考的先前位置；键 K 提供比较向量；值 V 提供被加权混合的向量。每个注意力头都有相应的计算。逐步生成时，可以保留旧位置的 K、V，避免重新计算。第 30 章会沿时间线展开这种缓存，本章先算架构要求它保存多少东西。

26.2 八个查询头，不一定要配八组键和值

例 26.1：先数要保存多少个数。 给定一层模型：查询头有 8 个，每个头宽 64 个坐标，保留 1,024 个位置的前缀，每个坐标占 2 字节。这里只算一个请求的键和值，不含填充、临时缓冲区和元数据。

若每个查询头各有一组 K、V，一个位置要保存 $2\times8\times64=1{,}024$ 个数。前面的 2 表示 K 和 V 两类对象，还不是字节数。再乘 1,024 个位置和每个数的 2 字节，得 $1{,}024\times1{,}024\times2=2{,}097{,}152$ 字节。若为 K、V 再乘一次 2，就会重复计数，因为每个位置的 1,024 个数已经包含了它们。

表 26.1：查询头相同，保存的上下文不同；单层字节数由规则推导

组织方式	查询头数	KV 头数	每个位置保存的数	1,024 个位置所需字节
每个查询头各有 K/V	8	8	1,024	2,097,152 = 2 MiB
每四个查询头共用一组 K/V	8	2	256	524,288 = 0.5 MiB
全部查询头共用一组 K/V	8	1	128	262,144 = 0.25 MiB

这里 1 MiB 为 $2^{20}=1{,}048{,}576$ 字节。第一行是 多头注意力（multi-head attention，MHA） 。所有查询头共用 K/V 投影，称为 多查询注意力（multi-query attention，MQA） ；介于两者之间、分组共享的方式，称为 分组查询注意力（grouped-query attention，GQA） 。中间一行可以把第 1—4 个查询头分为一组，第 5—8 个分为另一组。查询仍然各不相同，因此，共用键不代表注意力权重也必须相同。

存储减少多少，可以从这些数直接推出；速度和质量却不能。模型仍要产生八个查询头的输出，而投影方式和学到的行为已经改变。通常不能只改配置字段，就把已有 MHA 检查点变成质量不变的 GQA 模型。转换或训练之后仍须评测。MQA、GQA 也能与稠密前馈网络并存：“共享 KV”和“稀疏使用参数”回答的是两回事。

26.3 少调用几个专家，不等于少加载几个专家

假设一个很大的前馈阶段准备了多套变换，再由一个小型可学习路由器选择当前 token 走哪几套。每套变换叫作一个 专家（expert） 。它是一组数值运算，并非一个人，也不保证天然擅长某个学科。混合专家（mixture of experts，MoE） 按路由权重组合所选专家的输出，让参数按条件参与计算。

例 26.2：读一张稀疏模型卡。 给定的教学模型卡包含 2,000 万个共享参数，已计入路由器和其他全部共享组件；另有四个专家，每个 1,000 万参数。每个 token 恰好选择两个专家。总参数为 $20+4\times10=60$ 百万，当前 token 使用的参数为 $20+2\times10=40$ 百万。这里每组参数只计一次，算的不是 FLOPs。

若每个权重占 2 字节，全部权重的理想存储量是 120 MB；这里 MB 表示一百万字节。本次选中的权重对应 80 MB。只要全部专家仍驻留在加速器上，权重就仍需 120 MB，激活和缓存还没有算进去。若只从其他设备加载被选中的专家，就用传输和调度开销换取驻留空间。一个批次里的不同 token 也可能合起来用到全部四个专家。

假如路由器几乎把所有 token 都送给专家 1，其他设备可能闲着，专家 1 却超出分配给它的容量。负载均衡或容量规则可以缓解问题，但可能增加训练目标、数据搬运或超额 token 的处理。原始 Switch 论文研究的是每次选一个专家；本章选两个专家的数字是原创教学设定。两种计数都不能单独预测质量。模型卡只写“活跃参数 4,000 万”还不够，必须补上总权重、路由规则、批次条件和运行时的放置方式。

26.4 少做比较，也会少掉直接的信息通路

另一笔开销来自位置对。长度为 6 的因果序列中，完整注意力每个头允许 $1+2+3+4+5+6=21$ 个查询—键位置对。这个三角形计数包含当前位置自身，不包含未来位置。

例 26.3：远处的线索落在窗口之外。 现在每个位置只能看自己和紧邻的前一个位置。第一个位置有一个允许的位置对，其余五个各有两个，共 $1+5\times2=11$ 个。在位置 6，直接连到的只有位置 5 和 6；位置 1 的红钥匙不在其中。

信息仍可能跨层经中间位置传过来。这个窗口每经过一层，最多把信息向后传一个位置；从位置 1 到位置 6，至少需要五层这样的连接。这只说明路径存在，不保证训练会沿途保存颜色。若指定某些全局位置，路径可以缩短，但必须说清哪些位置是全局位置，并在生成时保留因果限制。

这种按规则省去部分位置对的做法叫作 稀疏注意力（sparse attention） 。Longformer 的局部／全局模式提供了一个研究实例，但不能把它的双向模式直接当成本书解码器的因果 Mask。“高效注意力”的范围更广：某种实现也可能不删任何逻辑位置对，只减少计算相同完整注意力所需的内存搬运。第 31 章的 FlashAttention 就属于这条思路。稀疏连接和高效执行不应混在一个没有说明口径的分数里。

26.5 用摘要状态代替逐位置保存

能不能把长前缀压进固定大小的状态？先看一个刻意简化、不同于 softmax 注意力的规则。每个新项目有非负标量键权重 $k_i$ 和标量值 $v_i$ ，查询权重取 1。只保留两个数：加权值之和 S，以及键权重之和 Z；开始都为零。

例 26.4：三个项目，两个累加量。 按顺序给定键 $[1,2,1]$ 和值 $[4,1,3]$ ，对应位置配成一对。第一个项目到来后，S 为 $1\times4=4$ ，Z 为 1，输出为 4。第二个项目给 S 增加 $2\times1=2$ ，给 Z 增加 2，输出变为 $6/3=2$ 。第三个项目之后，输出为 $9/4=2.25$ 。

表 26.2：按规则推导的递归状态

已处理项目	本次对 S 的贡献	S	Z	输出 S/Z
第一个	4	4	1	4
第二个	2	6	3	2
第三个	3	9	4	2.25

把这个标量例子写成一般更新式：

S_i=S_{i-1}+k_i v_i,\qquad Z_i=Z_{i-1}+k_i,\qquad y_i=S_i/Z_i.

下标 i 表示“处理完第 i 个项目之后”。S 累加的是值乘权重，Z 累加的是权重，因此分母取 Z。只有 Z 为正时，这里的输出才有定义。要得到这种加权平均，不必保存完整列表。线性注意力（linear attention） 的向量形式会对 Q、K 使用选定的特征映射，再保存类似的累加量；状态大小取决于特征和值的维度。维度固定时，递归状态不会随前缀长度增长。

代价也很具体。把第一对和第三对键值交换，最终 S、Z 都不变；仅凭这个标量状态，无法分辨它们谁先出现。这也不是普通的 softmax 注意力，因为我们有意换了相似度规则。其他特征映射和位置机制可以保留其他区别，但“线性”二字不能保证精确找回每一个旧 token。

状态空间模型（state-space model） 或其他递归路线，也会随着 token 到来更新状态。刚才的加权和已经做过这种事：保留旧值，再加上经过加权的新值。标量示意式 $h_i=a_i h_{i-1}+b_i x_i$ 中，h 是用一个数保存的状态，x 是本次输入的数，a 决定旧状态乘几倍，b 决定新输入乘几倍。把 a 设为 1、b 设为当前键权重、x 设为当前值，就得到前面的 S 更新；改变 a，则可以逐步减弱旧贡献。这个标量示意并不是完整的 Mamba 架构。Mamba 一类选择性状态空间模型让某些更新量随输入而变，从而影响哪些信息留下来。它们仍须在远距离钥匙和政策更新任务上接受检验。固定大小的状态既不是外部文档库，也不保证信息能够长期可靠地保存。

26.6 沿独立的比较轴读地图

共享保存的上下文

一组数字8 个查询头；8 个 KV 头
处理步骤查询头 1—4、5—8 分成两组
一组数字2 个 KV 头：2 MiB → 0.5 MiB

不同查询仍计算各自的混合结果。

按路由使用参数

模型保存的数值共享 2,000 万 + 4 × 1,000 万专家参数
处理步骤一个 token 选择两个专家
模型保存的数值活跃 4,000 万；总计 6,000 万

全部专家加载时，每个权重 2 字节仍需 120 MB。

限制位置对

处理步骤完整因果连接：1+2+3+4+5+6 = 21
处理步骤自己与前一位：1+5×2 = 11
一组数字位置 6：只直接读取位置 5、6

位置 1 的线索需要跨层传递或其他连接。

图 26.1：三种节省，留下不同要求。本书原创图示 · CC BY-SA 4.0。

图 26.1 把三种资源减少与仍需承担的要求放在一起。每条箭头都对应实际计数或信息约束，没有哪条箭头表示某条路线必然全面替代另一条。

表 26.3：读模型卡时，应当写在旁边的问题

模型卡的说法	定位了什么	还要查什么
使用 GQA 的稠密 Transformer	KV 共享，前馈网络不做专家路由	KV 维度、质量、可用运行实现
注意力模型配路由前馈专家	参数按条件参与	总参数／活跃参数口径、路由和传输
局部注意力	位置连接稀疏	窗口、全局位置、因果路径和任务质量
核函数形式的线性注意力	换用可累加状态的交互规则	特征／状态大小、数值稳定性和回忆能力
选择性递归／状态空间模块	这些模块更新状态，不保存完整 KV 历史	状态容量、训练行为、是否混入注意力层

混合架构可以同时占据几行。“稠密／稀疏”必须说明是在谈每个 token 使用的参数，还是位置对的连接；“注意力／替代路线”必须说明是在谈全部模块，还是其中一部分。活跃参数、存储字节和算术操作次数，应分别列出。

配套标量程序复现 data/part-v/mechanisms.json 中的计数，不实现或测试这些工业架构。完成模型卡标注，意味着读者能指出共享方式、路由方式和上下文表示，找出证据缺口，并拒绝把活跃参数数目直接当成内存保证。即使选好了架构，训练也远不止要装下权重。下一章继续列这份资源账。

26.7 本章小结

不同架构改动针对不同开销。MQA/GQA 共享 K/V，MoE 选择参数子集，稀疏注意力省去位置对，线性和递归路线以不同方式概括历史。每项节省都有自己的计数口径，也留下相应的信息或系统要求。比较时先分清它们，再用目标任务检验。

26.8 习题

把例 26.1 保留的位置数加倍，并改用四个 KV 头。计算单层字节数。保留八个查询头，是否还要把结果再乘 8？
MoE 模型卡改为每次选择一个专家。计算活跃参数，以及所有专家都加载时权重的理想驻留字节数。说明为什么这是两种问题。
例 26.3 改为可以看当前位置和之前两个位置。计算六个位置允许的全部位置对，以及从位置 1 连到位置 6 至少需要几层。
在例 26.4 后追加键 2、值 6，计算新状态和输出。某实现用项目数量作分母，错在哪里？
一张模型卡宣传“稀疏、活跃参数 4,000 万、长上下文”。要判断红钥匙任务是否放得下、做得对，还需提出哪三个问题？

26.9 参考解答

$2\times2{,}048\times4\times64\times2=2{,}097{,}152$ 字节，即 2 MiB。K、V 按 KV 头保存，八个不同的查询不会额外复制八份缓存。
活跃参数变成 $20+10=30$ 百万，总参数仍是 6,000 万。每个权重 2 字节时，理想驻留量仍为 120 MB。每个 token 少选一些权重，不代表其余权重已经卸载。
各行位置对数为 1、2、3、3、3、3，总计 15。每层最多前进两个位置，跨越五个位置的距离至少要三层，例如 1→3→5→6。这只建立了连通性。
S 变为 $9+2\times6=21$ ，Z 变为 $4+2=6$ ，输出为 3.5。若除以四个项目，就会得到 5.25，忽略了权重不同。Z 数的是权重，不是记录条数。
先问稀疏的是哪一部分、总共需要驻留多少权重；再问 KV／状态维度、精度和并发上下文长度；最后问在目标长度、这类远距离线索上有什么质量证据。允许输入多长，并不能回答是否做得对。

26.10 参考文献

Shazeer，Fast Transformer Decoding: One Write-Head is All You Need，2019。第 2.4—3 节解释 MQA 如何共享键和值。
Ainslie 等，GQA，2023。第 2.2 节定位 MHA、MQA 之间的分组方式，并说明模型转换与继续训练。
Fedus 等，Switch Transformers，2022。第 2 节说明专家路由和容量问题，其每次选择一个专家的设计与本章设定不同。
Beltagy 等，Longformer，2020。第 3 节介绍局部／全局注意力模式；本章因果窗口是原创比较例。
Katharopoulos 等，Transformers are RNNs，2020。第 3.2—3.4 节推导使用累加状态的线性注意力。
Gu、Dao，Mamba，v2。第 3.1—3.3 节联系选择性状态更新、信息压缩和实现开销。

第 26 章 高效、稀疏与替代架构地图