体系 › Transformerの仕組み
·無料·論文·11分で読めます

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

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

要点

  • KVキャッシュの大きさは「層数 × KVヘッド数 × ヘッドの幅 × 溜まったトークン数 × 同時実行数」の掛け算で決まる。
  • MQAはKとVのヘッドを1本だけにし、GQAはQueryヘッドをグループに分けてグループごとにKVを共有する。
  • スライディング窓注意は各トークンが直前 ww トークンだけを見るので、計算量は n×wn \times w になる。
  • 線形注意はsoftmaxをやめて過去を1つの箱に畳み込み、箱の大きさは文の長さによらない。
  • FlashAttentionは変種ではなく、出力は素の注意と数学的に完全に一致し、KVキャッシュは1バイトも減らない。
対象textタスクattention

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


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

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

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

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

高いのは2か所ある

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

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

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

MKV=2×L×hkv×dhead×n×B×pM_{\mathrm{KV}} = 2 \times L \times h_{kv} \times d_{\mathrm{head}} \times n \times B \times p
(1)

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

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

FIG 1文長が伸びたとき O(n²) と O(n) がどれだけ離れるか。縦軸を「線形」に切り替えると、n² の曲線だけが画面の上へ消えていく — 長文が高いのはこの離れ方のせいです

MQA — 名簿を1冊にする

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

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

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

GQA — 間を取る

2023年、Ainslieらが出した折衷案が Grouped-Query Attention(GQA) です。hh 本のQueryヘッドを gg 個のグループに分け、グループごとに1組のKVを共有します。g=1g = 1 ならMQA、g=hg = h ならMHA。両端に既存手法を持つ連続したつまみです。

もう一つの貢献は、既存のMHAチェックポイントからグループ内のKVヘッドを平均して初期化し、元の事前学習に比べればごくわずかな追加学習(uptraining)でGQA版に変換できると示したことです。ゼロから学習し直さなくてよいのが普及の決め手でした。Llama 2 の70Bは64本のQueryヘッドに対しKVヘッドが8本で、キャッシュはMHAの8分の1です。

スライディング窓 — 見る範囲を切る

ここまでは「誰の名札を持つか」の話でした。次は「どこまで見るか」を削ります。スライディング窓注意では、各トークンは直前 ww トークンだけを見ます。n×nn \times n の表のうち対角線に沿った幅 ww の帯だけを残し、外側は捨てる。計算量は n×wn \times w になり、nn に対して直線的にしか増えません。

そんなに切って文全体が分かるのか、と思うところですが、層を重ねると受容野が広がります。1層目で各トークンは ww 先まで見る。2層目が見る相手はすでに ww 先を含むので、実効的に 2w2w 先まで届く。kk 層あれば概ね k×wk \times w の範囲に伝わる、という伝言リレーです。

2020年のLongformerが長文向けに整備し、2023年のMistral 7Bが窓幅4096でデコーダに採用して一気に知られました。生成時のうまみが大きく、キャッシュに残すのは直近 ww トークンぶんだけなので、メモリが文脈長に比例して増えなくなります。弱点も明確で、窓の外を正確に引用する類のタスクは苦手です。Longformerは一部のトークンにだけ全域を見る権限(global attention)を与えて補いました。

線形注意 — softmaxを外して順番を入れ替える

最後は、いちばん大胆な変種です。n×nn \times n の表をそもそも作らない。

素の注意が n×nn \times n を経由するのは、softmaxが QK⊤QK^\top 全体にかかっているからです。行列積には結合則があるので、本来なら (QK⊤)V(QK^\top)V を Q(K⊤V)Q(K^\top V) と計算し直せます。前者は n×nn \times n の表を作り、後者が作るのは d×dd \times d の小さな箱です。ところが間に非線形なsoftmaxが挟まると、この入れ替えができません。

ならばsoftmaxをやめよう、というのが線形注意です。指数関数の代わりに、非負の特徴写像 ϕ\phi を通した内積で類似度を測ります。

outi=ϕ(qi)⊤∑j≤iϕ(kj)vj⊤ϕ(qi)⊤∑j≤iϕ(kj)\mathrm{out}_i = \frac{\phi(q_i)^\top \sum_{j \le i} \phi(k_j) v_j^\top}{\phi(q_i)^\top \sum_{j \le i} \phi(k_j)}
(2)

つまり、分子は「これまでの名札と発言の組を全部足し込んだ箱」に、いまの質問を当てているだけです。分母は重みの合計を1にするための正規化。肝心なのは ∑j≤i\sum_{j \le i} の部分で、過去が1つの箱に畳み込まれているという点です。箱の大きさは文の長さによりません。

S = np.zeros((d_k, d_v))   # 状態: 名札×発言 の箱
z = np.zeros(d_k)          # 正規化用の合計

for k, v, q in stream:               # 1トークンずつ
    S += np.outer(phi(k), v)         # 箱に足し込むだけ
    z += phi(k)
    out = (phi(q) @ S) / (phi(q) @ z + 1e-6)

これが Katharopoulos らの論文タイトル「Transformers are RNNs」の意味です。生成時の線形注意は、固定サイズの状態を1つ持って更新するRNNと同じ形になり、KVキャッシュが「長さに比例して伸びるリスト」から「大きさの決まった行列」に変わります。

では、なぜ標準にならなかったのか。捨てたものが大きすぎたからです。softmaxは、必要とあらばたった1つのトークンに重みを振り切れます。「この代名詞が指すのは3語前のあの名詞だ」という鋭い参照は、その振り切りで実現されている。非負の特徴写像の内積では同じ鋭さを作りにくく、実用モデルでの採用が限定的なのはこの一点によるところが大きいと考えられています。

FIG 2温度を下げると1本だけが突き出し、上げると横並びに潰れる。softmaxはこの「振り切り」を必要に応じて出せるのが強みで、線形注意が再現しにくいのがまさにこの尖った形です

混ぜて使う、という現在地

実際のモデルは、これらを排他的に選んではいません。GQAを土台に、層ごとにスライディング窓と全域注意を交互に配置する組み合わせが普通です。窓付きの層で近くを密に見て、時々挟まる全域の層で遠くを拾う分業になります。KVそのものを低ランクの潜在ベクトルに圧縮してキャッシュする MLA(Multi-head Latent Attention, DeepSeek-V2) のような手も出ています。

一つだけ線を引いておきます。FlashAttentionは変種ではありません。n×nn \times n の表を高速メモリの中でブロックごとに処理し、遅いメモリに書き出さずに済ませる実装上の工夫で、出力は素の注意と数学的に完全に一致します。近似ではないので品質は変わらず、そしてKVキャッシュは1バイトも減りません。

現場ではこう使う

この知識が要るのは、モデルを選ぶ人と、推論基盤を運用する人です。

モデル選定のとき。 config.json の num_attention_heads と num_key_value_heads を見比べます。同数ならMHA、KVが1ならMQA、その間ならGQA。キャッシュ量は num_key_value_heads に比例するので、この1行だけで「同じGPUで何並列さばけるか」の桁がおおよそ決まります。sliding_window があれば窓付きです。

サーバを立てるとき。 vLLMなら --max-model-len、--gpu-memory-utilization、--kv-cache-dtype が主な握りどころです。同時実行数はおおむね「KVキャッシュに回せるメモリ ÷ 1リクエストあたりのキャッシュ量」で決まるので、冒頭の掛け算式がそのまま見積もりに使えます。

設計レビューで問われる形。 「MQAとGQAの違いを、どこが速くなるかまで説明してください」。筋道は、削るのはKVヘッド数 → 減るのは生成のたびに読み出すバイト数 → 律速はメモリ帯域なので速くなる、そしてFLOPsはほとんど変わらない、まで言えると締まります。

落とし穴を3つ。

  1. テンソル並列とKVヘッド数。 KVヘッドが8本のモデルを8枚を超えるGPUで分割すると割り切れません。実装によってはKVヘッドを複製して辻褄を合わせるので、期待したメモリ削減が目減りします。並列度を決める前にKVヘッド数を見ること。
  2. 「128k対応」の読み方。 窓付きモデルの対応長は、全トークンを等しく参照できることを意味しません。長文の一箇所を正確に引用させる用途では、実データで確かめてから採用する。
  3. FlashAttentionへの期待違い。 入れてもKVキャッシュは減りません。OOMに効くのは多くの場合こちらではなく、文脈長・バッチ・KVのdtypeの側です。

まとめ

見積もりの手順はKVキャッシュを1から理解するにまとめてあります。変種を選ぶ前に、自分のワークロードでどの項が効いているかを数えるのが順序として正しいやり方です。

🎬 この内容の解説動画(Manim・字幕つき)

参考文献

  1. Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150論文ページ·PDF
  2. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245論文ページ·PDF
  3. Longformer: The Long-Document Transformer. arXiv:2004.05150論文ページ·PDF
  4. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. arXiv:2006.16236論文ページ·PDF

本記事は上記論文の本文にもとづいて執筆しています。数値・主張は原典を優先してください。

運営: TechAthletes · 編集方針 · 出典の付け方 · 訂正の記録

コメント

コメントにはログインが必要です