GQAとは
Grouped-Query Attention
複数のクエリヘッドをいくつかのグループに分け、グループ内でキー・バリューヘッドを共有するAttentionの構造
ひとことで言うと
質問役は大勢のまま、参照するメモ(キー・バリュー)役を少人数に共有させて、推論を軽くする工夫。
概要
GQA(Grouped-Query Attention)は、Multi-Head Attentionにおいて各クエリヘッドが個別のキー・バリューヘッドを持つ代わりに、複数のクエリヘッドをいくつかのグループへ分けるAttentionの設計。 同じグループ内のクエリヘッドは、キー・バリューヘッドを共有する。 通常のMulti-Head Attentionはクエリヘッドと同数のキー・バリューヘッドを保持するため、生成時に保持するKVキャッシュのサイズがヘッド数に比例して大きくなるが、GQAはキー・バリューヘッドの数をグループ数まで減らすことでKVキャッシュのメモリ使用量を削減する。 全クエリヘッドが単一のキー・バリューヘッドを共有するMulti-Query Attention(MQA)ほど大胆な共有ではないため、MQAで見られがちな精度低下を抑えつつ、通常のMulti-Head Attentionに近いKVキャッシュ削減効果を得られる。 Llama 2の70Bモデルをはじめ、多くの大規模モデルで推論効率と精度のバランスを取る設計として採用されている。
背景
LLMの生成時には各トークンのキー・バリューをKVキャッシュとしてGPUメモリに保持する必要があり、クエリヘッドの数が多いほどこのメモリ消費が増大する。 GQAは、精度への影響を抑えつつキー・バリューヘッドの数だけを減らすことで、KVキャッシュのメモリ消費と生成速度のトレードオフを改善する目的で考案された。
歴史
2023年: AinslieらがGQAを提案する論文「GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints」を発表。 2023年: Llama 2の70Bモデルなど大規模モデルの推論効率化手法として採用。 2024年以降: Llama 3やMistralなど多くの主要オープンウェイトモデルで標準的な構成として定着。
アーキテクチャ
H個のクエリヘッドをG個のグループへ均等に分割し、各グループが1組のキー・バリューヘッドを共有する。 G=1のとき全クエリヘッドが単一のキー・バリューヘッドを共有するMQAと一致し、G=Hのとき各クエリヘッドが個別のキー・バリューヘッドを持つ通常のMulti-Head Attentionと一致する。 この中間的なグループ数Gを調整することで、精度とKVキャッシュ削減効果のバランスを取る。
コード例
モデル設定でのGQAグループ数指定例
from transformers import LlamaConfig
config = LlamaConfig(
num_attention_heads=32,
num_key_value_heads=8, # GQA: 8グループでキー・バリューを共有
)利点
欠点
- グループ数Gの設計次第で精度とメモリ効率のバランスが変わり、タスクに応じた調整が必要になる
- 通常のMulti-Head Attentionと比べると、キー・バリューの表現力そのものは制限される
- アーキテクチャの変更を伴うため、既存モデルへ適用するには追加の学習コストがかかる
比較
- MQA — MQAは全クエリヘッドが単一のキー・バリューヘッドを共有する極端な形で、GQAはその中間的な構成にあたる
- Self-Attention — GQAは通常のMulti-Head Self-Attentionのキー・バリューヘッド構成を変更した効率化手法
- KVキャッシュ — GQAはキー・バリューヘッドの数を減らすことで、KVキャッシュのメモリ使用量を直接削減する
関連用語
よくある質問
GQAとMQAの違いは？
MQAは全クエリヘッドが単一のキー・バリューヘッドを共有するのに対し、GQAはクエリヘッドをいくつかのグループへ分け、グループごとにキー・バリューヘッドを共有する。GQAはMQAより精度を保ちやすい。
GQAはなぜKVキャッシュを削減できる？
生成時に保持するキー・バリューの数はヘッド数に比例するため、クエリヘッドと同数だったキー・バリューヘッドをグループ数まで減らすことで、保持すべきKVキャッシュのサイズを直接削減できる。