Pythonで理解するSelf-Attentionの仕組み：数式からスクラッチ実装まで

import numpy as np
import matplotlib.pyplot as plt

def softmax(x, axis=-1):
    """数値的に安定なsoftmax"""
    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):
    """
    Scaled Dot-Product Attention（式2）

    Parameters:
        Q: Query行列 (n, d_k)
        K: Key行列   (n, d_k)
        V: Value行列  (n, d_v)
    Returns:
        output: Attention出力 (n, d_v)
        weights: Attention重み (n, n)
    """
    d_k = Q.shape[-1]
    # ステップ1-2: 類似度計算とスケーリング
    scores = Q @ K.T / np.sqrt(d_k)
    # ステップ3: Attention重み
    weights = softmax(scores)
    # ステップ4: 重み付き和
    output = weights @ V
    return output, weights

np.random.seed(42)

# 入力: 4トークン、モデル次元8
n_tokens = 4
d_model = 8
d_k = d_v = 8

# ランダムな入力系列
X = np.random.randn(n_tokens, 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 の計算（式1）
Q = X @ W_Q
K = X @ W_K
V = X @ W_V

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

print("入力形状:", X.shape)
print("出力形状:", output.shape)
print("Attention重み:\n", np.round(weights, 3))

# Attention重みのヒートマップ
token_labels = ["Token 0", "Token 1", "Token 2", "Token 3"]
fig, ax = plt.subplots(figsize=(6, 5))
im = ax.imshow(weights, cmap="Blues", vmin=0, vmax=1)
ax.set_xticks(range(n_tokens))
ax.set_yticks(range(n_tokens))
ax.set_xticklabels(token_labels)
ax.set_yticklabels(token_labels)
ax.set_xlabel("Key position")
ax.set_ylabel("Query position")
ax.set_title("Attention Weights")
for i in range(n_tokens):
    for j in range(n_tokens):
        ax.text(j, i, f"{weights[i, j]:.2f}",
                ha="center", va="center", fontsize=11)
plt.colorbar(im, ax=ax)
plt.tight_layout()
plt.show()

class MultiHeadAttention:
    """Multi-Head Attention（式7, 8）"""

    def __init__(self, d_model, n_heads, seed=0):
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        rng = np.random.RandomState(seed)
        scale = 0.1

        # 各ヘッドの射影行列
        self.W_Q = rng.randn(n_heads, d_model, self.d_k) * scale
        self.W_K = rng.randn(n_heads, d_model, self.d_k) * scale
        self.W_V = rng.randn(n_heads, d_model, self.d_k) * scale
        # 出力射影行列
        self.W_O = rng.randn(n_heads * self.d_k, d_model) * scale

    def forward(self, X):
        """
        Parameters:
            X: 入力系列 (n, d_model)
        Returns:
            output: Multi-Head Attention出力 (n, d_model)
            all_weights: 各ヘッドのAttention重み (n_heads, n, n)
        """
        head_outputs = []
        all_weights = []

        for i in range(self.n_heads):
            Q = X @ self.W_Q[i]
            K = X @ self.W_K[i]
            V = X @ self.W_V[i]
            head_out, weights = scaled_dot_product_attention(Q, K, V)
            head_outputs.append(head_out)
            all_weights.append(weights)

        # ヘッドの結合と出力射影（式7）
        concat = np.concatenate(head_outputs, axis=-1)
        output = concat @ self.W_O

        return output, np.array(all_weights)

np.random.seed(42)

n_tokens = 6
d_model = 16
n_heads = 4

X = np.random.randn(n_tokens, d_model)
mha = MultiHeadAttention(d_model, n_heads, seed=42)
output, all_weights = mha.forward(X)

print("入力形状:", X.shape)
print("出力形状:", output.shape)
print("Attention重み形状:", all_weights.shape)

# 各ヘッドのAttention重みを可視化
fig, axes = plt.subplots(1, n_heads, figsize=(16, 4))
for h in range(n_heads):
    im = axes[h].imshow(all_weights[h], cmap="Blues", vmin=0, vmax=1)
    axes[h].set_title(f"Head {h}")
    axes[h].set_xlabel("Key")
    axes[h].set_ylabel("Query")

plt.suptitle("Multi-Head Attention Weights", fontsize=14)
plt.tight_layout()
plt.show()

def positional_encoding(max_len, d_model):
    """
    正弦波位置エンコーディング（式9, 10）

    Parameters:
        max_len: 最大系列長
        d_model: モデル次元
    Returns:
        PE: 位置エンコーディング行列 (max_len, d_model)
    """
    PE = np.zeros((max_len, d_model))
    position = np.arange(max_len)[:, np.newaxis]
    div_term = 10000 ** (2 * np.arange(d_model // 2) / d_model)

    PE[:, 0::2] = np.sin(position / div_term)
    PE[:, 1::2] = np.cos(position / div_term)
    return PE

# 位置エンコーディングの可視化
max_len = 50
d_model = 64
PE = positional_encoding(max_len, d_model)

fig, ax = plt.subplots(figsize=(10, 6))
im = ax.imshow(PE, cmap="RdBu", aspect="auto")
ax.set_xlabel("Dimension")
ax.set_ylabel("Position")
ax.set_title("Sinusoidal Positional Encoding")
plt.colorbar(im, ax=ax)
plt.tight_layout()
plt.show()

import torch
import torch.nn as nn

np.random.seed(42)
torch.manual_seed(42)

n_tokens = 4
d_model = 8
n_heads = 2
d_k = d_model // n_heads

# 入力データ
X_np = np.random.randn(n_tokens, d_model).astype(np.float32)

# --- スクラッチ実装 ---
mha_np = MultiHeadAttention(d_model, n_heads, seed=0)
out_np, _ = mha_np.forward(X_np)

# --- PyTorch実装 ---
mha_pt = nn.MultiheadAttention(d_model, n_heads, bias=False, batch_first=False)

# PyTorchのin_proj_weightにスクラッチ実装の重みをコピー
# PyTorchは[W_Q; W_K; W_V]を結合した形で保持 (3*d_model, d_model)
W_Q_cat = np.concatenate([mha_np.W_Q[i] for i in range(n_heads)], axis=1).T
W_K_cat = np.concatenate([mha_np.W_K[i] for i in range(n_heads)], axis=1).T
W_V_cat = np.concatenate([mha_np.W_V[i] for i in range(n_heads)], axis=1).T
in_proj_weight = np.concatenate([W_Q_cat, W_K_cat, W_V_cat], axis=0)

with torch.no_grad():
    mha_pt.in_proj_weight.copy_(torch.from_numpy(in_proj_weight))
    mha_pt.out_proj.weight.copy_(torch.from_numpy(mha_np.W_O.T))

# PyTorchはデフォルトで (seq_len, batch, d_model) の形式
X_pt = torch.from_numpy(X_np).unsqueeze(1)  # (n_tokens, 1, d_model)
out_pt, _ = mha_pt(X_pt, X_pt, X_pt)
out_pt = out_pt.squeeze(1).detach().numpy()

# 比較
print("スクラッチ実装の出力:\n", np.round(out_np, 4))
print("PyTorchの出力:\n", np.round(out_pt, 4))
print("最大誤差:", np.max(np.abs(out_np - out_pt)))