6 月 23 日登上 HuggingFace 当日论文的《Grouped Query Experts: Mixture-of-Experts on GQA Self-Attention》(arXiv:2606.20945,作者 Vishesh Tripathi、Abhay Kumar)提出了一个看起来朴素却切中要害的想法:把混合专家(MoE)做进注意力层本身,而不是只做进 FFN。具体做法是在分组查询注意力(GQA)的每个分组内部,把若干查询头当成「专家」,由一个路由器为每个 token 只挑选其中 top-k 个查询头去计算,而键值(KV)头全程保持稠密、始终参与计算。这样一来,GQA 最值钱的资产——共享的 KV 头、可压缩的 KV cache——被原样保留,被削减的只是激活的查询头算力。在 250M 参数 / 30B token 的固定预算下,完整版 GQE 在 HellaSwag、ARC-Easy、PIQA 三项上的平均分 56.04,略高于稠密 GQA 基线的 55.86;而预填充(prefill)阶段在 4k 到 1024k token 的长上下文区间提速约 1.67×–1.80×。论文以 CC0 公共领域协议发布,v1 于 6 月 18 日提交、v2 于 6 月 23 日修订。
这不是又一个万亿参数大模型,而是一篇小规模、单点验证、把话说得很克制的方法学论文。它值得一读的地方,在于它指出了一个长期被默认接受、却未必合理的事实:标准注意力对每个 token 都套用同一组注意力头,不论这个 token 是简单的虚词还是承载关键信息的实词——这里天然存在「按需计算」的空间。
发生了什么
过去两年里,降低大模型推理成本的主线之一,是把注意力的 KV cache 压小:多查询注意力(MQA)让所有查询头共享一组 KV,GQA 则是折中——把查询头分成若干组,每组共享一份 KV。GQA 已经是 Llama、Qwen 等主流模型的标配,因为它在几乎不掉点的前提下大幅缩小了 KV cache、缓解了长上下文的显存与带宽瓶颈。
GQE 在 GQA 之上又加了一层稀疏化,但稀疏的对象不是 KV,而是查询头。论文的核心机制可以拆成三件事:
- 组内路由(within-group routing):在每个固定的 GQA 分组里,设一个路由器,对该组内的查询头打分,softmax 之后选 top-k 个查询头专家来实际计算;未被选中的查询头这个 token 就不算。KV 路径完全不动,所有 KV 头照常计算——这是它和「直接砍注意力头」的根本区别,也是它能保住 KV cache 优势的原因。
- 可微的加权和槽(weighted-sum slot):除了被硬选中的几个专家头,论文额外保留一个经重归一化的加权和输出槽,给路由器提供可微的学习信号,避免 top-k 这种离散选择让路由器学不动。
- 常开共享头(always-on shared head):再设一个与 token 无关、永远参与计算的共享头,作为训练锚点,稳住整个结构的学习。
作者还用了负载均衡辅助损失防止路由器塌缩到总是选同一个专家,并用 ZClip 抑制训练中的 loss 尖峰。消融实验显示,光有硬选择(hard concat)效果会掉,只有同时加上「重归一化打分 + 常开共享头」这两味,GQE 才追平稠密基线——换句话说,给路由器一个像样的学习信号,是这套方法成立的前提,而不是可有可无的工程细节。
需要说清楚一个容易被标题简化的点:摘要里说「激活一半查询头」,指的是被路由的查询专家那部分(top-k 选中比例);若把常开的共享头也算进去,实际参与查询注意力计算的是 16 个里的 9 个(9/16 ≈ 56%),并非严格的 50%。下文的算力口径以论文正文为准。
关键数据 / 技术细节
论文主实验的配置是 16 个查询头 / 8 个 KV 头,即分组数 G=8、每组专家数 M=2、top-k=1(每组激活 1 个查询专家),训练 30B token。三项常识/推理基准(HellaSwag、ARC-Easy、PIQA)的结果如下,数字均为论文 Table 2 原文:
| 变体 | HellaSwag | ARC-E | PIQA | 平均 |
|---|---|---|---|---|
| GQA 基线(16 头全激活) | 41.31 | 61.36 | 64.90 | 55.86 |
| 加权拼接,无重归一化槽 | 40.16 | 60.52 | 64.85 | 55.18 |
| 仅硬拼接(hard concat) | 40.66 | 60.56 | 65.07 | 55.43 |
| 完整 GQE(重归一化打分 + 共享头) | 41.01 | 62.41 | 64.69 | 56.04 |
可以看到:无重归一化槽的加权拼接比稠密基线低 0.68,仅硬拼接低 0.43,而完整 GQE 反而高出基线 0.18(56.04 对 55.86),落在「与基线相当、误差不超过 0.2 个平均点」的范围内。作者非常克制地把这一结果定性为「匹配(match)而非提升(improvement)」,并明确写明应在多随机种子、更大规模下复现——这是一个加分项,而非可以拿来吹「超越稠密注意力」的结论。
预填充提速(GQE 相对 GQA 基线的延迟比)随上下文长度变化,论文只在两个锚点给出明确数值,其余长度以折线图呈现:
| 序列长度 | 预填充提速 | 说明 |
|---|---|---|
| 2k token | 约 1.15× | 短上下文下提速有限 |
| 4k → 1024k token | 约 1.67×–1.80× | 进入长上下文后趋于稳定 |
测试的序列长度覆盖 2k、4k、8k、16k、32k、64k、128k、256k、512k、1024k。直觉上这与机制一致:注意力计算随序列长度二次增长,序列越长、注意力占总开销的比例越大,削减查询头算力带来的收益就越能兑现成端到端的预填充加速;而短序列下注意力本身占比小,提速自然有限。
训练侧的关键超参(论文 Table 1):优化器为 Fused AdamW(β₁=0.9、β₂=0.95、ε=1e-7),WSD 学习率调度,最大学习率区间 1e-5 至 5e-4,预热 3B token,权重衰减 0.1,全局批大小 105 万 token,序列长度 2048,BFloat16 混合精度。训练数据取自 FineWeb-Edu(FineWeb2 的子集)。
架构口径与算力细节(展开)
- 在 W_O 之前的输出槽数(k=1 主设置):10 个 = 8 个硬拼接专家槽 + 1 个加权和槽 + 1 个共享头槽;稠密基线则投影 16 个头槽。
- 因此 W_O 从「16·d_h × d」缩为「10·d_h × d」。
- softmax 在 top-k 选择之前做(组内概率)。
- 被路由激活的查询专家数为 kG = 8;加上常开共享头,总查询注意力计算量为 9/16 ≈ 56%。
- 论文坦承每组专家池很小(M=2),「组内路由的余量有限」,更大的专家池没有扫过。
为何重要
把 MoE 引入 Transformer 不是新闻——主流大模型早已在 FFN 层用 MoE 把「总参数量」和「激活参数量」解耦。但 MoE 几乎都长在 FFN 上,注意力层一直是稠密的。GQE 的价值在于把「条件计算」的思路搬进了注意力,而且选了一个对 KV cache 友好的切口:只稀疏化查询头,不碰 KV。
这一点在长上下文时代尤其关键。当下推理成本的大头之一是长上下文预填充——把几万乃至上百万 token 的提示一次性吃进去,这一步的注意力计算是实打实的二次开销,且通常不像解码那样受 KV cache 显存主导。GQE 的提速恰恰落在这个被很多 KV 压缩方法忽略的环节:它不改善解码阶段的 KV 带宽(因为 KV 全程稠密),而是削减预填充阶段的查询头算力。两类优化方向正交,理论上可以叠加。
从方法学角度,GQE 也提供了一个有用的对照:它证明了「在注意力头粒度上做按需计算」至少在小规模上不掉点,且让稠密基线相形之下显得有冗余——每个 token 都用满全部查询头,可能确实是浪费。这给「token 难度自适应」这一长期命题添了一个可操作的结构。
但务必把分量摆正:这是 250M 参数 / 30B token 的单点实验,基准是 HellaSwag/ARC-E/PIQA 这类相对饱和的常识题,提升只有 0.18 个平均点,且每组只有 2 个专家、top-k=1,路由的「腾挪空间」本就很小。它离「可在前沿模型上落地的注意力 MoE」还有相当距离,作者自己也把话说得很满:这是一个匹配,不是突破。
接下来看什么
- 规模化能否站住:250M 是玩具规模,30B token 也只是小预算。这套方法在 1B、7B 乃至更大模型、更大数据预算下,「匹配稠密基线」的结论是否依然成立,是它能否被采纳的第一道关。论文已把这列为明确的 future work。
- 更大专家池的收益曲线:主实验每组仅 2 个专家,路由余量有限。若把每组专家数 M 加大、top-k 调高,精度与提速如何权衡,是判断这套机制天花板的关键。
- 多随机种子复现:0.18 个点的优势在单次实验里几乎落在噪声范围内。作者主动要求多种子复现——后续是否有人(包括作者自己)给出方差区间,直接决定「匹配」这个结论的可信度。
- 与 KV 压缩方法的叠加:GQE 削查询头算力、保 KV 稠密;MQA/GQA/KV 量化等削 KV 带宽。两者正交,能否叠加出「预填充更快 + 解码更省」的组合效果,值得后续工程验证。
- 与状态空间模型的对照:论文提到要与 Mamba 等架构比较。在长上下文这个 GQE 最占优的战场上,它相对线性注意力/SSM 的性价比如何,是更有意义的横向参照。
GQE 是一篇「想法漂亮、证据克制」的论文,值得记住的是它的问题切口而非它的跑分。它点破了一个被默认接受的浪费:稠密注意力对每个 token 都用满全部查询头。把 MoE 的条件计算搬进注意力、并刻意只稀疏化查询头而保住 KV 稠密——这个设计在工程上很「懂行」,因为它没有破坏 GQA 这两年好不容易换来的 KV cache 红利,而是去啃另一块没人动的肥肉:长上下文预填充的查询头算力。提速主要出现在 4k 以上的长序列(约 1.7–1.8×)、2k 时只有约 1.15×,这条曲线本身就把它的适用场景标得很清楚——它是给长上下文预填充省钱的,不是给短对话提速的。
但请不要被「在注意力里做 MoE、还超过稠密基线」这类叙述带节奏。第一,这是 250M/30B 的单点结果,优势仅 0.18 个平均点,作者自己都只敢说「匹配」并要求多种子复现;在如此小的 margin 下,换个种子结论可能就翻面。第二,「激活一半查询头」的省算力口径要打折扣——把常开共享头算进去,实际是 9/16(约 56%) 在算,不是 50%。第三,每组只有 2 个专家、top-k=1,路由几乎没有腾挪空间,真正考验这套机制的「大专家池 + 规模化」实验完全还没做。
更现实的判断是:GQE 当前的意义是学术信号而非可落地方案。它给「注意力层也能做条件计算」补了一个干净的小规模证据,和 FFN-MoE、KV 压缩在优化维度上正交,理论上可叠加。但要让人相信它能上前沿模型,需要看到三件事同时发生:规模放大后不掉点、专家池放大后提速进一步拉开、方差被压到 margin 之内。在那之前,它是一篇值得收藏的方法论文,而不是下一代注意力的既定答案。一个加分项值得点名:作者全程不夸张、主动标注局限、以 CC0 开放——在普遍「报喜不报忧」的论文风气里,这种克制本身就提高了结论的可信度。