Attention

実装: llm/attention/ / 実行: go test ./llm/attention/

各トークンが他のどのトークンにどれだけ注目するかを決めて、その割合で情報を混ぜる。式は行列積と softmax だけで書ける。だが並び順を入れ替えても結果は入れ替わるだけで中身は変わらない。attention は順序を見ていない。そして未来を見ないことは、マスク1枚で保証される。後ろのトークンを書き換えても、前の出力は1ビットも動かない。

この章で作るもの

行列演算で作った MatMul と SoftmaxRows だけを使って、self-attention を1ヘッド組む。次章で組み立てる層の中心部が、これだ。

やっていることは1文で言える。各トークンが他のどのトークンにどれだけ注目するかを決めて、その割合で情報を混ぜる。

そのために、各トークンは自分のベクトルから3つの役割のベクトルを作る。

• Q(query): 私は何を探しているか。そのトークンが出す問い合わせ
• K(key): 私は何を持っているか。各トークンが出す見出し
• V(value): 私の中身はこれ。各トークンが出す実体

トークン i が誰に注目するかは、i の Q と全トークンの K の内積で決まる。Q と K がよく合うトークンほど強く注目する。図書館に例えると、Q は探している本のテーマ、K は各本の背表紙、内積はどれだけ合致するかになる。

  注目の重み(4トークン、行ごとに合計 1)

         見られる →   t0    t1    t2    t3
    見る ↓
      t0             0.4   0.2   0.3   0.1      マスク無し:
      t1             0.1   0.5   0.2   0.2        全員が全員を見る
      t2             0.3   0.1   0.4   0.2        t0 が t3(未来)も見ている
      t3             0.2   0.2   0.3   0.3


  因果マスクをかける(上三角を -∞ にしてから softmax)

         見られる →   t0    t1    t2    t3
    見る ↓
      t0            1.00     0     0     0      自分しか見えない
      t1             0.2   0.8     0     0
      t2             0.5   0.1   0.4     0
      t3             0.2   0.2   0.3   0.3

  この下三角の形が、左から右へ生成できる理由そのものになる

注目の重みは正方の表になる。行が注目する側、列が注目される側。マスクは、その表の上半分を落とすだけの操作になる

順に見ていく。

1. 注目の重みを作って、その割合で混ぜる: 全体が行列積と softmax だけで書ける
2. 並びを見ていない: 入力を入れ替えると出力もそのまま入れ替わる。だから位置情報を別に足す
3. 未来を見ないことはマスク1枚で保証される: 後ろを変えても前の出力は1ビットも動かない

① 注目の重みを作って、その割合で混ぜる

まずヘッドを用意する。入力を Q・K・V に写す3つの重み行列を持つだけになる:

// Head は1つの attention ヘッド。Wq, Wk, Wv は入力を Q, K, V に変換する重み行列。
// (実物はこれらを学習で得るが、ここでは仕組みを見るため乱数 or 恒等で初期化する。)
type Head struct {
	Wq, Wk, Wv *tensor.Tensor
	dHead      int
}

// NewHead は dModel 次元の入力を dHead 次元の Q/K/V に写すヘッドを作る。
// 重みは決定的な擬似乱数で埋める(テストの再現性のため)。
func NewHead(dModel, dHead int) *Head {
	return &Head{
		Wq:    randMatrix(dModel, dHead, 1),
		Wk:    randMatrix(dModel, dHead, 2),
		Wv:    randMatrix(dModel, dHead, 3),
		dHead: dHead,
	}
}

// NewHeadIdentity は Wq=Wk=Wv=単位行列のヘッド(スケーリング等の検証用)。
func NewHeadIdentity(d int) *Head {
	return &Head{Wq: identity(d), Wk: identity(d), Wv: identity(d), dHead: d}
}

本体はこれで全部だ:

// rawScores は Q·Kᵀ / √dHead を返す(softmax 前のスコア = 注目の強さの生値)。
// √dHead で割るのは、次元が大きいと内積が大きくなりすぎて softmax が尖りすぎるのを防ぐため。
func (h *Head) rawScores(x *tensor.Tensor) *tensor.Tensor {
	q := tensor.MatMul(x, h.Wq) // (seq, dHead)
	k := tensor.MatMul(x, h.Wk) // (seq, dHead)
	kt := transpose(k)          // (dHead, seq)
	scores := tensor.MatMul(q, kt)

	scale := float32(1.0 / math.Sqrt(float64(h.dHead)))
	for i := range scores.Data {
		scores.Data[i] *= scale
	}
	return scores
}

// forwardWithWeights は attention の出力と、注目の重み(softmax後)を返す。
// causal=true なら因果マスクをかけ、各トークンが未来を見ないようにする。
func (h *Head) forwardWithWeights(x *tensor.Tensor, causal bool) (out, weights *tensor.Tensor) {
	scores := h.rawScores(x)

	if causal {
		// 上三角(j > i = 未来)を -Inf にする。softmax で 0 になり、注目が消える。
		negInf := float32(math.Inf(-1))
		for i := 0; i < scores.Rows; i++ {
			for j := i + 1; j < scores.Cols; j++ {
				scores.Set(i, j, negInf)
			}
		}
	}

	weights = tensor.SoftmaxRows(scores) // 各行を確率分布に
	v := tensor.MatMul(x, h.Wv)          // (seq, dHead)
	out = tensor.MatMul(weights, v)      // 重み付きで V を混ぜる
	return out, weights
}

// Weights は注目の重み(softmax 後)だけを返す。行が注目する側、列が注目される側。
func (h *Head) Weights(x *tensor.Tensor, causal bool) *tensor.Tensor {
	_, w := h.forwardWithWeights(x, causal)
	return w
}

// Scores は softmax 前の生のスコアを返す。scaled=false なら √dHead で割らない。
//
// 割らないとどうなるかを確かめるための入口になる。
func (h *Head) Scores(x *tensor.Tensor, scaled bool) *tensor.Tensor {
	q := tensor.MatMul(x, h.Wq)
	k := tensor.MatMul(x, h.Wk)
	scores := tensor.MatMul(q, transpose(k))
	if !scaled {
		return scores
	}
	scale := float32(1.0 / math.Sqrt(float64(h.dHead)))
	for i := range scores.Data {
		scores.Data[i] *= scale
	}
	return scores
}

// Forward は attention の出力だけを返す。
func (h *Head) Forward(x *tensor.Tensor, causal bool) *tensor.Tensor {
	out, _ := h.forwardWithWeights(x, causal)
	return out
}

順に見るとこうなる。

1. Q·Kᵀ: 全トークン対の注目スコアを一度に計算する(行列積)
2. ÷√d: 次元が大きいとスコアが大きくなりすぎるので割る
3. softmax: 各行を確率にする(トークン i が配る注目の合計が 1)
4. ·V: その確率で V を混ぜる。強く注目した相手の中身が多く混ざる

新しい道具は1つも要らなかった。行列演算の MatMul と SoftmaxRows、それに割り算だけで組めている。attention は行列積の塊だというのが、実装で確かめられたことになる。

÷√d を入れる理由も測れる。6トークンで、いちばん強い注目の重みがどうなるかを見る(一様に配れば 0.167):

次元	割らない	√d で割る
8	0.347	0.228
32	0.799	0.343
64	0.984	0.458
256	1.000	0.847

割らないと、次元が上がるほど注目が1点に集中していく。256次元では完全に1つの相手だけを見るようになり、他が0になる。内積は次元の数だけ足し算を重ねるので、次元に比例して大きくなる。softmax は入力の差を指数で拡大するので、差が開けば一気に尖る。

割ればこの進行は緩む。ただし表のとおり、割っても大きい次元では尖ってくる。ここでの重みは学習していない乱数なので、実物ほどきれいには収まらない。テストで、割らないと次元とともに尖ること、割ると同じ次元で尖りが弱まることを固定した。

② 並びを見ていない

ここが、この式のいちばん意外なところだ。

上の計算のどこにも「何番目のトークンか」は出てこない。Q と K の内積は中身だけで決まる。つまり入力の行を入れ替えると、出力の行も同じように入れ替わるだけで、中身は1つも変わらない。

テストで、4トークンのうち 1番目と3番目を入れ替えた入力を通し、出力が同じ入れ替えになることを固定した。誤差の範囲どころか、対応する行がぴたりと一致する。

だから attention だけでは「猫が犬を追う」と「犬が猫を追う」を区別できない。順序の情報は、attention の外から入れるしかない。あとの章で出てくる位置エンコーディングや RoPE が、そのための仕掛けだ。

式を書いたら、何を見ていないかも確かめる。見ていないものは、別に用意することになる。

③ 未来を見ないことはマスク1枚で保証される

GPT は次の単語を予測するモデルなので、予測するときにまだ書いていない未来の単語を見てはいけない。

保証の仕方は単純だ。注目スコアの表で「トークン i がトークン j > i を見る」マスは上三角にある。そこを -∞ にすると、softmax を通ったとき 0 になる。未来への注目が消える。

これが本当に効いているかは、確かめられる。最後のトークンだけを大きく書き換えて、前のトークンの出力が動くかどうかを見ればよい。

テストで、マスクありなら前の3行が1つも動かないこと、マスクを外すと動くことを、両方固定した。「未来を見ない」は方針ではなく、測れる性質になっている。

試す: 因果マスクの ON/OFF を切り替えると、表の上三角が消えたり現れたりする。マスク ON では各トークンは自分と過去にしか注目せず、いちばん上のトークンは過去が無いので自分に 100% 注目している。

デモself-attention の重み因果マスクON

因果マスク(GPT)マスクなし

行 = 見る側、列 = 見られる側。濃いほど強く注目

Thecatsatonmat

The100

cat1288

sat42671

on441181

mat10441073

「sat」は 、catに26%、satに71% 注目している(未来のトークンは見えない)。

この下三角の形には副産物がある。トークン i の出力は i 以前にしか依存しないので、生成のたびに前を計算し直す必要がない。K と V を取っておけば使い回せる。推論高速化の KV キャッシュは、この下三角の形から出てくる。

設計の観点

• 式を書いたら、何を見ていないかも確かめる: 順序を見ないと分かれば、位置情報を足す必要が見える
• 性質を測れる形にする: 「未来を見ない」を、後ろを書き換えて前が動くかで確かめる
• 数の大きさは次元とともに動く: 内積は次元だけ足し算を重ねる。割って抑えるのはそのため
• 制約が高速化を生む: 下三角という制約があるから、前の計算を使い回せる
• 同じ道具を積み上げる: 行列積と softmax だけで組める形にしておくと、速くする努力を1点に集められる
• 1ヘッドで足りないものを知る: 1つの見方しか持てないので、実物は複数を並べる

対照と実例

	見る範囲	順序を見るか	計算量	代表例
マスク無し attention	全トークン	見ない	系列長の2乗	BERT などの符号化器
因果マスクつき	自分と過去	見ない	系列長の2乗	GPT 系
窓つき	近くだけ	見ない	系列長に比例	Longformer、Mistral の sliding window
線形 attention	全トークン(近似)	見ない	系列長に比例	Performer、SSM 系
再帰(RNN)	直前の状態ごし	見る	系列長に比例	LSTM、GRU

裏どり:

• Attention Is All You Need(2017): Vaswani et al.。÷√d の理由も論文中に書かれている。内積の分散が次元に比例して増えることへの対処になる
• 順序を見ないこと: 論文が位置エンコーディングを足しているのは、この章で測った性質への直接の答えになる。attention は集合を扱う演算で、並びは外から与える
• RNN との違い: 再帰は前から順にしか進めないので順序が組み込まれているが、そのぶん並列に計算できない。attention は順序を捨てる代わりに全部を同時に計算できる。捨てたものと得たものが対になっている
• 系列長の2乗: 全対の内積を取るので、長い入力では計算量と記憶量がここで詰まる。attention の変種はここを削る工夫の系統になる
• KV キャッシュ: 因果マスクの下三角から出てくる。詳しくは推論高速化

簡略化したこと

• 1ヘッドのみ: 実物はマルチヘッド(8〜96個を並列に置いて、別々の関係を捉える)
• 重みは乱数か恒等: 実物は学習で得た Wq/Wk/Wv。ここは仕組みを見るのが目的
• 位置エンコーディングなし: この章で測ったとおり、順序は別に足す必要がある
• 出力射影なし: マルチヘッドを束ねる行列(Wo)は無い
• クロス attention なし: Q と K/V が別のところから来る形は扱わない
• 速くする工夫なし: 全対のスコアを素直に持つ。窓や線形化は変種の章

参考資料

• Vaswani et al., Attention Is All You Need(2017)
• The Illustrated Transformer — Q/K/V の図解
• nanoGPT (Karpathy) — 最小の GPT 実装
• 実装: llm/attention