分组查询专家:GQA自注意力中的专家混合模型
Grouped Query Experts: Mixture-of-Experts on GQA Self-Attention
摘要
自注意力机制是Transformer模型性能的关键所在,而在长上下文长度的情况下,它往往也是Transformer模型中成本最高的部分——因为其成对token之间的交互作用与序列长度的平方成正比。标准的密集注意力机制则会对每个token使用相同的注意力头结构,而不考虑该token的复杂程度或信息含量。这种统一的激活方式会浪费计算资源,尤其是当序列长度增加、注意力成本迅速上升时。我们提出了“分组查询专家”模型,该模型基于分组查询注意力机制而构建。在每个GQA组中,一个路由器会为每个token选择k个查询头专家,而所有键-值头则保持为密集结构且保持不变。因此,GQE能够保留GQA的键-值缓存优势,同时仅减少活跃查询头的计算量。在30B个token的固定预算下,以及250M参数的参数规模下,GQE在下游精度指标上与全活跃GQA模型相当,但每个token所需的查询头数量则减半。
English Abstract
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.