Bahdanau AttentionとLuong Attentionの違いを理論から実装まで完全解説

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

class BahdanauAttention(nn.Module):
    """Bahdanau Attention（加法的注意機構）

    スコア関数: e_ij = v^T tanh(W1 h_j + W2 s_{i-1})

    Args:
        encoder_dim: エンコーダ隠れ状態の次元
        decoder_dim: デコーダ隠れ状態の次元
        attention_dim: アライメント空間の次元
    """
    def __init__(self, encoder_dim, decoder_dim, attention_dim):
        super().__init__()
        # W1: エンコーダ隠れ状態をアライメント空間に射影
        self.W1 = nn.Linear(encoder_dim, attention_dim, bias=False)
        # W2: デコーダ隠れ状態をアライメント空間に射影
        self.W2 = nn.Linear(decoder_dim, attention_dim, bias=False)
        # v: アライメント空間からスカラーへ
        self.v = nn.Linear(attention_dim, 1, bias=False)

    def forward(self, encoder_outputs, decoder_hidden):
        """
        Args:
            encoder_outputs: (batch, src_len, encoder_dim) エンコーダの全隠れ状態
            decoder_hidden: (batch, decoder_dim) デコーダの前時刻の隠れ状態

        Returns:
            context: (batch, encoder_dim) コンテキストベクトル
            attn_weights: (batch, src_len) Attention重み
        """
        # encoder_outputs を射影: (batch, src_len, attention_dim)
        encoder_proj = self.W1(encoder_outputs)

        # decoder_hidden を射影: (batch, attention_dim)
        # unsqueeze で (batch, 1, attention_dim) にしてブロードキャスト
        decoder_proj = self.W2(decoder_hidden).unsqueeze(1)

        # 加法的スコアの計算: tanh(W1*h_j + W2*s_{i-1})
        # (batch, src_len, attention_dim)
        combined = torch.tanh(encoder_proj + decoder_proj)

        # v^T で各位置のスコアをスカラーに: (batch, src_len, 1) → (batch, src_len)
        scores = self.v(combined).squeeze(-1)

        # Softmaxで正規化して確率分布に
        attn_weights = F.softmax(scores, dim=-1)

        # コンテキストベクトル: 重み付き和
        # (batch, 1, src_len) × (batch, src_len, encoder_dim) → (batch, 1, encoder_dim)
        context = torch.bmm(attn_weights.unsqueeze(1), encoder_outputs).squeeze(1)

        return context, attn_weights

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

class LuongAttention(nn.Module):
    """Luong Attention（乗法的注意機構）

    3種のスコア関数:
        - dot:     score = h_t^T * h_s
        - general: score = h_t^T * W_a * h_s
        - concat:  score = v_a^T * tanh(W_a * [h_t; h_s])

    Args:
        encoder_dim: エンコーダ隠れ状態の次元
        decoder_dim: デコーダ隠れ状態の次元
        score_fn: スコア関数の種類 ("dot", "general", "concat")
    """
    def __init__(self, encoder_dim, decoder_dim, score_fn="dot"):
        super().__init__()
        self.score_fn = score_fn
        self.encoder_dim = encoder_dim
        self.decoder_dim = decoder_dim

        if score_fn == "general":
            # W_a: デコーダ次元 → エンコーダ次元 への線形変換
            self.W_a = nn.Linear(encoder_dim, decoder_dim, bias=False)
        elif score_fn == "concat":
            # 連結ベクトルを変換する行列とスカラー化するベクトル
            self.W_a = nn.Linear(encoder_dim + decoder_dim, decoder_dim, bias=False)
            self.v_a = nn.Linear(decoder_dim, 1, bias=False)
        # dot の場合は追加パラメータなし

    def _compute_score(self, decoder_hidden, encoder_outputs):
        """スコア関数を計算する

        Args:
            decoder_hidden: (batch, decoder_dim) 現在のデコーダ隠れ状態
            encoder_outputs: (batch, src_len, encoder_dim) エンコーダの全隠れ状態

        Returns:
            scores: (batch, src_len)
        """
        if self.score_fn == "dot":
            # h_t^T * h_s: 内積
            # (batch, decoder_dim) → (batch, decoder_dim, 1)
            # (batch, src_len, encoder_dim) × (batch, encoder_dim, 1) → (batch, src_len, 1)
            scores = torch.bmm(
                encoder_outputs, decoder_hidden.unsqueeze(-1)
            ).squeeze(-1)

        elif self.score_fn == "general":
            # h_t^T * W_a * h_s
            # まず W_a * h_s を計算: (batch, src_len, decoder_dim)
            encoder_proj = self.W_a(encoder_outputs)
            scores = torch.bmm(
                encoder_proj, decoder_hidden.unsqueeze(-1)
            ).squeeze(-1)

        elif self.score_fn == "concat":
            # v_a^T * tanh(W_a * [h_t; h_s])
            src_len = encoder_outputs.size(1)
            # decoder_hidden を src_len 分繰り返す
            decoder_expanded = decoder_hidden.unsqueeze(1).expand(-1, src_len, -1)
            # [h_t; h_s] を連結: (batch, src_len, encoder_dim + decoder_dim)
            concat = torch.cat([decoder_expanded, encoder_outputs], dim=-1)
            scores = self.v_a(torch.tanh(self.W_a(concat))).squeeze(-1)

        return scores

    def forward(self, encoder_outputs, decoder_hidden):
        """
        Args:
            encoder_outputs: (batch, src_len, encoder_dim) エンコーダの全隠れ状態
            decoder_hidden: (batch, decoder_dim) デコーダの現在の隠れ状態

        Returns:
            context: (batch, encoder_dim) コンテキストベクトル
            attn_weights: (batch, src_len) Attention重み
        """
        # スコアを計算
        scores = self._compute_score(decoder_hidden, encoder_outputs)

        # Softmaxで正規化
        attn_weights = F.softmax(scores, dim=-1)

        # コンテキストベクトル: 重み付き和
        context = torch.bmm(attn_weights.unsqueeze(1), encoder_outputs).squeeze(1)

        return context, attn_weights

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

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

# --- BahdanauAttention クラス定義 ---
class BahdanauAttention(nn.Module):
    def __init__(self, encoder_dim, decoder_dim, attention_dim):
        super().__init__()
        self.W1 = nn.Linear(encoder_dim, attention_dim, bias=False)
        self.W2 = nn.Linear(decoder_dim, attention_dim, bias=False)
        self.v = nn.Linear(attention_dim, 1, bias=False)

    def forward(self, encoder_outputs, decoder_hidden):
        encoder_proj = self.W1(encoder_outputs)
        decoder_proj = self.W2(decoder_hidden).unsqueeze(1)
        combined = torch.tanh(encoder_proj + decoder_proj)
        scores = self.v(combined).squeeze(-1)
        attn_weights = F.softmax(scores, dim=-1)
        context = torch.bmm(attn_weights.unsqueeze(1), encoder_outputs).squeeze(1)
        return context, attn_weights

# --- LuongAttention クラス定義 ---
class LuongAttention(nn.Module):
    def __init__(self, encoder_dim, decoder_dim, score_fn="dot"):
        super().__init__()
        self.score_fn = score_fn
        if score_fn == "general":
            self.W_a = nn.Linear(encoder_dim, decoder_dim, bias=False)
        elif score_fn == "concat":
            self.W_a = nn.Linear(encoder_dim + decoder_dim, decoder_dim, bias=False)
            self.v_a = nn.Linear(decoder_dim, 1, bias=False)

    def forward(self, encoder_outputs, decoder_hidden):
        if self.score_fn == "dot":
            scores = torch.bmm(
                encoder_outputs, decoder_hidden.unsqueeze(-1)
            ).squeeze(-1)
        elif self.score_fn == "general":
            encoder_proj = self.W_a(encoder_outputs)
            scores = torch.bmm(
                encoder_proj, decoder_hidden.unsqueeze(-1)
            ).squeeze(-1)
        elif self.score_fn == "concat":
            src_len = encoder_outputs.size(1)
            dec_exp = decoder_hidden.unsqueeze(1).expand(-1, src_len, -1)
            concat = torch.cat([dec_exp, encoder_outputs], dim=-1)
            scores = self.v_a(torch.tanh(self.W_a(concat))).squeeze(-1)

        attn_weights = F.softmax(scores, dim=-1)
        context = torch.bmm(attn_weights.unsqueeze(1), encoder_outputs).squeeze(1)
        return context, attn_weights

# --- 合成データの生成 ---
batch_size = 1
src_len = 6       # 入力系列の長さ（英語の単語数）
tgt_len = 5       # 出力系列の長さ（日本語の単語数）
hidden_dim = 32   # 隠れ状態の次元
attn_dim = 16     # アライメント空間の次元

# エンコーダの全隠れ状態（ランダム生成）
encoder_outputs = torch.randn(batch_size, src_len, hidden_dim)
# デコーダの各ステップの隠れ状態（ランダム生成）
decoder_hiddens = torch.randn(tgt_len, batch_size, hidden_dim)

# --- 各Attentionのインスタンス生成 ---
bahdanau = BahdanauAttention(hidden_dim, hidden_dim, attn_dim)
luong_dot = LuongAttention(hidden_dim, hidden_dim, score_fn="dot")
luong_general = LuongAttention(hidden_dim, hidden_dim, score_fn="general")
luong_concat = LuongAttention(hidden_dim, hidden_dim, score_fn="concat")

# 全モデルを評価モードに
for m in [bahdanau, luong_dot, luong_general, luong_concat]:
    m.eval()

# --- 各ステップの Attention重みを収集 ---
attention_maps = {}
models = {
    "Bahdanau\n(Additive)": bahdanau,
    "Luong\n(Dot)": luong_dot,
    "Luong\n(General)": luong_general,
    "Luong\n(Concat)": luong_concat,
}

with torch.no_grad():
    for name, model in models.items():
        weights_list = []
        for t in range(tgt_len):
            _, attn_w = model(encoder_outputs, decoder_hiddens[t])
            weights_list.append(attn_w.squeeze(0).numpy())
        attention_maps[name] = np.stack(weights_list)  # (tgt_len, src_len)

# --- ヒートマップの可視化 ---
src_tokens = ["The", "cat", "sat", "on", "the", "mat"]
tgt_tokens = ["猫が", "マットの", "上に", "座って", "いた"]

fig, axes = plt.subplots(1, 4, figsize=(20, 5))

for ax, (name, attn_map) in zip(axes, attention_maps.items()):
    im = ax.imshow(attn_map, cmap="YlOrRd", vmin=0, vmax=1, aspect="auto")
    ax.set_xticks(range(src_len))
    ax.set_xticklabels(src_tokens, fontsize=10)
    ax.set_yticks(range(tgt_len))
    ax.set_yticklabels(tgt_tokens, fontsize=10)
    ax.set_xlabel("Source (English)", fontsize=11)
    ax.set_ylabel("Target (Japanese)", fontsize=11)
    ax.set_title(name, fontsize=12, fontweight="bold")

    # 各セルに数値を表示
    for i in range(tgt_len):
        for j in range(src_len):
            val = attn_map[i, j]
            color = "white" if val > 0.4 else "black"
            ax.text(j, i, f"{val:.2f}", ha="center", va="center",
                    fontsize=8, color=color)

plt.colorbar(im, ax=axes, fraction=0.02, pad=0.04, label="Attention Weight")
plt.suptitle("Attention Weight Comparison (Random Initialization)", fontsize=14, y=1.02)
plt.tight_layout()
plt.savefig("bahdanau_luong_comparison.png", dpi=150, bbox_inches="tight")
plt.show()

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

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

# --- 簡易 Seq2Seq + Attention モデル ---
class BahdanauAttention(nn.Module):
    def __init__(self, enc_dim, dec_dim, attn_dim):
        super().__init__()
        self.W1 = nn.Linear(enc_dim, attn_dim, bias=False)
        self.W2 = nn.Linear(dec_dim, attn_dim, bias=False)
        self.v = nn.Linear(attn_dim, 1, bias=False)

    def forward(self, enc_out, dec_h):
        ep = self.W1(enc_out)
        dp = self.W2(dec_h).unsqueeze(1)
        scores = self.v(torch.tanh(ep + dp)).squeeze(-1)
        weights = F.softmax(scores, dim=-1)
        ctx = torch.bmm(weights.unsqueeze(1), enc_out).squeeze(1)
        return ctx, weights

class LuongAttention(nn.Module):
    def __init__(self, enc_dim, dec_dim, score_fn="dot"):
        super().__init__()
        self.score_fn = score_fn
        if score_fn == "general":
            self.W_a = nn.Linear(enc_dim, dec_dim, bias=False)

    def forward(self, enc_out, dec_h):
        if self.score_fn == "dot":
            scores = torch.bmm(enc_out, dec_h.unsqueeze(-1)).squeeze(-1)
        elif self.score_fn == "general":
            scores = torch.bmm(
                self.W_a(enc_out), dec_h.unsqueeze(-1)
            ).squeeze(-1)
        weights = F.softmax(scores, dim=-1)
        ctx = torch.bmm(weights.unsqueeze(1), enc_out).squeeze(1)
        return ctx, weights

class Seq2SeqWithAttention(nn.Module):
    """簡易Seq2Seq + Attention（コピータスク用）"""
    def __init__(self, vocab_size, embed_dim, hidden_dim, attn_type="bahdanau"):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.encoder = nn.GRU(embed_dim, hidden_dim, batch_first=True, bidirectional=True)
        self.decoder_cell = nn.GRUCell(embed_dim + hidden_dim * 2, hidden_dim)
        self.output_proj = nn.Linear(hidden_dim, vocab_size)

        enc_dim = hidden_dim * 2  # 双方向
        if attn_type == "bahdanau":
            self.attention = BahdanauAttention(enc_dim, hidden_dim, hidden_dim)
        elif attn_type == "luong_general":
            self.attention = LuongAttention(enc_dim, hidden_dim, score_fn="general")
        self.attn_type = attn_type

    def forward(self, src, tgt):
        """
        Args:
            src: (batch, src_len) 入力トークンID
            tgt: (batch, tgt_len) ターゲットトークンID（教師強制用）
        Returns:
            logits: (batch, tgt_len, vocab_size)
            all_weights: (batch, tgt_len, src_len)
        """
        # エンコード
        src_emb = self.embedding(src)
        enc_out, enc_h = self.encoder(src_emb)

        # デコーダ初期状態: 双方向の最終隠れ状態を加算
        dec_h = enc_h[0] + enc_h[1]  # (batch, hidden_dim)

        batch_size, tgt_len = tgt.size()
        logits_list = []
        weights_list = []

        for t in range(tgt_len):
            tgt_emb = self.embedding(tgt[:, t])  # (batch, embed_dim)

            # Attention計算
            ctx, attn_w = self.attention(enc_out, dec_h)
            weights_list.append(attn_w)

            # デコーダ更新
            dec_input = torch.cat([tgt_emb, ctx], dim=-1)
            dec_h = self.decoder_cell(dec_input, dec_h)

            logits_list.append(self.output_proj(dec_h))

        logits = torch.stack(logits_list, dim=1)
        all_weights = torch.stack(weights_list, dim=1)
        return logits, all_weights

# --- コピータスク用データ生成 ---
vocab_size = 20
seq_len = 8
num_samples = 2000
embed_dim = 32
hidden_dim = 64
num_epochs = 50
batch_size = 64

# ランダムなトークン列を生成（0はパディング用に予約）
data = torch.randint(1, vocab_size, (num_samples, seq_len))

# --- 2つのモデルを学習 ---
results = {}
for attn_type, label in [("bahdanau", "Bahdanau"), ("luong_general", "Luong (General)")]:
    model = Seq2SeqWithAttention(vocab_size, embed_dim, hidden_dim, attn_type)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    losses = []
    for epoch in range(num_epochs):
        # ミニバッチ学習
        perm = torch.randperm(num_samples)
        epoch_loss = 0.0
        for i in range(0, num_samples, batch_size):
            batch = data[perm[i:i+batch_size]]
            # コピータスク: 入力 = ターゲット
            logits, _ = model(batch, batch)
            loss = criterion(logits.reshape(-1, vocab_size), batch.reshape(-1))

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        losses.append(epoch_loss / (num_samples // batch_size))

    # 学習後の Attention 重みを取得
    model.eval()
    with torch.no_grad():
        test_seq = data[:1]  # テスト用に1サンプル
        _, attn_weights = model(test_seq, test_seq)
        results[label] = {
            "losses": losses,
            "weights": attn_weights.squeeze(0).numpy()
        }

print("学習完了")

import numpy as np
import matplotlib.pyplot as plt

# --- 学習曲線と Attention 重みの可視化 ---
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# (1) 学習曲線
for label, res in results.items():
    axes[0].plot(res["losses"], label=label, linewidth=2)
axes[0].set_xlabel("Epoch", fontsize=12)
axes[0].set_ylabel("Cross-Entropy Loss", fontsize=12)
axes[0].set_title("Training Loss (Copy Task)", fontsize=13, fontweight="bold")
axes[0].legend(fontsize=11)
axes[0].grid(True, alpha=0.3)

# (2) Bahdanau の Attention 重みヒートマップ
attn_b = results["Bahdanau"]["weights"]
im1 = axes[1].imshow(attn_b, cmap="YlOrRd", vmin=0, vmax=1, aspect="auto")
for i in range(attn_b.shape[0]):
    for j in range(attn_b.shape[1]):
        val = attn_b[i, j]
        axes[1].text(j, i, f"{val:.2f}", ha="center", va="center",
                     fontsize=8, color="white" if val > 0.4 else "black")
axes[1].set_xlabel("Source Position", fontsize=12)
axes[1].set_ylabel("Target Position", fontsize=12)
axes[1].set_title("Bahdanau Attention Weights", fontsize=13, fontweight="bold")
plt.colorbar(im1, ax=axes[1], fraction=0.046, pad=0.04)

# (3) Luong の Attention 重みヒートマップ
attn_l = results["Luong (General)"]["weights"]
im2 = axes[2].imshow(attn_l, cmap="YlOrRd", vmin=0, vmax=1, aspect="auto")
for i in range(attn_l.shape[0]):
    for j in range(attn_l.shape[1]):
        val = attn_l[i, j]
        axes[2].text(j, i, f"{val:.2f}", ha="center", va="center",
                     fontsize=8, color="white" if val > 0.4 else "black")
axes[2].set_xlabel("Source Position", fontsize=12)
axes[2].set_ylabel("Target Position", fontsize=12)
axes[2].set_title("Luong (General) Attention Weights", fontsize=13, fontweight="bold")
plt.colorbar(im2, ax=axes[2], fraction=0.046, pad=0.04)

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

import torch
import torch.nn as nn

# --- パラメータ数の比較 ---
d_enc = 512   # 双方向エンコーダの場合 256 * 2
d_dec = 256
d_attn = 256

# Bahdanau Attention のパラメータ数
W1_params = d_enc * d_attn       # W1: (d_attn, d_enc)
W2_params = d_dec * d_attn       # W2: (d_attn, d_dec)
v_params = d_attn * 1            # v: (1, d_attn)
bahdanau_total = W1_params + W2_params + v_params

# Luong Attention のパラメータ数
dot_params = 0                            # dot: パラメータなし
general_params = d_enc * d_dec            # general: W_a: (d_dec, d_enc)
concat_W = (d_enc + d_dec) * d_dec        # concat: W_a
concat_v = d_dec * 1                      # concat: v_a
concat_total = concat_W + concat_v

print("=== パラメータ数の比較 ===")
print(f"Bahdanau Attention:      {bahdanau_total:>10,} パラメータ")
print(f"  W1: {W1_params:,}, W2: {W2_params:,}, v: {v_params:,}")
print(f"Luong (dot):             {dot_params:>10,} パラメータ")
print(f"Luong (general):         {general_params:>10,} パラメータ")
print(f"Luong (concat):          {concat_total:>10,} パラメータ")
print(f"  W_a: {concat_W:,}, v_a: {concat_v:,}")

# PyTorchモデルで検証
class BahdanauAttention(nn.Module):
    def __init__(self, enc_dim, dec_dim, attn_dim):
        super().__init__()
        self.W1 = nn.Linear(enc_dim, attn_dim, bias=False)
        self.W2 = nn.Linear(dec_dim, attn_dim, bias=False)
        self.v = nn.Linear(attn_dim, 1, bias=False)

class LuongAttention(nn.Module):
    def __init__(self, enc_dim, dec_dim, score_fn="dot"):
        super().__init__()
        if score_fn == "general":
            self.W_a = nn.Linear(enc_dim, dec_dim, bias=False)

b = BahdanauAttention(d_enc, d_dec, d_attn)
l_g = LuongAttention(d_enc, d_dec, "general")

print(f"\n=== PyTorchで検証 ===")
print(f"Bahdanau: {sum(p.numel() for p in b.parameters()):,} パラメータ")
print(f"Luong (general): {sum(p.numel() for p in l_g.parameters()):,} パラメータ")