クエリ、キー、値
同一入力からの3つの線形写像
埋め込み(アクティビティ4)の後、各位置は 768 次元のベクトル x_t を保持します。アテンションは x の3つの異なる射影を生成することから始まります:
Q (query): この位置は何を知りたいのか?
K (key): この位置が他の位置に提供するものは何ですか?
V (value): この位置が注目された場合に提供する内容は何ですか?
各プロジェクションは学習された重み行列から来ています:
Q = x · W_Q # W_Q の形状: (d_model, d_k)
K = x · W_K # W_K の形状: (d_model, d_k)
V = x · W_V # W_V の形状: (d_model, d_k)
3つの行列で、すべてバックプロパゲーションで訓練されます。モデルは学習します:この位置で、役立つ過去のコンテキストを最もよく取得するクエリは何か?この位置のコンテンツをうまく宣伝するキーは何か?選択された場合に何を配信する価値か?
図書館のアナロジー
図書館のカードカタログを想像してください。あなたはクエリ(自分のトピック)を思い浮かべて入っていきます。すべてのカードにキーワード(キー)が記載されています。あなたのトピックがカードのキーワードと一致したら、本の内容(値)を取り出します。Attention はすべてのトークンに対して並列でこれを行います:すべての位置がすべての他の位置をクエリし、整合性をランク付けし、値ベクトルの重み付けされた組み合わせを取得します。
ANDREA-120M の次元
| 数量 | 値 | 備考 |
|---|---|---|
| d_model | 768 | 各位置でのベクトルサイズ |
| n_head | 12 | 並列アテンションヘッド |
| d_k | 64 | ヘッドごとの次元 (= d_model / n_head) |
| T | 1024 | コンテキスト長 |
d_k = d_model / n_head = 768 / 12 = 64。各ヘッドは、完全な768次元空間の64次元スライスを見ます。Activity 6 (grow_a_language_model_multi_head) でヘッドごとの分割を詳しく扱います。
d_k を計算する
なぜsqrt(d_k)で割るのか
スコア行列
QとKが存在すると(それぞれ形状(T, d_k))、アテンションはスコア行列を計算します:
scores = Q · K^T # shape: (T, T)
scores[i, j] = 位置 i のクエリが位置 j のキーどれだけ強く一致するかを示す。すべての (i, j) ペアに1つのスコアが割り当てられます:1024 × 1024 = 1,048,576 個のスコアが1回のフォワードパスあたりアテンションヘッドごとに。
なぜ除算するのか
d次元のランダムな単位ベクトルの2つのドット積の大きさは sqrt(d) のオーダーです。スケーリングなしでは、スコアが d_k に比例して大きくなります:
- d_k = 64: 典型的なドット積は8のオーダー。
- d_k = 256: 典型的なドット積は16のオーダー。
- d_k = 4096: 典型的なドット積は64のオーダー。
大きなスコアはピーキーなsoftmaxを生む(1つの位置が支配し、他の場所で勾配が消失)。トレーニングが停滞。スケーリングで大きさを修正:
scaled_scores = (Q · K^T) / sqrt(d_k)
ANDREA-120M では、sqrt(d_k) = sqrt(64) = 8 です。すべてのスコアが 8 で割り算されます。d_k に関係なく大きさがおおよそ単位スケールに保たれます。Softmax が良好に動作します。勾配が流れます。
Vaswani の元の正当化
Attention Is All You Need (2017) より: 'd_k の値が大きい場合、ドット積の大きさが大きくなり、softmax 関数を勾配が極めて小さい領域に押し込んでしまいます。' sqrt(d_k) による除算がその成長を相殺します。
コードのビュー
microgpt_cuda.cu の中で、このスケーリングはリテラルの除算として現れます:
scores[i][j] = dot(Q[i], K[j]) * (1.0f / sqrtf(d_k));
1つのスコアあたり1回の浮動小数点乗算。安価。重要。
d_model = 4096 でのスケーリング
位置 i が位置 j > i を見ることができない理由
生成から生まれた制約
ANDREA は1トークンずつ生成します。推論時、位置 0 が最初のトークンを生成し、次に位置 1 が位置 0 の出力を参照して2番目のトークンを生成し、以下同様です。モデルは生成中に未来のトークンにアクセスできません。
訓練もこれを反映する必要があります。訓練中に位置 5 が位置 6 に注目できる場合、モデルはショートカット「トークン 6 をトークン 6 を読んで予測する」を学習します。推論時、そのショートカットは消滅します(トークン 6 はまだ存在しません)。モデルの訓練時と推論時の動作が壊滅的に乖離します。
マスク
因果マスクは、任意の位置 i から j > i の位置への注意をブロックします。実装:j > i の箇所で scaled_scores[i][j] = -infinity に設定。softmax 後、それらのエントリは exp(-inf) = 0 になります。マスクは未来の位置への注意をきれいにゼロ化します。
for i in range(T):
for j in range(T):
if j > i:
scaled_scores[i][j] = -1e9 # 実質的に -inf
ソフトマックス(行ごと)を適用した後、各行の合計は1になりますが、[0, i] のエントリのものだけが確率質量を持ちます。位置 i は過去の位置からの情報のみを混合します。
マスクの可視化
マスクを適用したスコア行列の形状 (T, T) は、下三角構造のように見えます:
マスク後の scaled_scores、行ごとの softmax:
row 0: [1.0, 0, 0, 0, ...] # 自身のみを見ます
row 1: [0.4, 0.6, 0, 0, ...] # 位置 0, 1 を見ます
2行目: [0.2, 0.3, 0.5, 0, ...] # 0, 1, 2 が見える
3行目: [0.1, 0.2, 0.3, 0.4, ...] # 0, 1, 2, 3 が見える
...
各行ごとの厳密な下三角形確率分布。未来は見えないまま。
Decoder-Only Transformer がこれを必要とする理由
ANDREA、GPT、LLaMAなどのデコーダーオンリーモデルはすべて、過去から次のトークンを予測するという共通の目的を持っています。causal maskはこの目的を並列で訓練可能にします:すべての位置が同時に自身の次のトークン予測を計算し、どの位置も先を見ないようにします。
マスクとフレーバー
スコアから出力へ
ソフトマックス:スコアから確率へ
マスクされ、スケーリングされたスコアは依然として実数範囲です。ソフトマックスは各行を確率分布に変換します:
A[i][j] = exp(scaled_scores[i][j]) / sum_k exp(scaled_scores[i][k])
これにより3つの性質が生じます:
- すべての (i, j) に対して A[i][j] >= 0。
- 各行 i に対して sum_j A[i][j] = 1。
- より大きな生スコアがより大きな確率を生む(単調増加)。
行 i の確率ベクトルはモデルに伝える:出力計算時に位置 i が各前の位置にどれだけ注意を払うべきか?
加重 V の合計
位置 i の最終的な注意出力:
output[i] = sum_j A[i][j] · V[j]
各値ベクトル V[j] は注意確率 A[i][j] によって重み付けされ、その後合計されます。位置 i の出力は、すべての前の位置からの値ベクトルを関連性によって重み付けして組み合わせたものです。
行列形式で、一度にすべての位置に対して:
Attention(Q, K, V) = softmax(mask(Q · K^T / sqrt(d_k))) · V
1行。完全なattentionメカニズム。Vaswani et al. は2017年にこの行を書きました;transformerはその後根本的に変わっていません。
ヘッドごとの出力形状
1つのアテンションヘッドの出力: 形状 (T, d_k)。ANDREA-120M の場合: (1024, 64)。全12ヘッドが並列で計算され; その出力が (1024, 768) に連結され、最終線形投影 (W_O) に送られ、トランスフォーマーブロックの MLP に進む。
Activity 6 (grow_a_language_model_multi_head) はマルチヘッド分割を扱います。Activity 7 (grow_a_language_model_transformer_block) はアテンションの周囲のすべてを扱います: 残差接続、レイヤー正規化、MLP。