🤖 machine learning

Grouped Query Experts: Mixture-of-Experts on GQA Self-Attention

本論文は、Grouped Query Attentionに適用されるMixture-of-Experts層であるGrouped Query Experts（GQE）を提案しており、これはKey-Valueヘッドを密な状態に保ちつつ、トークンごとにクエリヘッドのエキスパートのサブセットを動的に選択することで、ダウンストリームの精度を損なうことなく、計算量を削減し効率を向上させるものである。

原著者： Vishesh Tripathi, Abhay Kumar

公開日 2026-06-23

📖 1 分で読めます☕ さくっと読める

CC0 1.0

原著者： Vishesh Tripathi, Abhay Kumar

原論文は CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/) のもとパブリックドメインに提供されています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

巨大な図書館（Transformer AIモデル）を想像してください。そこでは、物語を理解するために、すべての本（テキストのトークン）が他のすべての本と相互参照される必要があります。これが「自己注意（self-attention）」の仕組みです。

問題点：「一律対応」の司書
標準的なAIモデルでは、16人の専門司書（「アテンション・ヘッド」と呼ばれます）のチームが存在します。どんな本を取り出しても、16人全員がその本を読み、分析し、レポートを作成しなければなりません。

もしその本が「the」やカンマのような単純な単語であれば、16人もの専門家は必要ありません。1人か2人で十分でしょう。
しかし、もしそれが複雑な科学用語であれば、16人全員が必要になるかもしれません。
ところが現状では、システムはすべての単語に対して、強制的に16人全員に作業を行わせています。これは、たった一つの音符を奏でるために、オーケストラ全体を雇うようなものです。

既存の解決策：グループ化クエリ注意（GQA）
この論文が登場する前、研究者たちはコストを節約するために、司書をグループ化することを試みました。16のユニークなチームを作る代わりに、2人の司書が同じ「キーと値（KV）」のメモ（参照資料）を共有する8つのチームを作りました。これによりメモリは節約できましたが、16人の司書全員がいまだにすべての単語を読まなければなりませんでした。 それは、より小さなファイルキャビネットを用意したものの、従業員全員がいまだに建物全体を歩き回らなければならないような状態でした。

新しい解決策：グループ化クエリ・エキスパート（GQE）
著者らは、よりスマートなシステムである**グループ化クエリ・エキスパート（GQE）**を提案しています。これは、16人の司書を「混合エキスパート（Mixture of Experts）」システムに変えるようなものですが、少しひねりが加えられています。

セットアップ： 彼らは8つのグループの参照メモ（KVヘッド）をそのまま維持します。この部分は常に「高密度（dense）」であり、常にアクティブです。なぜなら、参照資料へのアクセスは安価だからです。
ルーター： 各グループ内には、複数の「エキスパート」司書（クエリ・ヘッド）が存在します。あらゆる単語に対して、スマートな「ルーター（交通整理員）」がその単語を見てこう問いかけます。「本当にこのグループ内の4人のエキスパート全員が、この単語を読む必要があるだろうか？」
選択： ルーターは、その特定の単語に対して、上位1つまたは2つのエキスパート（k=1またはk=2）を選び出し、実際に作業を行います。グループ内の他のエキスパートたちはコーヒーブレイクに入ります。
セーフティネット： システムが混乱したり怠慢になったりしないよう、彼らは2つの特別な機能を加えています。
- 共有ヘッド： チームの安定性を保つため、常に1人の司書が当番として、すべての単語を読んでいます。
- 加重要約： システムは、選択されたエキスパートたちの成果をブレンドした「合意レポート」を作成します。これが極めて重要です。なぜなら、これによりルーターが自分の仕事をどれほど上手くこなしたかという明確な信号が得られ、どのエキスパートがどの単語に最適かを学習できるからです。

結果：賢さを失わずに高速化を実現
著者らは、300億語で学習させた小規模なモデル（2億5千万パラメータ）を用いてテストを行いました。

精度： GQEモデルは、すべての単語に対して16人の司書をフル稼働させた従来のモデルと同等の性能を示しました。人をスキップしたことで「愚か」になることはありませんでした。
速度： クエリ・ヘッドの作業を約半分に省略したため、高速化されました。
- 短い物語の場合、わずかに高速化しました（1.15倍）。
- 長い物語（小説一冊のようなもの）の場合、1.7倍から1.8倍速くなりました。

なぜこれが重要なのか
この論文は、「読む」プロセスを条件付き（必要な時だけ作業を行う）にすることで、回答の質を損なうことなく、長い会話や文書分析のスピードを大幅に上げることができると主張しています。

落とし穴（限界）
著者らは、これは比較的規模の小さいモデルでテストされたものであることに注意を促しています。彼らは、これがまだ数兆パラメータを持つ巨大なモデルでも機能することを証明してはいません。また、「ルーター」は非常に慎重に訓練する必要があります。「セーフティネット（共有ヘッドと加重要約）」なしに、単にエキスパートをランダムに選ぶと、モデルの性能は actually（実際には）低下してしまいます。

要約すると、GQEは、特定のタスクに対して適切な専門家だけが出勤するチームを雇うようなものであり、仕事の質を落とすことなく、長い作業においてプロセスを大幅に高速化します。

技術要約：Grouped Query Experts (GQE)

問題提起

自己注意機構（Self-attention）は、Transformerアーキテクチャにおける計算上のボトルネックであり、特に長文コンテキストのシーケンスにおいて、トークン間のペアワイズな相互作用が二次関数的にスケールする際に顕著となります。標準的な高密度（dense）アテンションメカニズムは、トークンの情報量や難易度に関わらず、すべてのトークンに対して同一のセットのアテンションヘッドを適用します。この一様な活性化は、情報の少ないトークン（例：句読点、助詞など）が、内容のある単語と同じリソースを消費してしまうため、計算効率の低下を招きます。Mixture-of-Experts (MoE) は、Transformerのフィードフォワード（MLP）ブロックに対して条件付き計算を導入することに成功していますが、同様のスパース性をアテンションメカニズムに適用することは依然として未解決の課題です。既存のGrouped-Query Attention (GQA) は、複数のクエリヘッドが単一のKey-Value (KV) ヘッドを共有することでメモリ帯域幅とKVキャッシュのサイズを削減していますが、依然としてすべてのトークンに対してすべてのクエリヘッドを評価しています。

手法：Grouped Query Experts (GGQE)

著者らは、GQAモデルのアテンションブロックにMoEの原理を統合した手法であるGrouped Query Experts (GQE) を提案しています。コアとなる設計思想は、GQAの密なKVキャッシュの利点を維持しつつ、クエリヘッドの計算を条件付きにすることです。

アーキテクチャとルーティング

Grouped Expert Pool: 各固定GQAグループ内（複数のクエリヘッドが単一のKVヘッドを共有する構成）において、クエリヘッドは「エキスパート」のプールとして扱われます。
条件付き活性化: 各トークンに対して、ルーターが各GQAグループから $k$ 個のクエリヘッド・エキスパートのスパースなサブセットを選択します。
Dense KV Path: 極めて重要な点として、Key-Value投影およびKVキャッシュは密な状態のまま変更されません。スパース性はクエリ側の計算にのみ適用されます。
出力の構築: 学習の安定性を確保するために、特定のメカニズムを通じて出力が構築されます：
- ハード連結 (Hard Concatenation): 各グループで選択された $k$ 個のエキスパートの出力を連結します。
- 再正規化された加重和 (Renormalized Weighted Sum): ルーターの確率を用いて、選択されたエキスパート出力の加重和を計算します。これにより、離散的なtop- $k$ 選択に伴う非微分性の問題を解決し、言語モデルの損失がルーターに到達するための微分可能なパスを提供します。
- 常時稼働の共有ヘッド (Always-On Shared Head): ルーティングに関わらず、すべてのトークンに対して追加のアテンションヘッドが計算されます。これは学習中の安定したアンカーとして機能し、ルーターが単一のエキスパートに崩壊（collapse）することを防ぎ、ベースラインとなるアテンションパターンを常に利用可能にします。

学習と損失

標準的な言語モデルの損失に加え、GQEは、ルーターがグループ内のエキスパートプールにトークンを均等に分散させるよう促し、「エキスパートの崩壊」を防ぐためのロードバランシング補助損失を採用しています。

主な貢献

GQAにおけるMoEの定式化: 本論文は、すべてのKVヘッドを密な状態に保ちながら、各トークンに対して $k$ 個のエキスパートを活性化させるという、GQAグループに特化したMoEを定式化しました。これにより、GQAのメモリプロファイルとの整合性を保っています。
品質を損なわない計算削減: 著者らは、GQEが、すべてのクエリヘッドを活性化させるGQAベースラインと同等のダウンストリーム精度を実現しつつ、スパースなサブセットのクエリヘッド・エキスパートのみを活性化できることを示しました。
安定化メカニズム: 本論文は、アテンションにおける成功的なスパース・ルーティングには、適切な学習信号をルーターに提供するための「再正規化された加重和スロット」と、レイヤーを固定するための「常時稼働の共有ヘッド」という2つの具体的なコンポーネントが必要であることを特定しました。

実験結果

実験は、FineWeb-Eduデータセットを用い、250Mパラメータ規模の30Bトークン予算の条件下で行われました。

精度: メインの設定（16クエリヘッド、8 KVグループ、 $k=1$ ）において、GQEは8個のルーティングされたエキスパートと1個の共有ヘッドを活性化させ、合計16個のクエリ操作のうち9個を実行します。この削減（全クエリ操作の約56%を活性化）にもかかわらず、GQEはダウンストリームの精度（HellaSwag, PIQA, ARC-Easy）において、全活性化GQAベースラインと同等の性能を発揮し、ベースラインの55.86に対し、平均56.04を達成しました。
アブレーション研究: 再正規化された加重和スロット、または共有ヘッドを取り除くと、性能が低下する（それぞれ約55.18および55.43に低下）ことが示されており、スパース性だけでは、特定のアーキテクチャ上の安定化装置なしには不十分であることを裏付けています。
スループットと加速: GQEは、長文コンテキストの領域で大幅な高速化を示しています。
- 短文シーケンス（2kトークン）では、ルーティングのオーバーヘッドにより、高速化は緩やかです（約1.15倍）。
- 長文コンテキスト領域（4kから1024kトークン）において、GQEはGQAベースラインに対して1.67倍から1.80倍の高速化を達成します。これは、不要なクエリ計算をスキップするためであり、アテンションのコストが二次関数的に増大するにつれて、その節約効果がスケールするためです。

意義と主張

本論文は、GQEがモデルの品質を損なうことなく、長文コンテキストのTransformerにおける自己注意の計算コストを削減する実行可能な道筋を提供すると主張しています。その意義は以下の通りです：

効率性: テストされた構成において、性能を維持したまま、有効なクエリヘッドの計算量を約半分に削減します。
スケーラビリティ: 高速化の恩果はコンテキスト長が増すほど大きくなり、アテンションの二次関数的なスケーリングのボトルネックに対処します。
設計上の洞察: 本研究は、アテンションへのMoEの転用には、精度低下を伴うことが多いナイーブなスパースアテンションの試みとは異なり、慎重なアーキテクチャ設計（具体的には共有ヘッドと再正規化されたルーティング信号）が必要であることを浮き彫りにしました。

著者らは、結果が250Mスケールのものであることを認め、より大規模なスケールでの確認が必要であると謙虚に述べています。彼らは、この結果をベースラインに対する明確な「改善」ではなく、同等の性能で計算量を削減できたという「一致（match）」として位置づけています。