Grouped Query Experts: Mixture-of-Experts on GQA Self-Attention
作者：Vishesh Tripathi, Abhay Kumar
核心发表机构：FrontiersMind
论文链接：arXiv:2606.20945v2
发布于：arXiv 预印本（cs.LG）
| :— | :— | :— | :— |
| GQA Baseline (all 16 active) | 41.31 | 61.36 | 64.90 | 55.86 |
| Weighted concat, no renormalized slot | 40.16 | 60.52 | 64.85 | 55.18 |
| Hard concat only | 40.66 | 60.56 | 65.07 | 55.43 |
|GQE (renorm. scoring + shared head)|41.01|62.41|64.69|56.04|
训练过程中的损失曲线（下图）进一步显示了四种设置（GQA 基线、两个中间路由消融、最终 GQE）的收敛行为，GQE 的损失曲线与基线几乎重合。
下游准确率随训练 token 的变化曲线也证实了 GQE 与基线的稳定性。下图分别展示了 HellaSwag、ARC‑Easy 和 PIQA 上的准确率演进，GQE（最终配置）始终匹配或略优于基线。
在吞吐量方面，由于图片目录中未提供具体的加速比图，根据研究笔记所述，在序列长度 2K 时加速比约为 1.15x，从 4K 开始稳定在 1.7–1.8x，验证了长序列下查询侧计算节省的主导作用。
4.3 消融实验 / Ablation Study
消融实验系统地验证了 GQE 的两个核心设计。表 1 中的“Weighted concat, no renormalized slot”变体移除了重归一化加权和槽，仅将所有选中专家的输出进行简单加权拼接（权重不经重归一化），导致平均准确率下降至 55.18%，低于基线。“Hard concat only”变体完全取消加权和槽，仅使用硬路由输出的拼接，平均准确率为 55.43%，同样低于基线。而完整 GQE（重归一化加权和槽 + 共享头）平均准确率达到 56.04%，甚至轻微超越基线。这充分说明：重归一化加权和槽为路由器提供了有效的梯度信号，共享头则提供了稳定通路，两者缺一不可。
此外，负载均衡损失确保路由器不会坍缩到组内少数专家上，从而保持组内专家的公平利用。尽管当前实验只使用了k = 1 k=1k=1，但通过辅助损失，组内M MM个专家均能获得充分的训练。
五、相关工作 / Related Work
GQE 与以下方向紧密相关且具有明确区别：
MoE 在 FFN 中的应用：传统 MoE 将专家部署在 MLP 层，路由器专注于稀疏激活前馈计算。GQE 将 MoE 迁移至自注意力模块的查询投影，属于对注意力计算的稀疏化。
全局路由的头选择方法：如 Mixture of Attention (MoA) 和 MoH 等方法从所有注意力头中全局选择若干头。GQE 的独特之处在于它在每个固定的 GQA 分组内进行路由，确保每个 KV 组至少有一个选中的查询头，从而维持了分组结构的计算一致性，避免了某些 KV 组可能被完全跳过的问题。
稀疏化 KV 的方法：MoH 等方法同时稀疏化 KV 缓存和查询计算。GQE 则明确保持 KV 路径完全密集，因此保留了 GQA 的所有 KV 缓存优势（如减少内存占用、便于推理优化），只对查询侧进行稀疏化。
后训练剪枝/转换方法：如 LLaMA‑MoE v2 等从预训练密集模型通过剪枝或转换得到稀疏模型。GQE 是从零开始联合训练路由器与专家，模型可以更早地学习到适配的路由策略，避免了后训练方法可能带来的迁移损失。
六、局限性与展望 / Limitations & Future Work
当前实验仅限于 250M 参数规模和 30B token 预算，在多随机种子和更大规模（如 1B、7B）下的表现尚未验证。更大规模下，路由器可能学习到更精细的 token‑to‑expert 映射，但也可能面临更严重的训练稳定性问题。
与其它长上下文架构（如 Mamba、线性注意力）的对比未在本文中涉及，未来工作应系统评估 GQE 在不同架构中的相对效率。
专家池大小M MM（每个 GQA 组内的专家数）当前未系统研究。当M MM更大时，路由选择空间更广，可能进一步提升专业化收益，但也会增加路由器容量和负载均衡的难度。更优的k kk值选择（如k = 2 k=2k=2或动态k kk）值得探索。
GQE 仅稀疏化了查询侧计算，未减少 KV 缓存内存占用。未来可探索与 KV 稀疏化的联合设计，以进一步降低长上下文推理的资源需求。
七、总结 / Conclusion
本文提出 Grouped Query Experts (GQE)，将 MoE 机制引入 GQA 自注意力中的查询头，通过在每个 GQA 组内进行 top‑k 路由，实现了查询侧计算的稀疏化，同时保持 KV 缓存密集。GQE 通过重归一化加权和槽提供可微训练信号，并借助共享头稳定学习流程。在 250M 参数、30B token 预算的实验中，GQE 激活半数查询头却在下游准确率上匹配甚至超越全激活的 GQA 基线，并在长序列下获得 1.7–1.8 倍的预填充加速。该工作为 Transformer 注意力模块的高效化提供了一种简单而有效的稀疏化范式，其将条件计算引入注意力头的思路有望在更大规模模型和实际部署中发挥重要作用。
原文摘要:Self-attention is central to Transformer performance and is often the most expensive part of the Transformer at long context lengths because its pairwise token interactions scale quadratically with sequence length. Standard dense attention also applies the same set of attention heads to every token regardless of token difficulty or information content. This uniform activation can waste compute, especially as sequences grow longer and attention cost increases rapidly. We propose Grouped Query Experts (GQE), a mixture-of-experts layer on top of grouped-query attention (GQA). Within each GQA group, a router selects k query-head experts per token while all key-value (KV) heads remain dense and unchanged. Thus, GQE keeps the KV cache benefits of GQA and reduces only the active query-head computation. On a fixed 30B token budget at the 250M parameter scale, GQE matches the all-active GQA baseline in downstream accuracy while activating half the query heads per token.
PDF链接:https://arxiv.org/pdf/2606.20945v2
