Transformerをゼロから実装する — PyTorchで全コンポーネントを構築

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

def scaled_dot_product_attention(Q, K, V, mask=None):
    """Scaled Dot-Product Attention
    Args:
        Q: (batch, heads, seq_len, d_k)
        K: (batch, heads, seq_len, d_k)
        V: (batch, heads, seq_len, d_v)
        mask: (batch, 1, 1, seq_len) or (batch, 1, seq_len, seq_len)
    Returns:
        output: (batch, heads, seq_len, d_v)
        attn_weights: (batch, heads, seq_len, seq_len)
    """
    d_k = Q.size(-1)
    # スコア行列: (batch, heads, seq_len, seq_len)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

    # マスク適用（未来のトークンやパディングを無視）
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # softmaxで重みを正規化
    attn_weights = F.softmax(scores, dim=-1)
    # 重み付き和
    output = torch.matmul(attn_weights, V)
    return output, attn_weights

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

# テスト: バッチ2、ヘッド1、系列長4、次元8
batch, heads, seq_len, d_k = 2, 1, 4, 8
Q = torch.randn(batch, heads, seq_len, d_k)
K = torch.randn(batch, heads, seq_len, d_k)
V = torch.randn(batch, heads, seq_len, d_k)

# マスクなし
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
attn_weights = F.softmax(scores, dim=-1)
print("Attention weights shape:", attn_weights.shape)
print("各行の合計（1になるはず）:", attn_weights[0, 0].sum(dim=-1))

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0, "d_modelはnum_headsで割り切れる必要があります"
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # Q, K, V の線形射影（1つの大きな行列で一括計算）
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        # 出力の線形射影
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)

        # 線形射影: (batch, seq_len, d_model)
        Q = self.W_q(Q)
        K = self.W_k(K)
        V = self.W_v(V)

        # ヘッドに分割: (batch, seq_len, d_model)
        #   → (batch, seq_len, num_heads, d_k)
        #   → (batch, num_heads, seq_len, d_k)
        Q = Q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # Scaled Dot-Product Attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn_weights = F.softmax(scores, dim=-1)
        context = torch.matmul(attn_weights, V)

        # ヘッドを結合: (batch, num_heads, seq_len, d_k)
        #   → (batch, seq_len, num_heads, d_k)
        #   → (batch, seq_len, d_model)
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, -1, self.d_model)

        # 出力射影
        output = self.W_o(context)
        return output

import torch
import torch.nn as nn

# Multi-Head Attentionのテスト
d_model, num_heads = 512, 8
mha = MultiHeadAttention(d_model, num_heads)

# ダミー入力: (batch=2, seq_len=10, d_model=512)
x = torch.randn(2, 10, 512)
output = mha(x, x, x)  # Self-Attentionの場合、Q=K=V=x
print("入力形状:", x.shape)
print("出力形状:", output.shape)
print("形状が一致:", x.shape == output.shape)

import torch
import torch.nn as nn

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)
        self.activation = nn.GELU()

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        x = self.linear1(x)    # → (batch, seq_len, d_ff)
        x = self.activation(x) # GELU活性化
        x = self.dropout(x)
        x = self.linear2(x)    # → (batch, seq_len, d_model)
        return x

import torch

# FFNのテスト
ffn = PositionwiseFeedForward(d_model=512, d_ff=2048)
x = torch.randn(2, 10, 512)
output = ffn(x)
print("入力形状:", x.shape)
print("出力形状:", output.shape)
print("パラメータ数:", sum(p.numel() for p in ffn.parameters()))

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        # 位置エンコーディング行列を事前計算
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        # 分母の計算: 10000^(2i/d_model) = exp(2i * log(10000) / d_model)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数次元
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数次元
        pe = pe.unsqueeze(0)  # (1, max_len, d_model)

        # 学習対象ではないのでbufferとして登録
        self.register_buffer('pe', pe)

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)

import torch
import math
import matplotlib.pyplot as plt
import numpy as np

# 位置エンコーディングの計算
d_model = 128
max_len = 100
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(
    torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
)
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)

# ヒートマップ
plt.figure(figsize=(12, 5))
plt.imshow(pe.numpy(), cmap='RdBu', aspect='auto', vmin=-1, vmax=1)
plt.colorbar(label='Value')
plt.xlabel('Dimension')
plt.ylabel('Position')
plt.title('Positional Encoding (sin/cos)')
plt.tight_layout()
plt.savefig('positional_encoding_heatmap.png', dpi=150)
plt.show()

import torch
import math
import matplotlib.pyplot as plt

d_model = 128
max_len = 100
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(
    torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
)
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)

plt.figure(figsize=(12, 4))
for dim in [0, 10, 30, 60]:
    plt.plot(pe[:, dim].numpy(), label=f'dim {dim}')
plt.xlabel('Position')
plt.ylabel('Encoding Value')
plt.title('Positional Encoding for Different Dimensions')
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('positional_encoding_dims.png', dpi=150)
plt.show()

import torch
import torch.nn as nn

class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = PositionwiseFeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, src_mask=None):
        # Multi-Head Self-Attention + Add & Norm
        attn_output = self.self_attn(x, x, x, src_mask)
        x = self.norm1(x + self.dropout1(attn_output))

        # Feed-Forward Network + Add & Norm
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout2(ffn_output))
        return x

import torch
import torch.nn as nn

class Encoder(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, num_layers, dropout=0.1):
        super().__init__()
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x, src_mask=None):
        for layer in self.layers:
            x = layer(x, src_mask)
        return self.norm(x)

import torch

# Encoderのテスト
encoder = Encoder(d_model=512, num_heads=8, d_ff=2048, num_layers=6)
x = torch.randn(2, 20, 512)  # (batch=2, seq_len=20, d_model=512)
output = encoder(x)
print("Encoder入力形状:", x.shape)
print("Encoder出力形状:", output.shape)
total_params = sum(p.numel() for p in encoder.parameters())
print(f"Encoderパラメータ数: {total_params:,}")

import torch

def make_causal_mask(size):
    """因果マスク: 未来のトークンを参照できないようにする
    Args:
        size: 系列長
    Returns:
        mask: (1, 1, size, size) — 参照可能な位置が1、マスク位置が0
    """
    mask = torch.tril(torch.ones(size, size)).unsqueeze(0).unsqueeze(0)
    return mask  # (1, 1, size, size)

# 確認: 系列長5の因果マスク
print(make_causal_mask(5).squeeze())

import torch
import torch.nn as nn

class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        # 3つのサブレイヤー
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.cross_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = PositionwiseFeedForward(d_model, d_ff, dropout)

        # 3つのLayerNorm
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)

        # 3つのDropout
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, x, enc_output, src_mask=None, tgt_mask=None):
        # サブレイヤー1: Masked Self-Attention
        attn_output = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout1(attn_output))

        # サブレイヤー2: Cross-Attention
        # Query=x (Decoderの中間表現), Key=Value=enc_output
        attn_output = self.cross_attn(x, enc_output, enc_output, src_mask)
        x = self.norm2(x + self.dropout2(attn_output))

        # サブレイヤー3: Feed-Forward Network
        ffn_output = self.ffn(x)
        x = self.norm3(x + self.dropout3(ffn_output))
        return x

import torch
import torch.nn as nn

class Decoder(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, num_layers, dropout=0.1):
        super().__init__()
        self.layers = nn.ModuleList([
            DecoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x, enc_output, src_mask=None, tgt_mask=None):
        for layer in self.layers:
            x = layer(x, enc_output, src_mask, tgt_mask)
        return self.norm(x)

import torch

# Decoderのテスト
decoder = Decoder(d_model=512, num_heads=8, d_ff=2048, num_layers=6)
# Decoder入力とEncoder出力
tgt = torch.randn(2, 15, 512)   # 出力系列（長さ15）
enc_out = torch.randn(2, 20, 512)  # Encoder出力（長さ20）
# 因果マスク
tgt_mask = make_causal_mask(15)

output = decoder(tgt, enc_out, tgt_mask=tgt_mask)
print("Decoder入力形状:", tgt.shape)
print("Encoder出力形状:", enc_out.shape)
print("Decoder出力形状:", output.shape)

import torch
import torch.nn as nn
import math

class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512,
                 num_heads=8, d_ff=2048, num_layers=6,
                 max_len=5000, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        # 埋め込み層
        self.src_embedding = nn.Embedding(src_vocab_size, d_model)
        self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model)

        # 位置エンコーディング
        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)

        # Encoder と Decoder
        self.encoder = Encoder(d_model, num_heads, d_ff, num_layers, dropout)
        self.decoder = Decoder(d_model, num_heads, d_ff, num_layers, dropout)

        # 最終出力層（語彙サイズへの射影）
        self.generator = nn.Linear(d_model, tgt_vocab_size)

        # パラメータの初期化
        self._init_parameters()

    def _init_parameters(self):
        """Xavier初期化で学習を安定させる"""
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def make_src_mask(self, src, pad_idx=0):
        """パディングマスク: パディングトークンを無視する"""
        # src: (batch, src_len)
        src_mask = (src != pad_idx).unsqueeze(1).unsqueeze(2)
        return src_mask  # (batch, 1, 1, src_len)

    def make_tgt_mask(self, tgt, pad_idx=0):
        """パディングマスク + 因果マスクの組み合わせ"""
        batch_size, tgt_len = tgt.size()
        # パディングマスク
        pad_mask = (tgt != pad_idx).unsqueeze(1).unsqueeze(2)
        # 因果マスク
        causal_mask = torch.tril(
            torch.ones(tgt_len, tgt_len, device=tgt.device)
        ).unsqueeze(0).unsqueeze(0)
        # 両方のマスクを結合（論理AND）
        tgt_mask = pad_mask & (causal_mask.bool())
        return tgt_mask  # (batch, 1, tgt_len, tgt_len)

    def forward(self, src, tgt, pad_idx=0):
        # マスク生成
        src_mask = self.make_src_mask(src, pad_idx)
        tgt_mask = self.make_tgt_mask(tgt, pad_idx)

        # Encoder側: 埋め込み → 位置エンコーディング → Encoder
        src_emb = self.src_embedding(src) * math.sqrt(self.d_model)
        src_emb = self.pos_encoding(src_emb)
        enc_output = self.encoder(src_emb, src_mask)

        # Decoder側: 埋め込み → 位置エンコーディング → Decoder
        tgt_emb = self.tgt_embedding(tgt) * math.sqrt(self.d_model)
        tgt_emb = self.pos_encoding(tgt_emb)
        dec_output = self.decoder(tgt_emb, enc_output, src_mask, tgt_mask)

        # 最終出力: 語彙サイズへの射影
        output = self.generator(dec_output)
        return output  # (batch, tgt_len, tgt_vocab_size)

import torch

# Transformer全体のテスト
model = Transformer(
    src_vocab_size=1000,
    tgt_vocab_size=1000,
    d_model=512,
    num_heads=8,
    d_ff=2048,
    num_layers=6
)

# ダミー入力
src = torch.randint(1, 1000, (2, 20))  # (batch=2, src_len=20)
tgt = torch.randint(1, 1000, (2, 15))  # (batch=2, tgt_len=15)

output = model(src, tgt)
print("入力(src)形状:", src.shape)
print("入力(tgt)形状:", tgt.shape)
print("出力形状:", output.shape)
total_params = sum(p.numel() for p in model.parameters())
print(f"総パラメータ数: {total_params:,}")

import torch
import torch.nn as nn
import torch.optim as optim
import random
import numpy as np

# 再現性のためのシード固定
torch.manual_seed(42)
random.seed(42)
np.random.seed(42)

# 特殊トークン
PAD_IDX = 0
BOS_IDX = 1
EOS_IDX = 2
VOCAB_SIZE = 20   # 0-19のトークン（0,1,2は特殊トークン）
SEQ_LEN = 10      # コピーする系列の長さ

def generate_copy_data(batch_size, seq_len, vocab_size):
    """コピータスクのデータを生成する"""
    # データトークン: 3 ~ vocab_size-1 の範囲
    data = torch.randint(3, vocab_size, (batch_size, seq_len))

    # Encoder入力: [data tokens]
    src = data

    # Decoder入力: [BOS, data tokens]（右シフト）
    tgt_input = torch.cat([
        torch.full((batch_size, 1), BOS_IDX, dtype=torch.long),
        data
    ], dim=1)

    # 正解ラベル: [data tokens, EOS]
    tgt_output = torch.cat([
        data,
        torch.full((batch_size, 1), EOS_IDX, dtype=torch.long)
    ], dim=1)

    return src, tgt_input, tgt_output

import torch
import torch.nn as nn
import torch.optim as optim

# ハイパーパラメータ（コピータスク用の小さな設定）
D_MODEL = 64
NUM_HEADS = 4
D_FF = 128
NUM_LAYERS = 2
DROPOUT = 0.1
LR = 1e-3
NUM_EPOCHS = 50
BATCH_SIZE = 64

# モデル、損失関数、最適化器
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Transformer(
    src_vocab_size=VOCAB_SIZE,
    tgt_vocab_size=VOCAB_SIZE,
    d_model=D_MODEL,
    num_heads=NUM_HEADS,
    d_ff=D_FF,
    num_layers=NUM_LAYERS,
    dropout=DROPOUT
).to(device)

criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)
optimizer = optim.Adam(model.parameters(), lr=LR, betas=(0.9, 0.98), eps=1e-9)

print(f"デバイス: {device}")
total_params = sum(p.numel() for p in model.parameters())
print(f"パラメータ数: {total_params:,}")

import torch

# 学習ループ
losses = []
for epoch in range(NUM_EPOCHS):
    model.train()
    # バッチごとに新しいデータを生成
    src, tgt_input, tgt_output = generate_copy_data(BATCH_SIZE, SEQ_LEN, VOCAB_SIZE)
    src = src.to(device)
    tgt_input = tgt_input.to(device)
    tgt_output = tgt_output.to(device)

    # 順伝播
    output = model(src, tgt_input, pad_idx=PAD_IDX)
    # output: (batch, tgt_len, vocab_size)
    # CrossEntropyLossは (N, C) と (N,) を期待するのでreshape
    output = output.view(-1, VOCAB_SIZE)
    tgt_output = tgt_output.view(-1)

    # 損失計算と逆伝播
    loss = criterion(output, tgt_output)
    optimizer.zero_grad()
    loss.backward()
    # 勾配クリッピング（勾配爆発を防ぐ）
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()

    losses.append(loss.item())
    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1}/{NUM_EPOCHS}, Loss: {loss.item():.4f}")

import matplotlib.pyplot as plt

plt.figure(figsize=(10, 5))
plt.plot(losses, linewidth=2)
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Copy Task Training Loss')
plt.grid(True, alpha=0.3)
plt.yscale('log')
plt.tight_layout()
plt.savefig('copy_task_loss.png', dpi=150)
plt.show()

import torch

def greedy_decode(model, src, max_len, bos_idx, eos_idx, pad_idx, device):
    """貪欲法による自己回帰デコード"""
    model.eval()
    src = src.to(device)

    # Decoder入力を<BOS>で初期化
    tgt = torch.full((src.size(0), 1), bos_idx, dtype=torch.long, device=device)

    with torch.no_grad():
        for _ in range(max_len):
            output = model(src, tgt, pad_idx=pad_idx)
            # 最後の位置の予測を取得
            next_token = output[:, -1, :].argmax(dim=-1, keepdim=True)
            tgt = torch.cat([tgt, next_token], dim=1)
            # 全バッチでEOSが出たら終了
            if (next_token == eos_idx).all():
                break
    return tgt

import torch

# 推論テスト
src_test, _, _ = generate_copy_data(5, SEQ_LEN, VOCAB_SIZE)
print("入力系列:")
print(src_test.numpy())

predicted = greedy_decode(
    model, src_test, max_len=SEQ_LEN + 2,
    bos_idx=BOS_IDX, eos_idx=EOS_IDX, pad_idx=PAD_IDX, device=device
)
# <BOS>を除いた予測結果
pred_tokens = predicted[:, 1:SEQ_LEN+1].cpu().numpy()
print("\n予測系列:")
print(pred_tokens)

# 一致率の確認
correct = (src_test.numpy() == pred_tokens).mean()
print(f"\nトークン一致率: {correct * 100:.1f}%")

import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import math

# Attention重みを取得するための順伝播
model.eval()
src_vis = src_test[:1].to(device)
tgt_vis = torch.cat([
    torch.full((1, 1), BOS_IDX, dtype=torch.long),
    src_test[:1]
], dim=1).to(device)

with torch.no_grad():
    # Encoder
    src_mask = model.make_src_mask(src_vis, PAD_IDX)
    src_emb = model.src_embedding(src_vis) * math.sqrt(model.d_model)
    src_emb = model.pos_encoding(src_emb)
    enc_out = model.encoder(src_emb, src_mask)

    # Decoder（最初の層のCross-Attentionの重みを取得）
    tgt_mask = model.make_tgt_mask(tgt_vis, PAD_IDX)
    tgt_emb = model.tgt_embedding(tgt_vis) * math.sqrt(model.d_model)
    tgt_emb = model.pos_encoding(tgt_emb)

    # Decoder第1層のCross-Attentionの重みを手動計算
    dec_layer = model.decoder.layers[0]
    # Masked Self-Attention
    z1 = dec_layer.self_attn(tgt_emb, tgt_emb, tgt_emb, tgt_mask)
    z1 = dec_layer.norm1(tgt_emb + z1)
    # Cross-Attentionのスコアを直接計算
    Q = dec_layer.cross_attn.W_q(z1)
    K = dec_layer.cross_attn.W_k(enc_out)
    d_k = dec_layer.cross_attn.d_k
    num_heads = dec_layer.cross_attn.num_heads
    bs = Q.size(0)
    Q = Q.view(bs, -1, num_heads, d_k).transpose(1, 2)
    K = K.view(bs, -1, num_heads, d_k).transpose(1, 2)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    attn_w = F.softmax(scores, dim=-1)

import matplotlib.pyplot as plt

# 4つのヘッドのAttention重みを可視化
fig, axes = plt.subplots(1, 4, figsize=(16, 4))
for h in range(4):
    ax = axes[h]
    im = ax.imshow(
        attn_w[0, h].cpu().numpy(), cmap='Blues', aspect='auto'
    )
    ax.set_title(f'Head {h}')
    ax.set_xlabel('Source Position')
    ax.set_ylabel('Target Position')
fig.suptitle('Cross-Attention Weights (Decoder Layer 0)', fontsize=14)
plt.tight_layout()
plt.savefig('cross_attention_weights.png', dpi=150)
plt.show()