Self-Attention and QKV

Self-Attention は、各 token が同じ sequence 内の他の token を参照し、自分の表現を文脈化する仕組みです。Transformer の多くの能力は、各 token が「どの token から、どれくらい情報を受け取るか」を学習できることに支えられています。

自作概念図。Query は「何を探すか」、Key は「何に反応するか」、Value は「実際に混ぜ込む情報本体」と考えると直感的です。

Query、Key、Value の直感

Self-Attention では、入力 hidden state $X$ から三つの vector を作ります。

$$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$

名前	直感	役割
Query	自分が探している条件	どの token を参照したいかを問い合わせる
Key	自分が持っている見出し	Query と照合され、参照されやすさを決める
Value	実際に渡す中身	Attention weight によって加重平均される情報

Query と Key の内積は、「Query 側の token が Key 側の token をどれだけ参照したいか」を表す score になります。この score を softmax で正規化すると、Value を混ぜるための重みになります。

Scaled Dot-Product Attention

Transformer の基本式は次の通りです。

$$\mathrm{Attention}(Q,K,V) = \mathrm{softmax}\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)V$$

ここで、$M$ は mask です。Decoder-only LLM では causal mask を入れるため、未来 token を参照できません。

        key token
        t1  t2  t3  t4
query t1  ●   .   .   .
      t2  ●   ●   .   .
      t3  ●   ●   ●   .
      t4  ●   ●   ●   ●

なぜ $\sqrt{d_k}$ で割るのか

$QK^\top$ の内積は、dimension $d_k$ が大きくなるほど値の scale が大きくなりやすいです。そのまま softmax に入れると、分布が極端に尖り、gradient が不安定になります。そこで $\sqrt{d_k}$ で割り、softmax に入る値の scale を調整します。

Attention weight と Value の加重平均

Attention の出力は、単に「似ている token を探す」だけではありません。Query と Key で参照比率を作り、その比率で Value を混ぜます。

ある token の出力表現は、過去 token の Value を attention weight で加重平均したものです。そのため、同じ token でも文脈によって表現が変わります。

Self-Attention と Cross-Attention

種類	Q の元	K / V の元	例
Self-Attention	同じ sequence	同じ sequence	LLM の token 間 attention
Cross-Attention	生成側 hidden	別 modality / encoder output	Stable Diffusion の text conditioning、encoder-decoder translation

Self-Attention は同じ入力の中で情報を集約します。Cross-Attention は、ある情報源から別の情報源を参照します。Stable Diffusion では、U-Net の latent feature が Query を出し、text embedding が Key / Value として使われます。

Multi-Head Attention の意味

一つの head だけでは、すべての関係を一種類の attention pattern で表すことになります。Multi-Head Attention では、head ごとに異なる subspace で attention を計算します。

$$\mathrm{MHA}(X)=\mathrm{Concat}(\mathrm{head}_1,\ldots,\mathrm{head}_h)W^O$$

ある head は近い token を見て、別の head は long-range dependency を見て、さらに別の head は構文的関係や entity の対応を見ます。このように複数の見方を並列に持てることが Multi-Head Attention の強みです。

KV Cache との関係

Autoregressive decoding では、次 token を生成するたびに過去 token の Key と Value が必要になります。しかし、過去 token の Key / Value は一度計算すれば変わりません。そのため、推論時には Key / Value を cache して再利用します。詳しくは KV Cache を参照してください。

関連ページ

• Transformer Architecture
• KV Cache
• Long Context and Position Encoding
• LLM Inference Optimization
• Stable Diffusion

主なソース

• Attention Is All You Need: https://arxiv.org/abs/1706.03762
• Zenn, 30分で完全理解するTransformerの世界: https://zenn.dev/zenkigen_tech/articles/2023-01-shimizu
• Qiita, TransformerのSelf AttentionのQKVを直感的に解説する: https://qiita.com/kenmatsu4/items/1b3853a3314ab66eb2a3
• The Illustrated Transformer: https://jalammar.github.io/illustrated-transformer/