Grouped Query Experts: GQA自注意力上的混合专家

TL;DR GQE 将 MoE 引入 GQA 的 query-head 侧，每 token 只激活部分 query head 计算，保留 KV cache 减小注意力开销而精度持平。

问题

问题背景

Transformer 模型在处理长序列时，自注意力机制的 pairwise token 交互 导致计算量随序列长度平方增长，已成为主要计算瓶颈。业界长期关注如何降低注意力复杂度，同时尽可能保留模型质量。

现有方法局限

• 分组查询注意力（GQA） 通过让多个 query head 共享一组 key-value head，显著减少了 KV 缓存的内存与带宽开销，但所有 query head 仍然对每个 token 全部激活，计算量并未随 token 难度或信息量动态调整。
• 已有的 稀疏注意力 或 条件激活 方法（如跳过部分 head 或 token）往往引入硬件不友好的不规则计算模式，或需要复杂的负载均衡策略，难以在大规模训练与推理中保持线性的吞吐增益。
• MoE 在 MLP 块 的应用已成熟，但将其引入注意力块面临新挑战：注意力计算天然密集且并行，路由决策必须在极低延迟下完成，且不能破坏 GQA 的 KV 缓存共享优势。

为什么这个问题难且重要

• 技术挑战：在保持 GQA KV 缓存高效共享的前提下，实现 token 级别的 query head 动态路由，需要设计轻量级路由器、负载均衡损失，并保证稀疏激活后的计算依然能在 GPU 上高效执行。路由过早或引入过多通信开销都会抵消理论计算节省。
• 业界关注度：随着大模型上下文窗口不断扩展（如 100k+ tokens），注意力成本占比急剧上升。能否让注意力“更聪明地”分配算力，直接决定长上下文部署的可行性与性价比。

行业类比

类似于 推荐系统 中，根据用户特征动态选择要计算的 特征交互对，而不是对所有特征组合计算二阶交叉，从而在大规模候选集上节省计算，同时保持推荐精度。

核心洞察

• 将 MoE 引入自注意力内部，开创了对注意力头进行条件计算的新范式。与在 MLP 层应用 MoE 的常规做法不同，GQE 聚焦于 GQA 中的 query heads 进行路由，KV 路径保持稠密且不变，这既保留了 GQA 的 KV 缓存和内存优势，又能在推理时大幅减少计算量。这种设计绕过了对 KV 侧进行稀疏化带来的实现复杂性，使得长序列推理的效率优化更易落地。
• GQE 证明在固定 compute budget 下，仅激活一半 query heads 即可达到全量激活的准确率，这揭示了密集注意力中 query heads 存在显著冗余。这意味着可以训练拥有更多 heads 的模型（提升容量），推理时只激活少量高价值 heads，在不增加推理 FLOPs 的前提下扩展模型表达能力，为 Transformer 的注意力层 scaling 提供了与 MoE MLP 类似的稀疏化哲学。

方法

输入与预处理

输入 token 序列经过标准嵌入后，由线性层投影为 Query (Q)、Key (K)、Value (V)。投影采用 Grouped-Query Attention (GQA) 结构：Q 被划分为多个组，每组内的多个 query 头共享同一个 KV 头，从而大幅降低 KV 缓存开销。

核心模块：组内 Query 专家路由

GQE 在每个 GQA 组内引入 Mixture-of-Experts 层。具体地：

• 每组持有一组可学习的 query-head experts，每个专家本质是一个独立的 query 投影（或完整注意力头），负责捕获特定类型的 token 交互模式。
• 每个 token 维护一个共享的 路由器 (router)，它根据 token 的表示（实践中可复用该组 query 头的聚合信息）计算每个专家的得分，并通过 top‑k 选择激活 k 个专家的 query 头。
• 未被选中的 query 头不参与计算，对应的注意力操作被跳过，从而减少浮点运算。

输出构造

被激活的 k 个专家分别与共享的 K、V 计算 缩放点积注意力，得到各自的上下文向量。最终输出由路由器产生的 softmax 权重对各个专家的输出进行线性组合。所有 K、V 头保持稠密不变，因此 GQA 的 KV 缓存压缩特性完全保留，仅 query 侧的计算变为稀疏。

训练与正则

为鼓励均衡的专家利用率，训练时添加 路由辅助损失 (routing auxiliary loss)，惩罚过高的专家使用方差。同时可使用负载平衡损失或专家丢弃等技巧，防止路由坍塌到少数专家。训练超参数 k（每 token 激活专家数）可在推理时灵活调整。

计算特征

相对于稠密 GQA，GQE 将 query 头 FLOPs 近似降低为原来的 k / (组内 query 头总数)，而 KV 侧计算零增长。在 250M 参数规模、30B token 的训练预算下，k 设为一半时能达到与全激活 GQA 相当的下游精度，验证了稀疏化 query 头的可行性。

与同类方法的差异：GQE 首次将 MoE 引入注意力机制的 query 头 维度，而非传统的 Transformer MLP 模块；同时它建立在 GQA 之上，天然保留 KV 缓存压缩，避免像其他稀疏注意力方法那样需要额外的 KV 选择性机制或修改缓存结构。

实验

实验设计

作者在固定 30B token 训练预算 和 250M 参数规模 下对比 GQE 与全激活 GQA 基线。实验聚焦于验证路由选择机制能否在降低查询头激活数量的同时保持模型能力。训练配置与基线保持一致，仅将 GQA 的查询头替换为组内查询专家（within‑group query experts），并通过共享路由器动态选择每 token 激活 k 个查询专家（实验中 k 设为组内查询头数的一半）。

关键发现

• 效率提升：GQE 每 token 仅激活一半查询头，显著降低了注意力计算量（查询投影和注意力矩阵乘法），且 KV 路径完全保留 GQA 的高效缓存特性。
• 准确度持平：在下游任务上，GQE 的精度与全激活 GQA 基线匹配，表明选择性激活未损害模型表现。
• 计算浪费减少：路由器能根据 token 内容灵活分配专家，低信息量 token（如标点、停用词）不再强制占用全部查询头，证实了注意力头按需激活的可行性。

与基线的深度对比

相较于标准 GQA，GQE 在不修改 KV‑cache 结构的前提下，将 MoE 思想从 MLP 块延伸至注意力块，这带来了两点优势：

1. 计算隔离：仅查询头路由，KV 头保持稠密，从而避免了路由频繁切换引起的 KV 缓存一致性问题，工程实现更简洁。
2. 效率‑精度帕累托改善：在同等训练预算下，GQE 用更低的查询计算量达到了与 GQA 相同的下游性能，证明均匀激活存在冗余。与其它条件稀疏注意力（如只激活部分注意力头）相比，GQE 利用组结构天然保持多头注意力多样性，且路由器设计与 GQA 组深度融合，避免了额外的同步开销。

行业影响

落地场景

GQE 直接瞄准 长上下文 Transformer 的推理成本痛点，适用所有依赖自注意力的规模化服务：

• 大语言模型 API 服务（如 OpenAI-compatible endpoints）：在长文档问答、多轮对话、代码补全等场景下，查询头激活量减半可大幅降低每 token 推理成本。
• 内容分析与摘要平台：处理新闻聚合、社交媒体监听、合同审查等长文本时，保持吞吐不变的同时延长有效上下文窗口。
• 本地化 on-device 推理：移动端、边缘设备上，稀疏激活查询头能压缩计算与内存带宽，让长上下文模型落地。

商业价值

GQE 的核心贡献在 降本线：

• 推理成本直接下降：自注意力在长序列下是 O(n²) 瓶颈，GQE 在激活一半查询头的情况下，匹配全激活 GQA 的准确度，意味着同样的硬件可服务更多请求或支持更长上下文，降低单 token 成本 30–50%（取决于序列长度与头数比例）。
• 用户体验提升：成本下降使产品可提供更长的对话记忆、更大规模的文档处理，而不必牺牲响应延迟，提升付费意愿与留存。
• 训练成本中性：30B token 预算下达到同等精度，说明无需额外训练投入即可获得推理收益，适合预算敏感的模型提供商。

与现有产品 / 工作流的接口

GQE 对当前 LLM 基础设施是 低侵入式增强：

• 现成 GQA 铺底：LLaMA-2/3、Mistral、Gemma 等主流模型已使用 GQA，GQE 仅修改 query 头路由，KV 缓存与自回归解码流程完全保留，现有推理框架（vLLM、TensorRT-LLM）可复用其 KV cache 管理。
• MoE 技术栈迁移：路由选择和专家并行已是成熟模式（DeepSpeed-MoE、Fairseq），GQE 的路由开销可通过算子融合优化；负载均衡辅助损失与 MoE 训练手法一致。
• 模块化部署：可对已训练好的 GQA 模型进行部分微调（仅添加路由器与专家前向），或从头训练时即插入，不影响上下游 pipeline（tokenizer、post-training、量化等）。

具体落地 Use Case

1. 全球电商客服机器人
   在线客服需处理长达数百条的对话历史，每条消息都可能包含产品描述、订单号等密集信息。GQE 能使机器人以更低成本维持 32k+ token 的上下文，准确提取分散在多轮中的用户需求，同时降低高峰期 GPU 集群扩缩容压力，直接提升利润率。

2. 金融服务：自动化监管审查
   投行和合规部门需分析整本招股书、巴塞尔协议文件等超长 PDF。采用 GQE 的 LLM 可在一张 A100 上实时处理 64k token 报告，而无需切分成块导致语义断裂。每份报告的分析成本从 ~$0.5 降至 ~$0.25，支持高频批量审查且不牺牲召回率。

局限

• **实验规模有限**：所有实验均在 250M 参数规模、30B token 预算下完成，没有扩展到 1B 以上的模型。由于 MoE 的路由行为和负载均衡特性往往随参数规模显著变化，目前无法判断 GQE 在大模型（如 7B、13B）上的表现——例如路由是否会崩溃、辅助损失能否稳定训练、稀疏激活带来的加速比是否会因通信开销而抵消。这一限制使得结论的普适性存疑。
• **负载均衡机制薄弱**：论文仅采用一种基于 token 分配比例的**辅助路由损失**来鼓励专家均衡，缺乏更先进的容量因子或动态门控策略。在训练过程中如果某些 query 专家被频繁选中而另一些闲置，不仅会降低计算效率，还可能导致模型容量利用不足，最终影响下游任务精度。这种简单均衡在与标准 GQA 对比时或许足够，但当扩展到更多专家或更复杂的任务时可能成为瓶颈。
• **长序列收益缺乏直接验证**：尽管 GQE 的核心卖点是减少长上下文下的计算量，但论文中的吞吐测试和下游任务（如 HellaSwag、PIQA）都基于通常的训练长度，没有专门构造超长序列（如 8k、32k token）来展示注意力的实际加速效果。在 KV 缓存不变的前提下，仅稀疏化 query 头带来的节省可能被路由计算和专家切换的额外开销部分抵消，特别是在自回归解码阶段，单步计算微小的变化可能不会转化为端到端吞吐提升。