Transformerの解説記事を読むと、ほぼ必ず出てくる図がある。“The animal didn’t cross the street because it was too tired.” という文で、“it” から “animal” に向かって太い矢印(高いattention重み)が伸びている、あの図だ。「Attentionは文脈の中で本当に重要な単語を自動的に見つけ出す」という説明とセットで語られることが多い。

これを見るたびに、素朴な疑問があった。この「itがanimalに注目する」という賢い振る舞いは、学習済みモデルの中で起きていることなのか、それともQuery・Key・Valueという仕組みそのものが持っている性質なのか。 実際に手を動かして、QKVを自分で実装し、本物の単語ベクトルで試してみることにした。ついでに、Attentionの説明では必ず出てくるのに理由を深く考えたことがなかった「1/√d_kでスケーリングする」という一手間と、GPTで使われる「causal mask」についても実験で確かめた。

Self-Attentionが計算していること

Self-Attentionは、系列の各トークンについて、Query $Q$・Key $K$・Value $V$ の3つのベクトルを作り、次の式で「どのトークンにどれだけ注目するか」を計算する。

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

$QK^\top$ でQueryとKeyの内積(類似度)を全トークンのペアについて計算し、softmaxで正規化して「注目の重み」にする。この重みでValueを加重平均したものが出力になる。式はシンプルだが、$\sqrt{d_k}$ で割るという部分の意味は、実際に試すまで実感が持てていなかった。

実験1: 本物の単語ベクトルで「it」の注目先を見る

まず、あの有名な例文をそのままGloVe(学習済み単語埋め込み、50次元)に通し、Query=Key=単語埋め込みそのものとして(射影なしの)attentionを計算してみた。

import numpy as np
from gensim.models import KeyedVectors

def softmax(x, axis=-1):
    x = x - np.max(x, axis=axis, keepdims=True)
    e = np.exp(x)
    return e / np.sum(e, axis=axis, keepdims=True)

kv = KeyedVectors.load("glove50.kv")  # glove-wiki-gigaword-50

sentence = ["the", "animal", "did", "not", "cross", "the",
            "street", "because", "it", "was", "too", "tired"]
X = np.stack([kv[w] for w in sentence])  # (12, 50) 本物のGloVeベクトル

d_k = X.shape[1]
scores = X @ X.T / np.sqrt(d_k)
attn = softmax(scores)

“it” の行を見てみると、結果は次の通りだった。

it         0.1802
not        0.1452
because    0.1198
too        0.0989
the        0.0922
did        0.0810
was        0.0696
street     0.0358
animal     0.0310   ← 12語中11位
cross      0.0276
tired      0.0266

予想が外れた。 “animal” への注目は12語中で下から2番目、ほぼ最下位だった。次に、実際のTransformer層と同じように、ランダムに初期化した(まだ学習していない)Q・K射影行列を通してみたが、結果はさらに悪化し、“animal” は完全に最下位になった。

行列のヒートマップを並べるより、単語同士の位置関係で見た方が直感的だと思ったので、12個の単語ベクトルをPCAで3次元に落とし、“it” を中心に、各単語への実際のattention重み(ランダム射影後)を点の大きさ・色・“it”からの線の太さで表した3D散布図を作った(マウスでドラッグして回転できる)。

回転させてみるとよく分かるのだが、“tired”・“too”・“did”・“because” のような単語が “it” の近くに大きく暗い赤の点として現れる一方、肝心の “animal” は画面の隅の方に、小さく薄い黄色の点としてぽつんと浮いている。空間的な距離(=埋め込みの類似度)という観点で見ても、“it” と “animal” はそもそも近くにいない。 つまり、「itがanimalに注目する」というあの有名な図は、QKVという仕組みが自動的にもたらす性質ではなく、実際にその文脈解決タスクを解けるようにQ/K/Vの射影行列を学習した結果だった。生の単語ベクトルの類似度も、ランダムな射影も、“it” を “animal” に結びつける手がかりを何も持っていない。むしろ “not” や “because” のような機能語、あるいは自分自身への注目の方が強い。教科書の図は「学習が終わった後の到達点」を見せているのであって、「仕組みの初期状態」や「教師なしの自然な振る舞い」ではないという、当たり前だけど実際にやってみるまで区別できていなかった点に気づかされた。

なぜ1/√d_kで割るのか

Attentionの式にある $\sqrt{d_k}$ でのスケーリングは、原論文(Attention Is All You Need)の脚注で「次元数が大きいとき、内積の分散が大きくなりすぎてsoftmaxが勾配のほぼない領域に入ってしまうため」と説明されている。文章で読むだけでは実感が湧かなかったので、これも実際に数値で確かめた。

d_ks = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
for d_k in d_ks:
    Q = rng.normal(0, 1, size=(seq_len, d_k))
    K = rng.normal(0, 1, size=(seq_len, d_k))
    scores = Q @ K.T
    # scores.var() の実測値と、softmax後の最大重みを記録

d_k    var(QK^T)  理論値d_k   最大重み(無scale)  最大重み(scale後)
   4      3.87        4        0.4190             0.2305
   8      7.74        8        0.5649             0.2401
  16     15.74       16        0.6791             0.2393
  64     63.63       64        0.8404             0.2447
 256    254.02      256        0.9245             0.2458
1024   1032.39     1024        0.9557             0.2413

数字の表より、実際に同じQuery1つの注目先が$d_k$を増やすごとにどう変わっていくかをそのままGIFにした方が伝わりやすいので、$d_k$を4→1024までスケールアップさせながら、スケーリングなし(左・赤)とスケーリングあり(右・青)のattention重みの棒グラフを並べてアニメーションさせた。

$d_k=4$ の時点ではどちらもそれなりに重みが分散しているが、$d_k$ を増やすアニメーションが進むにつれて、スケーリングなし(左)の棒グラフはみるみる1本だけが伸びていき、$d_k=1024$ では文字通り高さ1.0の棒が1本立つだけになる。他の15本の棒はほぼ0に潰れてしまう。一方スケーリングあり(右)は、$d_k$がどれだけ増えても複数の棒に重みが分散したままで、形がほとんど変わらない。内積スコアの分散は理論通りぴったり $d_k$ に比例して増えていくので、対策なしではどんな入力でもいずれ「たまたま」1箇所だけ内積が突出し、softmaxがそこだけを選んでしまう。これでは他のトークンの情報がValueにほとんど反映されず、勾配もほぼ流れなくなる。$1/\sqrt{d_k}$ というたった一行の割り算が、この「次元が増えるほど不可避に起きる一極集中」を打ち消していることが、静止画のグラフより動きで見るとずっと腹落ちした。

causal maskで何が変わるか

GPTのような自己回帰モデルでは、未来のトークンを見てしまうとカンニングになるので、各位置は自分より前のトークンにしか注目できないよう「causal mask」をかける。BERTのような双方向モデルとの違いを、実際にattention行列で可視化してみた。

mask = np.triu(np.ones((seq_len, seq_len)), k=1).astype(bool)
scores_masked = scores.copy()
scores_masked[mask] = -np.inf
attn_causal = softmax(scores_masked)

これも静止画のヒートマップ1枚より、GPTが実際に1トークンずつ生成していく過程をそのまま再現した方が伝わると思い、「生成済みトークン数」を1から10まで増やしながら、その時点までのattention行列がどう埋まっていくかをGIFにした。グレーはまだ生成されておらず参照不可能な部分、色付きが実際にattentionが計算された部分。

生成が進むごとに、色のついた領域が左下から階段状に広がっていき、右上の三角形部分は最後まで一切色がつかない。特に面白かったのは1トークン目(一番左上のマス)の挙動で、自分より前のトークンが1つも存在しないため、attention重みは [1, 0, 0, ..., 0] と自分自身に100%張り付くしかない。系列の最初のトークンは、原理的に他の情報を一切参照できないまま出力を作らざるを得ない、という制約が、アニメーションの最初のコマにそのまま表れていた。

手を動かして意外だったこと

今回一番印象に残ったのは、Attentionの数式そのものよりも、「よく見る説明図」と「仕組みの素の状態」のギャップだった。“it”が”animal”に注目するという図は、仕組みが賢いからではなく、その振る舞いを学習した結果にすぎない。同じことは、人が何かに「注目」するときにも言えるかもしれない。何が重要かを見極める力は、生まれつき備わっているものではなく、経験を通じて「何と何を結びつけるべきか」を学習した結果として身についていく。逆に言えば、学習(経験)を積む前の状態では、重要なはずの情報(animal)より、目立つだけの情報(not、because)に注目が引っ張られてしまうのは、人でもモデルでも同じなのかもしれない。

スケーリングの実験も、次元(情報量)が増えるほど、何も対策しなければ注目は一点に極端に偏っていくという構造を見せてくれた。情報が増えたときほど、意識的に正規化する一手間を挟まないと、視野は勝手に狭くなっていく。$1/\sqrt{d_k}$ というたった1行の割り算が、次元の呪いによる極端な偏りを打ち消していたのは、地味だが効いている工夫の好例だと思う。

まとめ

• 有名な「itがanimalに注目する」という説明図は、QKVという仕組みが自動的にもたらす性質ではなく、そのタスクを解けるようQ/K/Vを学習した結果。生のGloVe埋め込みでもランダムな射影でも、“animal”への注目は12語中最下位クラスだった
• $QK^\top$ の分散は理論通り次元数$d_k$に比例して増加し、スケーリングなしでは$d_k=1024$で最大attention重みが0.96(ほぼワンホット)まで偏るが、$1/\sqrt{d_k}$でスケーリングすると$d_k$によらず0.24前後で安定する
• causal maskをかけると綺麗な下三角行列になり、系列の先頭トークンは参照できる過去が存在しないため、attention重みが自分自身に100%集中する

• 実践記
• 機械学習

関連記事

• 二重降下を実験で再現した。補間閾値でテスト誤差が76.8倍に跳ね上がり、そこから64分の1まで下がったが、シンプルな次数5モデルには最後まで届かなかった
• 「破局的忘却」を実際に起こしてみたら、Task A精度は0.9834から0.0000まで消えた。10%のリハーサルでも一度は9.9%まで沈んでから戻ってきた
• 自分のブログ53記事をUMAPで並べてみたら、「実践記」と「Claude入門」はきれいには分かれなかった