Grouped Query Experts: Mixture-of-Experts on GQA Self-Attention作者Vishesh Tripathi, Abhay Kumar核心发表机构FrontiersMind论文链接arXiv:2606.20945v2发布于arXiv 预印本cs.LG| :— | :— | :— | :— || GQA 基线全部 16 头激活 |41.31| 61.36 | 64.90 |55.86|| 加权拼接无重归一化槽 | 40.16 | 60.52 | 64.85 | 55.18 || 仅硬拼接 | 40.66 | 60.56 |65.07| 55.43 || GQE重归一化评分 共享头 | 41.01 |62.41| 64.69 |56.04|从表中可以看到完整的 GQE 配置在平均准确率上达到 56.04不仅没有低于全激活的 GQA 基线55.86甚至略高 0.18 个百分点。考虑到 GQE 只激活了一半的路由查询头总查询计算量为9 / 16 9/169/16这一结果说明查询侧的稀疏化在固定计算预算下是可行的模型确实可以通过条件路由用更少的注意力计算达到相当的模型质量。除了准确率论文还报告了 GQE 相对 GQA 基线的预填充加速比。在序列长度为 2k 时加速比较温和约 1.15 倍因为此时路由和调度开销相对较大。从 4k token 开始GQE 进入长上下文优势区加速比稳定在 1.67 倍到 1.80 倍之间并随序列长度增加缓慢上升在 1024k 长度时接近 1.80 倍。这一趋势符合设计预期随着序列长度增加查询侧注意力计算逐渐占据主导跳过非活跃查询专家带来的收益超过了固定的路由开销。4.3 消融实验 / Ablation Study论文对路由和输出构建的关键设计进行了消融分析表 1 中已包含四种变体的下游任务准确率。论文通过这种消融说明稀疏路由要成功至少需要两个条件路由器需要获得归一化的、可微的学习信号通过重归一化槽并且需要一个稳定的共享路径来锚定训练。具体来看加权拼接但无重归一化槽这一变体仅对各组被选专家输出做加权拼接可能直接使用组内 softmax 概率作为权重但没有额外的路由器监督槽。它的平均分为 55.18比 GQE 低 0.86 分比 GQA 基线低 0.68 分。这表明如果路由器只能通过隐式路径获得梯度其学习信号不足路由质量会明显下降。仅硬拼接不使用任何加权平均或路由器监督槽完全通过离散的 top-k kk选择产生输出。其平均分为 55.43虽然比“无重归一化槽”略高但仍低于 GQA 基线和 GQE。由于硬选择本身不可微路由器的学习只能依赖弱的代理信号因此路由策略不够准确。GQE重归一化评分 共享头在硬拼接的基础上加入重归一化的路由器监督槽并添加共享头。这一完整配置达到 56.04 的平均分匹配并略微超越 GQA 基线。图 2 展示了表 1 中四种变体在训练过程中的损失曲线。可以看到不同路由机制在训练动态上存在明显差异GQA 基线的损失下降相对平稳而无重归一化槽的变体在早期训练中表现出更不稳定的学习信号GQE 与基线的损失曲线最终收敛到相近水平但中间过程存在偏移。这说明路由器的学习信号设计直接影响训练动态而 GQE 的最终配置在训练稳定性和收敛质量之间取得了更好的平衡。图 3 至图 5 分别给出了 HellaSwag、ARC-Easy、PIQA 三个下游任务上准确率随训练 token 变化的曲线对比了 GQA 基线、中间路由消融变体和最终 GQE 模型。从这些曲线中可以看出GQE 与 GQA 基线在大部分训练阶段保持接近中间路由消融变体则在不同任务上表现出不同程度的差距。尤其需要注意的是在 ARC-Easy 上最终 GQE 模型表现出超越基线的趋势这与其在下游任务上的最终平均准确率优势一致。综合消融结果论文的关键论证是单纯的硬路由或单纯的加权拼接都不能在保持计算节省的同时达到全激活基线的精度只有将硬路由头与重归一化加权槽、共享头结合起来才能让路由器在可微信号驱动下学到有效的专家选择策略。五、相关工作 / Related WorkGQE 的研究处于两条工作线的交叉点一是针对注意力机制效率优化的 GQA 及其变体二是将 MoE 思想引入注意力计算的若干尝试。与GQA的关系最直接。GQA 通过分组共享 KV 头减少了 KV 缓存和带宽但其查询头计算仍是全激活的。GQE 在 GQA 基础上增加查询头的条件路由使 KV 侧的优势保持不变同时进一步裁剪查询侧的计算量。因此GQE 可以被视为 GQA 的“稀疏查询头扩展”。与将 MoE 应用于 FFN 层的常规做法相比GQE 是少见的将 MoE 应用于注意力头的方案。传统 MoE 通常放置在前馈网络中对 token 的隐藏表示进行条件计算GQE 则直接在注意力头级别做路由粒度更细并且特殊的难点在于离散选择与注意力输出的不可微性。GQE 通过路由器监督槽解决这一问题。与MoA、MoH 等注意力头路由方法相比GQE 的区别体现在两个层面。首先路由粒度不同MoA 类方法通常在整个模型范围内选择 top-k kk个注意力头因此某些组可能完全没有被选中GQE 则在每个固定的 GQA 组内选择 top-k kk保证每组都参与计算从而维持了 GQA 的结构性归纳偏置。其次KV 缓存处理方式不同MoH 等稀疏化注意力头的方法可能会同时稀疏化 KV 缓存而 GQE 明确保持所有 GQA KV 头密集且不变。Token 可能跳过某个查询专家但仍会针对该组对应的 KV 头做注意力计算因此不会损失 KV 侧的上下文信息。与LLaMA-MoE v2 等先训练后转换的方法不同GQE 是端到端从头联合训练路由器和专家的方法而不是对预训练模型进行专家剪枝或转换。这意味着 GQE 的路由器与查询头专家在整个预训练过程中共同适应可以更好地协调专家分工。此外训练过程中使用的 ZClip 技术与算术运算无关其主要作用是通过梯度裁剪抑制训练损失尖峰属于训练稳定性技术而 FineWeb2/FineWeb-Edu 是数据侧的处理策略。GQE 本身是架构层面的创新与上述两项工作解决的问题不同但可以相互配合。六、局限性与展望 / Limitations Future Work当前研究仍存在若干明确的局限需要在后续工作中关注。第一验证规模有限。论文仅在 250M 参数模型和 30B token 预算下进行实验该规模下观察到的“匹配甚至略超基线”需要谨慎解释。论文本身将这一结果定性为“匹配”而非“提升”。在更大规模的模型上路由器的行为、专家特化程度以及训练稳定性都可能发生变化小规模上的收益不一定能线性外推。因此未来需要在多种子、更多参数规模和更长训练预算下进行验证。第二专家池规模较小。主要实验使用了有限的每组专家数量限制了路由器的选择空间。增大M N / G M N/GMN/G有望带来更丰富的专家候选和更细粒度的 token 专化但论文未对N NN做广泛的扫描。更大的专家池也可能加剧负载不平衡和路由开销因此需要权衡。第三路由和调度开销。在较短的上下文长度下例如 2k tokenGQE 的加速比并不高只有约 1.15 倍说明路由器计算和稀疏调度本身存在不可忽略的固定成本。在更长序列中这些开销会被查询侧节省的计算量覆盖但对于经常使用短上下文的场景GQE 的收益有限。未来的工作可以探索更轻量的路由器或更高效的稀疏调度算法。第四评估范围有限。论文报告的下游任务限于 HellaSwag、ARC-Easy 和 PIQA 三个常识推理基准缺乏语言建模、代码、多任务指令跟随等更广泛的评测。此外虽然论文报告了预填充阶段的加速但没有详细分析解码decode阶段或 KV 缓存带宽受限场景下的性能而这些在实际部署中同样重要。第五负载均衡的潜在风险。虽然 GQE 使用了负载均衡辅助损失但不同 token 的查询头使用频率可能仍然存在长尾分布。如果某些专家在训练中始终只被少数 token 选择其参数可能无法得到充分优化最终影响模型质量。更大规模上的 router collapse 现象需要进一步研究未来可尝试更适合注意力路由的负载均衡损失设计。未来研究方向包括将 GQE 扩展到更大的 Transformer 架构并验证其通用性与 Mamba、线性注意力等长上下文架构进行直接对比探索非均匀的k kk值即每个组选择不同数量的专家以及将 GQE 与其他稀疏注意力方法结合进一步压缩长上下文场景下的计算和内存开销。七、总结 / ConclusionGQE 提出了一种在 GQA 自注意力内部进行查询头条件路由的方法。它通过在每个固定 GQA 组内使用路由器选择 top-k kk查询头专家同时保持 KV 路径完全密集成功地在保留 GQA KV 缓存优势的同时将查询侧的计算量减半。为了让离散的专家选择能够顺利训练论文设计了“先 softmax 再 top-k”的组内路由、硬路由头拼接、重归一化路由器监督槽以及共享头等一系列关键模块并通过负载均衡辅助损失维持专家利用率。在固定 30B token 和 250M 参数的设置下GQE 达到或略微超过全激活 GQA 基线的平均下游准确率同时将每个 token 的活跃查询注意力计算量降至9 / 16 9/169/16。在长上下文预填充阶段GQE 相对基线实现了 1.7–1.8 倍的加速且加速比随序列长度增加而提高。这些结果表明基于 token 内容动态选择查询头是一条可行且有潜力的稀疏注意力路径为后续在更大模型和更长上下文中探索条件计算与注意力稀疏化提供了基础。原文摘要: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部分平台可能图片显示异常请以我的博客内容为准