Transformer扩展：架构改进

可视化演示

1. 注意力瓶颈

聚会规模还小的时候，“每个人和所有人聊一遍”的规则很温馨——信息最全，谁都不会被冷落。可当人数从几十涨到几千，这个规则就成了灾难：每个人要发出几千次对话请求，总对话次数按人数的平方爆炸式增长。Transformer 的标准注意力机制恰好就是这样一种“全聊”规则：每个词都要计算与所有词的关系，文章越长，代价越高。

在技术世界里，这个代价被写成一个公式。假设一篇文章有 \(n\) 个词，每个词用长度为 \(d\) 的向量表示，那么标准注意力需要做的计算量大约正比于 \(n^2 \times d\)。为什么是平方？因为每个词都要和其余 \(n\) 个词逐一配对打分，n 个词就有 \(n^2\) 个配对。维度 \(d\) 只是给每次配对打分时的工作量。

更麻烦的是内存：计算过程中需要把 \(n^2\) 个打分结果全部存下来。文章从 1 千词涨到 1 万词，光是这份“打分表”就要膨胀 100 倍。显存被占满后，模型要么跑不动，要么只能缩短输入，长序列建模便无从谈起。

回到聚会：如果一场聚会注定要到场几千人，“全员两两聊天”的原始规则一定撑不住。要想把聚会办下去，必须修改聊天的规则——这正是后面所有改进方案的出发点。

2. 标准注意力回顾

在改良规则之前，先弄明白原始的“全员聊天”到底在聊什么。班级聚会里，每位同学见面时会带着三样东西：自己要找的人的特征（称为查询 Q）、自己身上贴着的标签（称为键 K）、以及自己真正想说的内容（称为值 V）。两个人是否投缘，由 Q 和 K 的匹配程度决定；聊完之后带走多少信息，则由 V 决定。

Transformer 的注意力机制把这套流程数学化了。第一步，把 A 的 Q 与 B 的 K 做点积，得到一个相似度分数，代表“你们两个有多相关”。第二步，把所有分数除以 \(\sqrt{d}\) 做缩放，再经过 softmax 归一化成一组和为 1 的权重——就像把“好感度”换算成“我该花多少比例的时间听你说话”。第三步，用这些权重对所有人的 V 做加权平均，得到融合后的新表示。

这个“全员聊天”机制是后续所有改进的基线。它的优点是一视同仁、信息无损，缺点是开销太大——接下来的每一种改进，本质上都是在“砍掉不必要的聊天”。

3. 稀疏注意力

有经验的组织者很快发现：聚会上真正值得深聊的人并不多。大多数人只需要和两类人说话就够了——“班委”（少数全局重要的关键角色）和“同桌”（位置邻近的人）。于是他们制作了一份“通讯录名单”：名单上列了每个同学允许和谁聊天，名单外的一律免谈。这份名单，在数学上就叫掩码矩阵（Mask）。

稀疏注意力就是在标准注意力的打分矩阵上，套一个掩码。掩码中为 1 的位置允许计算注意力，为 0 的位置直接屏蔽。这样，原本 \(n^2\) 个配对中，只有少数位置需要真正计算，其余都被跳过，计算量随之大幅下降。

稀疏模式通常是结构化的：一部分 token 与全局固定的少数 token 相连（相当于认识班委），一部分与局部邻域相连（相当于认识同桌），还有一部分按固定规律做随机采样（相当于随缘认识几个新朋友）。这三种模式组合起来，就是经典的“全局 + 局部 + 随机”稀疏注意力。

回到聚会：有了“通讯录名单”，每个人只需要和自己名单上的人聊天，对话总数立刻从“全员互聊”降为“重点社交”。只要名单设计得合理，该聊的重要人物一个都没落下，聚会质量几乎不受影响。

4. 滑动窗口注意力

另一种更省事的做法是：干脆只和座位前后左右的几位同学聊天，离得远的就暂时不理。这就像排队买奶茶，你只需要和前面几个人确认一下排队顺序，不必和队尾的人喊话。这种规则被称为滑动窗口注意力：每个词只关注相邻窗口大小为 \(w\) 的 \(2w+1\) 个词。

这样，每个词的配对数量从 \(n\) 降到了 \(2w+1\)，总计算量从 \(n^2\) 降为 \(n \times (2w+1)\)，变成了随序列长度线性增长。窗口 \(w\) 是常数时，无论文章多长，开销都只是成比例地增加，不再平方爆炸。

但有人会担心：只看邻居，远方的信息不就永远传不过来了吗？答案藏在“多层堆叠”里。Transformer 有很多层，每一层只和邻居交流，但信息会像“传话游戏”一样逐层扩散——第一层传给邻居，第二层邻居再传给邻居的邻居。层数足够多时，序列两端的信息照样能汇合到同一个词上。

滑动窗口把“全员聊天”改成了“邻居聊天 + 层层传话”。虽然单次对话范围变窄了，但依靠层数堆叠，整场聚会依然能信息互通，而代价却从平方级降到了线性级，真正让超长文本成为可能。

5. 多查询注意力（MQA）

聚会开到一半又出现新问题：每个人都想记笔记，于是教室里坐着几十位“速记员”同时奋笔疾书，场面混乱且纸张浪费严重。多查询注意力（MQA）的想法很干脆——全班只保留一位速记员，所有人的谈话要点都交给这一个人记录，大家共用这一份笔记。

在 Transformer 里，标准注意力有 \(h\) 个“头”，每个头都维护自己独立的 K（键）和 V（值），相当于 \(h\) 位速记员各记各的。MQA 则让所有头共用同一份 K 和 V，只保留各自独立的 Q。这样，K/V 的数量从 \(h\) 份变成了 1 份，KV 缓存的内存占用直接缩为原来的 \(1/h\)。

为什么省内存能带来加速？因为在自回归生成时，模型要反复读取已经算好的 K/V 缓存。缓存越大，从显存搬运数据的时间就越长。MQA 把缓存大幅压缩后，显存带宽压力骤减，推理吞吐显著提升。代价是：所有头共享一套信息，表达多样性变弱，模型质量可能轻微下降。

全班共用一位速记员，纸张立刻省下了 \(h\) 倍，大家读笔记也快得多。但一份笔记要满足所有人的口味，难免丢失某些细节——这就是 MQA 用“质量换速度”的取舍。

6. 分组查询注意力（GQA）

既然一位速记员记全班容易遗漏细节，那能不能多配几位速记员，又不至于给每个人都配一个？分组查询注意力（GQA）给出的答案是：把全班分成若干小组，每组配一位速记员。组内成员共享同一份 K/V，组与组之间各自独立。

GQA 引入一个分组数 \(g\)：一共有 \(g\) 份 K/V 和 \(h\) 个 Q 头，每组有 \(h/g\) 个 Q 头共享同一份 K/V。当 \(g=1\) 时，GQA 就退化成 MQA（全班一位速记员）；当 \(g=h\) 时，每组只有一个 Q 头，GQA 又回到标准注意力（每人一位速记员）。因此 GQA 是夹在 MQA 与标准注意力之间的连续调节旋钮。

实际使用中，\(g\) 通常取 4 或 8，在“笔记够细致”和“纸张够省”之间取得平衡。例如 Llama 2 的 70B 模型就采用 GQA，用多组 K/V 保留了足够的信息多样性，同时把 KV 缓存压缩到标准注意力的几分之一，兼顾生成质量与推理吞吐。

分组速记员既避免了 MQA“一份笔记太粗”的问题，又比“人手一份笔记”省下大量纸张。GQA 的火爆正是因为它在质量与速度之间找到了一个性价比极高的甜蜜点。

7. 方案综合对比

现在，我们手里有了好几套聚会规则：全员互聊（标准注意力）、通讯录名单（稀疏注意力）、邻居聊天（滑动窗口）、一位速记员（MQA）、分组速记员（GQA）。到底该选哪套？答案是：看聚会规模和信息需求。

• 质量优先：标准注意力信息最全，适合短文本、对精度要求极高的任务。

• 超长文本：稀疏注意力和滑动窗口把计算量从平方级降为线性级，适合处理长文档。

• 高并发推理：MQA 把 KV 缓存压到最小，吞吐最高，适合部署在流量巨大的线上服务。

• 均衡之选：GQA 介于标准注意力与 MQA 之间，常用于需要“质量不差、速度要快”的大模型。

如果量化 KV 缓存的影响，可以定义一个缓存缩减比：标准注意力的 KV 头数为 \(h\)，GQA 为 \(g\)，缓存大小与头数成正比，缩减比为 \(h/g\)。

这场“聚会规则改良”的旅程告诉我们：没有任何一种方案是绝对最优的。标准注意力像一场坦诚而昂贵的大型聚会，稀疏与窗口方案像懂得取舍的社交高手，MQA 是极致效率的独行侠，GQA 则是大多数场景下最务实的组织者。理解这些权衡，正是 AI 工程师做模型设计与部署决策的核心能力。

AI 应用场景

Python 演示

"""Compare standard, sliding-window, MQA and GQA attention costs."""import math  def full_attn_flops(n: int, d: int) -> int:    """Standard full attention: QK^T and AV, each ~ n^2 * d."""    return 2 * n * n * d  def window_attn_flops(n: int, d: int, w: int) -> int:    """Sliding window attention: each row attends to 2w+1 positions."""    return 2 * n * (2 * w + 1) * d  def kv_cache_mb(layers: int, seq: int, d_head: int, kv_heads: int) -> float:    """KV cache = 2 K/V * layers * seq * d_head * kv_heads."""    total = 2 * layers * seq * d_head * kv_heads    return total * 4 / 1024 / 1024  def main() -> None:    n, d, w = 4096, 128, 64    full = full_attn_flops(n, d)    win = window_attn_flops(n, d, w)    print(f"seq_len={n}, d_model={d}, window={w}")    print(f"full attention flops : {full:>12,}")    print(f"window attention flops: {win:>12,}")    print(f"flops ratio          : {win / full:.1%}")    print()    print("KV cache per layer (seq=4096, d_head=128):")    for name, heads in [("standard(32)", 32), ("GQA(4)", 4), ("MQA(1)", 1)]:        mb = kv_cache_mb(1, n, d, heads)        print(f"  {name:<14} {mb:8.1f} MB")  if __name__ == "__main__":    main()

1. def full_attn_flops(n: int, d: int) -> int:定义标准注意力计算量函数，接收序列长度与维度。
2. return 2 * n * n * dQK^T 与 AV 两次矩阵乘法，各约 \(n^2d\) 次运算。
3. return 2 * n * (2 * w + 1) * d每行只看 \(2w+1\) 个邻居，复杂度降为线性。
4. def kv_cache_mb(...)计算 KV 缓存大小并换算为 MB，用于对比三种方案。
5. total = 2 * layers * seq * d_head * kv_headsK 和 V 各一份，乘以层数、序列长度、维度与头数。
6. for name, heads in [("standard(32)", 32), ...]用 32 头、4 组、1 头代表标准、GQA、MQA。

常见误区

面试题