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 k1k1但通过辅助损失组内M MM个专家均能获得充分的训练。五、相关工作 / Related WorkGQE 与以下方向紧密相关且具有明确区别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 k2k2或动态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部分平台可能图片显示异常请以我的博客内容为准