グループ化クエリエキスパート：GQA自己注意機構における混合エキスパート

要旨

自己注意はTransformerの性能の中心であり、長いコンテキスト長においては、トークン間のペアワイズ相互作用がシーケンス長の二乗でスケールするため、Transformerの中で最も計算コストの高い部分となることが多い。標準的な密な注意は、トークンの難易度や情報量に関わらず、すべてのトークンに同一の注意ヘッドセットを適用する。この一様な活性化は、特にシーケンスが長くなり注意コストが急速に増大するにつれて、計算を無駄にする可能性がある。我々は、グループ化クエリ注意（GQA）の上に混合エキスパート層を重ねたGrouped Query Experts（GQE）を提案する。各GQAグループ内で、ルーターはトークンごとにk個のクエリヘッドエキスパートを選択し、すべてのキー・バリュー（KV）ヘッドは密なまま変更されない。これにより、GQEはGQAのKVキャッシュの利点を維持しつつ、アクティブなクエリヘッドの計算のみを削減する。250Mパラメータ規模で30Bトークンの固定予算において、GQEはトークンあたりのクエリヘッドを半分だけ活性化しながら、すべてのヘッドを活性化するGQAベースラインと同等の下流タスク精度を達成する。

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.