注意機構（Attention）

注意機構の誕生

Seq2Seq モデルの課題

機械翻訳などのSequence-to-Sequence（Seq2Seq）タスクでは、エンコーダが入力系列を 固定長のコンテキストベクトルに圧縮していた。これには問題があった。

Attention の解決策

Attention機構（Bahdanau et al., 2014）は、デコーダが各ステップで エンコーダのすべての隠れ状態を参照できるようにした。

Attention の計算

スコア関数の種類

Self-Attention

Self-Attention（自己注意）は、入力系列が自分自身に対して Attentionを計算する仕組みである。系列内の任意の位置間の関係を直接モデル化できる。

なぜスケーリングが必要か

Multi-Head Attention

単一のAttentionでは表現力が限られるため、複数の「ヘッド」で 異なる観点からAttentionを計算し、結果を結合する。

ヘッドはどう分割する？

1. 入力（各トークン \(d_{model}\) 次元）を重み行列 \(W^Q, W^K, W^V\) で線形変換し、全体の Q, K, V（\(d_{model}\) 次元）を作る。
2. その \(d_{model}\) 次元ベクトルを \(h\) 個に等分し、各ヘッドが \(d_k = d_{model}/h\) 次元の小さな Q, K, V を担当する（実装上は reshape：\((\text{系列長},\, d_{model}) \to (\text{系列長},\, h,\, d_k)\)）。別々の重みを \(h\) 組用意するのではなく、1 つの大きな射影を \(h\) 分割しているだけで、数式の \(W_i^Q\) は \(W^Q\) のうち各ヘッドが担当する列に対応する。
3. 各ヘッドで独立に Scaled Dot-Product Attention を計算し、得られた出力（各 \(d_k\) 次元）を連結して \(d_{model}\) 次元へ戻し、最後に \(W^O\) で混ぜ合わせる。

ヘッド数 \(h\) の決め方：\(h\) は人が決めるハイパーパラメータで、\(d_k = d_{model}/h\) が整数になるよう \(d_{model}\) の約数から選ぶ。ヘッドを増やしても各ヘッドが細くなるぶん計算量はほぼ一定だが、増やしすぎると 1 ヘッドあたりの次元が小さくなり表現力が落ちる。経験的に 8〜16 が定番である。

典型的な設定

• ヘッド数 h：8〜16
• 各ヘッドの次元：\(d_k = d_v = d_{model} / h\)
• 例：\(d_{model} = 512, h = 8 \Rightarrow d_k = 64\)

位置エンコーディング

Self-Attentionは順序を考慮しない（並べ替え不変）。 そのため、位置情報を明示的に追加する必要がある。

位置エンコーディングの特性

• 任意の固定オフセット \(k\) に対して、\(PE_{pos+k}\) は \(PE_{pos}\) の線形関数で表現可能
• これにより相対位置の情報を学習しやすい
• 学習済みモデルより長い系列にも対応可能（外挿可能）

PyTorchでの実装

Scaled Dot-Product Attention

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

def scaled_dot_product_attention(query, key, value, mask=None):
    """
    Scaled Dot-Product Attention

    Args:
        query: (batch, heads, seq_len, d_k)
        key: (batch, heads, seq_len, d_k)
        value: (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)
        attention_weights: (batch, heads, seq_len, seq_len)
    """
    d_k = query.size(-1)

    # QK^T / sqrt(d_k)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    # マスク適用（オプション）
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # Softmax
    attention_weights = F.softmax(scores, dim=-1)

    # Attention × V
    output = torch.matmul(attention_weights, value)

    return output, attention_weights

Multi-Head Attention

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # 線形変換層
        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 split_heads(self, x, batch_size):
        """(batch, seq_len, d_model) -> (batch, heads, seq_len, d_k)"""
        x = x.view(batch_size, -1, self.num_heads, self.d_k)
        return x.transpose(1, 2)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        # 線形変換
        Q = self.W_q(query)
        K = self.W_k(key)
        V = self.W_v(value)

        # ヘッドに分割
        Q = self.split_heads(Q, batch_size)
        K = self.split_heads(K, batch_size)
        V = self.split_heads(V, batch_size)

        # Scaled Dot-Product Attention
        attn_output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)

        # ヘッドを結合
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, -1, self.d_model)

        # 最終線形変換
        output = self.W_o(attn_output)

        return output, attn_weights


# 使用例
d_model = 512
num_heads = 8
seq_len = 100
batch_size = 32

mha = MultiHeadAttention(d_model, num_heads)
x = torch.randn(batch_size, seq_len, d_model)

output, attention = mha(x, x, x)  # Self-Attention
print(f"Output shape: {output.shape}")  # (32, 100, 512)
print(f"Attention shape: {attention.shape}")  # (32, 8, 100, 100)

位置エンコーディング

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

        # 位置エンコーディングを事前計算
        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)  # 奇数次元

        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
        self.register_buffer('pe', pe)

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


# 使用例
pe = PositionalEncoding(d_model=512)
x = torch.randn(32, 100, 512)
output = pe(x)
print(f"Output shape: {output.shape}")  # (32, 100, 512)

Attentionの可視化

Attention重みを可視化することで、モデルが「どこに注目しているか」を理解できる。

まとめ

よくある質問（FAQ）