Attention Mechanism

アテンション機構とは、各入力が処理対象の要素にどれほど関連するかに基づいて重みを計算し、入力特徴量の加重平均として各出力を構成するニューラルネットワークの演算です。クエリベクトルとキー・ベクトルを比較して関連度を評価し、そのスコアで対応するバリューベクトルに重みを付けます。アテンションは2014年に機械翻訳向けに導入され、リカレントニューラルネットワーク（RNN）のデコーダーが、文全体を固定長の1つのベクトルに圧縮せず、すべての入力語を参照できるようにしました。2017年のTransformerはアテンションだけでアーキテクチャ全体を構成しました。アテンションは現在、大規模言語モデル、ビジョンTransformer、多くの物体検出モデルの中核的な演算です。

注意機構の仕組み#

アテンションでは、入力に対して学習可能な3つの線形射影を使用します。クエリ（Q）は現在のトークンが探しているものを表し、各キー（K）はトークンが提供する情報を表し、各バリュー（V）は伝達される内容を保持します。モデルはクエリをすべてのキーと比較し、比較スコアを重みに変換して、バリューの加重和を返します。翻訳モデルが「The black cat」に対して「noir」を生成する場合、そのステップのクエリは、フランス語では形容詞が名詞の後に来るにもかかわらず、「black」のキーに対して最も高いスコアを示します。

マルチヘッドアテンションでは、複数の計算をそれぞれ別の低次元射影上で並列実行し、その結果を連結します。これにより、異なるヘッドが同じレイヤー内の異なる関係に特化できます。

スケールド・ドット積アテンション#

TransformerのアテンションはAttention(Q, K, V) = softmax(QKᵀ / √d_k) Vと表されます。各クエリをすべてのキーとドット積で比較し、そのスコアをキーの次元数d_kの平方根で割ります。これにより、ベクトルサイズに応じてスコアが増大し、softmaxが飽和するのを防ぎます。Softmaxは各クエリのスコアを、合計が1になる重みに変換します。出力はバリューベクトルの加重和です。

import torch
import torch.nn.functional as F

torch.manual_seed(0)

# それぞれ8次元の4トークンからなるシーケンスのクエリ、キー、バリュー
q, k, v = torch.randn(3, 4, 8).unbind(0)

# 各クエリをすべてのキーと照合し、sqrt(d_k)でスケーリングして、各行をsoftmaxで正規化する
weights = (q @ k.T / k.shape[-1] ** 0.5).softmax(dim=-1)

# 各出力トークンはバリューベクトルの加重平均
output = weights @ v

print("Attention weights of token 0:", [round(w, 3) for w in weights[0].tolist()])
print("Matches PyTorch SDPA:", torch.allclose(output, F.scaled_dot_product_attention(q, k, v), atol=1e-6))

[0.121, 0.741, 0.106, 0.032]とTrueが出力されます。トークン0は、出力の74%をトークン1のバリューベクトルから取得します。また、4つの重みの合計は1となり、手動で計算した結果はPyTorchの融合scaled_dot_product_attentionと一致します。この処理は、対応するGPUでFlashAttentionなどのメモリー効率に優れたカーネルを呼び出します。

アテンションの種類#

• セルフアテンション： クエリ、キー、バリューはすべて同じシーケンスから得られるため、各トークンは他のすべてのトークンからコンテキストを集めます。TransformerエンコーダーやViT内部のアテンションです。
• クロスアテンション： クエリは一方のシーケンスから、キーとバリューはもう一方のシーケンスから得られます。翻訳デコーダーは入力文にアテンションを適用し、テキストから画像を生成する拡散モデルはプロンプトにアテンションを適用します。また、DETRデコーダーの物体クエリは画像特徴量にアテンションを適用します。
• 因果アテンション： マスクによって各トークンが後続のトークンにアテンションを適用できないようにするため、言語モデルはテキストを左から右へ生成できます。また、先行するトークンのキーとバリューはKVキャッシュに保存できます。
• 効率的なアテンション： スパースアテンション、スライディングウィンドウアテンション、線形アテンションは、すべてのトークンを他のすべてのトークンと比較する際の二次計算量を削減します。また、変形可能アテンションは、各基準位置の周囲にある、学習可能な少数の点だけをサンプリングします。
• CNNのチャネルアテンション： Squeeze-and-Excitationネットワークは、特徴マップのグローバル平均を用いてチャネルの重みを調整します。これは計算コストの低いアテンション形式であり、ILSVRC 2017分類コンテストの優勝モデルの基盤となりました。

注意機構と畳み込みの比較#

畳み込みでは、同じ小さなフィルターをすべての位置に適用するため、各出力が参照する固定の局所ウィンドウがあり、その重みは画像によって変わりません。一方、アテンションは入力内容から重みを計算し、1つのレイヤー内の任意の2つの位置を結び付けられます。その代償として、トークン数の2乗に比例して計算コストが増大します。640×640の画像をストライド32で処理すると、20×20のグリッドに400トークンができ、クエリとキーの組み合わせは160,000通りになります。同じ画像のストライド8のグリッドでは6,400トークンとなり、組み合わせは約4,100万通りに達します。YOLO26は、アテンションブロックを適切な位置に配置しています。バックボーンの末尾にあるC2PSAブロックと、ネック内のアテンション対応C3k2ブロックの1つが、20×20のストライド32の特徴マップ上でマルチヘッドセルフアテンションを実行します。それより高解像度のレイヤーはすべて畳み込みのままです。