A-2：Self-Attentionの詳細

A-1では、TransformerがEmbedding、Self-Attention、FFN、残差接続と正規化から成ることを確認しました。その中でSelf-Attentionは、トークン同士の情報を交換する中心部品です。しかし「重要な単語に注目する仕組み」とだけ捉えると、実際に何を計算しているのかが曖昧なままになります。

今回のA-2では、Self-Attentionが文脈をどのように集めるのかを、Query・Key・Value（Q・K・V）の役割から追います。数式を先に暗記するのではなく、どのトークンが何を問い合わせ、どの情報が運ばれるのかを理解することが目標です。そのうえで、GPTのような生成モデルに必要なCausal Mask、複数の視点を持つマルチヘッド、長い文脈で問題になる計算量までを整理します。

この記事を読み終える頃には、「Attentionが文脈を見る」という説明を一歩進めて、「各トークンのQとKの類似度から参照先の重みを決め、Vの加重平均で表現を更新する」と説明できるようになるはずです。

1. Self-Attentionは何をしているのか？

Self-Attentionは、各トークンが系列内の他のトークンを参照し、その結果を使って自分の表現を更新する仕組みです。ここでいう「Self」は、外部の別系列ではなく、同じ入力系列の中を参照することを意味します。

たとえば「太郎は、疲れていたので、早く寝た」という文では、「寝た」を解釈する際に「太郎」や「疲れていた」が手掛かりになります。一方で、読点や助詞がすべて同じ重要度とは限りません。Self-Attentionは、トークンごとに文脈から参照すべき情報の重みを計算します。

重要なのは、重みが固定のルールではない点です。同じ語でも周囲の文が変われば、参照する相手と重みは変わります。モデルが学習するのは、特定の単語を常に見る規則ではなく、入力に応じて参照関係を作るための変換です。

2. Q・K・Vはそれぞれ何を表すのか？

Self-Attentionでは、各入力ベクトルからQuery、Key、Valueという3種類のベクトルを作ります。名前は検索にたとえると理解しやすくなります。

入力を $X$、学習する重み行列を $W_Q$、$W_K$、$W_V$ とすると、3つのベクトルは次のように作られます。

$$ Q = XW_Q, \quad K = XW_K, \quad V = XW_V $$

同じ入力 $X$ から作るものの、Q・K・Vは同じベクトルではありません。どの特徴を「問い合わせ」に使うか、どの特徴を「照合」に使うか、どの情報を「渡す値」に使うかを、別々の重みで学習できるようにしているためです。

この分離があることで、Self-Attentionは単なる単語の一致ではなく、文脈に応じた関係を学習できます。ただしQ・K・Vの各次元が人間にとって一対一に解釈できるとは限りません。ここでの検索は、説明のための有用な比喩です。

コードの実装背景： 実装上x,x,xと同じ入力を3回渡すのは、同一のの系列からこの3つのベクトルを生成するため。 独立した学習： 「問い合わせる特徴」「照合される特徴」「渡す情報」 を別々の重みで学習できるように分離している。

3. Attentionの重みはどのように計算されるのか？

まず、あるトークンのQueryと、系列内の各トークンのKeyとの内積を取ります。内積が大きいほど、そのQueryにとってKeyが関連すると判断しやすくなります。値が大きくなりすぎないように $\sqrt{d_k}$ で割り、Softmaxで合計が1になる重みへ変換します。

計算の流れは次の通りです。

1. $QK^\mathsf{T}$ で、各トークンの組み合わせごとの関連度を計算する
2. $\sqrt{d_k}$ で割り、内積のスケールを調整する
3. Softmaxで、各QueryがどのKeyをどの程度参照するかという重みに変換する
4. その重みでValueを加重平均し、各トークンの新しい表現を作る

Softmax後の各行は確率のように合計1になりますが、これを人間にとっての「重要度」や「因果関係」とそのまま見なすことには注意が必要です。Attentionの重みは、次トークン予測に役立つよう最適化された計算上の重みであり、モデルの説明として十分かどうかは別問題です。

4. 小さな例で、情報が集まる流れを追う

「猫が寝る」という3トークンの入力を考えます。「寝る」の表現を更新するとき、モデルは「猫」を比較的強く参照し、「が」を弱く参照するかもしれません。すると「寝る」の出力は、「猫」のValueを多く含むベクトルになります。

これは、モデルが文法を明示的なルールとして保存していることを意味しません。どのトークンをどれだけ参照するかは、重みと入力から計算され、層やヘッドによっても異なります。複数層を通るうちに、主語・述語のような関係、局所的なつながり、より抽象的な意味の関係を扱える表現が形成されると考えられます。

実際に確認するときは、「Attentionマップがきれいに見えるか」だけで判断しないことが大切です。出力の品質、損失、タスク指標、他の層との相互作用も合わせて見る必要があります。

5. GPT型の生成モデルでCausal Maskが必要なのはなぜか？

GPTのようなDecoder-onlyモデルは、過去のトークンから次のトークンを予測します。学習中に答えとなる未来のトークンを見られると、推論時には使えない情報を利用してしまいます。この不整合を防ぐ仕組みがCausal Maskです。

Causal Maskでは、位置 $i$ のトークンが位置 $j > i$ を参照するスコアを、Softmaxの前に非常に小さな値へ置き換えます。Softmax後の重みはほぼ0となり、未来のValueは出力に混ざりません。

この違いは、モデルが優れているかどうかではなく、目的の違いです。穴埋めや分類では前後の文脈を使えることが有利な場合があり、左から右への生成では未来を隠す必要があります。

6. マルチヘッドに分けるのはなぜか？

1つのAttentionだけでも関連度を計算できますが、すべての関係を単一の表現空間で扱うことになります。Multi-Head Attentionは、Q・K・Vを複数のヘッドに分割し、それぞれでAttentionを計算してから結合する仕組みです。

ヘッドごとに、近くのトークンを主に参照する傾向や、離れた位置を参照する傾向が観察されることがあります。ただし、すべてのヘッドに人間が理解しやすい役割が一つずつあるとは限りません。似た振る舞いをするヘッドや、削除しても影響が小さいヘッドが報告されていることもあります。

「ヘッド数を増やせば必ず性能が上がる」とは言えません。隠れ次元、計算予算、データ量、実装方式との組み合わせで決まる設計上の選択です。

7. 系列長を伸ばすと、なぜ計算量が増えるのか？

Self-Attentionでは、系列長を $n$ とすると、基本的に $n \times n$ 個のトークン対のスコアを作ります。そのため、Attention行列の計算量とメモリ使用量は、系列長に対して概ね二乗で増えます。

これはAttention部分だけを切り出した概算です。実際のメモリ使用量は、バッチサイズ、ヘッド数、隠れ次元、精度形式、勾配や中間活性の保持などにも左右されます。それでも、コンテキスト長を2倍にするとAttention行列の要素数が4倍になる、という感覚は設計上重要です。

この課題に対しては、FlashAttentionのようにメモリ入出力を減らす実装、局所Attentionや疎なAttention、生成時のKVキャッシュを削減するMQA・GQAなど、複数の手法があります。各手法は同じ問題を異なる条件で緩和するため、万能な置き換えではありません。実装面の詳細はB-5で扱います。

8. 最小コードで構造を確認する

以下は、Causal Self-Attentionの流れを確認するための最小例です。学習済みモデルを再現するコードではなく、入力形状とMaskの役割をつかむことを目的にしています。

import torch
import torch.nn as nn

attention = nn.MultiheadAttention(
    embed_dim=128,
    num_heads=4,
    batch_first=True,
)

# (batch_size, sequence_length, hidden_size)
x = torch.randn(2, 5, 128)

# True の位置は参照禁止。上三角部分を未来として隠す。
causal_mask = torch.triu(
    torch.ones(5, 5, dtype=torch.bool), diagonal=1
)

output, weights = attention(
    x, x, x,
    attn_mask=causal_mask,
    need_weights=True,
)

print(output.shape)   # torch.Size([2, 5, 128])
print(weights.shape)  # torch.Size([2, 5, 5])

x, x, xと同じ入力を3回渡しているのは、Self-Attentionが同一系列からQ・K・Vを作るためです。nn.MultiheadAttentionの内部では、それぞれ別の学習可能な線形変換が適用されます。また、need_weights=Trueで返る重みは、既定ではヘッド方向を平均したものです。個々のヘッドを観察したい場合は、設定を追加する必要があります。

コードを動かす際に、Maskの向きを逆にしてしまう失敗はよくあります。未来の位置が確実に0へ近い重みになっているかを小さな入力で確認すると、生成時の情報漏れを早い段階で見つけやすくなります。

9. よくある誤解と確認ポイント

9.1 Attentionの重みは、そのまま説明ではない

Attentionの重みは、モデルの出力に関わる重要な中間値です。しかし、重みが高いトークンだけが最終予測の理由である、と断定することはできません。Value、出力射影、残差接続、FFN、後続層なども結果に影響します。可視化は観察の入口として有用ですが、説明の結論として扱うには追加の検証が必要です。

9.2 Attentionだけで、すべての知識を保存するわけではない

Attentionはトークン間の情報経路を作ります。一方で、各トークン表現の変換にはFFNが関わり、入力の順序には位置情報が必要です。A-1で見た通り、Transformerは複数の部品が組み合わさって機能します。

9.3 長いコンテキストを扱えることと、内容を正しく使えることは別である

最大コンテキスト長を大きくしても、モデルが長い入力の中の必要な情報を常に正しく利用できるとは限りません。位置情報、学習データ、Attentionの実装、評価方法などが影響します。長さの上限と、長い文脈に対する実際の性能を分けて評価しましょう。

10. まとめ

今回の要点を整理します。

• Self-Attentionは、同じ系列内のトークンを参照し、各トークンの表現を文脈に応じて更新する仕組みである
• Qは問い合わせ、Kは照合の手掛かり、Vは参照後に集める情報として理解できる
• QとKから重みを作り、その重みでVを加重平均する
• GPT型の生成モデルでは、未来の情報を隠すCausal Maskが必要である
• Multi-Head Attentionは、複数の部分空間で参照関係を並行して扱う
• 系列長が伸びると、Attention行列は概ね二乗で大きくなる

Self-Attentionは、文脈を集める強力な仕組みです。ただし、その出力を安定して深い層へ渡すには、正規化、活性化関数、残差接続の設計も欠かせません。

11. 今回のブログの考察：A-3への橋渡し

A-2では、Transformerが文脈を集める方法を、Q・K・V、Attentionの重み、Causal Maskという流れで確認しました。Self-Attentionは、どの情報を参照するかを決める重要な部品ですが、集めた情報をそのまま積み重ねるだけでは、深いモデルを安定して学習させることは難しくなります。

そこで次に重要になるのが、各層の表現のスケールを整える正規化と、非線形な変換を加える活性化関数です。Self-Attentionが「文脈のどこを見るか」を担うなら、正規化と活性化関数は「集めた情報をどう安定して変換するか」を担います。

次回A-3では、Transformerを深く積み重ねるための正規化と活性化関数を扱います。LayerNormとRMSNormは何を安定させるのか？ GELUやSwiGLUはどのように表現を変換するのか？を、Self-Attentionの出力が次の処理へ渡る流れに沿って見ていきましょう。

参考文献

1. Vaswani, A., Shazeer, N., Parmar, N., et al. (2017). “Attention Is All You Need.” NeurIPS 2017. https://arxiv.org/abs/1706.03762
2. Clark, K., Khandelwal, U., Levy, O., & Manning, C. D. (2019). “What Does BERT Look at? An Analysis of BERT’s Attention.” https://arxiv.org/abs/1906.04341
3. Dao, T., Fu, D. Y., Ermon, S., et al. (2022). “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” NeurIPS 2022. https://arxiv.org/abs/2205.14135
4. Shazeer, N. (2019). “Fast Transformer Decoding: One Write-Head is All You Need.” https://arxiv.org/abs/1911.02150
5. Ainslie, J., Lee-Thorp, J., de Jong, M., et al. (2023). “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.” https://arxiv.org/abs/2305.13245

このシリーズの案内

ブログA「Transformer構成」（全7回）では、現代のLLMを支えるTransformerの構成要素を、全体像から実装への接続まで段階的に扱います。

1. 前の記事

• A-1：Transformer構成の基本概念と重要性

2. 次に読む

• A-3：正規化と活性化関数