Grouped Query Attention (GQA)

Grouped-query attention（GQA）は、Transformer向けのマルチヘッドアテンションの一種で、クエリヘッドのグループが1つのキーヘッドと1つのバリューヘッドを共有します。たとえばLlama 3.1 8Bでは、各層で32個のクエリヘッドと8個のキー・バリューヘッドを組み合わせます。これは、各クエリヘッドが専用のキーと値を持つマルチヘッドアテンション（MHA）と、すべてのクエリヘッドが単一のキーヘッドとバリューヘッドを共有するマルチクエリアテンション（MQA）の中間に位置します。Google ResearchのJoshua Ainslieらは、2023年のEMNLPで発表した「GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints」でこの手法を紹介し、GQAモデルはMHAに近い品質を保ちながら、MQAに近いデコード速度を実現すると示しました。Llama 2 70B、Llama 3ファミリー、Mistral 7B、Gemma 2などのオープンな大規模言語モデル（LLM）がGQAを採用しています。また、Gemma 2を基盤とするGoogleのPaliGemma 2など、それらを基盤とする視覚言語モデル（VLM）も、言語デコーダーでGQAを継承しています。

GQAの仕組み#

自己アテンションでは、各トークンを各ヘッドのクエリ、キー、値に射影し、各ヘッドがクエリとキーを比較します。GQAでは、クエリヘッドの数はそのままに、キーヘッドとバリューヘッドの数を減らします。クエリヘッドの各グループは、グループで共有するキーと値を使ってアテンションを計算します。クエリヘッドが8個、キー・バリューヘッドが2個の場合、各キー・バリューヘッドは4個のクエリヘッドを受け持ちます。

各ヘッド内のアテンション計算は、従来どおりのスケールド・ドット積アテンションです。グループが1つなら層はMQAになり、グループ数がクエリヘッド数と同じならMHAになります。PyTorchでは、scaled_dot_product_attentionのenable_gqa引数で、このレイアウトを直接サポートしています。

import torch
import torch.nn.functional as F

torch.manual_seed(0)
q = torch.randn(1, 8, 16, 64)  # 8 query heads, 16 tokens, 64 dimensions per head
k = torch.randn(1, 2, 16, 64)  # only 2 key heads
v = torch.randn(1, 2, 16, 64)  # and 2 value heads

out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)

# Equivalent multi-head attention: copy each key-value head to its group of 4 query heads
k_full, v_full = k.repeat_interleave(4, dim=1), v.repeat_interleave(4, dim=1)
reference = F.scaled_dot_product_attention(q, k_full, v_full, is_causal=True)

print(f"Output {tuple(out.shape)}, same as copied heads: {torch.allclose(out, reference, atol=1e-6)}")
print(f"Cached values per token per layer: {k[0, :, 0].numel() * 2} with GQA, {k_full[0, :, 0].numel() * 2} with MHA")

Output (1, 8, 16, 64), same as copied heads: Trueと256 with GQA, 1024 with MHAが出力されます。MHAが保存するキーと値の4分の1だけを使って、8個すべてのクエリヘッドが出力を生成します。

GQAがKVキャッシュを縮小する仕組み#

テキスト生成では、モデルは以前のすべてのトークンのキーと値をKVキャッシュに保持し、新しいトークンごとにキャッシュ全体を読み取ります。そのため、デコード速度は演算よりもメモリ帯域幅に制限されます。キャッシュサイズはキー・バリューヘッドの数に比例します。Llama 3.1 8Bは32層で、各層に128次元のキー・バリューヘッドを8個持つため、16ビット精度でトークンあたり128 KiBを保存します。最大131,072トークンのコンテキストウィンドウでは、合計16 GiBになります。MHAでは同じモデルにトークンあたり512 KiB、最大長で64 GiBが必要です。空いたメモリを使ってコンテキストを長くしたり、連続バッチ処理のバッチを大きくしたり、トークンあたりの推論レイテンシを下げたりできます。Ultralytics LLMインターフェースは、セルフホスト型Llama 3.1 8Bなど、OpenAI互換サーバー向けのbase_urlを受け付けます。キー・バリューヘッドが8個のこのサーバーは、32ヘッドのMHAレイアウトと比べ、同じキャッシュメモリで会話トークンを4倍保持できます。

モデルをGQAに変換する#

GQAでは、ゼロからトレーニングする必要はありません。著者らは、既存のMHAチェックポイントを変換する際、各グループ内のキーとバリューヘッドを平均プーリングし、その後、元の計算量の約5%を使って事前学習を継続しました。この工程をアップトレーニングと呼んでいます。共有ヘッドが1つだけの極端な形であるMQAは、Noam Shazeerが2019年に発表しました。最速のレイアウトですが、品質が低下し、トレーニングが不安定になることもありました。中間のグループ数を使うと、品質の大部分を取り戻せます。

GQAと関連手法との比較#

• マルチヘッド潜在アテンション（MLA）: DeepSeek-V2とともに導入され、各トークンのキーと値を、小さな共有潜在ベクトルに圧縮します。
• PagedAttention: KVキャッシュを固定サイズのブロックに割り当て、キャッシュを埋めるヘッドの数に関係なく、メモリの断片化を防ぎます。
• スライディングウィンドウアテンション: 層ごとにキャッシュするトークン数を制限します。Mistral 7BはGQAと組み合わせて使用しています。
• Flash attention: 厳密なアテンションをより高速に計算し、GQAレイアウトをネイティブにサポートするGPUカーネルです。