GQA：マルチヘッド事前学習チェックポイントからの汎用マルチクエリアテンション構築

大規模言語モデルが一語ずつ文章を生成する背後では、次の単語を順に予測する自己回帰デコードが行われています。モデルは新しい単語と過去の全単語との関連度をアテンション機構で計算します。過去の単語について毎回再計算するのを避けるため、計算済みのKey（参照用の索引見出し）とValue（取り出すべき要点メモ）を「KVキャッシュ」としてメモリに保持し再利用します。

仕事の現場として、複数の調査員が過去の議事録を調べながら報告書を一文ずつ執筆する場面を考えてみましょう。文を書き加えるたび、過去の議事録全体を読み直す代わりに、手元の検索用見出し（Key）と要点メモ（Value）を参照して作業を進めます。これがKVキャッシュの役割です。

しかし文章が長くなり同時処理数（バッチサイズ）が増えるほど、保持すべきメモの量は膨大になります。単語を1つ出すたびに過去の全ステップのKVキャッシュをプロセッサへ読み込む必要があり、このメモリ転送の待ち時間が推論全体の深刻なボトルネックとなっていました。

本論文発表当時、主流だったのはMulti-Head Attention（MHA）でした。MHAでは複数の調査員（Queryヘッド）が専用の索引見出し（Key）と専用メモ（Value）を独立して保持します。精密な文脈把握により高い品質を維持できる一方、調査員の数に比例してメモが増加し、推論時にメモリ帯域を激しく消費する弱点がありました。

これに対し推論時のメモリ負担を解消するために考案されたのがMulti-Query Attention（MQA）です。MQAではすべての調査員が一組の共通メモを参照します。保持するメモ量が激減し推論速度は向上しますが、全調査員に同一メモを強制するため情報の多様性が損なわれ、特定のタスクで品質が低下する代償を抱えていました。

当時の現場は、高品質だがメモリ消費が重いMHAか、高速だが品質低下を招きうるMQAかという両極端の二者択一に直面していました。

さらに実務上、モデル再学習にかかる莫大な計算コストという壁がありました。当時すでに膨大な計算資源を投じて事前学習された大規模MHAモデルが多数存在していましたが、推論を高速化するためだけにMQAモデルを一から事前学習し直すことは極めて非現実的でした。

こうした背景のもと、Google ResearchのJoshua Ainslie氏らが挑んだ問いは明快でした。第一に、MHAの高品質を損なわずにMQAに迫る推論速度を引き出す中間的なアテンション構造を設計できないか。そして第二に、手元の学習済みMHAモデル資産を活用し、その効率的な構造へ低コストで適応させることができないか、という実践的課題でした。

著者らが提案したGrouped-Query Attention（GQA）の核心は、KeyとValueの共有度を「全個別（MHA）」か「単一共有（MQA）」かの二元論ではなく、中間の設計変数として定式化した点にあります。

GQAではQueryヘッドを均等なグループに分割し、同一グループ内でのみ一組のKeyとValueを共有させます。調査チームの例で言えば、全員に個別メモを持たせるのでも全員で一冊を共有するのでもなく、「事実確認班」「トーン調整班」といったグループを編成し、班ごとに一組の専用索引と専用メモを共有して作業を進める形態です。

ただし、この例えが実際の研究と乖離する境界にも注意が必要です。人間の調査員と異なり、アテンションのKeyとValueは多次元ベクトル空間における線形射影数値であり、内積計算により機械的に重みが求まる数学的手続きに過ぎません。比喩は共有の空間的範囲を直感的に捉えるためのものです。

なおGQAには、全体のQueryヘッド数がグループ数で割り切れる整数比でなければならない設計制約があり、約数の範囲でメモリ削減量と表現力のバランスを選択します。

推論効率化のために既存の学習済みMHAモデルを一から再学習する必要があっては、実用的な普及は望めません。そこで著者らは、既存MHAの重みパラメータを活用してGQAモデルを初期化するチェックポイント変換手法を考案しました。

MHAでは各Queryヘッドに対応する独立したKeyとValueの射影行列が存在します。著者らは同一グループに割り当てられる複数のKeyヘッドの射影重みを行列レベルで算術平均し、Valueヘッドも同様に平均化して、新しい共有Key・Valueヘッドの初期パラメータを生成しました。

調査チームの例で言えば、各自の個別メモを持ち寄り班のベースメモを作る作業に似ています。ただし重みの平均化は人間による文章の要約ではなく、高次元行列の数値を機械的に足し合わせて割る演算です。そのため平均化直後のパラメータは元の精緻な関係を完全には再現できず、そのままでは性能低下が生じます。

重みの平均化による表現のズレを解消し、新しいグループ共有構造にモデルを順応させるため、著者らは「アップトレーニング（Uptraining）」と呼ばれる追加学習を導入しました。

アップトレーニングでは、重みを平均化して初期化したGQAモデルに対し、元の事前学習計算量のわずか5パーセントに相当する計算リソースを用いて追加学習を行います。この適応ステップにより、Queryヘッドは共有されたKeyとValueの表現空間を適切に参照できるように調整され、低下していた精度を回復させます。

ゼロからの事前学習に比べ95パーセントもの計算量を削減しながら既存の高価なモデル資産を推論効率の高い構造へ再構築できるこの手法は、極めて実用性の高い工学的成果となりました。GQAは、重みの平均化による初期化とわずかなアップトレーニングによる適応という一連の変換手順が確立されたことで、産業界で即座に活用できる実践的技術となったのです。

著者らは提案手法の有効性を検証するため、確立されたエンコーダ・デコーダ型モデルであるT5をベースに実験を行いました。規模による影響を把握するため、T5-Base、T5-Large、そして当時最大級のT5-XXLという3つのパラメータスケールで体系的な比較を実施しています。

比較基準には元のMHAと単一ヘッドを共有するMQAを設定しました。その上でKey/Valueヘッドのグループ数を段階的に変えた複数のGQA構成を用意し、同一データセットと学習条件で統制実験を行いました。モデル初期化には既存MHAの重み平均化を用い、5パーセントの計算量でアップトレーニングを行う手順を一貫して適用しています。

評価指標には、言語モデルの予測の不確実性を測るPerplexity（パープレキシティ）、学術ベンチマークの下流タスク精度、そして実際の推論レイテンシを測定しました。

著者らの実験結果において注目すべきは、既存MHAからアップトレーニングされたGQAが、MHAに迫る高い品質を維持しながらMQAに匹敵する高速な推論処理速度を達成した点です。

とりわけ最大規模のT5-XXLにおける検証では、適切なグループ数を選択したGQAは、元のMHAと同等水準のPerplexityおよび学術ベンチマークスコアを達成しつつ、推論レイテンシにおいてMQAとほぼ同等の処理時間を記録したと著者らは報告しています。MQAで顕著だった品質劣化が、GQAの中間的なグループ共有によって効果的に抑制された形です。

この結果は、ヘッド共有の効果が単純な直線的関係ではないことを示しています。KeyとValueのヘッド数を全数からある程度絞り込んでもモデルの表現力は急激には失われず、KVキャッシュの読み出し負荷はヘッド削減に比例して大きく軽減されます。著者らの報告は、品質と速度の間に有利な実用バランスが存在することを実証しました。

原著論文の評価を読み解く上では、実験が立証できた範囲と前提を厳密に区別する必要があります。

第一に、著者らの定量的成果はT5という特定の構造と設定で得られた結果です。そのため、あらゆる構造、異なるコーパス、極端に長い文脈長、多様なハードウェア構成の下で常に同じグループ数が最適になるとは結論づけられません。

第二に、本実験では重みの平均化と5パーセントのアップトレーニングが一体で機能しています。そのため記録された情報のみから、性能回復の何割が初期化方法により、何割が追加学習によるものかを個別に切り分けることはできません。

第三に、推論速度の向上幅は、デコード時のKVキャッシュ読み出しがどれほど支配的なボトルネックになっているかに依存します。演算器の計算能力が限界の状況や他のオーバーヘッドが大きい環境では、示された速度向上がそのまま再現されるとは限りません。

本論文発表以降、大規模言語モデルの実用化は商用サービスの大規模展開へと急速に移行しました。この運用局面で最大の課題となったのが、多数のリクエストを同時に処理し長い文脈を扱う際のGPUメモリ枯渇とメモリアクセス帯域の限界でした。

GQAの提案は、この実務上の壁に対する直接的な処方箋となりました。KVキャッシュを従来の数分の一に圧縮できるため、同一ハードウェアで保持できる文脈長を伸ばし、同時処理可能なバッチサイズを拡大することが可能になりました。Hugging Faceの論文コミュニティでも260件の賛同票（upvotes）を集め、KVキャッシュを数分の一に圧縮しながらMHAと同等のベンチマークスコアを達成できる実用的な着地点として、高い評価が寄せられました。

ハードウェアの現実に即した解決策を提示したことが、本研究を推論最適化の中心技術へと押し上げる原動力となりました。

本研究の影響は、その後に登場した世界的な主要オープンモデル群の設計に明確に刻まれています。普及の大きな契機となったのは、Metaが公開したLlama 2の最大規模モデルであるLlama 2（70B）において、GQAが中核アーキテクチャとして採用されたことでした。これにより巨大モデルでもGQAが安定して高い性能と優れた推論効率を発揮することが広く認知されました。

Llama 2（70B）での成功を皮切りに、後継モデルのLlama 3では全サイズでGQAが基本設計として組み込まれました。さらにMistralやGemma、Qwenといった主要オープンモデル群でも採用が相次ぎ、現代のモデルを構築・運用する上での必須アーキテクチャとして定着しました。

後続モデルが取り入れたのは曖昧な類似性ではなく、Key/Valueヘッド数をグループ共有するというGQAの数理的メカニズムの直接の継承です。注意機構の設計はGQAによって事実上の標準設計へと更新されました。

現代の視点からGQAを見つめ直したとき、まず理解すべき限界は、すべての計算フェーズが一様に高速化されるわけではないという点です。

LLMの推論には、入力プロンプト全体を一括処理するPrefillフェーズと、一語ずつ順に出力するデコード（自己回帰生成）フェーズが存在します。入力プロンプトを一括処理するPrefill段階は並列行列演算が中心で、ボトルネックは主にプロセッサの計算能力にあります。これに対しデコード段階では、単語を出すたびに過去のKVキャッシュを呼び出すため、データの転送速度がボトルネックになります。

GQAの恩恵は主に後者のデコード段階におけるキャッシュ読み出し削減にあります。そのためプロンプト処理が中心のタスクや、入力が長く出力が極めて短い用途では速度向上効果は限定的です。またモデルの初期事前学習にかかる計算量そのものを削減するわけではない点も混同してはなりません。

GQAの構造には、実務上の適用において留意すべき固有の制約も存在します。

第一に、Queryヘッド数とKey/Valueヘッド数は整数比でなければならないため、選択可能なグループ数は元のヘッド数の約数に限られます。連続的に微調整できるわけではなく、離散的な候補から選ぶ必要があります。

第二に、もともとのヘッド数が少ない小規模モデルでは、削減できるKey/Valueヘッドの絶対数が限られるため、KVキャッシュ圧縮や速度向上の恩恵が相対的に小さくなります。一方で共有を極端に強めれば、表現力の低下や品質劣化のリスクが再燃します。

第三に、極限の大規模モデルにおける調整の難しさや、特定のハードウェアおよび分散並列構成（テンソル並列など）への依存性が挙げられます。GPU間でアテンションヘッドをどう分散配置するかという実装上の都合で採用可能なグループ数が制約されるケースもあり、原著の設定を異なる基盤へ無条件に当てはめることはできません。

これらの制約を踏まえた上で引き継ぐべき教訓は、推論時の物理的ボトルネック（メモリ帯域とデータ転送量）を起点として内部表現を逆算・最適化するという設計思想です。

現代の実務者がGQA採用モデルの導入や調整を検討する際は、評判を鵜呑みにせず自らの運用環境で客観的証拠を確認しなければなりません。具体的には、対象業務のプロンプト長と生成長、想定される同時リクエスト数（バッチサイズ）、稼働ハードウェア環境において、Prefill時間とデコード時間を個別に実測することが不可欠です。

さらにPerplexityのような平均的指標だけでなく、業務で要求される具体的な下流タスクの精度が許容水準を満たしているかを検証し、推論コスト削減と出力品質のバランスが真に成立しているかを実証データに基づいて判断する姿勢が求められます。