Attention は過去を読み直している — Q/K/V と O(T²) の壁を最小実装で覗く

Part 1 で、LLM は次のトークンを 1 つずつ予測して繋げているだけだ、という土台を確かめました。そのとき最後にひとつ宿題を残しました。「1 トークンずつ・逐次ということは、過去のトークンを毎回参照しているということです」と。

今回はその「参照」の正体を開けます。モデルが過去をどう読み直しているのか、その仕組みが self-attention です。そして、この読み直し方には避けられない代償があります。文脈が長くなるほど、計算量とメモリが系列長の 二乗 で膨らむのです。これを数式の暗記ではなく、numpy だけで書いた 50 行ほどの実装を動かして、自分の手元の数字で見ていきます。

この連載の現在地 — Attention は「過去への重みづけ」

• ここまで分かっていること — Part 1 で、LLM は次トークン予測器で、生成はその関数を 1 トークンずつ回す逐次ループだと確かめました
• 今回新しく説明できること — その「過去を毎回参照する」中身である self-attention（Q/K/V）、未来を見ないための causal mask、そして系列長の二乗 O(T²) の壁
• 今回はやらないこと — 二乗の再計算を緩める KV キャッシュは Part 3、Q/K/V の重みを学習で獲得する話は Part 5 へ送ります
• PHOTON のどの主張につながるか — 「attention は系列長の二乗で重い」という事実が、最終回「なぜ最大 475 倍が成立し、なぜ品質が落ちるのか」を分解するときの土台の半分になります（もう半分は Part 3 のメモリの話です）

「過去を読み直す」とはどういう操作か

自己回帰のループの中で、モデルは次のトークンを決めるたびに、それまでに出てきた全部のトークンをもう一度見ています。ただし「見る」と言っても、すべてを平等に見るわけではありません。今の文脈にとって関係の深いトークンを重く、関係の薄いトークンを軽く見ます。

この「関係の深さに応じた重みづけ平均」が attention の本質です¹。たとえば「その犬は疲れていた。なぜなら一日中走っていたから。」という文で「走っていた」の主語を決めるとき、モデルは「犬」を重く、「疲れて」を中ぐらいに、「。」をほぼ無視して見る、という配分を学びます。問題は、この配分をどうやって計算するか、です。

self-attention の仕組み — Q・K・V を最小実装で

attention は 3 つの登場人物で動きます。各トークンの埋め込みベクトルに、3 つの別々の重み行列を掛けて作る Q（Query）、K（Key）、V（Value） です²。直観としては、Q が「今のトークンが何を探しているか」という問い合わせ、K が「各トークンが何を提供できるか」という見出し、V が「実際に渡す中身」にあたります。

手順はこうです。あるトークンの Q を、全トークンの K と内積して「どれだけ合うか」のスコアを出します。そのスコアを softmax で和が 1 になる重みに変え、その重みで全トークンの V を平均します。式で書けば softmax(Q Kᵀ / √d) · V の一行です²。√d で割っているのは、次元 d が大きいと内積が大きくなりすぎて softmax が尖り、勾配が消えやすくなるのを防ぐためです²。

numpy だけで書くと、本当にこれだけです。

def causal_self_attention(x, w_q, w_k, w_v):
    seq_len, d = x.shape
    q = x @ w_q                       # 各トークンの「問い合わせ」
    k = x @ w_k                       # 各トークンの「見出し」
    v = x @ w_v                       # 各トークンの「中身」
    scores = (q @ k.T) / np.sqrt(d)   # (T, T) ← ここが二乗の発生源
    future = np.triu(np.ones((seq_len, seq_len), dtype=bool), k=1)
    scores = np.where(future, -np.inf, scores)   # 未来を見えなくする
    attn = softmax(scores, axis=-1)   # 行ごとに和が 1 になる
    out = attn @ v
    return out, attn                  # (T, d) と重み行列 (T, T)

q @ k.T の一行に注目してください。T 個の Q と T 個の K の総当たりなので、ここで T × T のスコア行列ができます。各トークンが他の全トークンとの相性を計算する、という構造そのものが、この正方形の行列です。後半で重くなる犯人は、ここにいます。

なお、実際の Transformer は 1 つの attention ではなく、Q/K/V を複数の小さな部分空間に分けて並列に計算する multi-head attention を使います²。「複数の視点を同時に持つ」ための工夫ですが、計算量のオーダーは単一ヘッドと同じなので、本記事では 1 ヘッドで話を進めます。

未来を見てはいけない — causal mask

上のコードに future という行があります。これが生成にとって決定的に重要です。

文章を左から生成するとき、3 番目のトークンを予測している最中に 5 番目のトークンが見えてしまったら、それはカンニングです。学習も生成も成り立ちません。そこで、各トークンが自分より後ろのトークンを見ないようにします。実装は単純で、スコア行列の上三角（未来側）を −∞ にしてから softmax をかけるだけです¹。−∞ は softmax を通すと 0 になるので、未来への重みがちょうど消えます。nanoGPT などの実装も、下三角だけを残すこの方法を使っています³。

本当に未来が見えていないかは、小さい行列を表示すれば一目です。attention_from_scratch.py --demo を走らせると、T=8 のトークン列でこんな重み行列が出ます。

# causal attention 重み (T=8, d=16)  ※各行の和=1・上三角(未来)=0
[[1.   0.   0.   0.   0.   0.   0.   0.  ]
 [0.72 0.28 0.   0.   0.   0.   0.   0.  ]
 [0.16 0.81 0.03 0.   0.   0.   0.   0.  ]
 [0.12 0.1  0.5  0.28 0.   0.   0.   0.  ]
 [0.03 0.05 0.11 0.01 0.79 0.   0.   0.  ]
 [0.1  0.7  0.02 0.09 0.03 0.06 0.   0.  ]
 [0.08 0.12 0.18 0.18 0.36 0.03 0.05 0.  ]
 [0.17 0.07 0.1  0.02 0.1  0.03 0.02 0.5 ]]

行 i が「トークン i がどこをどれだけ見るか」です。右上の三角形がすべて 0 になっているのが、未来を見ていない証拠です。最初の行は自分しか見るものがないので 1.0、行が下がるほど見られる過去が増え、重みが分散していきます。どの行も足すと 1 になります（softmax なので当然です）。これが「過去への重みづけ平均」の正体です。

なぜ系列長の二乗で重くなるのか — O(T²) の壁

ここからが本題です。先ほどの T × T のスコア行列は、系列長 T が 2 倍になれば 4 倍の大きさになります。計算量で言うと、attention の中心部分は時間が O(T²·d)、メモリが O(T²) です⁴。T に対して二乗で効く、というのがポイントです。

理屈は分かったので、本当にそうなるか測ります。attention_from_scratch.py --bench は、系列長を変えながら「Q/K/V の射影にかかる時間」と「attention コア（スコア計算と重みづけ平均）にかかる時間」を分けて計測します。Apple M5 Pro・numpy 2.5.0・次元 d=128・7 回の中央値での結果です。

メモリの列は理論値そのものです。スコア行列は T × T の float32 なので、T=2048 では 2048 × 2048 × 4 バイト = ちょうど 16 MiB になります。これは 1 ヘッド・1 層ぶんの話で、実際のモデルでは層の数とヘッドの数とバッチサイズを掛けた量が必要になります。長い文脈が一気にメモリを食う理由がここにあります。

時間のほうも、T が 2 倍になるたびに attention コアの時間がどう変わるかを見ると、二乗がはっきり出ます。

attention コアは T が 2 倍になると、大きい T では約 4 倍（= 2²）に近づいています。きれいな二乗です。一方、射影（Q/K/V を作る掛け算）は約 2 倍（= 2¹）で、こちらは系列長に対して線形です。二乗で効く部分と線形で効く部分が、同じ計算の中に同居しているのが見えます。

ちなみに、この「近似なしの attention は二乗から逃げられない」というのは、経験則ではなく理論的にも示されています。強指数時間仮説（SETH）が崩れない限り、self-attention の計算量は系列長に対して必然的に二乗になる、という下限が知られています⁴。

「O(T²) だから常に重い」は、少し不正確です

ここで誇張を 1 つ取り除いておきます。「attention は O(T²) だから二乗で重い」という言い方は、大づかみには正しいのですが、小さい文脈ではそのとおりにはなりません。

さっきの倍率表をもう一度見てください。T が 64 → 128 のとき、attention コアの倍率は 2.23 でした。二乗なら 4 のはずなのに、半分ほどです。二乗の倍率（約 4）にきちんと近づくのは、だいたい T が 512 を超えたあたりからです。理由は、小さい T では二乗の項よりも、Q/K/V の射影（線形）や、関数呼び出し・メモリ確保といった一定のオーバーヘッドのほうが効いているからです。O(T²) はあくまで T が十分大きいときの漸近的な振る舞いで、短いプロンプトでは二乗項はまだ顔を出していません。

「計算量」と「実際の速度」が別物だという例は、ほかにもあります。FlashAttention という有名な手法は、計算量を O(T²) のままにしておきながら、巨大な T × T 行列を遅いメモリに書き出さない工夫（タイリング）で、attention の実速度を大きく改善しました⁵。オーダーが同じでも、メモリの動かし方次第で実速度はまったく変わるのです。だからこそ、オーダーの議論だけで終わらせず、自分の環境で測る価値があります。なお、ここで出した ms の値も M5 Pro・numpy という条件での数字で、機種や BLAS が変われば変動します。無料で動くことと速いことは別、という Part 1 の注意はここでも有効です。

O(T²) の壁の先へ

二乗の壁が見えたところで、その先に何があるかだけ予告します。

ひとつは KV キャッシュです。自己回帰生成では、1 トークン進むたびに過去の K と V を作り直すのは無駄です。一度計算した K/V を取っておいて使い回せば、毎ステップの計算は減ります。ただし、取っておく量は文脈の長さ・層の数・ヘッドの数に比例して増えるので、今度はメモリが新しいボトルネックになります。この「decode はメモリで詰まる」という話が Part 3 の主題で、Part 1 で generate_step が裏でやっていた cached decode の正体でもあります。

もうひとつは、二乗そのものを近似で崩すアプローチ（疎な attention や線形 attention）ですが、これは exact ではなくなります⁴。本連載では深追いせず、名前だけ置いておきます。

そして最終回 Part 6 の伏線です。富士通の PHOTON が掲げる「最大 475 倍」⁶は、ひとことで言えば実スループットを約 44 倍、メモリ効率を約 11 倍にして、その積が 475 という合成値です。今回見た O(T²) の計算量と、次回見るメモリの線形増加は、その「なぜ速く・軽くできるのか」を自分の言葉で分解するための、ちょうど 2 つの軸になります。

手元に残るもの と 次回予告

今回手を動かして残ったものを確認します。

• 動くコード — attention_from_scratch.py（causal self-attention の最小実装・causal 行列の表示・系列長スイープの実測）
• 測った数字 — attention コアは系列長が 2 倍で約 4 倍（O(T²)）、射影は約 2 倍（線形）、スコア行列は T=2048 で 16 MiB。そして「小さい T では二乗はまだ見えない」という観察
• メンタルモデル — attention は過去への重みづけ平均で、その「全 × 全」の構造が系列長の二乗を生む

次回 Part 3 では、この二乗の計算を生成のたびに繰り返すのがいかに無駄かを見て、KV キャッシュでそれを避けます。そして、避けた先に待っているメモリの壁を測ります。「過去を毎回作り直すのはもったいない」と感じられたら、次に進む準備ができています。

コードは サンプルコード（companion repo）の part-02 タグで再現できます（attention_from_scratch.py と、MLX も外部モデルも不要で通る --selftest）。clone と各 Part のタグは、上で挙げた companion repo か、この記事末尾のシリーズナビからたどれます。

参考文献