MQAとは
Multi-Query Attention
全てのクエリヘッドが単一のキー・バリューヘッドを共有し、KVキャッシュを最小限に抑えるAttentionの構造
ひとことで言うと
質問役は複数いても、参照するメモ(キー・バリュー)役は1人だけにして、推論をとても軽くする工夫。
概要
MQA(Multi-Query Attention)は、Multi-Head Attentionの各クエリヘッドがそれぞれ専用のキー・バリューヘッドを持つ代わりに、全てのクエリヘッドが単一の共有されたキー・バリューヘッドを参照するように設計されたAttentionの構造。 キー・バリューヘッドを1組だけに減らすことで、生成時にGPUメモリへ保持しておくKVキャッシュのサイズをヘッド数分の1まで削減でき、特にメモリ帯域幅がボトルネックになりやすい自己回帰生成のスループットを改善できる。 一方で、複数のクエリヘッドが単一のキー・バリューヘッドしか参照できないため、通常のMulti-Head Attentionと比べてモデルの表現力がやや制限され、タスクによっては精度がわずかに低下することが報告されている。 この精度低下を緩和しつつメモリ削減効果もある程度保つ折衷案として、後にGrouped-Query Attention(GQA)が提案された。近年の大規模モデルでは、このGQAを採用することが多い。
背景
自己回帰的なテキスト生成では、トークンを1つ生成するたびに過去のキー・バリューをGPUメモリから読み込む必要があり、このメモリ帯域幅の制約が生成速度のボトルネックになりやすい。 MQAは、キー・バリューヘッドの数を極限まで減らすことで、このメモリ帯域幅のボトルネックを解消する目的で考案された。
歴史
2019年: ShazeerがMQAを提案する論文「Fast Transformer Decoding: One Write-Head is All You Need」を発表。 2020年代前半: PaLMなど一部の大規模モデルが、生成速度を優先してMQAを採用。 2023年以降: 精度と効率のバランスに優れたGQAへ主流が移行。
アーキテクチャ
H個のクエリヘッドはそれぞれ独立して学習されるが、キーとバリューの射影行列は全ヘッドで共通の1組のみを使う。 各クエリヘッドは、この共有されたキー・バリューに対してAttentionスコアを計算し、それぞれ異なる出力を得る。 生成時にキャッシュすべきキー・バリューは1組分のみとなり、ヘッド数分のキャッシュを持つ通常のMulti-Head Attentionと比べてメモリ使用量が大きく減る。
コード例
モデル設定でのMQA指定例
from transformers import LlamaConfig
config = LlamaConfig(
num_attention_heads=32,
num_key_value_heads=1, # MQA: 全クエリヘッドで1組のキー・バリューを共有
)利点
欠点
比較
- GQA — GQAはMQAとMulti-Head Attentionの中間的な構成で、精度低下を抑えつつメモリ削減効果を得る折衷案として提案された
- Self-Attention — MQAは通常のMulti-Head Self-Attentionのキー・バリューヘッドを1組に絞った効率化手法
- KVキャッシュ — MQAはキー・バリューヘッドを1組に減らすことで、KVキャッシュのメモリ使用量を最大限に削減する
関連用語
よくある質問
MQAで精度が下がりやすいのはなぜ？
複数のクエリヘッドが単一のキー・バリューヘッドしか参照できないため、通常のMulti-Head Attentionと比べて表現できるパターンの多様性が制限されるため。
MQAは現在も使われている？
生成速度を最優先する一部の用途では使われるが、多くの新しいモデルは精度低下を抑えられるGQAを採用する傾向にある。