import torch import torch.nn as nn import torch.nn.functional as F import math # torch: PyTorchの基本機能 # torch.nn: ニューラルネットワークのモジュール # torch.nn.functional: Softmaxなどの関数 # math: 平方根の計算（√d_k）に使用

def scaled_dot_product_attention(Q, K, V, mask=None): “”” Scaled Dot-Product Attentionの計算 Args: Q: Query (batch, seq_len, d_k) K: Key (batch, seq_len, d_k) V: Value (batch, seq_len, d_v) mask: マスク（オプション） Returns: output: 出力 (batch, seq_len, d_v) attention_weights: 注目度 (batch, seq_len, seq_len) “”” # d_kを取得（スケーリング用） d_k = Q.size(-1) # ステップ1: QK^T を計算 # Q: (batch, seq_len, d_k) # K.transpose(-2, -1): (batch, d_k, seq_len) # scores: (batch, seq_len, seq_len) scores = torch.matmul(Q, K.transpose(-2, -1)) # ステップ2: √d_k でスケーリング scores = scores / math.sqrt(d_k) # マスク適用（Decoderで使用、未来の単語を見えなくする） if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) # ステップ3: Softmaxで正規化 # dim=-1: 最後の次元（各行）に対してSoftmax attention_weights = F.softmax(scores, dim=-1) # ステップ4: Valueとの重み付け和 # attention_weights: (batch, seq_len, seq_len) # V: (batch, seq_len, d_v) # output: (batch, seq_len, d_v) output = torch.matmul(attention_weights, V) return output, attention_weights

# テスト batch_size = 2 seq_len = 4 d_k = 8 # ダミーデータを作成 Q = torch.randn(batch_size, seq_len, d_k) K = torch.randn(batch_size, seq_len, d_k) V = torch.randn(batch_size, seq_len, d_k) # Attention計算 output, weights = scaled_dot_product_attention(Q, K, V) print(f”Output shape: {output.shape}”) print(f”Attention weights shape: {weights.shape}”) print(f”\n1文目のAttention Weights:”) print(weights[0].detach().numpy().round(3)) print(f”\n各行の合計: {weights[0].sum(dim=-1).detach().numpy()}”)

class MultiHeadAttention(nn.Module): “””Multi-Head Attention””” def __init__(self, d_model, n_heads, dropout=0.1): “”” Args: d_model: 入力の次元 (例: 512) n_heads: ヘッド数 (例: 8) dropout: ドロップアウト率 “”” super(MultiHeadAttention, self).__init__() # d_modelがn_headsで割り切れることを確認 assert d_model % n_heads == 0, “d_model must be divisible by n_heads” self.d_model = d_model # 512 self.n_heads = n_heads # 8 self.d_k = d_model // n_heads # 64 # Q, K, V用の線形変換（全ヘッド分まとめて） self.W_Q = nn.Linear(d_model, d_model) # (512, 512) 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) self.dropout = nn.Dropout(dropout) def forward(self, Q, K, V, mask=None): batch_size = Q.size(0) # 線形変換 Q = self.W_Q(Q) # (batch, seq_len, d_model) K = self.W_K(K) V = self.W_V(V) # ヘッドに分割 # (batch, seq_len, d_model) → (batch, n_heads, seq_len, d_k) Q = Q.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) K = K.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) V = V.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) # Attention計算（全ヘッド同時） d_k = Q.size(-1) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) attention_weights = F.softmax(scores, dim=-1) attention_weights = self.dropout(attention_weights) output = torch.matmul(attention_weights, V) # output: (batch, n_heads, seq_len, d_k) # ヘッドを結合 # (batch, n_heads, seq_len, d_k) → (batch, seq_len, d_model) output = output.transpose(1, 2).contiguous() output = output.view(batch_size, -1, self.d_model) # 出力の線形変換 output = self.W_O(output) return output, attention_weights

# Multi-Head Attentionのテスト d_model = 512 n_heads = 8 batch_size = 2 seq_len = 10 # モデル作成 mha = MultiHeadAttention(d_model, n_heads) # 入力データ X = torch.randn(batch_size, seq_len, d_model) # Self-Attention（Q = K = V = X） output, weights = mha(X, X, X) print(f”Input shape: {X.shape}”) print(f”Output shape: {output.shape}”) print(f”Attention weights shape: {weights.shape}”) print(f”\nHead 0のAttention Weights (最初の5×5):”) print(weights[0, 0, :5, :5].detach().numpy().round(3))