Self-Attentionの理論と実装 — Query・Key・Valueの計算

前回の記事では、エンコーダとデコーダの間でAttentionを計算する仕組みを学びました。しかしTransformerの革新的な点は、Attentionを系列内部に適用したことです。入力系列の各要素が、同じ系列の他の全要素との関係を動的に計算する — これがSelf-Attentionです。

たとえば「The animal didn’t cross the street because it was too tired」という文で、「it」が何を指すかを理解するには、文全体の文脈を考慮する必要があります。Self-Attentionでは、「it」というトークンが「animal」や「street」との関連度を計算し、「animal」への重みが高くなることで、「it = animal」という参照関係を捉えます。

RNNは逐次的に情報を処理するため、離れた位置のトークン間の依存関係の学習が困難でしたが、Self-Attentionは全ペアを一度に計算するため、長距離依存を直接的に捉えることができます。

Self-Attentionを理解すると、以下のことが可能になります。

Transformerの完全理解: BERT、GPT、Vision Transformerの核心部分
Multi-Head Attentionの設計: なぜ複数のAttentionヘッドが必要なのか
計算量の分析: Self-Attentionの $O(n^2)$ 計算量とその対策
Attentionパターンの解釈: モデルが何に注目しているかの可視化

本記事の内容

Self-Attentionの動機と定式化
Query・Key・Valueの線形射影
Scaled Dot-Product Attention の導出
Multi-Head Attention の理論と実装
Masked Attention（因果的Attention）
Pythonでのスクラッチ実装

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

Attention機構の基礎 — Attention の基本概念
ニューラルネットワークの基礎 — 行列演算の基礎

Self-Attentionとは

基本的なアイデア

通常のAttentionでは、Query（デコーダ）とKey/Value（エンコーダ）が異なる系列から来ます。Self-Attentionでは、Query、Key、Valueが全て同じ系列から計算されます。

入力系列 $\bm{X} = (\bm{x}_1, \bm{x}_2, \ldots, \bm{x}_n) \in \mathbb{R}^{n \times d}$ に対して、各要素 $\bm{x}_i$ が他の全要素との関連度を計算し、文脈を考慮した表現を生成します。

なぜ線形射影が必要なのか

入力ベクトル $\bm{x}_i$ をそのままQuery、Key、Valueとして使うと、内積 $\bm{x}_i^\top \bm{x}_j$ は単に2つのベクトルの類似度を測るだけです。しかし、「何を探しているか」（Query）と「何が提供されているか」（Key）は異なる概念であり、「取り出す情報」（Value）も異なる変換が必要です。

そこで、入力を3つの異なる線形変換で射影します。

$$ \begin{align} \bm{Q} &= \bm{X}\bm{W}_Q \quad (\text{Query行列}: n \times d_k) \\ \bm{K} &= \bm{X}\bm{W}_K \quad (\text{Key行列}: n \times d_k) \\ \bm{V} &= \bm{X}\bm{W}_V \quad (\text{Value行列}: n \times d_v) \end{align} $$

$\bm{W}_Q \in \mathbb{R}^{d \times d_k}$、$\bm{W}_K \in \mathbb{R}^{d \times d_k}$、$\bm{W}_V \in \mathbb{R}^{d \times d_v}$ は学習可能なパラメータです。

この射影により、同じ入力から「検索クエリ」「検索キー」「取得する値」という3つの異なる役割を持つ表現を生成できます。

Scaled Dot-Product Attention

定式化

Query行列 $\bm{Q}$、Key行列 $\bm{K}$、Value行列 $\bm{V}$ が得られたら、Scaled Dot-Product Attentionは次のように計算されます。

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

この式を分解して理解しましょう。

ステップ1: スコア行列の計算 $\bm{S} = \bm{Q}\bm{K}^\top \in \mathbb{R}^{n \times n}$

$S_{ij} = \bm{q}_i^\top \bm{k}_j$ は、位置 $i$ のQueryと位置 $j$ のKeyの内積（類似度）です。$n \times n$ のスコア行列は、全ペアの関連度を一度に計算しています。

ステップ2: スケーリング $\bm{S} / \sqrt{d_k}$

$d_k$ が大きいとき、内積の値は $O(\sqrt{d_k})$ のオーダーで大きくなります。$\bm{q}$ と $\bm{k}$ の各要素が独立に平均0・分散1の場合、$\bm{q}^\top \bm{k}$ の分散は $d_k$ になります。$\sqrt{d_k}$ で割ることで分散を1に正規化し、Softmaxの入力が極端に大きくなるのを防ぎます。

Softmaxの入力が大きすぎると、最大値に対応する位置の重みがほぼ1になり、他の位置の重みがほぼ0になります（ハードAttention化）。勾配もほぼ0になるため、学習が停滞します。

ステップ3: Softmaxで正規化 $\bm{A} = \text{softmax}(\bm{S} / \sqrt{d_k})$

各行が確率分布になります。$A_{ij}$ は「位置 $i$ が位置 $j$ にどれだけ注目するか」を表します。

ステップ4: Valueの重み付き和 $\text{Output} = \bm{A}\bm{V}$

Attention重みに基づいてValueを集約します。出力の $i$ 行目は $\sum_j A_{ij} \bm{v}_j$ であり、位置 $i$ の文脈を考慮した表現です。

計算量の分析

Self-Attentionの計算量は、主にスコア行列の計算 $\bm{Q}\bm{K}^\top$（$O(n^2 d_k)$）とValue の重み付き和 $\bm{A}\bm{V}$（$O(n^2 d_v)$）から、$O(n^2 d)$ となります。

系列長 $n$ に対して二乗のコストがかかるため、非常に長い系列（数千〜数万トークン）では計算がボトルネックになります。これに対して、Linear Attention、Sparse Attention、FlashAttentionなどの効率化手法が提案されています。

Self-Attentionの基本を理解したところで、1つのAttentionだけでは不十分な理由と、その解決策であるMulti-Head Attentionを見ていきましょう。

Multi-Head Attention

なぜ複数のヘッドが必要なのか

1つのAttentionヘッドは、入力の1つの側面（たとえば構文的な関係）に注目する傾向があります。しかし、言語の理解には複数の側面（構文、意味、参照関係など）を同時に捉える必要があります。

Multi-Head Attentionは、入力を複数の異なる部分空間に射影し、各部分空間で独立にAttentionを計算してから結果を連結します。

$$ \begin{align} \text{head}_i &= \text{Attention}(\bm{X}\bm{W}_Q^{(i)}, \bm{X}\bm{W}_K^{(i)}, \bm{X}\bm{W}_V^{(i)}) \\ \text{MultiHead}(\bm{X}) &= \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\bm{W}_O \end{align} $$

$h$ はヘッド数、$\bm{W}_Q^{(i)} \in \mathbb{R}^{d \times d_k}$、$d_k = d / h$ です。$\bm{W}_O \in \mathbb{R}^{d \times d}$ は連結後の射影行列です。

パラメータ数

各ヘッドの次元を $d_k = d_v = d/h$ とすると、Multi-Head Attentionの全パラメータ数は

$$ \underbrace{3 \times h \times d \times (d/h)}_{\text{Q, K, Vの射影}} + \underbrace{d \times d}_{\text{出力射影}} = 4d^2 $$

つまり、ヘッド数 $h$ を変えてもパラメータ数は変わりません。各ヘッドの次元が $d/h$ に小さくなるだけです。

各ヘッドの役割

学習されたMulti-Head Attentionでは、異なるヘッドが異なるパターンを捉えることが観察されています。

あるヘッドは隣接するトークンに注目（局所的なパターン）
あるヘッドは文頭や文末に注目（構造的なパターン）
あるヘッドは特定の品詞（名詞、動詞など）に注目（意味的なパターン）
あるヘッドは参照関係（代名詞とその先行詞）に注目

Masked Attention（因果的Attention）

自己回帰モデルでの必要性

GPTのような自己回帰言語モデルでは、位置 $i$ のトークンは位置 $1, \ldots, i-1$ の過去のトークンのみを参照でき、未来のトークン $i+1, \ldots, n$ は見えてはいけません。

これを実現するのがMasked Attention（Causal Attention）です。スコア行列の上三角部分を $-\infty$ に設定します。

$$ \text{Mask}_{ij} = \begin{cases} 0 & (j \leq i) \\ -\infty & (j > i) \end{cases} $$

$$ \text{MaskedAttention}(\bm{Q}, \bm{K}, \bm{V}) = \text{softmax}\left(\frac{\bm{Q}\bm{K}^\top}{\sqrt{d_k}} + \text{Mask}\right)\bm{V} $$

$-\infty$ はSoftmax後に0になるため、未来の情報が遮断されます。

Pythonでの実装

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

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

def softmax(x, axis=-1):
    e = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e / np.sum(e, axis=axis, keepdims=True)

def scaled_dot_product_attention(Q, K, V, mask=None):
    """Scaled Dot-Product Attention"""
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)

    if mask is not None:
        scores = scores + mask

    attn_weights = softmax(scores, axis=-1)
    output = attn_weights @ V
    return output, attn_weights

class MultiHeadAttention:
    """Multi-Head Attention"""

    def __init__(self, d_model, n_heads):
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        # 射影行列の初期化
        scale = np.sqrt(2.0 / d_model)
        self.W_Q = np.random.randn(d_model, d_model) * scale
        self.W_K = np.random.randn(d_model, d_model) * scale
        self.W_V = np.random.randn(d_model, d_model) * scale
        self.W_O = np.random.randn(d_model, d_model) * scale

    def forward(self, X, mask=None):
        """X: (seq_len, d_model)"""
        n = X.shape[0]

        # Q, K, V の計算
        Q = X @ self.W_Q  # (n, d_model)
        K = X @ self.W_K
        V = X @ self.W_V

        # ヘッドに分割
        Q_heads = Q.reshape(n, self.n_heads, self.d_k)
        K_heads = K.reshape(n, self.n_heads, self.d_k)
        V_heads = V.reshape(n, self.n_heads, self.d_k)

        # 各ヘッドでAttentionを計算
        head_outputs = []
        head_weights = []
        for h in range(self.n_heads):
            out, weights = scaled_dot_product_attention(
                Q_heads[:, h, :], K_heads[:, h, :], V_heads[:, h, :],
                mask=mask)
            head_outputs.append(out)
            head_weights.append(weights)

        # ヘッドの連結
        concat = np.concatenate(head_outputs, axis=-1)  # (n, d_model)

        # 出力射影
        output = concat @ self.W_O
        return output, head_weights

# --- デモ ---
d_model = 16
n_heads = 4
seq_len = 6

# 模擬的な入力（文のトークン埋め込み）
tokens = ["The", "cat", "sat", "on", "the", "mat"]
X = np.random.randn(seq_len, d_model)

# Multi-Head Attention
mha = MultiHeadAttention(d_model, n_heads)
output, head_weights = mha.forward(X)

print(f"入力: {X.shape}")
print(f"出力: {output.shape}")
print(f"ヘッド数: {n_heads}, 各ヘッドの次元: {d_model // n_heads}")

# --- Attention重みの可視化 ---
fig, axes = plt.subplots(1, 4, figsize=(18, 4))

for h in range(n_heads):
    ax = axes[h]
    im = ax.imshow(head_weights[h], cmap="YlOrRd", vmin=0, vmax=0.5)
    ax.set_xticks(range(seq_len))
    ax.set_xticklabels(tokens, fontsize=9, rotation=45)
    ax.set_yticks(range(seq_len))
    ax.set_yticklabels(tokens, fontsize=9)
    ax.set_title(f"Head {h+1}", fontsize=12)
    plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

plt.suptitle("Multi-Head Self-Attention Weights", fontsize=14, y=1.05)
plt.tight_layout()
plt.savefig("multi_head_attention.png", dpi=150, bbox_inches="tight")
plt.show()

この可視化から、Multi-Head Attentionの動作が確認できます。

4つのヘッドが異なるパターンを示している: ランダム初期化の段階でも、各ヘッドのAttention重みの分布は異なっています。学習後には、各ヘッドが構文的・意味的に異なるパターンを捉えるように特化します
行ごとの重みの合計は1: 各行はSoftmaxで正規化されているため、各トークンが他のトークンに割り当てる注目の合計は1になります。これは「注目の予算」を全トークンに配分するイメージです
出力の次元は入力と同じ: Multi-Head Attentionの出力は入力と同じ $(n, d_\text{model})$ の形状を持ちます。これにより、残差接続で入力と出力を直接加算できます

Causal Masked Attentionの実装

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

def softmax(x, axis=-1):
    e = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e / np.sum(e, axis=axis, keepdims=True)

# --- Causal Mask の効果 ---
seq_len = 6
d_k = 8
tokens = ["I", "love", "machine", "learning", "very", "much"]

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

# Causal Mask: 上三角を-inf
causal_mask = np.triu(np.ones((seq_len, seq_len)) * (-1e9), k=1)

# Attentionの計算
scores = Q @ K.T / np.sqrt(d_k)
scores_masked = scores + causal_mask

attn_no_mask = softmax(scores, axis=-1)
attn_masked = softmax(scores_masked, axis=-1)

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

# マスクなし
ax = axes[0]
im = ax.imshow(attn_no_mask, cmap="YlOrRd", vmin=0, vmax=0.5)
ax.set_xticks(range(seq_len))
ax.set_xticklabels(tokens, fontsize=9, rotation=45)
ax.set_yticks(range(seq_len))
ax.set_yticklabels(tokens, fontsize=9)
ax.set_title("Without Mask (Bidirectional)", fontsize=13)
plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

# マスクあり
ax = axes[1]
im = ax.imshow(attn_masked, cmap="YlOrRd", vmin=0, vmax=0.5)
ax.set_xticks(range(seq_len))
ax.set_xticklabels(tokens, fontsize=9, rotation=45)
ax.set_yticks(range(seq_len))
ax.set_yticklabels(tokens, fontsize=9)
ax.set_title("With Causal Mask (Autoregressive)", fontsize=13)
plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

plt.tight_layout()
plt.savefig("causal_attention.png", dpi=150, bbox_inches="tight")
plt.show()

この比較から、Causal Maskの効果が明確に確認できます。

マスクなし（左図）: 各トークンが全てのトークン（自分自身と過去・未来の両方）に注目できます。BERTのようなエンコーダモデルで使用されます
マスクあり（右図）: 上三角部分（未来のトークン）の重みが0になっています。各トークンは自分自身と過去のトークンのみに注目できます。GPTのようなデコーダモデル（自己回帰モデル）で使用されます。これにより、学習時に並列計算しながらも、推論時の自己回帰生成と整合性を保てます

まとめ

本記事では、Self-Attentionの理論をQuery・Key・Valueの枠組みで導出し、Multi-Head Attentionまで実装しました。

Self-Attention: 系列内部で全ペアの関連度を計算し、文脈を考慮した表現を生成する
Q, K, V の射影: 同じ入力から異なる役割の表現を生成。「何を検索するか」「何で検索されるか」「何を取り出すか」
Scaled Dot-Product: $\sqrt{d_k}$ でスケーリングしてSoftmaxの飽和を防ぐ
Multi-Head Attention: 複数の部分空間で独立にAttentionを計算し、異なるパターンを同時に捉える
Causal Mask: 自己回帰モデルで未来の情報を遮断。上三角を $-\infty$ に設定

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

位置エンコーディングの理論 — Self-Attentionに位置情報を与える方法
BERTとGPTの違い — エンコーダ vs デコーダのAttention
Attention機構の基礎 — Attentionの原点

Self-Attentionの理論と実装 — Query・Key・Valueの計算

前提知識

Self-Attentionとは

基本的なアイデア

なぜ線形射影が必要なのか

Scaled Dot-Product Attention

定式化

計算量の分析

Multi-Head Attention

なぜ複数のヘッドが必要なのか

パラメータ数

各ヘッドの役割

Masked Attention（因果的Attention）

自己回帰モデルでの必要性

Pythonでの実装

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

Causal Masked Attentionの実装

まとめ

方形導波管のTE/TMモードとカットオフ周波数を理解して実装する

マイクロストリップ線路の特性インピーダンスと実効誘電率の理論と導出と実装

Self-Attentionの理論と実装 — Query・Key・Valueの計算

前提知識

Self-Attentionとは

基本的なアイデア

なぜ線形射影が必要なのか

Scaled Dot-Product Attention

定式化

計算量の分析

Multi-Head Attention

なぜ複数のヘッドが必要なのか

パラメータ数

各ヘッドの役割

Masked Attention（因果的Attention）

自己回帰モデルでの必要性

Pythonでの実装

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

Causal Masked Attentionの実装

まとめ

関連記事

Multi-head Latent Attention（MLA）とは？KVキャッシュを圧縮するDeepSeekの注意機構を完全解説

Context Window拡張技術 — 学習済みモデルのコンテキスト長を後から伸ばす手法を徹底解説

BERTology — Probingタスクで事前学習モデルの内部知識を探る

方形導波管のTE/TMモードとカットオフ周波数を理解して実装する

マイクロストリップ線路の特性インピーダンスと実効誘電率の理論と導出と実装