Grouped-Query Attentionとは?KVキャッシュを減らして推論を速くする仕組みと使い道

Grouped-Query Attentionは、TransformerのAttentionで使うKey/Valueヘッドを共有し、品質を保ちながら推論を軽くしやすくする技術です。GQAの仕組み、既存手法との違い、実務での使い道を日本語で整理します。

参考文献

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, Sumit Sanghai

論文を見る

今回の論文

今回取り上げるのは、Joshua Ainslie らによる 2023 年の論文「GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints」です。公開元は EMNLP 2023 で、研究分野は Transformer アーキテクチャ、LLM 推論最適化、Attention 設計です。URL は https://aclanthology.org/2023.emnlp-main.298/ です。

この論文を選んだ理由は、LLM の実運用でずっと効いてくる「KV キャッシュが重い」「応答速度を上げたい」という問題に、かなり実装しやすい形で答えているからです。しかも完全に新しいモデルを一から学習するのではなく、既存の multi-head attention のチェックポイントから変換し、少ない追加学習で使える点が実務的です。

どんな技術か

Grouped-Query Attention、略して GQA は、Transformer の attention で使う Query ヘッド数は多く保ちつつ、Key と Value のヘッド数だけを減らして共有する技術です。

通常の multi-head attention では、各 Query ヘッドがそれぞれ独立した Key/Value ヘッドを持ちます。これは表現力が高い一方で、推論時には過去トークンぶんの Key/Value を大量に保持する必要があり、KV キャッシュのメモリ使用量が大きくなります。長い文脈や大規模バッチを扱うと、ここがかなり重くなります。

GQA は、この Key/Value を複数の Query ヘッドで共有します。極端に言えば、全 Query が 1 組の Key/Value を共有するのが Multi-Query Attention です。GQA はその中間で、たとえば 32 個の Query ヘッドに対して 4 組や 8 組の Key/Value ヘッドを持たせます。これにより、品質を落としにくくしながら、キャッシュ量とメモリアクセス量を減らしやすくなります。

つまり GQA は、「attention の柔軟さをできるだけ残しつつ、推論の重さを下げるためのバランス設計」です。いまの LLM 推論基盤を理解するうえでかなり重要な技術です。

課題

この技術が解決しようとしている課題は、LLM のデコード時に attention がメモリ帯域と KV キャッシュに強く縛られることです。

何が難しいのかというと、自己回帰生成ではトークンを 1 つずつ増やすたびに、過去全トークンの Key と Value を参照する必要があるからです。モデルサイズが大きくなり、コンテキストが長くなるほど、保持すべき KV キャッシュも増えます。計算器そのものより、重いキャッシュの読み書きが遅さの原因になる場面も少なくありません。

既存の方法では、multi-head attention は品質面で安定していても重く、逆に multi-query attention はかなり軽くできる一方で、品質低下が出ることがありました。特に大規模モデルでは、「速さのために全部の Key/Value を 1 つにまとめる」と精度や学習済み能力を削りやすい、というトレードオフがあります。

なぜこの課題を解く必要があるかというと、チャット、コード生成、RAG、エージェント、社内検索補助のようなシステムでは、推論コストとレイテンシがそのまま UX と利益率に響くからです。特に長い会話履歴や長文文書を扱う場合、KV キャッシュがメモリを圧迫して同時実行数を下げたり、GPU の利用効率を悪くしたりします。

実際の AI システムでは、モデルの重みだけでなく、リクエストごとに持つ KV キャッシュがボトルネックになります。GQA はこの「attention の中でも、特に Key/Value の持ち方が重い」という問題に正面から手を入れた技術です。

用語解説

Multi-Head Attention
Transformer の標準的な attention 方式で、複数のヘッドが別々の視点で関連性を計算します。GQA はこの構造を完全に捨てるのではなく、Query は多ヘッドのまま保ち、Key/Value だけ共有化する点が重要です。
KVキャッシュ
自己回帰生成で、過去トークンの Key と Value を保存して再利用する仕組みです。LLM 推論を速くする基本技術ですが、コンテキスト長や同時実行数が増えるとメモリを大きく消費します。GQA が直接削減したい対象です。
Multi-Query Attention(MQA)
全 Query ヘッドで 1 組の Key/Value を共有する方式です。推論はかなり軽くなりますが、品質が落ちやすいことがあります。GQA は MQA と通常の multi-head attention の中間に位置する設計です。
アップトレーニング(uptraining)
既存の学習済みチェックポイントを別構造へ変換し、少量の追加学習で性能を回復させる考え方です。この論文では、multi-head のモデルを GQA や MQA に変換し、元の学習コストの約 5% で再適応させる点が実務上の肝になります。
デコーダ推論
LLM が 1 トークンずつ出力する生成フェーズです。学習時よりも KV キャッシュとメモリ帯域の影響が大きく、attention の設計変更が体感速度や同時実行性能に直結しやすい場面です。

技術の仕組み

この論文のポイントは、単に「Key/Value を減らしたら軽い」という話ではありません。既存の multi-head モデル資産を活かしながら、どこまで共有しても品質を保てるかを、構造と学習の両面から整理しているところが本質です。

基本アイデア

通常の multi-head attention では、h 個の Query ヘッドに対して、Key と Value も h 個あります。MQA では Key/Value を 1 組にします。GQA では、その中間として h / g 個の Query ヘッドを 1 グループにまとめ、各グループが同じ Key/Value ヘッドを共有します。

たとえば 32 ヘッドの attention で 8 グループに分ける場合、Query は 32 のままですが、Key/Value は 8 組になります。これだけで、各トークンについて保存する KV キャッシュ量を大きく減らせます。

重要なのは、Query 側の多様性を完全には捨てないことです。各 Query ヘッドは依然として別々に attention を計算できるので、MQA よりも表現力を残しやすくなります。GQA は、この「共有しすぎない」バランスによって、推論効率と品質の中間点を狙っています。

モデル構造

論文では、T5 系の encoder-decoder モデルを対象に、特にデコーダ側 attention の Key/Value ヘッド共有を扱っています。自己注意でもクロス注意でも同じ考え方を適用できますが、実務的に効くのはまずデコード時の KV キャッシュです。

GQA への変換では、複数の Key ヘッドと Value ヘッドをそれぞれグループ単位で平均し、新しい共有ヘッドとして初期化します。つまり完全な再学習ではなく、既存の multi-head attention が持っていた重みをなるべく活かして出発します。

この初期化が重要なのは、いきなりランダムな共有ヘッドへ置き換えると、学習済み能力が大きく壊れるからです。平均を取ることで、近い役割を持つヘッド群の情報をある程度保ったまま圧縮できます。

学習方法

この論文の実務的な価値は、GQA 用に最初から大規模事前学習しなくてよい点です。著者らは、既存の multi-head チェックポイントを GQA あるいは MQA に変換したあと、元の事前学習計算量の約 5% の追加学習で性能回復を試しています。論文ではこれを uptraining と呼んでいます。

やっていることはシンプルで、変換後のモデルをそのまま追加学習するだけです。ただし意味は大きく、推論に有利な attention 形状へモデルを寄せながら、元の性能を実用レベルまで戻せることを示しています。

これはプロダクトの観点だとかなり重要です。既存モデルを全部捨てず、推論コスト改善のために後から構造変更できる可能性を示しているからです。

推論時に何が軽くなるのか

GQA の利点は、主に Key/Value の保存量とメモリアクセス量が減ることです。生成中は新しいトークンが出るたびに、そのトークンの Key/Value をキャッシュへ追加し、次のトークン生成では過去のすべての KV を読みます。

Key/Value ヘッド数を減らすと、各トークンあたりに持つキャッシュが小さくなります。その結果、長文生成や多数同時リクエストでメモリ使用量を抑えやすくなり、デコーダ推論のスループット改善につながります。

特に最近の LLM サービングでは、計算量そのものより HBM 帯域やキャッシュサイズが詰まりやすいため、GQA のような構造変更はかなり効きます。これは論文外の実装知見も含む一般化ですが、GQA が現在の推論系モデルで広く採用される背景でもあります。

なぜ MQA より安定しやすいのか

MQA は全 Query が同じ Key/Value を使うため、非常に軽い代わりにヘッドごとの差が消えやすいです。attention の多様性が必要なタスクでは、これが品質低下につながることがあります。

GQA は、Query の多様性は維持しつつ、Key/Value だけをグループ単位で共有します。そのため、「全部独立」ほど重くなく、「全部共有」ほど乱暴でもない設計になっています。論文では、この中間設計が perplexity と速度のバランスを最も取りやすいことが示されています。

実験と結果

論文では、GQA が本当に multi-head attention と MQA の中間としてうまく機能するかを、事前学習後のアップトレーニングと下流タスク評価の両方で検証しています。

何を検証したのか

主な検証は 2 つです。1 つ目は、既存の multi-head チェックポイントを MQA や GQA に変換したあと、少ない追加学習で性能をどこまで回復できるかです。2 つ目は、推論速度の改善と品質低下のバランスです。

つまり論文の問いは、「GQA は本当に軽いのか」だけではなく、「既存資産を活かしながら、どの程度まで品質を維持できるのか」です。

使ったモデルと評価設定

著者らは T5 系モデルを使い、C4 で事前学習されたチェックポイントから出発しています。評価では pre-training の perplexity に加え、SuperGLUE、自然言語生成、要約、翻訳など複数の下流タスクを見ています。

推論面では、デコーダ速度を token/s で比較し、multi-head attention、MQA、複数設定の GQA を並べて評価しています。これにより、精度と速度の両面でどの構成が実務向きかを比較しやすくしています。

どのような結果が出たのか

Table 1 では、T5-XXL の multi-head attention が平均性能 47.2、推論時間 1.51 秒/サンプルだったのに対し、5% uptraining した MQA-XXL は平均性能 46.6、推論時間 0.24 秒/サンプルでした。かなり速くなる一方で、品質は少し落ちています。

これに対して GQA-8-XXL は、平均性能 47.1、推論時間 0.28 秒/サンプルでした。つまり MHA-XXL にかなり近い品質を保ちながら、推論時間は MQA にかなり近い水準まで短縮できています。比較対象として MHA-Large は 0.37 秒/サンプル、平均性能 46.0 なので、GQA-8-XXL はそれより高品質で、しかも速いという結果です。

著者らは、GQA が「速度と品質の良い折衷」だと結論づけています。グループ数を増やせば品質は multi-head に寄り、減らせば速度は MQA に寄ります。この連続的な調整ができる点も重要です。

結果から何が言えるのか

この結果から言えるのは、推論最適化は単純な二択ではないということです。従来は「品質重視なら multi-head」「速度重視なら MQA」と見えがちでしたが、GQA によってその間に実用的な設計空間があるとわかりました。

アブレーションでも、Key/Value ヘッドの初期化は単純平均が最も安定し、GQA は変換直後でも比較的まともに動き、5% の uptraining で大きく改善、10% では改善が逓減する傾向が示されています。また 1 グループから 8 グループへ増やす程度なら速度低下は比較的小さく、8 グループ前後が良い中間点だと報告されています。

既存チェックポイントから変換して少ない追加学習で実用域へ持ち込めるなら、研究用だけでなくプロダクト側でも採用しやすくなります。これはアーキテクチャ研究としてだけでなく、運用コスト改善手段としても価値があります。

何に使える?

GQA は、長い文脈や高い同時実行数を扱う LLM アプリで特に使いやすい技術です。品質をなるべく維持しながら、KV キャッシュの重さを減らしたい場面に向いています。

チャットや社内ナレッジ検索

社内アシスタントやサポートチャットでは、会話履歴や検索結果を長くプロンプトへ入れがちです。このとき GQA を使ったモデルは、長文脈でもキャッシュの圧迫を抑えやすく、同時接続数を増やしやすくなります。RAG の品質以前に、そもそも遅くて使われない問題を減らせる可能性があります。

コード生成と開発支援

コード補完やレビュー支援では、プロジェクト文脈や差分、ログを長く渡したくなることがあります。GQA はこうした長めのコンテキストで推論基盤を軽くしやすいため、応答時間と GPU コストの改善に効きます。特に IDE 補完や CI 上の自動レビューのようにリクエスト数が多い用途で有効です。

エージェントの長履歴処理

AI エージェントは、ユーザー対話だけでなく、ツール実行ログ、観測、計画、中間メモを抱えます。こうした長履歴ワークロードでは、モデル精度と同じくらい、履歴をどれだけ安く保持できるかが重要です。GQA はエージェント基盤そのものを変えるというより、長い履歴を扱う LLM レイヤーの効率化として効きます。

GPU メモリが限られる環境

オンプレミス GPU や小規模クラスタで LLM を回す場合、最大の制約は計算性能よりメモリであることが多いです。GQA はこの制約に対してかなり素直に効くため、同じハードウェアでより長いコンテキストや多い同時実行を取りにいくときの選択肢になります。

開発や事業へのヒント

この論文から得られるヒントは、LLM の差別化はモデルサイズだけでなく、attention の持ち方でも作れるということです。特に推論コストが収益に直結するサービスでは、GQA 的な設計思想はかなり重要です。

「精度を上げる」前に「KV を減らす」を考える

AI アプリ開発では、遅い・高いという問題に対して、つい小さいモデルへ変えるか量子化を先に考えがちです。ただ、品質低下が困るなら、まず KV キャッシュの構造を軽くする方向を検討する価値があります。GQA はその代表例で、モデル本体の知識量を大きく落とさずに推論基盤を改善しやすいです。

既存資産を捨てずにアーキテクチャ改善できる

この論文の uptraining という発想は、小規模チームにも参考になります。全部を一から作り直すのではなく、既存モデルや既存重みを活かしながら、構造変更に少量追加学習で追従させる考え方です。独自ドメインモデルや社内向けモデルでも、完全再学習より現実的な改善策になる可能性があります。

長文脈プロダクトでは原価改善の余地が大きい

RAG、議事録要約、コードベース検索、問い合わせ自動化のように、長いコンテキストを扱うサービスでは、GQA の効果が積み上がりやすいです。1 リクエストごとの差が小さく見えても、同時実行数や月間トークン量で見ると収益差になりやすいです。

今後注目すべき方向性

今後見るべきなのは、GQA 単体というより「attention のどの部分を圧縮しても品質が落ちにくいか」という方向性です。KV キャッシュ量の削減、head sharing、sliding window、speculative decoding、state space 系モデルはすべて、推論ボトルネックを別角度から削っています。GQA はその中でも、比較的実装しやすい基礎技術として押さえておく価値があります。

限界

GQA にも限界はあります。まず、どのグループ数が最適かはモデルサイズやタスクに依存します。共有を強くしすぎると MQA に近づき、品質低下が出やすくなります。逆にグループ数を増やしすぎると、multi-head attention に近づいて軽量化の利点が薄れます。

また、この論文は T5 系モデルでの検証が中心で、すべてのモデル系列で同じ最適点になるとは限りません。現在のデコーダ専用 LLM へ応用する考え方は自然ですが、どの程度効くかはモデル設計や実装依存です。

実装面でも注意があります。単にヘッド数を減らせば終わりではなく、推論エンジン側がその構造を効率よく扱える必要があります。KV キャッシュ管理、カーネル実装、テンソル並びが噛み合わないと、理論ほど速度が出ないことがあります。

さらに、GQA は attention 設計の改善であって、モデル知識の不足や hallucination を直接解決する技術ではありません。RAG やエージェントの品質問題を直すものではなく、あくまで推論基盤を軽くする技術だと切り分けて考える必要があります。

よくある質問

Q. GQA は MQA と何が違うのですか?

A. MQA は全 Query ヘッドが 1 組の Key/Value を共有しますが、GQA は Query ヘッドを複数グループに分け、各グループごとに Key/Value を共有します。つまり GQA は MQA より共有を弱めた中間設計で、品質を保ちやすいのが違いです。

Q. GQA を使うと必ず速くなりますか?

A. KV キャッシュ量とメモリアクセス量は減りやすいので、長文脈や高同時実行では有利になりやすいです。ただし実際の速度は推論エンジン、GPU、バッチサイズ、カーネル実装に左右されます。モデル構造だけでなくサービング実装も重要です。

Q. 既存モデルをそのまま GQA 化できますか?

A. 論文では、multi-head チェックポイントを変換し、少量の追加学習で性能を回復させています。理屈としては既存モデル資産を活かせますが、実際には対象モデルの構造や学習環境に応じた実装が必要です。

Q. GQA はどんなプロダクトで特に効果がありますか?

A. 長い履歴を持つチャット、RAG、コード生成、エージェントのように、KV キャッシュが膨らみやすい用途です。1 回の性能差より、同時実行数や運用コストの差として効いてくることが多いです。

Q. GQA があれば量子化や speculative decoding は不要ですか?

A. 不要にはなりません。GQA は KV キャッシュ由来の重さを減らす技術で、量子化は重みの軽量化、speculative decoding はデコード逐次性の削減です。ボトルネックが違うので、実運用では組み合わせて効かせることが多いです。

今日の学び

この論文は、LLM 推論で attention の Key/Value キャッシュが重いという課題を扱いました。そこに対して、Query は多ヘッドのまま保ちつつ、Key/Value だけをグループ単位で共有する GQA を提案し、少量の uptraining で品質と速度のバランスが良いことを示しました。

ここから得られるヒントは、推論最適化は量子化だけではないということです。長文脈や多同時実行のプロダクトを作るなら、モデル精度だけでなく KV キャッシュの構造そのものを設計対象として見るべきだとわかります。

関連記事

推論最適化

LayerSkipとは?追加モデルなしでLLM推論を高速化する早期退出と自己Speculative Decodingの仕組み

LayerSkipは、LLMの中間層から早めに出力を作れるよう訓練し、同じモデルだけでspeculative decodingを行う推論最適化技術です。追加ドラフトモデルを持たずに高速化する考え方、仕組み、実務での使い道を整理します。

参照論文:LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

推論最適化

Speculative Decodingとは?LLM推論を品質そのままで高速化する仕組みと使い道

Speculative Decodingは、小さなドラフトモデルで先読みした候補を大きな本命モデルがまとめて検証し、出力分布を変えずにLLM推論を高速化する技術です。仕組み、評価結果、実務での使い道まで日本語で整理します。

参照論文:Fast Inference from Transformers via Speculative Decoding

推論最適化

PagedAttentionとは?LLMサービングのKVキャッシュ断片化を減らし高スループット化する仕組みと使い道

PagedAttentionは、LLM推論で肥大化するKVキャッシュをページ単位で管理し、メモリ断片化と重複を抑える技術です。vLLMの中核になったこの仕組みを、課題、実装アイデア、実験結果、開発への応用まで日本語で整理します。

参照論文:Efficient Memory Management for Large Language Model Serving with PagedAttention