Self-Attention機構の理論と実装を完全解説

文章を読んでいるとき、私たちは無意識に「注目すべき場所」を切り替えています。たとえば、”The cat sat on the mat because it was tired.” という英文に出会ったとき、”it” が何を指すのかを理解するために、脳は瞬時に文中の他の単語へ意識を向けます。”cat” と “mat” のどちらが疲れるのか――文脈から “cat” だと判断できるのは、私たちが「itという単語と、文中の他の全ての単語との関連度」を暗黙的に計算しているからです。

この「どこに注目するか」を数学的にモデル化したのが Self-Attention（自己注意機構）です。2017年の論文 Attention Is All You Need で提案され、Transformerアーキテクチャの中核として、自然言語処理から画像認識まで幅広い分野に革命をもたらしました。

Self-Attentionを理解すると、以下のような応用先が見えてきます。

自然言語処理: GPTやBERTといった大規模言語モデルの基盤。文中の単語同士の関係を捉え、翻訳・要約・質問応答などの高度なタスクを実現します
画像認識: Vision Transformer（ViT）は画像パッチ間のSelf-Attentionで物体認識を行い、CNNに匹敵・凌駕する性能を発揮しています
音声認識・生成: Whisperなどの音声モデルでは、音声フレーム間のSelf-Attentionで時系列パターンを捉えています
タンパク質構造予測: AlphaFold2はアミノ酸残基間のSelf-Attentionで立体構造を高精度に予測しました

では、なぜSelf-Attentionが必要になったのでしょうか。それを理解するには、それ以前の手法が抱えていた限界を知る必要があります。

本記事の内容

RNN/LSTMの限界とSelf-Attentionの動機
「図書館の検索」アナロジーによる直感的理解
Query, Key, Valueの役割と生成方法
Scaled Dot-Product Attentionの数式と導出
なぜ $\sqrt{d_k}$ でスケーリングするのか
マスク付きAttentionの仕組みと必要性
計算量 $O(n^2)$ の意味と実務的影響
Pythonでのスクラッチ実装とAttention重みの可視化

前提知識

この記事を読む前に、以下の概念を理解しておくと理解が深まります。

内積 — Query・Key間の類似度計算の基盤です
ソフトマックス関数 — Attentionスコアを確率分布に変換する際に使います

なぜSelf-Attentionが必要なのか — RNN/LSTMの限界

Self-Attentionの意義を理解するために、まずそれ以前の主役であったRNN（再帰型ニューラルネットワーク）とLSTMの限界を振り返りましょう。

逐次処理のボトルネック

RNNは文中の単語を先頭から順に処理します。単語 $t$ の隠れ状態 $\bm{h}_t$ は、直前の隠れ状態 $\bm{h}_{t-1}$ と現在の入力 $\bm{x}_t$ から計算されます。

$$ \bm{h}_t = f(\bm{h}_{t-1}, \bm{x}_t) $$

この逐次的な構造は2つの問題を引き起こします。

第一に、長距離依存の困難です。 文頭の単語の情報が文末に届くまでに、何十回もの非線形変換を通過しなければなりません。この過程で勾配が指数的に減衰（勾配消失）し、遠く離れた単語同士の関連を学習することが極めて困難になります。LSTMはゲート機構で勾配消失をある程度緩和しましたが、本質的な「逐次処理の壁」は残りました。

第二に、並列化の不可能性です。 $\bm{h}_t$ を計算するには $\bm{h}_{t-1}$ が必要なので、GPUの並列計算能力を十分に活かせません。シーケンスが長くなるほど訓練時間が増大し、大規模データでの学習が非効率になります。

Self-Attentionが解決したこと

Self-Attentionは、これら2つの問題を根本的に解決しました。

長距離依存: 任意の2つの単語の関連度を「1ステップ」で直接計算します。文頭と文末の単語であっても、情報がバケツリレーで薄まることはありません
並列化: 全ての単語ペアの関連度を行列積で一括計算できるため、GPUの並列処理と非常に相性が良い設計です

この「全ての要素が互いに直接参照できる」というSelf-Attentionのアイデアを、日常的なアナロジーで深く理解していきましょう。

Attentionの直感的理解 — 図書館の検索アナロジー

あなたが図書館で調べ物をするとき

Self-Attentionの仕組みは、図書館での調べ物に似ています。このアナロジーを丁寧に追いかけることで、Query・Key・Valueという3つの概念が自然に理解できます。

あなたが「量子力学の歴史」について調べたいとしましょう。図書館には何千冊もの本が並んでいます。全ての本を1ページずつ読むのは非現実的です。そこで、次のようなプロセスで情報を探します。

検索キーワードを決める（Query）: あなたの頭の中に「量子力学歴史起源」という検索意図があります
本棚のラベルを確認する（Key）: 各本の背表紙やタイトル、目次を見て、自分の検索意図との関連度を判断します
関連する本の中身を読む（Value）: 関連度が高いと判断した本の実際の内容を取り出して読みます

ここで重要なのは、全ての本を同じ程度に読むわけではないという点です。「量子力学入門」には高い注目を払い、「園芸の手引き」にはほぼ注目しません。関連度の高い本の内容を重点的に取り出し、関連度の低い本の内容は無視する――これがまさにAttentionの本質です。

アナロジーをSelf-Attentionに対応させる

Self-Attentionでは、1つの文の中で各単語が同時に「調べ物をする人」になります。 つまり、全ての単語が他の全ての単語に対して「あなたと私はどれくらい関連しているか？」と問いかけ、関連度に応じて情報を集約します。

具体的に “The cat sat on the mat” という文で考えてみましょう。

“cat” という単語が Query（検索する側）になったとき、”sat” や “mat” といった単語の Key（見出し）との関連度を計算します。”sat” はcatの動作を表すので高い関連度、”the” は冠詞なので低い関連度になるでしょう
関連度に応じて各単語の Value（実際の意味情報）を重み付けして集約した結果が、”cat” の新しい表現になります

この処理を文中の全ての単語について同時に行うのがSelf-Attentionです。処理後の各単語の表現は、もはや単独の意味ではなく、文脈を踏まえた意味を持つようになります。

アナロジーの対応表

図書館の調べ物	Self-Attention	数学的表現
あなたの検索意図	Query（問い合わせ）	$\bm{Q} = \bm{X}\bm{W}^Q$
本棚の各本のタイトル	Key（見出し）	$\bm{K} = \bm{X}\bm{W}^K$
本の中身	Value（情報）	$\bm{V} = \bm{X}\bm{W}^V$
検索意図とタイトルの合致度	Attentionスコア	$\bm{Q}\bm{K}^T / \sqrt{d_k}$
関連度に応じた読み込み	重み付き和	$\text{softmax}(\cdot)\bm{V}$

この図書館アナロジーで1点だけ異なるのは、QueryとKeyとValueが全て同じ入力シーケンスから作られるということです。「自分自身を検索する」から “Self”-Attention なのです。通常のAttention（Cross-Attention）では、QueryとKey/Valueが異なるシーケンスから来ますが、Self-Attentionでは同一のシーケンスが3つの異なる役割を同時に担います。

では、具体的にQ・K・Vはどのように生成されるのでしょうか。次のセクションで、その数学的な仕組みを見ていきます。

Query, Key, Valueの生成 — なぜ3つの異なる表現が必要なのか

同じ入力から3つの「視点」を作る

Self-Attentionでは、入力シーケンス $\bm{X} \in \mathbb{R}^{n \times d_{\text{model}}}$（$n$ 個のトークン、各 $d_{\text{model}}$ 次元）から、3つの学習可能な重み行列を使って Query, Key, Value を生成します。

$$ \begin{align} \bm{Q} &= \bm{X}\bm{W}^Q \quad (\bm{W}^Q \in \mathbb{R}^{d_{\text{model}} \times d_k}) \\ \bm{K} &= \bm{X}\bm{W}^K \quad (\bm{W}^K \in \mathbb{R}^{d_{\text{model}} \times d_k}) \\ \bm{V} &= \bm{X}\bm{W}^V \quad (\bm{W}^V \in \mathbb{R}^{d_{\text{model}} \times d_v}) \end{align} $$

ここで自然な疑問が生まれます。なぜ同じ入力 $\bm{X}$ を3回変換するのか？ 1つの表現ではだめなのか？

役割の分離が必要な理由

この疑問に答えるために、再び図書館のアナロジーを使います。

ある本が「量子力学の教科書」だとします。この本には少なくとも3つの異なる「側面」があります。

「この本が欲しがる情報」 — この本が参考文献として必要とする情報（例：線形代数、確率論の知識）。これがQueryに対応します
「この本が提供するトピック」 — 他の本から見たときの「見出し」（例：量子力学、波動関数、シュレーディンガー方程式）。これがKeyに対応します
「この本の実際の中身」 — 他の本に対して提供する具体的な情報内容。これがValueに対応します

もしQ・K・Vが全く同じ表現だったら、「何を検索するか」と「どういう見出しとして見られるか」と「何の情報を提供するか」が全て同一になってしまい、表現力が著しく制限されます。3つの重み行列 $\bm{W}^Q$, $\bm{W}^K$, $\bm{W}^V$ を別々に学習することで、各単語は「問う側」「応える側」「情報を渡す側」として異なる最適な表現を持てるのです。

射影の幾何学的意味

幾何学的に見ると、$\bm{W}^Q$ や $\bm{W}^K$ は $d_{\text{model}}$ 次元空間から $d_k$ 次元の部分空間への射影です。元の $d_{\text{model}}$ 次元の表現には、構文情報・意味情報・位置情報など様々な特徴が混在しています。$\bm{W}^Q$ はその中から「何を求めているか」に関連する軸を抽出し、$\bm{W}^K$ は「何を提供できるか」に関連する軸を抽出します。

この射影により、高次元空間で直接内積を取るよりも、関連度の計算に本当に必要な情報だけを使った効率的な比較が可能になります。

Q・K・Vの生成方法がわかったところで、次はこれらを使ってAttentionスコアを計算し、情報を集約する一連の処理――Scaled Dot-Product Attention――の詳細を見ていきましょう。

Scaled Dot-Product Attention — 数式の全体像

入力の定義

Scaled Dot-Product Attentionへの入力は以下の3つの行列です。

$\bm{Q} \in \mathbb{R}^{n \times d_k}$: Query行列。$n$ 個のクエリベクトル（各 $d_k$ 次元）を行方向に並べたもの
$\bm{K} \in \mathbb{R}^{m \times d_k}$: Key行列。$m$ 個のキーベクトル（各 $d_k$ 次元）を行方向に並べたもの
$\bm{V} \in \mathbb{R}^{m \times d_v}$: Value行列。$m$ 個の値ベクトル（各 $d_v$ 次元）を行方向に並べたもの

Self-Attentionでは $n = m$ です。つまり、同一シーケンスの各要素が互いを参照します。一方、Cross-Attentionでは $n \neq m$ となることもあります（例えばデコーダのクエリがエンコーダの出力を参照する場合）。

4つのステップ

Scaled Dot-Product Attentionは4つの明確なステップから成り立ちます。各ステップが何をしているのかを、直感と数式の両面から追いかけていきましょう。

ステップ1: Attentionスコアの計算（内積による類似度）

まず、全てのQueryと全てのKeyの間の類似度を一括で計算します。内積は2つのベクトルがどれだけ同じ方向を向いているかを測る尺度であり、ここでは「QueryとKeyの関連度」を表します。

$$ \bm{S} = \bm{Q}\bm{K}^T \in \mathbb{R}^{n \times m} $$

行列 $\bm{S}$ の要素 $S_{ij}$ は、$i$ 番目のQueryベクトル $\bm{q}_i$ と $j$ 番目のKeyベクトル $\bm{k}_j$ の内積です。

$$ S_{ij} = \bm{q}_i \cdot \bm{k}_j = \sum_{\ell=1}^{d_k} q_{i\ell} \, k_{j\ell} $$

値が大きいほど関連度が高く、小さい（または負の）ほど関連度が低いことを意味します。

ステップ2: スケーリング

内積の値を $\sqrt{d_k}$ で割ります。

$$ \bm{S}_{\text{scaled}} = \frac{\bm{Q}\bm{K}^T}{\sqrt{d_k}} $$

なぜこのスケーリングが必要なのかは、次のセクションで詳しく解説します。ここでは「内積の値が大きくなりすぎるのを防ぐための正規化」と理解しておいてください。

ステップ3: Softmaxによる正規化

スケーリングされたスコアの各行にソフトマックス関数を適用し、確率分布（Attention重み $\bm{A}$）に変換します。

$$ \bm{A} = \text{softmax}\left(\frac{\bm{Q}\bm{K}^T}{\sqrt{d_k}}\right) $$

Softmaxにより、各行の要素は非負で合計が1になります。

$$ A_{ij} \geq 0, \quad \sum_{j=1}^{m} A_{ij} = 1 $$

$A_{ij}$ は「$i$ 番目のトークンが $j$ 番目のトークンにどれだけ注目しているか」の割合を表します。

ステップ4: Valueの重み付き和

最後に、Attention重みを使ってValueを重み付き和で集約します。

$$ \begin{equation} \text{Attention}(\bm{Q}, \bm{K}, \bm{V}) = \text{softmax}\left(\frac{\bm{Q}\bm{K}^T}{\sqrt{d_k}}\right) \bm{V} \end{equation} $$

出力行列の $i$ 行目を見ると、次のようになっています。

$$ \text{output}_i = \sum_{j=1}^{m} A_{ij} \, \bm{v}_j $$

これは、全てのValueベクトル $\bm{v}_j$ の重み付き和です。関連度の高いトークンのValueが大きな重みで反映され、関連度の低いトークンのValueはほぼ無視されます。つまり、各トークンは文脈に応じて他のトークンから必要な情報だけを選択的に集めるのです。

なぜ内積で類似度を測れるのか

ここで一歩引いて、「なぜ内積が類似度になるのか」を確認しておきましょう。2つのベクトル $\bm{q}$ と $\bm{k}$ の内積は、幾何学的には次のように分解できます。

$$ \bm{q} \cdot \bm{k} = \|\bm{q}\| \, \|\bm{k}\| \cos\theta $$

ここで $\theta$ は2つのベクトルのなす角です。同じ方向を向いているベクトル（$\theta \approx 0$）は大きな正の内積を、直交するベクトル（$\theta = 90°$）はゼロの内積を、反対方向のベクトル（$\theta \approx 180°$）は大きな負の内積を持ちます。したがって、内積はベクトルの「方向の近さ」を測る自然な尺度なのです。

ステップ1〜4で全体像がわかりました。しかし、ステップ2のスケーリングについてはまだ「なぜ必要か」を詳しく説明していませんでした。次のセクションで、この $\sqrt{d_k}$ スケーリングの背後にある統計的な理由を掘り下げます。

なぜ $\sqrt{d_k}$ でスケーリングするのか — 内積の分散の問題

直感的な説明

スケーリングの理由を直感的に理解するために、次のような状況を想像してください。

サイコロを1個振ると、出る目の合計は1〜6です。しかしサイコロを100個振って合計すると、合計値は100〜600の範囲に広がり、「たまたま大きな値が出る」と「たまたま小さな値が出る」の差が非常に大きくなります。

内積 $\bm{q} \cdot \bm{k} = \sum_{i=1}^{d_k} q_i k_i$ はまさにこの「足し合わせ」であり、次元数 $d_k$ が大きいほど値のばらつき（分散）が大きくなります。値のばらつきが大きいと、Softmaxの入力に極端に大きな値と小さな値が混在し、出力がほぼ one-hot（1つだけ1に近く、他は0に近い）になってしまいます。

統計的な導出

これを数式で確認しましょう。$\bm{q}$ と $\bm{k}$ の各成分 $q_i$, $k_i$ が互いに独立で、平均0、分散1の確率変数であると仮定します。

内積は $d_k$ 個の積 $q_i k_i$ の和なので、まず各項の分散を求めます。

$q_i$ と $k_i$ が独立で平均0のとき、積 $q_i k_i$ の分散は次のようになります。

$$ \text{Var}(q_i k_i) = E[q_i^2 k_i^2] – (E[q_i k_i])^2 $$

独立性より $E[q_i k_i] = E[q_i]E[k_i] = 0$、また $E[q_i^2 k_i^2] = E[q_i^2]E[k_i^2] = 1 \cdot 1 = 1$ なので、

$$ \text{Var}(q_i k_i) = 1 $$

各項が独立であるため、$d_k$ 個の和の分散は各項の分散の和になります。

$$ \text{Var}(\bm{q} \cdot \bm{k}) = \sum_{i=1}^{d_k} \text{Var}(q_i k_i) = d_k $$

つまり、内積の分散は次元数 $d_k$ に正比例します。 $d_k = 64$ なら分散は64、標準偏差は8です。Transformerの原論文では $d_k = 64$ が使われており、スケーリングしないと内積の典型的な値が $\pm 16$ 程度にまで広がります。

Softmaxへの影響

Softmaxの出力は $\text{softmax}(z_i) = e^{z_i} / \sum_j e^{z_j}$ です。入力の値の差が大きいと、最大値の指数関数が他を圧倒的に上回り、出力がほぼ one-hot になります。

one-hot に近い出力は、ほぼ1つのトークンだけに注目していることを意味し、Attentionの「複数の関連トークンから情報を柔軟に集約する」という利点が失われます。さらに、Softmaxのone-hot 領域は勾配が極めて小さく（飽和領域）、勾配消失が発生して学習が進まなくなります。

スケーリングの効果

$\sqrt{d_k}$ で割ることで、内積の分散を $d_k / d_k = 1$ に正規化します。

$$ \text{Var}\left(\frac{\bm{q} \cdot \bm{k}}{\sqrt{d_k}}\right) = \frac{1}{d_k} \text{Var}(\bm{q} \cdot \bm{k}) = \frac{d_k}{d_k} = 1 $$

分散が1に正規化されることで、Softmaxの入力値は適度な範囲に収まり、出力も適度に分散した確率分布になります。これにより、勾配が健全に保たれ、学習が安定して進むのです。

この「分散の爆発をスケーリングで抑える」テクニックは、深層学習の初期化戦略（Xavier初期化、He初期化など）にも通じる、重要な設計原則です。

ここまでで、Scaled Dot-Product Attentionの全貌が明らかになりました。しかし実際のTransformerでは、特にデコーダにおいて「未来の情報を見てはいけない」という制約があります。次に、この制約をどのように実現するかを見ていきましょう。

マスク付きAttention — 未来を覗かせない仕組み

なぜマスクが必要なのか

GPTのようなデコーダモデルは、テキストを左から右へ順に生成します。「I love」の次の単語を予測するとき、モデルが見てよいのは「I」と「love」だけです。もし正解の「cats」を事前に見てしまったら、カンニングと同じで、予測タスクとして意味をなしません。

このような「自己回帰（autoregressive）生成」の文脈では、各トークンが未来のトークンを参照できないようにする必要があります。これが因果マスク（causal mask） の役割です。

仕組み — $-\infty$ マスク

マスクは、Softmaxの前にスコア行列に加算する行列 $\bm{M}$ として実装されます。

$$ \text{Attention}(\bm{Q}, \bm{K}, \bm{V}) = \text{softmax}\left(\frac{\bm{Q}\bm{K}^T}{\sqrt{d_k}} + \bm{M}\right) \bm{V} $$

マスク行列 $\bm{M}$ は、参照してよい位置には $0$ を、参照してはいけない位置には $-\infty$ を設定した行列です。因果マスクの場合は上三角部分（$i < j$ の位置）が $-\infty$ になります。

$$ M_{ij} = \begin{cases} 0 & (i \geq j: \text{現在以前のトークン}) \\ -\infty & (i < j: \text{未来のトークン}) \end{cases} $$

$-\infty$ が加算された位置は、$e^{-\infty} = 0$ となるため、Softmax後のAttention重みが厳密に0になります。つまり、未来のトークンからは一切の情報が流入しません。

因果マスクの例

4トークンのシーケンスに対する因果マスクを具体的に見てみましょう。

$$ \bm{M} = \begin{pmatrix} 0 & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0 \end{pmatrix} $$

1行目（トークン1）: 自分自身のみ参照可能
2行目（トークン2）: トークン1と自分自身を参照可能
3行目（トークン3）: トークン1, 2と自分自身を参照可能
4行目（トークン4）: 全トークンを参照可能

このマスクにより、Softmax後のAttention重み行列は下三角行列になり、各トークンは自分自身と過去のトークンだけから情報を集約します。

エンコーダ vs デコーダ

Transformerアーキテクチャのエンコーダでは、通常マスクは使いません。エンコーダは入力シーケンス全体が与えられた状態で処理するため、全てのトークンが互いを自由に参照してよいからです。BERTはエンコーダ型であり、マスクなしの双方向Self-Attentionを使っています。

一方、デコーダでは因果マスクが必須です。GPTはデコーダ型であり、因果マスク付きSelf-Attentionで自己回帰生成を行います。

なお、因果マスク以外にもパディングマスク（異なる長さのシーケンスをバッチ処理する際に、パディングトークンを無視するためのマスク）もあります。

マスクの仕組みがわかったところで、次にSelf-Attentionの計算コストについて考えてみましょう。全てのトークンペアの関連度を計算するというのは、直感的にもコストが高そうですが、具体的にどの程度なのでしょうか。

計算量 $O(n^2)$ の意味と実務的影響

各演算の計算量

Scaled Dot-Product Attentionの計算量を、ステップごとに分解して見ていきましょう。Self-Attentionの場合 $n = m$ とします。

演算	計算量	説明
$\bm{Q}\bm{K}^T$ の計算	$O(n^2 d_k)$	$n \times d_k$ と $d_k \times n$ の行列積
スケーリング	$O(n^2)$	$n \times n$ 行列の各要素を $\sqrt{d_k}$ で割る
Softmax	$O(n^2)$	$n \times n$ 行列の各行にSoftmaxを適用
$\bm{A}\bm{V}$ の計算	$O(n^2 d_v)$	$n \times n$ と $n \times d_v$ の行列積
合計	$O(n^2 d)$	$d = \max(d_k, d_v)$

なぜ $O(n^2)$ が問題なのか

シーケンス長 $n$ に対して計算量が2乗で増加するということは、具体的には以下のような影響を持ちます。

$n = 512$ のとき: 約26万の要素ペア
$n = 2048$ のとき: 約420万の要素ペア（$n = 512$ の16倍）
$n = 8192$ のとき: 約6700万の要素ペア（$n = 512$ の256倍）

シーケンス長を4倍にすると計算量は16倍になります。このため、長文の処理や高解像度画像のViTでは、計算量とメモリが深刻なボトルネックになります。

メモリの問題

計算量だけでなく、メモリ使用量も大きな課題です。$n \times n$ のAttention重み行列をGPUメモリ上に保持する必要があり、$n = 8192$ では32ビット浮動小数点で約256MBのメモリを1つのAttentionレイヤー・1つのヘッドだけで消費します。Multi-Head Attentionで複数ヘッドを使い、さらにバッチ処理をすると、メモリ使用量は急速に増大します。

効率化の研究

この $O(n^2)$ の壁を突破するために、多くの研究が行われています。

Flash Attention: メモリアクセスパターンを最適化し、Attention重み行列を明示的にメモリ上に構築することなく計算を行います。計算量のオーダーは変わりませんが、実行時間とメモリ使用量を大幅に削減します
KVキャッシュ: 自己回帰生成時に過去のKey/Valueを再計算せずキャッシュすることで、推論時の計算量を削減します
Sparse Attention: 全トークンペアではなく、局所的・ストライド的なパターンのみでAttentionを計算し、$O(n\sqrt{n})$ や $O(n \log n)$ に削減します
Linear Attention: カーネル近似などでSoftmaxを置き換え、$O(n)$ の線形計算量を実現する手法です

これらの手法により、現在のLLMでは $n = 100{,}000$ を超えるコンテキスト長が実用化されています。

計算量の理解が深まったところで、ここまでの理論をPythonコードで実装し、実際にAttention重みがどのように振る舞うかを確認していきましょう。

Pythonでの実装

Scaled Dot-Product Attentionのスクラッチ実装

まず、Scaled Dot-Product Attentionをゼロから実装します。NumPyのみを使い、各ステップがどの数式に対応しているかをコメントで明示します。

import numpy as np


def softmax(x, axis=-1):
    """数値的に安定なSoftmax実装

    最大値を引いてからexpを取ることで、オーバーフローを防ぎます。
    """
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e_x / np.sum(e_x, axis=axis, keepdims=True)


def scaled_dot_product_attention(Q, K, V, mask=None):
    """Scaled Dot-Product Attention

    Parameters:
        Q: Query行列 (n, d_k)
        K: Key行列 (m, d_k)
        V: Value行列 (m, d_v)
        mask: マスク行列 (n, m) or None。禁止位置に-infを設定

    Returns:
        output: Attention出力 (n, d_v)
        attention_weights: Attention重み (n, m)
    """
    d_k = Q.shape[-1]

    # ステップ1: QとKの内積でAttentionスコアを計算
    scores = Q @ K.T  # (n, m)

    # ステップ2: sqrt(d_k)でスケーリング
    scores = scores / np.sqrt(d_k)

    # マスクがある場合、禁止位置に-infを加算
    if mask is not None:
        scores = scores + mask

    # ステップ3: Softmaxで確率分布に変換
    attention_weights = softmax(scores, axis=-1)  # (n, m)

    # ステップ4: Valueの重み付き和で出力を計算
    output = attention_weights @ V  # (n, d_v)

    return output, attention_weights

上のコードの各ステップは、前セクションの数式にそのまま対応しています。softmax 関数では、数値的な安定性のために入力から最大値を引いてから指数関数を計算しています。これは $\text{softmax}(\bm{z}) = \text{softmax}(\bm{z} – c)$（任意の定数 $c$ を引いてもSoftmaxの値は変わらない）という性質を利用したテクニックです。

Q, K, Vの生成と動作確認

次に、入力シーケンスからQ, K, Vを生成し、Self-Attentionの計算が正しく動作するか確認します。

import numpy as np

np.random.seed(42)

# パラメータ設定
n = 4           # シーケンス長（4トークン）
d_model = 8     # 入力の次元数
d_k = 4         # Query/Keyの次元数
d_v = 4         # Valueの次元数

# 入力シーケンス（4トークン、各8次元）
X = np.random.randn(n, d_model)

# 学習可能な射影行列（実際には学習で最適化される）
W_Q = np.random.randn(d_model, d_k) * 0.1
W_K = np.random.randn(d_model, d_k) * 0.1
W_V = np.random.randn(d_model, d_v) * 0.1

# Q, K, V の生成
Q = X @ W_Q
K = X @ W_K
V = X @ W_V

# Self-Attentionの計算
output, weights = scaled_dot_product_attention(Q, K, V)

print("入力 X の形状:", X.shape)
print("Q の形状:", Q.shape)
print("K の形状:", K.shape)
print("V の形状:", V.shape)
print("出力の形状:", output.shape)
print()
print("Attention重み行列（各行の和 = 1）:")
print(weights.round(3))
print()
print("各行の和:", weights.sum(axis=1).round(6))

このコードの出力から、いくつかの重要な点を確認できます。まず、出力の形状が $(4, 4)$ であり、入力シーケンス長は保ったまま次元が $d_v$ になっていることがわかります。次に、Attention重み行列の各行の和が厳密に1.0であり、Softmaxが正しく確率分布を作っていることが確認できます。さらに、重み行列は一様分布（$1/n = 0.25$）からやや偏りがあり、これはランダムなQ, K間でも若干の類似度の差が生じていることを示しています。学習が進むと、この偏りが意味のあるパターン（同じ主語−動詞の組を高い重みで結ぶなど）へと変化します。

スケーリングの効果を実験で確認

スケーリングの重要性を実験的に確認してみましょう。次元数 $d_k$ を変化させたとき、スケーリングの有無がSoftmax出力にどう影響するかを観察します。

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

d_k_values = [4, 16, 64, 256, 1024]
n = 6  # シーケンス長

fig, axes = plt.subplots(2, len(d_k_values), figsize=(20, 8))

for idx, d_k in enumerate(d_k_values):
    Q = np.random.randn(n, d_k)
    K = np.random.randn(n, d_k)

    # スケーリングなし
    scores_raw = Q @ K.T
    weights_raw = softmax(scores_raw)

    # スケーリングあり
    scores_scaled = Q @ K.T / np.sqrt(d_k)
    weights_scaled = softmax(scores_scaled)

    # 上段: スケーリングなし
    im1 = axes[0, idx].imshow(weights_raw, cmap='Blues', vmin=0, vmax=1)
    axes[0, idx].set_title(f'd_k={d_k}\n(no scaling)', fontsize=11)
    for i in range(n):
        for j in range(n):
            axes[0, idx].text(j, i, f'{weights_raw[i,j]:.2f}',
                              ha='center', va='center', fontsize=7,
                              color='white' if weights_raw[i,j] > 0.4 else 'black')

    # 下段: スケーリングあり
    im2 = axes[1, idx].imshow(weights_scaled, cmap='Blues', vmin=0, vmax=1)
    axes[1, idx].set_title(f'd_k={d_k}\n(with scaling)', fontsize=11)
    for i in range(n):
        for j in range(n):
            axes[1, idx].text(j, i, f'{weights_scaled[i,j]:.2f}',
                              ha='center', va='center', fontsize=7,
                              color='white' if weights_scaled[i,j] > 0.4 else 'black')

axes[0, 0].set_ylabel('Without scaling', fontsize=12)
axes[1, 0].set_ylabel('With scaling', fontsize=12)

plt.suptitle('Effect of scaling on Attention weights', fontsize=14, y=1.02)
plt.tight_layout()
plt.show()

この可視化から、スケーリングの効果が明瞭に読み取れます。上段（スケーリングなし）では、$d_k$ が大きくなるにつれてAttention重みが急速にone-hotに近づいていく様子が確認できます。$d_k = 1024$ ではほぼ全ての行で1つの要素が1.00に近くなり、他が0.00に潰れています。一方、下段（スケーリングあり）では、$d_k$ の値によらずAttention重みが適度に分散した分布を保っています。これは、$\sqrt{d_k}$ で割ることで内積の分散が常に1に正規化されている効果です。この実験は、スケーリングが学習の安定性にとって不可欠であることを実証的に裏付けています。

Attention重みの可視化 — 文の意味構造を読み取る

次に、具体的な文を使ってAttention重みを可視化し、Self-Attentionが文の意味構造をどのように捉えるかを見てみましょう。ここでは学習済みの重みを模擬するために、特定の単語ペアの類似度を手動で調整しています。

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(0)
words = ["The", "cat", "sat", "on", "the", "mat"]
n = len(words)
d_k = 8

# ダミーのQ, K を生成
Q = np.random.randn(n, d_k) * 0.5
K = np.random.randn(n, d_k) * 0.5

# 学習後の状態を模擬: 意味的に関連する単語ペアの類似度を高める
# "cat" (idx=1) と "sat" (idx=2) は主語-動詞の関係
Q[1] = K[2] * 0.7 + np.random.randn(d_k) * 0.2
# "sat" (idx=2) と "cat" (idx=1) — 動詞から主語への逆参照
Q[2] = K[1] * 0.6 + np.random.randn(d_k) * 0.2
# "sat" (idx=2) と "mat" (idx=5) — 動詞と場所の関係
Q[2] += K[5] * 0.4
# "on" (idx=3) と "mat" (idx=5) — 前置詞と目的語
Q[3] = K[5] * 0.8 + np.random.randn(d_k) * 0.15

scores = Q @ K.T / np.sqrt(d_k)
weights = softmax(scores)

fig, ax = plt.subplots(figsize=(8, 6))
im = ax.imshow(weights, cmap='Blues', aspect='auto', vmin=0, vmax=0.5)

ax.set_xticks(range(n))
ax.set_yticks(range(n))
ax.set_xticklabels(words, fontsize=12)
ax.set_yticklabels(words, fontsize=12)
ax.set_xlabel('Key (referenced token)', fontsize=12)
ax.set_ylabel('Query (attending token)', fontsize=12)
ax.set_title('Self-Attention weights: "The cat sat on the mat"', fontsize=14)

for i in range(n):
    for j in range(n):
        ax.text(j, i, f'{weights[i,j]:.2f}',
                ha='center', va='center', fontsize=10,
                color='white' if weights[i,j] > 0.3 else 'black')

plt.colorbar(im, label='Attention weight')
plt.tight_layout()
plt.show()

この可視化結果から、Self-Attentionが文の構文構造を反映していることが読み取れます。”cat”（主語）は “sat”（動詞）に高い注目を向けており、主語−動詞の関係が重みに表れています。同様に、”on”（前置詞）は “mat”（目的語）に強く注目しています。一方、冠詞 “The” や “the” は特定の単語に対して突出した注目を見せず、比較的均等な分布になっています。実際の学習済みTransformerでも、異なるAttentionヘッドが構文関係・意味関係・位置関係など異なるパターンを学習することが知られています。この「複数の視点からAttentionを計算する」のがMulti-Head Attentionです。

因果マスク付きAttentionの実装と可視化

デコーダで使われる因果マスク付きAttentionを実装し、マスクの有無によるAttention重みの違いを比較します。

import numpy as np
import matplotlib.pyplot as plt


def causal_mask(n):
    """因果マスクを生成

    下三角部分は0（参照可能）、上三角部分は-inf（参照禁止）。
    """
    mask = np.full((n, n), -np.inf)
    mask[np.tril_indices(n)] = 0  # 下三角（対角含む）を0にする
    return mask


np.random.seed(42)
n, d_k = 5, 8
tokens = ["I", "love", "deep", "learning", "!"]

Q = np.random.randn(n, d_k)
K = np.random.randn(n, d_k)
V = np.random.randn(n, d_k)

# マスクなし（エンコーダ的）
_, weights_no_mask = scaled_dot_product_attention(Q, K, V)

# 因果マスクあり（デコーダ的）
mask = causal_mask(n)
_, weights_masked = scaled_dot_product_attention(Q, K, V, mask=mask)

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

titles = ['Encoder-style (no mask)', 'Decoder-style (causal mask)']
weight_list = [weights_no_mask, weights_masked]

for ax, w, title in zip(axes, weight_list, titles):
    im = ax.imshow(w, cmap='Blues', aspect='auto', vmin=0, vmax=0.6)
    for i in range(n):
        for j in range(n):
            ax.text(j, i, f'{w[i,j]:.2f}', ha='center', va='center',
                    fontsize=10, color='white' if w[i,j] > 0.3 else 'black')
    ax.set_xticks(range(n))
    ax.set_yticks(range(n))
    ax.set_xticklabels(tokens, fontsize=11)
    ax.set_yticklabels(tokens, fontsize=11)
    ax.set_xlabel('Key position', fontsize=12)
    ax.set_ylabel('Query position', fontsize=12)
    ax.set_title(title, fontsize=13)
    plt.colorbar(im, ax=ax)

plt.tight_layout()
plt.show()

この比較から、因果マスクの効果が明確に見えます。左側（マスクなし）では全てのトークンが全てのトークンを参照でき、Attention重みが行列全体に分布しています。右側（因果マスクあり）では上三角部分が完全にゼロになっており、各トークンが自分自身と過去のトークンだけを参照しています。注目すべきは、マスクにより参照可能なトークン数が制限されることで、各行の重みが少数のトークンに集中する傾向がある点です。特に最初のトークン “I” は自分自身しか参照できないため、重みが1.00になっています。自己回帰生成においてモデルが未来の情報に依存しないことを、この可視化で確認できます。

エントロピーによるAttention分布の鋭さの定量化

最後に、Attention重み分布の「鋭さ」をエントロピーで定量化するコードを紹介します。エントロピーが低いほど少数のトークンに集中（鋭い分布）、高いほど均等に分散した分布を意味します。

import numpy as np

def attention_entropy(weights):
    """Attention重みのエントロピーを各行について計算

    H = -sum(p * log(p))
    p=0 のとき p*log(p)=0 と定義（情報理論の慣例）。
    """
    # 0の対数を避けるためクリップ
    w_clipped = np.clip(weights, 1e-12, 1.0)
    entropy = -np.sum(weights * np.log(w_clipped), axis=-1)
    return entropy


np.random.seed(42)
n = 6
d_k_values = [4, 64, 256]

print("=== Attention重み分布のエントロピー ===")
print(f"（均等分布のエントロピー: {np.log(n):.3f}）\n")

for d_k in d_k_values:
    Q = np.random.randn(n, d_k)
    K = np.random.randn(n, d_k)

    # スケーリングなし
    scores_raw = Q @ K.T
    w_raw = softmax(scores_raw)
    h_raw = attention_entropy(w_raw)

    # スケーリングあり
    scores_scaled = Q @ K.T / np.sqrt(d_k)
    w_scaled = softmax(scores_scaled)
    h_scaled = attention_entropy(w_scaled)

    print(f"d_k = {d_k}:")
    print(f"  スケーリングなし — 平均エントロピー: {h_raw.mean():.3f}")
    print(f"  スケーリングあり — 平均エントロピー: {h_scaled.mean():.3f}")
    print()

この出力から、スケーリングの効果を定量的に確認できます。均等分布のエントロピーは $\ln(6) \approx 1.791$ です。スケーリングなしの場合、$d_k$ が大きくなるとエントロピーが急激に低下し、one-hotに近い分布になっていることがわかります。一方、スケーリングありの場合は $d_k$ に関わらずエントロピーが安定しており、適度に分散した分布が保たれています。これは先ほどの可視化で視覚的に確認した結果と完全に一致しています。

Self-Attentionの出力が持つ意味

ここまでの内容を振り返り、Self-Attentionの出力が何を表しているのかを改めて整理します。

入力シーケンス $\bm{X}$ の各トークンは、もともと単語埋め込みや位置エンコーディングによる固定的な表現を持っています。Self-Attentionを通すと、各トークンの表現は「文脈を考慮した表現」に変換されます。

具体的には、”bank” という単語は、単語埋め込みの段階では「銀行」と「川岸」の両方の意味を含んだ汎用的なベクトルです。しかし Self-Attention を通すと、周囲に “money” や “account” がある場合は「銀行」に近い表現に、”river” や “water” がある場合は「川岸」に近い表現に変化します。これは、Attention重みを通じて周囲のトークンの情報が混合されるためです。

この「文脈に応じた動的な表現」こそが、Transformerが多義語の曖昧性解消や長距離の照応関係の把握に優れている理由です。

Multi-Head Attentionへの発展

ここまで見てきたSelf-Attentionは「Single-Head」のAttentionです。実際のTransformerでは、これを複数並列に実行するMulti-Head Attentionが使われます。各ヘッドが異なる $\bm{W}^Q$, $\bm{W}^K$, $\bm{W}^V$ を持つことで、「構文関係に注目するヘッド」「意味的類似性に注目するヘッド」「近接位置に注目するヘッド」など、多様な観点からの情報集約が可能になります。

また、Self-Attentionは入力トークンの順序情報を持たない（全ての位置ペアを等しく扱う）ため、位置エンコーディングやRotary Position Embedding（RoPE）で位置情報を補う必要があります。さらに、Self-Attentionの出力はLayer Normalizationと残差接続を経て次の層に渡されます。

まとめ

本記事では、Self-Attention機構の動機・理論・実装を、直感的な説明から数式、そしてPythonコードまで一貫して解説しました。

Self-Attentionの動機: RNN/LSTMの長距離依存の困難と並列化の限界を解決するために生まれた。全てのトークンペアの関連度を1ステップで直接計算できる
図書館アナロジー: Query（検索意図）・Key（本のタイトル）・Value（本の中身）の3つの役割で理解できる。関連度の高い情報を選択的に集約する仕組み
Q, K, Vの生成: 同じ入力から3つの異なる射影を作ることで、「問う側」「応える側」「情報を渡す側」の役割を分離し、表現力を確保する
Scaled Dot-Product Attention: $\text{Attention}(\bm{Q}, \bm{K}, \bm{V}) = \text{softmax}\left(\bm{Q}\bm{K}^T / \sqrt{d_k}\right)\bm{V}$ の4ステップで構成される
$\sqrt{d_k}$ スケーリング: 内積の分散が $d_k$ に比例するため、スケーリングしないとSoftmaxが飽和して勾配消失が起きる。$\sqrt{d_k}$ で割ることで分散を1に正規化する
因果マスク: デコーダの自己回帰生成で未来の情報参照を防ぐために、上三角部分に $-\infty$ を設定する
計算量 $O(n^2 d)$: シーケンス長の2乗に比例し、長いシーケンスでボトルネックとなる。Flash AttentionやSparse Attentionなどの効率化手法が研究されている

Self-Attentionは、現代の深層学習における最も重要な構成要素の1つです。本記事で扱った Single-Head の仕組みを理解していれば、Multi-Head AttentionやTransformer全体のアーキテクチャもスムーズに理解できるでしょう。

次のステップとして、以下の記事も参考にしてください。

Multi-Head Attention — 複数の視点からAttentionを並列実行し、表現力を高める仕組み
位置エンコーディング — Self-Attentionが持たない位置情報を補う手法
Transformerアーキテクチャ — Self-Attentionを組み込んだ全体設計
Flash Attention — $O(n^2)$ メモリの壁を実用的に突破する手法