Attentionの変種図鑑 — MQA・GQA・スライディング窓・線形注意

MQA・GQA・スライディング窓・線形注意は、それぞれ勝手に生まれた別々の工夫ではありません。KVキャッシュの大きさを決める一本の掛け算式があり、変種は「どの項を叩いたか」で分類できます。系譜を式の上に並べ直すと、何を捨てて何を得たのかが一目で見えます。

Fast Transformer Decoding: One Write-Head is All You Need

「注意」は一種類ではない

Attention機構を1から理解するで見たのは、教科書に載っている素の自己注意でした。全単語が全単語を見渡し、内積で重みを決め、Valueを混ぜる。美しい仕組みですが、実際に動いているモデルの設定ファイルを開くと、そのままの形はむしろ少数派です。

会議の比喩を続けます。あなたは参加者全員の名札（Key）を見比べ、発言（Value）をメモに混ぜる係でした。参加者が10人ならこれで困りません。ところが実際のLLMは、会話が進むほど参加者が増え続ける会議室にいます。1万トークンの文脈とは、1万人分の名札と発言を捨てずに机に積み上げたまま、次の一言を考えるということです。

変種と呼ばれるものたちは、この積み上がった机をどう畳むかの答えです。そして重要なのは、それぞれが好き勝手に生まれた別々の工夫ではないことです。畳める場所は数えられるほどしかなく、変種はそのどれを叩いたかで整理できます。

高いのは2か所ある

まず「何が高いのか」を分けます。混同されがちですが、コストの正体は場面によって別物です。

学習時と、プロンプトを一度に読み込むプリフィルの段階では、$n \times n$（$n$ は文の長さ）の点数表そのものが重荷になります。文が2倍なら表は4倍。ここは計算量とメモリの両方の話です。

一方、1トークンずつ吐き出す生成の段階では、点数表は1行しかできません。新しいトークンのQueryと過去全部のKeyとの照合、それだけです。にもかかわらず生成は遅い。理由は計算ではなく、過去のKとVを丸ごと読み出すことにあります。これがKVキャッシュで、その大きさは全部が掛け算で決まります。

要するに、KVキャッシュの大きさは「層数 × KVヘッド数 × ヘッドの幅 × 溜まったトークン数 × 同時実行数」の掛け算で、それにKとVの2種類ぶんと、数値1個あたりのバイト数を掛けただけです。順に読みます。先頭の $2$ はKとVの2種類ぶん。$L$ は層の数（層ごとに別のキャッシュを持つ）。$h_{kv}$ はKVヘッドの本数。$d_{\mathrm{head}}$ は1ヘッドあたりの幅。$n$ は溜まったトークン数。$B$ は同時に捌いている会話の数。$p$ は数値1個のバイト数（FP16なら2）。

全部が掛け算だ、という点が肝心です。どれか1つを $1/8$ にすれば、全体がそのまま $1/8$ になります。逆に言えば、削れる場所は式の中にしかありません。だから変種の系譜は、この式のどこを叩いたかで分類できます。

• $h_{kv}$ を減らす → MQA / GQA（そしてMLA）
• 実効的な $n$ を減らす → スライディング窓
• そもそも $n \times n$ の表を作らない → 線形注意
• $p$ を減らす → KVキャッシュの量子化

MQA — 名簿を1冊にする

最初の一手は、いちばん乱暴な引き算でした。2019年にNoam Shazeerが提案した Multi-Query Attention（MQA） は、Queryヘッドは複数のまま、KとVのヘッドを1本だけにします。質問は人それぞれでよいが、参照する名簿は1冊を全員で回し読みする設計です。式の $h_{kv}$ が $1$ になるので、キャッシュはヘッド数ぶんの1に縮みます。

なぜ速度に直結するのか。生成では1トークン作るたびにKVキャッシュを丸ごと読み出します。演算自体は軽く、読み出しの帯域が律速です。読むバイト数が $1/h$ になれば、そのまま待ち時間が縮む。FLOPsはほとんど変わらないのに速くなる、という一見奇妙な現象の正体はこれです。

代償は表現力です。ヘッドごとに違う見方を持てるのが利点だったのに、見る対象を全ヘッドで共有してしまう。品質の低下は小さいと報告されましたが、ゼロではありません。