分組查詢專家：基於GQA自注意力的混合專家模型

摘要

自注意力機制是 Transformer 效能的關鍵，但在長上下文情境下，它常是 Transformer 中最耗費計算的部分，因為其逐對 token 互動的規模與序列長度成二次方增長。標準的密集注意力還會對每個 token 套用相同的一組注意力頭，而不考慮 token 的難度或資訊內容。這種均勻活化會浪費計算資源，尤其在序列變長、注意力成本快速增加時尤為明顯。我們提出「分組查詢專家」（Grouped Query Experts, GQE），這是一種基於分組查詢注意力（Grouped-Query Attention, GQA）之上的混合專家層。在每個 GQA 群組中，路由器會為每個 token 選取 k 個查詢頭專家，而所有鍵值頭（KV heads）保持密集且不變。如此一來，GQE 保留了 GQA 的 KV 快取優勢，並僅減少活躍的查詢頭計算量。在固定 30B token 預算、250M 參數規模下，GQE 在下游準確率上與全活躍的 GQA 基準相當，同時每個 token 僅活化一半的查詢頭。

English

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.