アテンション機構

「明月」（明るい月）の「明」と「明天」（明日）の「明」は、意味が違います。一つの文字をどう理解すべきかは、周りにどんな文字があるかで決まります。

私たちが学習させるモデルは、各位置で次の文字を予測するのが仕事です。「床前明月」（寝台の前の明るい月）の後に「光」（光）を書くには、前に「明月」が出てきたことを知っていなければなりません。ですから、どの位置にも前の文字から情報を集める方法が必要です。

アテンション（attention）がその方法で、Transformer の核心です。この課では最も単純なやり方から始めて、一歩ずつ組み立てていきます。

python attention.py

最も単純な方法：平均を取る

各文字はまずベクトルで表します（モジュール 01 第 5 課で埋め込みを扱いました）。最も単純な「前の情報を集める」方法は、自分と前のすべての文字のベクトルの平均を取ることです。

これは 1 回の行列の掛け算でできます。下三角の行列を作り、各行は前のいくつかの位置にだけ値があって、しかも合計が 1 になるようにします。

x = torch.randn(T, 4)  # 5 个词，每个词一个 4 维向量（先随便取）
weights = torch.ones(T, T).tril()  # 下三角：第 i 行只有前 i+1 个位置是 1
weights = weights / weights.sum(dim=1, keepdim=True)  # 每行除以个数，变成平均
out = weights @ x

== 1. 用一个矩阵乘法，让每个位置得到自己和前面所有位置的平均
tensor([[1.00, 0.00, 0.00, 0.00, 0.00],
        [0.50, 0.50, 0.00, 0.00, 0.00],
        [0.33, 0.33, 0.33, 0.00, 0.00],
        [0.25, 0.25, 0.25, 0.25, 0.00],
        [0.20, 0.20, 0.20, 0.20, 0.20]])
  第 3 个位置（'明'）的结果 = 前三个向量的平均？True

i 行目は、i 番目の文字が各位置からどれだけ取るかを表しています。1 番目の文字は自分しか見られず、3 番目の文字は前の三つを均等に見て、5 番目の文字は五つすべてを均等に見ます。weights @ x で、すべての位置の結果が一度に計算されます。

この行列の掛け算の書き方を覚えておいてください。後のアテンションも同じ形です。重みの行列 × ベクトルの束。違うのは、重みがどこから来るかだけです。

平均の問題ははっきりしています。すべての文字を同じに扱ってしまうのです。「光」が知りたいのは前に「月」があるかどうかで、「床」や「前」はそれほど重要ではありません。重みは内容によって決まるべきです。

クエリとキー：各文字に誰を見るべきかを決めさせる

方法は、各文字に二つのベクトルを用意することです。

• クエリ（query、q）：自分はどんな情報を探しているか。
• キー（key、k）：自分のところにはどんな情報があるか。

ある文字が別の文字にどれだけ「注目」するかは、自分の q と相手の k の内積で採点します。二つのベクトルの向きが近いほど、内積は大きくなります。

まず手作りのベクトルで効果を見てみましょう。ベクトルは 2 次元で、1 次元目は「空に関係する」、2 次元目は「位置に関係する」を表すとします。「光」のクエリは [1, 0] で、空に関係するものを探していることを表します。

k = torch.tensor([[0.1, 0.9],   # 床：一个地点
                  [0.0, 1.0],   # 前：一个方位
                  [0.9, 0.1],   # 明：和天空、光有关
                  [1.0, 0.0],   # 月：天空里的东西
                  [0.8, 0.2]])  # 光
q_guang = torch.tensor([1.0, 0.0])  # "光"想找的是：和天空有关的东西
scores = k @ q_guang

== 2. 点积打分：查询（q）和每个键（k）越像，分数越高
  '光' 对每个词的分数： {'床': 0.1, '前': 0.0, '明': 0.9, '月': 1.0, '光': 0.8}
  softmax 之后的权重：   {'床': 0.12, '前': 0.11, '明': 0.26, '月': 0.29, '光': 0.23}
  分数放大 5 倍再 softmax： {'床': 0.01, '前': 0.0, '明': 0.3, '月': 0.5, '光': 0.18}

スコアは softmax で重みに変わり（モジュール 08 第 5 課で softmax を扱いました）、合計は 1 になります。「月」のスコアが最も高く、最も大きな重みを得ています。

最後の行に注目してください。スコアを全体で 5 倍に拡大すると、softmax 後の重みはずっと集中し、「月」一つで半分を持っていきます。スコアの大きさが、アテンションが「まんべんなく行き渡る」か「一つだけを見つめる」かを決めるのです。この点は後でスケーリングを扱うときに重要になります。

本物のモデルでは、q と k は手作りではなく、二つの行列で各文字のベクトルから計算され、行列の中のパラメータは学習で得られます。どんな文字がどんな情報を探すべきかを、モデルは自分で学ぶのです。

値：実際に伝える内容

採点して、各位置からどれだけ取るかが決まりました。では取るのは何でしょうか。三つ目のベクトル、値（value、v）です。「私に注目するなら、この情報をあげる」という意味です。

なぜ元のベクトルをそのまま取らないのでしょうか。「照合に使う情報」と「伝えるための情報」は同じとは限らないからです。二つを分ければ、モデルはより柔軟になります。

三つのベクトルを合わせたものが、アテンションのすべてです。

权重 = softmax(q 和每个 k 的点积)
输出 = 用这些权重，把每个位置的 v 加起来

因果マスク：後ろを覗いてはいけない

学習のとき、モデルは各位置で次の文字を予測しなければなりません。3 番目の位置が 4 番目の文字を見られれば、「答えを写す」だけで済み、何も学べません。ですから各位置は自分と前の位置しか見られないようにします。

やり方は、softmax の前に、後ろの位置のスコアを負の無限大にすることです。負の無限大の指数は 0 なので、softmax 後の重みは 0 になります。

scores = q @ k.T / math.sqrt(8)
mask = torch.ones(T, T, dtype=torch.bool).tril()
scores = scores.masked_fill(~mask, float("-inf"))

== 3. 因果掩码：把后面位置的分数设成负无穷，softmax 之后权重就是 0
tensor([[1.00, 0.00, 0.00, 0.00, 0.00],
        [0.96, 0.04, 0.00, 0.00, 0.00],
        [0.24, 0.26, 0.50, 0.00, 0.00],
        [0.09, 0.39, 0.43, 0.09, 0.00],
        [0.02, 0.14, 0.02, 0.09, 0.73]])

最初の節の平均の行列と比べてみてください。形は同じで、どちらも下三角で、各行の合計はどちらも 1 です。違いは、各行の重みがもう均等ではなく、q と k によって決まっていることです。ここの q と k はランダムなので、重みには規則性がないように見えます。学習した後には、意味を持つようになります。

このマスクがあるので、この種のアテンションは因果的自己アテンションと呼ばれます。自己アテンションとは q、k、v がすべて同じ文章から来ること、因果的とは過去しか見られず未来は見られないことを指します。GPT のような「1 文字ずつ後ろへ書いていく」モデルは、どれもこれを使っています。

なぜ次元の平方根で割るのか

上のコードに / math.sqrt(8) がありました。8 は q と k の次元です。このステップをスケーリングと呼びます。しないとどうなるか見てみましょう。

== 4. 为什么要除以 √d：维度越大，点积的数值越大，softmax 会变得非常极端
  d=   16  点积的标准差   3.89（√d =  4.00）  每行最大权重的平均：不缩放 0.79，除以 √d 后 0.53
  d=   64  点积的标准差   8.17（√d =  8.00）  每行最大权重的平均：不缩放 0.86，除以 √d 后 0.47
  d=  256  点积的标准差  16.08（√d = 16.00）  每行最大权重的平均：不缩放 0.95，除以 √d 后 0.47
  d= 1024  点积的标准差  32.53（√d = 32.00）  每行最大权重的平均：不缩放 0.92，除以 √d 后 0.39

二つのランダムなベクトルの内積は、d 個の積を足し合わせたものです。足す項が多いほど結果の揺らぎは大きくなり、標準偏差はおよそ d の平方根になります。d=256 のとき、内積の標準偏差は 16 です。

スコアの揺らぎが大きくなると、2 番目の節の「5 倍に拡大」と同じように、softmax が極端になります。スケーリングしないと、各行で最も大きな重みは平均 0.8 から 0.95 になり、ほぼすべての注目が一つの位置に向かいます。これは学習を始めたばかりのときにはとてもまずいことです。重みが 0 や 1 に近いところでは softmax の勾配がとても小さく、モデルが学習しにくくなるからです。

d の平方根で割れば、スコアの標準偏差は 1 前後に戻り、各行で最も大きな重みは 0.4 から 0.5 の間になって、注目は分散し、勾配も正常で、モデルは誰に注目すべきかを少しずつ学べます。

完全なアテンションヘッド

これまでのものを合わせます。

d_model, d_head = 16, 8
x = torch.randn(T, d_model)
W_q, W_k, W_v = (torch.randn(d_model, d_head) / math.sqrt(d_model) for _ in range(3))
q, k, v = x @ W_q, x @ W_k, x @ W_v
scores = (q @ k.T / math.sqrt(d_head)).masked_fill(~mask, float("-inf"))
out = F.softmax(scores, dim=-1) @ v

== 5. 一个完整的注意力头：q、k、v 都由同一个输入经过不同的矩阵得到
  输入 (5, 16) → q、k、v 各 (5, 8) → 输出 (5, 8)
  和 PyTorch 自带的 scaled_dot_product_attention 比，最大差别 2.4e-07

5 行のコードです。三つの行列で入力を q、k、v に変え、採点し、スケーリングし、マスクし、softmax を取り、重み付きで足し合わせる。結果は PyTorch に組み込まれた scaled_dot_product_attention と一致します（差は小数点以下 7 桁目で、浮動小数点の計算順序の違いによる誤差です）。

W_q、W_k、W_v が、このアテンションヘッドが学ぶパラメータです。学習のときは、ほかのパラメータと同じように、誤差逆伝播法と勾配降下法で更新されます。

この式の出典は 2017 年の論文『Attention Is All You Need』で、数式で書くと softmax(QKᵀ/√d)V です。今ではその中の記号一つ一つが何をしているのかわかるはずです。

アテンションの代償

どの位置も、前のすべての位置と採点しなければなりません。系列の長さが T なら、スコアの表は T×T の大きさです。長さが倍になれば、計算量とこの表が占めるメモリは 4 倍になります。

これが長いコンテキストが高くつく根本的な理由の一つです（モジュール 01 第 4 課でコンテキストとコストを扱いました）。各社のモデルは数十万、数百万トークンのコンテキストに対応するため、アテンションの計算方法に大量の最適化を施していますが、基本的な考え方はやはりこの課の数行です。

練習問題

1. 2 番目の節で、「光」のクエリベクトルを [0, 1]（位置に関係するものを探す）に替えてください。重みはどう変わりますか。
2. 5 番目の節で因果マスクを取り除き、scaled_dot_product_attention(..., is_causal=False) と比べてください。
3. 5 番目の節の系列の長さを 5 から 1000、2000、4000 に変え、time.time() でそれぞれの長さにかかる時間を測ってください。おおよそ二乗で増えていますか。

確認テスト

1. クエリ、キー、値はそれぞれどんな役割を果たしていますか？

クエリは「自分は何を探しているか」、キーは「自分のところには何があるか」を表します。ある位置のクエリとすべての位置のキーの内積を取ると、各位置への注目スコアが得られ、softmax で重みに変わります。値は「私に注目したときに得られる情報」を表し、出力はそれらの重みですべての位置の値を足し合わせたものです。

2. 因果マスクはどう実装されていますか？なぜ必要なのですか？

softmax の前に、各位置より後ろのスコアを負の無限大にし、softmax 後にそれらの位置の重みが 0 になるようにします。必要なのは、モデルが次の文字の予測を学ぼうとしているからで、後ろの文字が見えてしまえば答えを写すだけで済み、何も学べません。

3. アテンションのスコアを次元の平方根で割るのはなぜですか？

d 次元の二つのベクトルの内積は、標準偏差がおよそ d の平方根になり、次元が大きいほどスコアの揺らぎが大きくなって、softmax 後の重みが一つの位置に極端に集中し、勾配がとても小さくなって、モデルが学習しにくくなります。d の平方根で割ると、スコアの標準偏差が 1 前後に戻ります。