Self-Attentionとは？3トークンの重みをPythonで計算する
埋め込みと位置情報によって、各トークンは順番を含むベクトルになりました。次に必要なのは、ある位置の表現を作るとき、文中のどの位置からどれだけ情報を集めるかです。Self-Attentionは、同じ系列の各位置どうしの関係を点数化し、情報を重み付きで混ぜます。
「Attentionは重要な単語を見つける仕組み」とだけ覚えると、計算の一部が抜けます。重みは文全体に一つではなく、情報を受け取る位置ごとに作られます。同じ単語でも、どの位置から問い合わせるかによって参照の配分は変わります。
Query、Key、Valueの役割
各トークンの入力ベクトルから、学習可能な線形変換でQuery、Key、Valueを作ります。
Queryは「この位置がどんな情報を探すか」、Keyは「各位置がどんな照合用の特徴を持つか」、Valueは「選ばれたときに渡す情報」に相当します。これは理解のための言い換えで、各次元に人が役割を指定するわけではありません。
あるQueryとすべてのKeyの内積を取ると、参照先ごとのスコアができます。スコアをsoftmaxで合計1の重みに変え、その重みでValueを平均します。
Attention(Q, K, V) = softmax(QK^T / √d_k) Vd_k はKeyの次元数です。内積を √d_k で割るのは、次元が大きいと内積の大きさが増え、softmaxが極端な領域に入りやすくなるためです。
3トークンを一度に計算する
説明を短くするため、すでに作られたQ、K、Vを直接使います。実際のモデルでは、入力行列へ別々の重み行列を掛けて作ります。
import numpy as np
Q = np.array([ [1.0, 0.0], [0.0, 1.0], [1.0, 1.0],])K = np.array([ [1.0, 0.0], [0.0, 1.0], [1.0, 1.0],])V = np.array([ [10.0, 0.0], [0.0, 10.0], [5.0, 5.0],])
scores = Q @ K.T / np.sqrt(K.shape[1])shifted = scores - scores.max(axis=1, keepdims=True)weights = np.exp(shifted)weights /= weights.sum(axis=1, keepdims=True)output = weights @ V
print("スコア:\n", np.round(scores, 3))print("Attention重み:\n", np.round(weights, 3))print("出力:\n", np.round(output, 3))print("行ごとの重み合計:", np.round(weights.sum(axis=1), 3))出力:
スコア: [[0.707 0. 0.707] [0. 0.707 0.707] [0.707 0.707 1.414]]Attention重み: [[0.401 0.198 0.401] [0.198 0.401 0.401] [0.248 0.248 0.503]]出力: [[6.017 3.983] [3.983 6.017] [5. 5. ]]行ごとの重み合計: [1. 1. 1.]重み行列の各行は、一つのQueryが3個のValueをどの比率で混ぜるかを示します。1行目は1番目と3番目へ約0.401ずつ、2番目へ約0.198を割り当てました。その結果、1行目の出力は3個のValueの重み付き平均 [6.017, 3.983] になります。
softmaxを行ごとに計算するため、各行の合計は1です。ただしAttention重みをそのまま、人間が説明する「重要度」や因果関係と同一視はできません。Valueの内容、後続の線形変換、残差接続、ほかの層も最終出力へ影響します。
なぜSelfと呼ぶのか
Query、Key、Valueの元が同じ系列だからSelf-Attentionと呼びます。原論文のエンコーダーでは、一つ前の層の出力からQ、K、Vを作り、各位置が同じ系列の全位置を参照します。
一方、エンコーダーとデコーダーの間のAttentionでは、QueryとKey・Valueの出どころが異なります。Attentionという計算がすべてSelf-Attentionなのではありません。
文章生成では未来を見せない
次トークンを予測するデコーダー型の処理で、現在位置が後ろの正解トークンを参照できると、学習課題の答えを先に見てしまいます。そのため未来側のスコアをsoftmax前に無効化する因果マスクを使います。
原論文のデコーダーも、現在位置より後ろへAttentionを向けられないようにマスクしています。推論時にはまだ未来のトークン自体が存在しませんが、訓練では系列をまとめて計算するため、この制約が必要です。
一つのAttentionだけではない
実際のTransformerでは、複数のAttentionヘッドを並行して計算するMulti-Head Attentionが使われます。ヘッドごとに異なる射影を持つため、同じ入力から異なる関係を捉えられます。各ヘッドの出力を結合し、もう一度線形変換して層の出力にします。
今回の例は一つのヘッドと、すでに用意されたQ、K、Vだけに絞りました。次にQuery・Key・Valueを入力から作る過程を分けて見ると、Attention重みが固定の規則ではなく、学習された射影によって変わることが分かります。