🤖 machine learning

Grouped Query Experts: Mixture-of-Experts on GQA Self-Attention

该论文提出了分组查询专家（Grouped Query Experts, GQE），这是一种应用于分组查询注意力机制的混合专家层，它通过在每个 token 层面动态选择查询头专家的子集，同时保持键值头（key-value heads）的稠密性，从而在不牺牲下游准确性的情况下，减少了活跃计算量并提高了效率。

原作者： Vishesh Tripathi, Abhay Kumar

发布于 2026-06-23

📖 1 分钟阅读☕ 轻松阅读

CC0 1.0

原作者： Vishesh Tripathi, Abhay Kumar

原始论文根据 CC0 1.0（http://creativecommons.org/publicdomain/zero/1.0/）发布到公有领域。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一座巨大的图书馆（一个 Transformer AI 模型），其中的每一本书（一个文本 token）都需要与其他每一本书进行交叉引用，才能理解整个故事。这就是“自注意力机制”（self-attention）的工作方式。

问题：“一刀切”的图书管理员
在标准的 AI 模型中，有一组由 16 名专业图书管理员组成的团队（称为“注意力头”，attention heads）。无论你从书架上取下哪本书，这 16 名专家都必须阅读、分析并撰写报告。

如果这本书是一个简单的词，比如“the”或者一个逗号，你并不需要 16 位专家；一两个就足够了。
如果这本书是一个复杂的科学术语，你可能确实需要全部 16 位专家。
但目前的系统强制要求这 16 名专家对每一个词都进行处理。随着故事变得越来越长（长上下文），这会变得极其缓慢且昂贵，就像雇佣了一整个管弦乐团来演奏一个单音符。

现有的解决方案：分组查询注意力 (GQA)
在此论文之前，研究人员尝试通过对图书管理员进行分组来节省成本。他们没有使用 16 个独特的团队，而是有了 8 个团队，其中每两个图书管理员共享同一份“键与值”（KV heads）笔记（即参考资料）。这节省了一些内存，但所有面的 16 名图书管理员仍然必须阅读每一个词。这就像是换了一个更小的文件柜，但仍然让每一位员工都要走遍整个大楼去完成工作。

新的解决方案：分组查询专家 (GQE)
作者提出了一种更聪明的系统，称为分组查询专家 (GQE)。可以将其想象为将这些 16 名图书管理员变成了一个“混合专家”（Mixture of Experts）系统，但带有一个转折。

设置： 他们保持 8 组参考笔记（KV heads）完全不变。这部分始终是“稠密”的（全量激活），因为获取参考资料的成本较低。
路由 (The Router)： 在每个组内部，有多个“专家”图书管理员（查询头，query heads）。对于每一个单词，一个聪明的“路由”（交通警察）会观察这个词并询问：“我们真的需要这个组里的全部 4 位专家来阅读这个词吗？”
选择： 路由会为这个特定的词挑选出前 1 或 2 位专家（k=1 或 k=2）来执行实际工作。组内的其他专家则去喝杯咖啡休息一下。
安全网： 为了确保系统不会变得混乱或偷懒，他们增加了两个特殊功能：
- 共享头 (The Shared Head)： 会有一位图书管理员始终在岗，阅读每一个词，以保持团队的稳定性。
- 加权摘要 (The Weighted Summary)： 系统会创建一个“共识报告”，将所选专家的工作进行融合。这至关重要，因为它为路由提供了一个清晰的信号，让它知道自己的表现如何，从而让它能够学习哪些专家最适合处理哪些词。

结果：在不损失智能的情况下实现提速
论文在了一个拥有 2.5 亿参数、并在 300 亿个词汇上训练的小型模型上进行了测试。

准确度： GQE 模型的表现与旧有的、对每个词都使用全部 16 名专家的模型一样出色。通过跳过部分专家，它并没有变得“更笨”。
速度： 由于跳过了大约一半的查询头工作，它变得更快了。
- 对于短篇故事，它稍微快了一些（1.15 倍）。
- 对于非常长的故事（如一整本小说），它的速度提升了 1.7 到 1.8 倍。

为什么这很重要
该论文声称，通过使“阅读”部分变为条件性的（仅在需要时才进行工作），同时保持“参考”部分保持恒定，可以在不牺牲答案质量的前提下，显著提高长对话或文档分析的速度。

局限性 (The Catch)
作者谨慎地指出，这是在一个相对较小的模型上进行的测试。他们尚未证明这在万亿级参数的巨型模型上同样有效。此外，“路由”需要经过非常精心的训练；如果你只是随机挑选专家而没有“安全网”（共享头和加权摘要），模型的表现实际上会变差。

简而言之：GQE 就像是聘请了一支专家团队，只有最合适的专家才会出现在特定的任务面前，这使得处理长任务的过程更加高效，同时又不会降低工作的质量。

技术摘要：分组查询专家 (Grouped Query Experts, GQE)

问题陈述

自注意力机制是 Transformer 架构中的计算瓶颈，特别是在长上下文序列中，成对的 Token 交互呈二次方增长。标准的稠密注意力机制会对每个 Token 应用相同的注意力头集合，而不考虑该 Token 的信息含量或复杂度。这种统一的激活方式导致了计算效率的低下，因为低信息量的 Token（例如标点符号、停用词）消耗了与承载内容的单词相同的资源。虽然混合专家模型（Mixture-of-Experts, MoE）已成功将条件计算引入 Transformer 的前馈（MLP）块，但将类似的稀疏性应用于注意力机制仍是一个尚未解决的挑战。现有的分组查询注意力（Grouped-Query Attention, GQA）通过在若干查询头组内共享一对键值（KV）头来减少内存带宽和 KV 缓存（KV-cache）大小，但它仍然会对每个 Token 计算所有的查询头。

方法论：分组查询专家 (GQE)

作者提出了 分组查询专家 (GQE)，这是一种将 MoE 原理集成到 GQA 模型注意力模块中的方法。其核心设计理念是在保留 GQA 稠密 KV 缓存优势的同时，使查询头（query-head）的计算具有条件性。

架构与路由

分组专家池 (Grouped Expert Pool)： 在每个固定的 GQA 组内（其中多个查询头共享一个 KV 头），查询头被视为一个“专家”池。
条件激活 (Conditional Activation)： 对于每个 Token，路由器会从每个 GQA 组中选择 $k$ 个查询头专家组成的稀疏子集。
稠密 KV 路径 (Dense KV Path)： 至关重要的是，键值投影（Key-Value projections）和 KV 缓存保持稠密且保持不变。稀疏性仅应用于查询侧（query-side）的计算。
输出构建： 通过特定的机制构建输出，以确保训练的稳定性：
- 硬拼接 (Hard Concatenation)： 每个组内被选中的 $k$ 个专家的输出被进行拼接。
- 重归一化加权求和 (Renormalized Weighted Sum)： 使用路由概率对选中的专家输出进行加权求和。这提供了一个可微路径，使语言模型损失能够传递给路由器，从而解决了离散 top- $k$ 选择的不可微问题。
- 常驻共享头 (Always-On Shared Head)： 无论路由结果如何，每个 Token 都会计算一个额外的注意力头。该头在训练期间充当稳定的锚点，防止路由器坍缩到单个专家，并确保始终存在一个基准注意力模式。

训练与损失

除了标准的语言模型损失外，GQE 还采用了负载均衡辅助损失（load-balancing auxiliary loss），以鼓励路由器在每个组内的专家池中均匀分布 Token，从而防止“专家坍缩”。

核心贡献

在 GQA 中实现 MoE 形式化： 论文专门在 GQA 组内制定了 MoE 方案，在激活每个 Token 的 $k$ 个专家同时保持所有 KV 头为稠密状态。这使其与底层的 GQA 配置保持一致，并保留了 GQA 的内存特性。
在不损失质量的情况下降低计算量： 作者证明了 GQE 可以匹配全激活 GQA 基准模型的下游准确度，同时仅激活查询头的稀疏子集。
稳定机制： 论文指出，成功的注意力稀疏路由需要两个特定组件：用于提供适当学习信号的重归一化加权求和槽位，以及用于锚定层的常驻共享头。

实验结果

实验是在使用 FineWeb-Edu 数据集的 30B-token 预算 和 250M 参数规模 下进行的。

准确度： 在主要设置下（16 个查询头，8 个 KV 组， $k=1$ ），GQE 激活 8 个路由专家加 1 个共享头，即在 16 个查询注意力计算中总共激活 9 个。尽管存在这种减少（激活了约 56% 的总查询操作），GQE 仍能匹配全激活 GQA 基准模型的下游准确度（HellaSwag, PIQA, ARC-Easy），其平均得分达到 56.04，而基准模型为 55.86。
消融实验： 论文表明，移除重归一化加权求和槽位或共享头会导致性能下降（分别降至 ~55.18 和 55.43），证实了仅靠稀疏性而不具备特定的架构稳定器是不足以维持性能的。
吞吐量与加速比： GQE 在长上下文场景中表现出显著的加速。
- 在短序列（2k tokens）时，由于路由开销，加速效果较为温和（~1.15×）。
- 在长上下文场景（4k 到 1024k tokens）中，GQE 相比 GQA 基准实现了 1.67× 到 1.80× 的加速。这是因为该方法跳过了不需要的查询计算，且随着序列长度增加，注意力成本呈二次方增长，其节省的收益也会随之扩大。

重要性与主张

论文声称 GQE 提供了一条在不牺牲模型质量的前提下，降低长上下文 Transformer 自注意力计算成本的可行路径。其重要性在于：

效率： 在测试配置中，它减少了大约一半的活跃查询头计算，同时保持了性能对等。
可扩展性： 随着上下文长度的增加，加速效果变得越来越有价值，解决了注意力机制的二次方缩放瓶颈。
设计洞察： 该工作强调，将 MoE 转移到注意力机制需要精细的架构设计（特别是共享头和重归一化路由信号），以避免在朴素的稀疏注意力尝试中常见的准确度损失。

作者保持了谦逊的态度，指出结果是在 250M 规模下的，应当在大规模场景下进一步验证。他们将结果定位为对基准模型的“匹配”而非明显的“超越”，强调其成就的核心在于以相等的性能实现了更低的计算量。