第11章: PyTorchで作るAttention — 50行で動くSelf-Attention

# PyTorchをインストール（CPU版でOK）
pip install torch

# Jupyter Notebook で実行するなら
pip install jupyter
jupyter notebook

import math
import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q, K, V: (B, H, T, d_k)  ※ Bはバッチ、Hはヘッド数、Tは系列長、d_kは次元
    mask   : (B, 1, T, T)、True が「無視したい位置」
    return : 出力 (B, H, T, d_k), 注意重み (B, H, T, T)
    """
    d_k = Q.size(-1)
    # Step 1: QK^T 内積を計算 → (B, H, T, T)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    # Step 2: マスクがあれば -inf で埋める
    if mask is not None:
        scores = scores.masked_fill(mask, float("-inf"))
    # Step 3: softmaxで重みに変換
    attn = F.softmax(scores, dim=-1)
    # Step 4: 重みでVを加重平均
    out = torch.matmul(attn, V)
    return out, attn

import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0):
        super().__init__()
        assert d_model % n_heads == 0, "d_model は n_heads で割り切れること"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        # Q, K, V を作る重み行列（bias なし）
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        # 最後の出力射影
        self.W_o = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    def _split_heads(self, x):
        """(B, T, d_model) -> (B, H, T, d_k)"""
        B, T, _ = x.shape
        # view で分割し、transpose でヘッド次元を前に
        return x.view(B, T, self.n_heads, self.d_k).transpose(1, 2)

    def _merge_heads(self, x):
        """(B, H, T, d_k) -> (B, T, d_model)"""
        B, H, T, d_k = x.shape
        return x.transpose(1, 2).contiguous().view(B, T, H * d_k)

    def forward(self, x_q, x_k, x_v, mask=None):
        """
        x_q, x_k, x_v: (B, T, d_model)  自己注意なら全部同じテンソル
        mask: (B, 1, T_q, T_k) ブロードキャスト可能、True = 無視
        """
        # 1. 射影してヘッド分割
        Q = self._split_heads(self.W_q(x_q))  # (B, H, T, d_k)
        K = self._split_heads(self.W_k(x_k))
        V = self._split_heads(self.W_v(x_v))

        # 2. Scaled Dot-Product Attention（さきほどの関数を使う）
        out, attn = scaled_dot_product_attention(Q, K, V, mask=mask)
        out = self.dropout(out)

        # 3. ヘッドを連結して出力射影
        out = self._merge_heads(out)  # (B, T, d_model)
        return self.W_o(out), attn

torch.manual_seed(0)
B, T, d_model, H = 2, 6, 32, 4  # バッチ2、系列長6、d_model=32、4ヘッド
x = torch.randn(B, T, d_model)

# 因果マスク作成: 上三角を True にする (GPT流)
causal = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
causal = causal.view(1, 1, T, T)  # ブロードキャスト用

mha = MultiHeadAttention(d_model=d_model, n_heads=H)
out, attn = mha(x, x, x, mask=causal)

print("out shape :", out.shape)         # torch.Size([2, 6, 32])
print("attn shape:", attn.shape)        # torch.Size([2, 4, 6, 6])
print("行和=1 ?  :", attn.sum(-1)[0, 0])  # 各行ほぼ 1.0
print("未来は0 ? :", attn[0, 0, 0, 1:].sum().item())  # t=0 から未来は 0

import torch.nn.functional as F

# 公式の Scaled Dot-Product Attention（is_causal=True で因果マスク内蔵）
out_official = F.scaled_dot_product_attention(Q, K, V, is_causal=True)

# 自作版との一致確認
out_mine, _ = scaled_dot_product_attention(Q, K, V, mask=causal)
print("差:", (out_official - out_mine).abs().max().item())  # ~1e-6 程度

# pip install transformers

from transformers import pipeline

# BERT で穴埋め
fill = pipeline("fill-mask", model="bert-base-uncased")
print(fill("Paris is the [MASK] of France."))
# → [{'sequence': 'paris is the capital of france.', ...}]

# GPT-2 で文章生成
gen = pipeline("text-generation", model="gpt2")
print(gen("Once upon a time"))