Self-Attention

自己注意は、クエリ、キー、値がすべて同じシーケンスから得られる注意機構の一種です。そのため、各トークンは、シーケンス内の他のすべてのトークンから新しい表現を構築します。1つの自己注意層は、段落の最初と最後の単語や画像の対角にある隅同士を1ステップで結び付けられます。一方、再帰型ニューラルネットワーク (RNN)では位置ごとに1ステップが必要であり、畳み込みニューラルネットワーク (CNN)では視野を広げるために多くの層を積み重ねる必要があります。フィードフォワード層と積み重ねることで、2017年の論文Attention Is All You Needで紹介されたTransformerが構成されます。これは、大規模言語モデル (LLM)、Vision Transformer (ViT)、Transformerベースの検出器の基盤となっています。

自己注意機構の仕組み#

埋め込みとして表現される各入力トークンに、学習可能な3つの行列を乗算し、クエリ、キー、値を得ます。トークンiの出力は、すべてのトークンの値ベクトルの加重平均です。その重みは、トークンiのクエリとすべてのキーとのスケーリング済み内積に対するsoftmaxで与えられ、注意機構のページで説明されているスケーリング済みドット積アテンションに相当します。クエリ、キー、値は同じトークンの射影であるため、n個のトークンからなるシーケンスではn × nの重み行列が生成されます。「その動物は疲れすぎていたため、通りを渡らなかった」という文では、「それ」に対応する行が「動物」に大きな重みを割り当て、その単語の意味を自身の表現に取り込むことができます。

マルチヘッド自己注意では、これらの計算を小さな射影上で複数並列に実行し、結果を連結します。そのため、あるヘッドは構文を、別のヘッドは照応関係を追跡するなど、異なる関係を捉えられます。元のTransformerでは、512次元モデル内でそれぞれ64次元のヘッドを8個使用していました。

順序の認識と因果マスク#

自己注意は入力を集合として扱います。トークンを並べ替えると出力も同じように並べ替わるだけで、ほかに変化はありません。そのため、元のTransformerで使われた正弦波埋め込みや、現在のほとんどのLLMで使われる回転位置埋め込み (RoPE)などを用いて位置情報を追加します。BERTなどのエンコーダーでは、すべてのトークンが双方向に注意を向けます。デコーダーのみのLLMは因果マスクを適用し、各トークンが自身とそれ以前のトークンだけに注意を向けるようにします。これにより、モデルは左から右へテキストを生成し、以前のトークンのキーと値をKVキャッシュから再利用できます。同じテンソルをクエリ、キー、値として渡すと、PyTorchのMultiheadAttentionでこの2つの特性を確認できます。

import torch

torch.manual_seed(0)
tokens = torch.randn(1, 5, 16)  # それぞれ16個の特徴量を持つ5トークンのシーケンス1つ
attn = torch.nn.MultiheadAttention(embed_dim=16, num_heads=2, batch_first=True)

# 自己注意: クエリ、キー、値はすべて同じトークンの射影です
out, weights = attn(tokens, tokens, tokens)
print("Attention weights shape:", tuple(weights.shape))

# トークンを並べ替えると出力も並べ替わるだけです。そのため、位置情報がなければ単語の順序は認識されません
perm = torch.tensor([3, 0, 4, 1, 2])
out_shuffled, _ = attn(tokens[:, perm], tokens[:, perm], tokens[:, perm])
print("Same outputs after shuffling:", torch.allclose(out_shuffled, out[:, perm], atol=1e-6))

# 因果マスクは、デコーダーのみのLLMと同様に、後続トークンへの注意を遮断します
mask = torch.triu(torch.ones(5, 5, dtype=torch.bool), diagonal=1)
_, causal_weights = attn(tokens, tokens, tokens, attn_mask=mask)
print(causal_weights[0].detach().round(decimals=2))

出力には、(1, 5, 5)の重みの形状、シャッフルテスト用のTrue、下三角行列が表示されます。最初のトークンは自身に重み1.00を割り当て、2番目のトークンは最初の2つのトークンに0.31と0.69を分けて割り当てます。また、対角線より上のすべての要素は0.00です。

自己注意と関連用語の比較#

• 注意機構: クエリとキーの類似度に基づいて値に重み付けする一般的な演算です。自己注意は、3つすべてが1つのシーケンスから得られる場合に相当します。
• クロスアテンション: クエリは1つのシーケンスから、キーと値は別のシーケンスから得られます。翻訳デコーダーはエンコードされた原文に注意を向け、DETRデコーダーのオブジェクトクエリは画像特徴に注意を向けます。また、テキストから画像を生成する拡散モデルはプロンプトに注意を向けます。
• 畳み込み: 各位置に同じ小さなフィルターを適用するため、各出力が見る範囲は固定された局所ウィンドウであり、重みは内容に依存しません。自己注意では内容から重みを計算し、入力全体を対象にします。

元のエンコーダー・デコーダー型Transformerは、双方向自己注意、因果自己注意、クロスアテンションを組み合わせています。

計算コストと効率的なバリエーション#

トークンのすべての組み合わせにスコアを付けるため、計算時間とメモリ使用量はシーケンス長の2乗に比例して増加します。1,000トークンではヘッドごとに100万個、100,000トークンでは100億個のスコアが必要です。FlashAttentionは行列全体を保存せず、タイル単位で正確な自己注意を計算します。グループ化クエリアテンション (GQA)はヘッド間でキーと値を共有してKVキャッシュを縮小し、スライディングウィンドウ、スパース、線形アテンションは組み合わせの数を減らします。コンピュータビジョンではトークン数が画像解像度に応じて増えるため、多くの検出器は特徴マップが最も小さい場所に限って自己注意を適用します。YOLO26は、バックボーン末端のC2PSAブロックと、ネック内の注意機構を備えたC3k2ブロック1つだけでマルチヘッド自己注意を実行します。どちらも、640ピクセル画像のストライド32に対応する20×20の特徴マップ上で処理されるため、トークン数は400です。その注意モジュールは、そのマップに1×1畳み込みを1回適用してクエリ、キー、値を計算し、値に適用する3×3の深さ方向畳み込みで位置情報を追加します。