PyTorch アテンション機構
アテンション機構（Attention Mechanism）は、深層学習において最も重要な概念の1つです。
アテンション機構は、モデルが入力の中で最も関連性の高い部分に「注目」することを学習させ、自然言語処理やコンピュータビジョンなどの分野で大きな成功を収めています。
この節では、アテンション機構の核となる原理、PyTorch の実装、およびさまざまなアテンション変種について詳しく説明します。
対応バージョン：本記事のコードは PyTorch 2.0+ を前提に書かれています。
nn.MultiheadAttentionのbatch_firstパラメータは PyTorch 1.9 で導入され、初期バージョンでは次元を手動で処理する必要があります。
1. アテンション機構の基礎
1.1 なぜアテンションが必要か
従来のシーケンス・ツー・シーケンス（Seq2Seq）モデルには根本的な問題があります：Encoder はすべての情報を固定長のベクトルに圧縮する必要があります。長い系列の場合、このベクトルが情報のボトルネックになります。文が長いほど、情報の損失は深刻になります。
アテンション機構の核となる考え方は、Decoder が各出力を生成する際に、Encoder のすべての隠れ状態を「見る」ことができ、現在のコンテキストに応じて異なるアテンション重みを動的に割り当てられるようにすることです。これは、人間が翻訳するときに、単語を訳すたびに原文の対応する部分を振り返って見るのと似ています。
1.2 アテンション機構の本質
アテンション機構は一種の加重和操作と見なすことができます。クエリ（Query）、キー（Key）、バリュー（Value）が与えられたとき、Query と各 Key の類似度を計算して重みを割り当て、Value の加重和を取ります：
\[\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V\]$$
ここで：
- Q（Query）：クエリベクトル。「私が探している情報」を表します
- K（Key）：キーベクトル。「ここにどんな情報があるか」を表す（マッチング用）
- V（Value）：バリューベクトル。「実際に提供できる内容」を表す
- $$d_k$$：Key の次元。$$\sqrt{d_k}$$ はスケーリングに使用され、ドット積の値が大きくなりすぎて softmax の勾配が消失するのを防ぎます
なぜ $$\sqrt{d_k}$$ で割るのか？ $$d_k$$ が大きい場合、Q と K のドット積の分散は次元に比例して増加し、softmax の入力値が大きくなりすぎて、勾配がゼロに近づきます。スケーリングにより分散は 1 に戻り、勾配が正常に流れることが保証されます。
例
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):
"""
スケーリングドット積アテンション
パラメータ:
Q: クエリテンソル [batch, n_heads, seq_len_q, d_k]
K: キーテンソル [batch, n_heads, seq_len_k, d_k]
V: バリューテンソル [batch, n_heads, seq_len_v, d_v]
mask: マスクテンソル。マスクされた位置は False/0 に設定
戻り値:
output: アテンション出力 [batch, n_heads, seq_len_q, d_v]
attention_weights: アテンション重み [batch, n_heads, seq_len_q, seq_len_k]
"""
d_k = Q.size(-1)
# 1. QK^T / sqrt(d_k) を計算
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
# 2. マスクを適用（オプション）
# マスクされた位置は極小値に設定し、softmax 後は 0 に近づく
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
# 3. Softmax 正規化 → アテンション重み
attention_weights = F.softmax(scores, dim=-1)
# 4. 加重和
output = torch.matmul(attention_weights, V)
return output, attention_weights
# テスト
batch_size, n_heads = 2, 4
seq_len_q, seq_len_k = 5, 6
d_k, d_v = 8, 16
Q = torch.randn(batch_size, n_heads, seq_len_q, d_k)
K = torch.randn(batch_size, n_heads, seq_len_k, d_k)
V = torch.randn(batch_size, n_heads, seq_len_k, d_v)
output, attn_weights = scaled_dot_product_attention(Q, K, V)
print(f“出力形状: {output.shape}”) # [2, 4, 5, 16]
print(f“アテンション重み形状: {attn_weights.shape}”) # [2, 4, 5, 6]
print(f“重み行の合計（1.0 になるはず）: {attn_weights[0, 0, 0].sum().item():.4f}”)
import torch.nn as nn
import torch.nn.functional as F
import math
def scaled_dot_product_attention(Q, K, V, mask=None):
"""
スケーリングドット積アテンション
パラメータ:
Q: クエリテンソル [batch, n_heads, seq_len_q, d_k]
K: キーテンソル [batch, n_heads, seq_len_k, d_k]
V: バリューテンソル [batch, n_heads, seq_len_v, d_v]
mask: マスクテンソル。マスクされた位置は False/0 に設定
戻り値:
output: アテンション出力 [batch, n_heads, seq_len_q, d_v]
attention_weights: アテンション重み [batch, n_heads, seq_len_q, seq_len_k]
"""
d_k = Q.size(-1)
# 1. QK^T / sqrt(d_k) を計算
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
# 2. マスクを適用（オプション）
# マスクされた位置は極小値に設定し、softmax 後は 0 に近づく
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
# 3. Softmax 正規化 → アテンション重み
attention_weights = F.softmax(scores, dim=-1)
# 4. 加重和
output = torch.matmul(attention_weights, V)
return output, attention_weights
# テスト
batch_size, n_heads = 2, 4
seq_len_q, seq_len_k = 5, 6
d_k, d_v = 8, 16
Q = torch.randn(batch_size, n_heads, seq_len_q, d_k)
K = torch.randn(batch_size, n_heads, seq_len_k, d_k)
V = torch.randn(batch_size, n_heads, seq_len_k, d_v)
output, attn_weights = scaled_dot_product_attention(Q, K, V)
print(f“出力形状: {output.shape}”) # [2, 4, 5, 16]
print(f“アテンション重み形状: {attn_weights.shape}”) # [2, 4, 5, 6]
print(f“重み行の合計（1.0 になるはず）: {attn_weights[0, 0, 0].sum().item():.4f}”)
アテンション重みの本質は確率分布です。softmax を経た後、各行の合計は 1 になります。重みが大きいほど、モデルがその位置に「注目」していることを示します。
2. PyTorch アテンションモジュール
2.1 Multi-Head Attention
Multi-Head Attention（マルチヘッドアテンション）モデルが、異なる位置の異なる表現部分空間の情報を同時に注目できるようにします。Q、K、V をそれぞれ異なる線形投影で複数の部分空間にマッピングし、各部分空間で独立にアテンションを計算し、最後に結果を連結してもう一度線形変換を行います。
これは、複数の人が同じテキストを異なる視点から同時に調べるようなものです。ある人は文法構造に注目し、ある人は意味的な関連に注目し、ある人は長距離依存に注目します。マルチヘッド機構により、モデルはより豊かなパターンを捉えることができます。
$$\text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W^O$$
$$\text{where head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$$
例
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, n_heads, dropout=0.1):
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 の線形投影（すべてのヘッドを一度に計算するため高効率）
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)
self.dropout = nn.Dropout(dropout)
def forward(self, query, key, value, mask=None):
"""
パラメータ:
query: [batch, seq_len_q, d_model]
key: [batch, seq_len_k, d_model]
value: [batch, seq_len_k, d_model]
mask: [batch, 1, 1, seq_len_k] または [batch, 1, seq_len_q, seq_len_k]
"""
batch_size = query.size(0)
# 1. 線形投影し、ヘッドに分割
# [batch, seq_len, d_model] → [batch, seq_len, n_heads, d_k]
# → [batch, n_heads, seq_len, d_k]
Q = self.w_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
K = self.w_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
V = self.w_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
# 2. スケーリングドット積アテンション
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)
# 3. 加重和
context = torch.matmul(attn_weights, V)
# 4. マルチヘッドの結合: [batch, n_heads, seq_len, d_k] → [batch, seq_len, d_model]
context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
# 5. 出力投影
output = self.w_o(context)
return output, attn_weights
# テスト
d_model, n_heads = 128, 8
seq_len, batch = 10, 4
layer = MultiHeadAttention(d_model, n_heads)
query = torch.randn(batch, seq_len, d_model)
key = torch.randn(batch, seq_len, d_model)
value = torch.randn(batch, seq_len, d_model)
output, attn_weights = layer(query, key, value)
print(f“出力形状: {output.shape}”) # [4, 10, 128]
print(f“アテンション重み形状: {attn_weights.shape}”) # [4, 8, 10, 10]
print(f“各ヘッドの次元 d_k = {d_model // n_heads}”)
import torch.nn as nn
import torch.nn.functional as F
import math
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, n_heads, dropout=0.1):
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 の線形投影（すべてのヘッドを一度に計算するため高効率）
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)
self.dropout = nn.Dropout(dropout)
def forward(self, query, key, value, mask=None):
"""
パラメータ:
query: [batch, seq_len_q, d_model]
key: [batch, seq_len_k, d_model]
value: [batch, seq_len_k, d_model]
mask: [batch, 1, 1, seq_len_k] または [batch, 1, seq_len_q, seq_len_k]
"""
batch_size = query.size(0)
# 1. 線形投影し、ヘッドに分割
# [batch, seq_len, d_model] → [batch, seq_len, n_heads, d_k]
# → [batch, n_heads, seq_len, d_k]
Q = self.w_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
K = self.w_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
V = self.w_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
# 2. スケーリングドット積アテンション
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)
# 3. 加重和
context = torch.matmul(attn_weights, V)
# 4. マルチヘッドの結合: [batch, n_heads, seq_len, d_k] → [batch, seq_len, d_model]
context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
# 5. 出力投影
output = self.w_o(context)
return output, attn_weights
# テスト
d_model, n_heads = 128, 8
seq_len, batch = 10, 4
layer = MultiHeadAttention(d_model, n_heads)
query = torch.randn(batch, seq_len, d_model)
key = torch.randn(batch, seq_len, d_model)
value = torch.randn(batch, seq_len, d_model)
output, attn_weights = layer(query, key, value)
print(f“出力形状: {output.shape}”) # [4, 10, 128]
print(f“アテンション重み形状: {attn_weights.shape}”) # [4, 8, 10, 10]
print(f“各ヘッドの次元 d_k = {d_model // n_heads}”)
2.2 PyTorch 組み込み MultiheadAttention
PyTorch は高度に最適化されたnn.MultiheadAttentionモジュールを提供しています。内部では融合演算子（fused kernels）を使用しており、ほとんどの場面で手動実装より高速です。本番環境での使用をお勧めします。
例
import torch
import torch.nn as nn
class TransformerAttention(nn.Module):
"""PyTorch 組み込み MultiheadAttention を使用した自己アテンション層"""
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
self.attention = nn.MultiheadAttention(
embed_dim=d_model,
num_heads=n_heads,
dropout=dropout,
batch_first=True, # 入力形式は [batch, seq, features]
# PyTorch 2.0+ では Flash Attention バックエンドを有効化できます：
# attn_implementation="flash_attention_2" # flash-attn のインストールが必要
)
self.layernorm = nn.LayerNorm(d_model)
def forward(self, x, key_padding_mask=None):
"""
自己注意: Q、K、V はすべて x
パラメータ:
x: [batch, seq_len, d_model]
key_padding_mask: [batch, seq_len]、True はその位置が padding であることを示す
"""
attn_output, attn_weights = self.attention(
x, x, x, # self-attention
key_padding_mask=key_padding_mask
)
# Pre-Norm 残差接続（Post-Norm よりも訓練が安定）
output = self.layernorm(x + attn_output)
return output, attn_weights
# テスト
d_model, n_heads = 128, 8
seq_len, batch = 10, 4
model = TransformerAttention(d_model, n_heads)
x = torch.randn(batch, seq_len, d_model)
# padding mask を作成: True はその位置が無視されるべきことを示す
key_padding_mask = torch.zeros(batch, seq_len, dtype=torch.bool)
key_padding_mask[0, 7:] = True # 最初のサンプルの後ろ 3 つの位置は padding
key_padding_mask[1, 5:] = True # 2 番目のサンプルの後ろ 5 つの位置は padding
output, attn_weights = model(x, key_padding_mask=key_padding_mask)
print(f"出力形状: {output.shape}") # [4, 10, 128]
print(f"注意重み形状: {attn_weights.shape}") # [4, 10, 10]
import torch.nn as nn
class TransformerAttention(nn.Module):
"""PyTorch 組み込み MultiheadAttention を使用した自己アテンション層"""
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
self.attention = nn.MultiheadAttention(
embed_dim=d_model,
num_heads=n_heads,
dropout=dropout,
batch_first=True, # 入力形式は [batch, seq, features]
# PyTorch 2.0+ では Flash Attention バックエンドを有効化できます：
# attn_implementation="flash_attention_2" # flash-attn のインストールが必要
)
self.layernorm = nn.LayerNorm(d_model)
def forward(self, x, key_padding_mask=None):
"""
自己注意: Q、K、V はすべて x
パラメータ:
x: [batch, seq_len, d_model]
key_padding_mask: [batch, seq_len]、True はその位置が padding であることを示す
"""
attn_output, attn_weights = self.attention(
x, x, x, # self-attention
key_padding_mask=key_padding_mask
)
# Pre-Norm 残差接続（Post-Norm よりも訓練が安定）
output = self.layernorm(x + attn_output)
return output, attn_weights
# テスト
d_model, n_heads = 128, 8
seq_len, batch = 10, 4
model = TransformerAttention(d_model, n_heads)
x = torch.randn(batch, seq_len, d_model)
# padding mask を作成: True はその位置が無視されるべきことを示す
key_padding_mask = torch.zeros(batch, seq_len, dtype=torch.bool)
key_padding_mask[0, 7:] = True # 最初のサンプルの後ろ 3 つの位置は padding
key_padding_mask[1, 5:] = True # 2 番目のサンプルの後ろ 5 つの位置は padding
output, attn_weights = model(x, key_padding_mask=key_padding_mask)
print(f"出力形状: {output.shape}") # [4, 10, 128]
print(f"注意重み形状: {attn_weights.shape}") # [4, 10, 10]
mask のタイプに注意:
nn.MultiheadAttentionのkey_padding_mask使用True = 無視という約定は、一部のカスタム実装で 0 = 無視 という約定と逆である。使用時には必ず確認すること。
3. アテンション機構の変種
3.1 自己アテンション（Self-Attention）
自己注意は注意機構の特殊な形式の一つであり、Q、K、V はすべて同じ入力シーケンスから得られる。これにより、シーケンス内の各位置が他のすべての位置に直接注意を向けることができ、グローバルな依存関係を捕捉できる。これは Transformer の中核コンポーネントであり、RNN よりも優れている鍵でもある: RNN は段階的に情報を伝達する必要があるが、自己注意は一度に行う。
例
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class SelfAttention(nn.Module):
"""マルチヘッド自己注意層、因果マスクをサポート（自己回帰生成用）"""
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
assert d_model % n_heads == 0
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_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.out_proj = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)
def _split_heads(self, x, batch_size):
"""[batch, seq, d_model] → [batch, n_heads, seq, d_k]"""
return x.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
def forward(self, x, causal=False):
"""
自己注意: Q、K、V はすべて同じ入力 x から得られる
パラメータ:
x: [batch, seq_len, d_model]
causal: 因果マスクを使用するかどうか（未来の位置を見るのを防ぐ）
"""
batch_size, seq_len, _ = x.size()
# プロジェクション + ヘッド分割
Q = self._split_heads(self.w_q(x), batch_size)
K = self._split_heads(self.w_k(x), batch_size)
V = self._split_heads(self.w_v(x), batch_size)
# 注意スコアを計算
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
# 因果マスク: 上三角行列を -inf に設定し、未来のトークンを見るのを防ぐ
if causal:
causal_mask = torch.triu(
torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool),
diagonal=1
)
scores = scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))
attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)
# 加重和 + マルチヘッドの結合
context = torch.matmul(attn_weights, V)
context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
return self.out_proj(context), attn_weights
# テスト
d_model, n_heads = 256, 8
seq_len, batch = 20, 2
attn = SelfAttention(d_model, n_heads)
x = torch.randn(batch, seq_len, d_model)
# 因果マスクを使用しない（Encoder シナリオ）
output_enc, weights_enc = attn(x, causal=False)
print(f"[Encoder] 出力: {output_enc.shape}, 重み: {weights_enc.shape}")
# 因果マスクを使用（Decoder / 自己回帰生成シナリオ）
output_dec, weights_dec = attn(x, causal=True)
print(f"[Decoder] 出力: {output_dec.shape}, 重み: {weights_dec.shape}")
# 因果性を確認: 最初のトークンは後続のトークンに注意を向けるべきではない
print(f"causal=True の場合、weights[0,0,0,5:]（すべて 0 になるはず）: {weights_dec[0, 0, 0, 5:].tolist()[:5]}")
import torch.nn as nn
import torch.nn.functional as F
import math
class SelfAttention(nn.Module):
"""マルチヘッド自己注意層、因果マスクをサポート（自己回帰生成用）"""
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
assert d_model % n_heads == 0
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_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.out_proj = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)
def _split_heads(self, x, batch_size):
"""[batch, seq, d_model] → [batch, n_heads, seq, d_k]"""
return x.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
def forward(self, x, causal=False):
"""
自己注意: Q、K、V はすべて同じ入力 x から得られる
パラメータ:
x: [batch, seq_len, d_model]
causal: 因果マスクを使用するかどうか（未来の位置を見るのを防ぐ）
"""
batch_size, seq_len, _ = x.size()
# プロジェクション + ヘッド分割
Q = self._split_heads(self.w_q(x), batch_size)
K = self._split_heads(self.w_k(x), batch_size)
V = self._split_heads(self.w_v(x), batch_size)
# 注意スコアを計算
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
# 因果マスク: 上三角行列を -inf に設定し、未来のトークンを見るのを防ぐ
if causal:
causal_mask = torch.triu(
torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool),
diagonal=1
)
scores = scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))
attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)
# 加重和 + マルチヘッドの結合
context = torch.matmul(attn_weights, V)
context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
return self.out_proj(context), attn_weights
# テスト
d_model, n_heads = 256, 8
seq_len, batch = 20, 2
attn = SelfAttention(d_model, n_heads)
x = torch.randn(batch, seq_len, d_model)
# 因果マスクを使用しない（Encoder シナリオ）
output_enc, weights_enc = attn(x, causal=False)
print(f"[Encoder] 出力: {output_enc.shape}, 重み: {weights_enc.shape}")
# 因果マスクを使用（Decoder / 自己回帰生成シナリオ）
output_dec, weights_dec = attn(x, causal=True)
print(f"[Decoder] 出力: {output_dec.shape}, 重み: {weights_dec.shape}")
# 因果性を確認: 最初のトークンは後続のトークンに注意を向けるべきではない
print(f"causal=True の場合、weights[0,0,0,5:]（すべて 0 になるはず）: {weights_dec[0, 0, 0, 5:].tolist()[:5]}")
3.2 交差アテンション（Cross-Attention）
交差注意では、Q はあるシーケンス（例: Decoder）から、K と V は別のシーケンス（例: Encoder）から得られる。これは Encoder-Decoder アーキテクチャの橋渡しであり、Decoder が各トークンを生成する際に、Encoder のすべての出力を動的に「参照」できるようにする。
例
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class CrossAttention(nn.Module):
"""交差注意: Q はターゲット系列から、K/V はソース系列から"""
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
assert d_model % n_heads == 0
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_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.out_proj = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, target, source, mask=None):
"""
パラメータ:
target: ターゲット系列 [batch, target_len, d_model] → Q を提供
source: ソース系列 [batch, source_len, d_model] → K, V を提供
mask: オプションのマスク
"""
batch_size = target.size(0)
target_len = target.size(1)
source_len = source.size(1)
# プロジェクション + ヘッド分割
Q = self.w_q(target).view(batch_size, target_len, self.n_heads, self.d_k).transpose(1, 2)
K = self.w_k(source).view(batch_size, source_len, self.n_heads, self.d_k).transpose(1, 2)
V = self.w_v(source).view(batch_size, source_len, self.n_heads, self.d_k).transpose(1, 2)
# 注意の計算
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)
attn_weights = self.dropout(attn_weights)
# 加重和 + マルチヘッドの結合
context = torch.matmul(attn_weights, V)
context = context.transpose(1, 2).contiguous().view(batch_size, target_len, self.d_model)
return self.out_proj(context), attn_weights
# テスト
d_model, n_heads = 256, 8
batch = 2
target_len, source_len = 10, 15
cross_attn = CrossAttention(d_model, n_heads)
target = torch.randn(batch, target_len, d_model) # Decoder 出力
source = torch.randn(batch, source_len, d_model) # Encoder 出力
output, weights = cross_attn(target, source)
print(f"ターゲット系列形状: {target.shape}")
print(f"ソース系列形状: {source.shape}")
print(f"出力形状: {output.shape}") # [2, 10, 256]
print(f"重み形状: {weights.shape}") # [2, 8, 10, 15]
# 重みの意味: 各 target 位置の各 source 位置に対する注意の度合い
import torch.nn as nn
import torch.nn.functional as F
import math
class CrossAttention(nn.Module):
"""交差注意: Q はターゲット系列から、K/V はソース系列から"""
def __init__(self, d_model, n_heads, dropout=0.1):
super().__init__()
assert d_model % n_heads == 0
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_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.out_proj = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, target, source, mask=None):
"""
パラメータ:
target: ターゲット系列 [batch, target_len, d_model] → Q を提供
source: ソース系列 [batch, source_len, d_model] → K, V を提供
mask: オプションのマスク
"""
batch_size = target.size(0)
target_len = target.size(1)
source_len = source.size(1)
# プロジェクション + ヘッド分割
Q = self.w_q(target).view(batch_size, target_len, self.n_heads, self.d_k).transpose(1, 2)
K = self.w_k(source).view(batch_size, source_len, self.n_heads, self.d_k).transpose(1, 2)
V = self.w_v(source).view(batch_size, source_len, self.n_heads, self.d_k).transpose(1, 2)
# 注意の計算
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)
attn_weights = self.dropout(attn_weights)
# 加重和 + マルチヘッドの結合
context = torch.matmul(attn_weights, V)
context = context.transpose(1, 2).contiguous().view(batch_size, target_len, self.d_model)
return self.out_proj(context), attn_weights
# テスト
d_model, n_heads = 256, 8
batch = 2
target_len, source_len = 10, 15
cross_attn = CrossAttention(d_model, n_heads)
target = torch.randn(batch, target_len, d_model) # Decoder 出力
source = torch.randn(batch, source_len, d_model) # Encoder 出力
output, weights = cross_attn(target, source)
print(f"ターゲット系列形状: {target.shape}")
print(f"ソース系列形状: {source.shape}")
print(f"出力形状: {output.shape}") # [2, 10, 256]
print(f"重み形状: {weights.shape}") # [2, 8, 10, 15]
# 重みの意味: 各 target 位置の各 source 位置に対する注意の度合い
自己注意 vs 交差注意:自己注意では Q/K/V の形状が同じ（$$seq \times seq$$ の重み行列）。交差注意では Q の長さと K の長さを異なるものにできる（$$target \times source$$ の重み行列）。
3.3 位置エンコーディング（Positional Encoding）
注意機構自体は置換不変性——入力の順序を入れ替えても出力は変わらない。しかし言語や画像は順序を持つため、位置エンコーディングによって位置情報を明示的に注入する必要がある。元の Transformer はサイン/コサイン関数を使用して固定の位置エンコーディングを生成する:
$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right), \quad PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$
例
import torch
import torch.nn as nn
import math
class PositionalEncoding(nn.Module):
"""サイン/コサイン位置エンコーディング（Transformer オリジナル方式）"""
def __init__(self, d_model, max_len=5000, dropout=0.1):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
# 位置エンコーディング行列を事前計算 [max_len, d_model]
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
# div_term: 異なる周波数次元の減衰速度を制御
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):
"""
x: [batch, seq_len, d_model]
戻り値: x + positional_encoding
"""
x = x + self.pe[:, :x.size(1), :]
return self.dropout(x)
# テスト
d_model = 128
seq_len = 50
batch = 4
pe = PositionalEncoding(d_model)
x = torch.randn(batch, seq_len, d_model)
output = pe(x)
print(f"入力形状: {x.shape}")
print(f"出力形状: {output.shape}")
print(f"位置エンコーディング形状: {pe.pe.shape}") # [1, 5000, 128]
import torch.nn as nn
import math
class PositionalEncoding(nn.Module):
"""サイン/コサイン位置エンコーディング（Transformer オリジナル方式）"""
def __init__(self, d_model, max_len=5000, dropout=0.1):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
# 位置エンコーディング行列を事前計算 [max_len, d_model]
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
# div_term: 異なる周波数次元の減衰速度を制御
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):
"""
x: [batch, seq_len, d_model]
戻り値: x + positional_encoding
"""
x = x + self.pe[:, :x.size(1), :]
return self.dropout(x)
# テスト
d_model = 128
seq_len = 50
batch = 4
pe = PositionalEncoding(d_model)
x = torch.randn(batch, seq_len, d_model)
output = pe(x)
print(f"入力形状: {x.shape}")
print(f"出力形状: {output.shape}")
print(f"位置エンコーディング形状: {pe.pe.shape}") # [1, 5000, 128]
4. アテンション機構の視覚応用
注意機構はシーケンスデータだけでなく、コンピュータビジョンでも大いに活躍している。以下では、3 つの古典的な視覚注意モジュールを紹介する:
4.1 チャネルアテンション（Channel Attention）
チャネル注意（SENet など）ネットワークに「どのチャネルがより重要か」を学習させる。各チャネルに対してグローバルプーリングを行い空間情報を圧縮し、2 層の全結合ネットワークを通じてチャネル間の依存関係を学習する。
例
import torch
import torch.nn as nn
class SEBlock(nn.Module):
"""Squeeze-and-Excitation Block（チャネル注意）"""
def __init__(self, channels, reduction=16):
super().__init__()
self.squeeze = nn.AdaptiveAvgPool2d(1) # グローバル平均プーリング
self.excitation = nn.Sequential(
nn.Linear(channels, channels // reduction, bias=False),
nn.ReLU(inplace=True),
nn.Linear(channels // reduction, channels, bias=False),
nn.Sigmoid() # 0〜1 のチャネル重みを出力
)
def forward(self, x):
# x: [batch, channels, H, W]
b, c, _, _ = x.size()
# Squeeze: [B, C, H, W] → [B, C, 1, 1] → [B, C]
y = self.squeeze(x).view(b, c)
# Excitation: チャネル重みを学習 [B, C]
y = self.excitation(y).view(b, c, 1, 1)
# チャネル重み付け: 各チャネルに対応する重みを掛ける
return x * y.expand_as(x)
# テスト
se = SEBlock(channels=256, reduction=16)
x = torch.randn(4, 256, 32, 32)
output = se(x)
print(f"入力: {x.shape} → 出力: {output.shape}")
import torch.nn as nn
class SEBlock(nn.Module):
"""Squeeze-and-Excitation Block（チャネル注意）"""
def __init__(self, channels, reduction=16):
super().__init__()
self.squeeze = nn.AdaptiveAvgPool2d(1) # グローバル平均プーリング
self.excitation = nn.Sequential(
nn.Linear(channels, channels // reduction, bias=False),
nn.ReLU(inplace=True),
nn.Linear(channels // reduction, channels, bias=False),
nn.Sigmoid() # 0〜1 のチャネル重みを出力
)
def forward(self, x):
# x: [batch, channels, H, W]
b, c, _, _ = x.size()
# Squeeze: [B, C, H, W] → [B, C, 1, 1] → [B, C]
y = self.squeeze(x).view(b, c)
# Excitation: チャネル重みを学習 [B, C]
y = self.excitation(y).view(b, c, 1, 1)
# チャネル重み付け: 各チャネルに対応する重みを掛ける
return x * y.expand_as(x)
# テスト
se = SEBlock(channels=256, reduction=16)
x = torch.randn(4, 256, 32, 32)
output = se(x)
print(f"入力: {x.shape} → 出力: {output.shape}")
4.2 空間アテンション（Spatial Attention）
空間注意ネットワークに「どの位置がより重要か」を学習させる。チャネル次元で圧縮（平均と最大値を取得）し、その後コンボリューションで空間重みマップを学習する。
例
import torch
import torch.nn as nn
class SpatialAttention(nn.Module):
"""空間注意モジュール"""
def __init__(self, kernel_size=7):
super().__init__()
padding = kernel_size // 2
# 2 チャネルの入力（平均 + 最大値）、1 チャネルの出力（注意マップ）
self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
# x: [batch, channels, H, W]
# チャネル圧縮: 平均と最大値を取得 → [B, 1, H, W] をそれぞれ 1 つ
avg_out = torch.mean(x, dim=1, keepdim=True)
max_out, _ = torch.max(x, dim=1, keepdim=True)
# 連結 → [B, 2, H, W]
scale = torch.cat([avg_out, max_out], dim=1)
# コンボリューション + Sigmoid → 空間重みマップ [B, 1, H, W]
scale = self.sigmoid(self.conv(scale))
# 空間重み付け
return x * scale
# テスト
spatial_attn = SpatialAttention(kernel_size=7)
x = torch.randn(4, 256, 32, 32)
output = spatial_attn(x)
print(f"入力: {x.shape} → 出力: {output.shape}")
import torch.nn as nn
class SpatialAttention(nn.Module):
"""空間注意モジュール"""
def __init__(self, kernel_size=7):
super().__init__()
padding = kernel_size // 2
# 2 チャネルの入力（平均 + 最大値）、1 チャネルの出力（注意マップ）
self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
# x: [batch, channels, H, W]
# チャネル圧縮: 平均と最大値を取得 → [B, 1, H, W] をそれぞれ 1 つ
avg_out = torch.mean(x, dim=1, keepdim=True)
max_out, _ = torch.max(x, dim=1, keepdim=True)
# 連結 → [B, 2, H, W]
scale = torch.cat([avg_out, max_out], dim=1)
# コンボリューション + Sigmoid → 空間重みマップ [B, 1, H, W]
scale = self.sigmoid(self.conv(scale))
# 空間重み付け
return x * scale
# テスト
spatial_attn = SpatialAttention(kernel_size=7)
x = torch.randn(4, 256, 32, 32)
output = spatial_attn(x)
print(f"入力: {x.shape} → 出力: {output.shape}")
4.3 CBAM（Convolutional Block Attention Module）
CBAM はチャネル注意と空間注意を直列に適用する: 最初にチャネル注意で重要なチャネルを選択し、次に空間注意で重要な位置を選択する。両者は相補的であり、単独で使用するよりも効果が優れている。
例
import torch
import torch.nn as nn
class CBAM(nn.Module):
"""CBAM: チャネル注意 + 空間注意"""
def __init__(self, channels, reduction=16, kernel_size=7):
super().__init__()
# チャネル注意
self.channel_attn = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(channels, channels // reduction, bias=False),
nn.ReLU(inplace=True),
nn.Linear(channels // reduction, channels, bias=False),
nn.Sigmoid()
)
# 空間注意
self.spatial_attn = nn.Sequential(
nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False),
nn.Sigmoid()
)
def forward(self, x):
# 1. チャネル注意
b, c, _, _ = x.size()
ca = self.channel_attn(x).view(b, c, 1, 1)
x = x * ca
# 2. 空間注意
avg_out = torch.mean(x, dim=1, keepdim=True)
max_out, _ = torch.max(x, dim=1, keepdim=True)
sa = self.spatial_attn(torch.cat([avg_out, max_out], dim=1))
x = x * sa
return x
# テスト
cbam = CBAM(channels=256)
x = torch.randn(4, 256, 32, 32)
output = cbam(x)
print(f"入力: {x.shape} → 出力: {output.shape}")
import torch.nn as nn
class CBAM(nn.Module):
"""CBAM: チャネル注意 + 空間注意"""
def __init__(self, channels, reduction=16, kernel_size=7):
super().__init__()
# チャネル注意
self.channel_attn = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(channels, channels // reduction, bias=False),
nn.ReLU(inplace=True),
nn.Linear(channels // reduction, channels, bias=False),
nn.Sigmoid()
)
# 空間注意
self.spatial_attn = nn.Sequential(
nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False),
nn.Sigmoid()
)
def forward(self, x):
# 1. チャネル注意
b, c, _, _ = x.size()
ca = self.channel_attn(x).view(b, c, 1, 1)
x = x * ca
# 2. 空間注意
avg_out = torch.mean(x, dim=1, keepdim=True)
max_out, _ = torch.max(x, dim=1, keepdim=True)
sa = self.spatial_attn(torch.cat([avg_out, max_out], dim=1))
x = x * sa
return x
# テスト
cbam = CBAM(channels=256)
x = torch.randn(4, 256, 32, 32)
output = cbam(x)
print(f"入力: {x.shape} → 出力: {output.shape}")
5. アテンション機構の NLP への応用
5.1 Transformer Encoder
Transformer Encoder は N 個の同一レイヤーを積み重ねて構成され、各レイヤーには 2 つのサブレイヤーが含まれる:マルチヘッド自己注意とフィードフォワードネットワーク。各サブレイヤーは残差接続とレイヤー正規化を使用する。
例
import torch
import torch.nn as nn
import math
class TransformerEncoderLayer(nn.Module):
"""単層 Transformer Encoder"""
def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
super().__init__()
# マルチヘッド自己注意
self.self_attn = nn.MultiheadAttention(
d_model, n_heads, dropout=dropout, batch_first=True
)
# フィードフォワードネットワーク（2 層の線形 + 活性化）
self.ffn = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(), # GELU は ReLU より一般的に使われる
nn.Dropout(dropout),
nn.Linear(d_ff, d_model)
)
# 層正規化（Pre-Norm スタイル、学習がより安定）
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x, key_padding_mask=None):
# サブレイヤー 1: 自己注意
residual = x
x = self.norm1(x) # Pre-Norm
attn_out, _ = self.self_attn(x, x, x, key_padding_mask=key_padding_mask)
x = residual + self.dropout(attn_out)
# サブレイヤー 2: フィードフォワードネットワーク
residual = x
x = self.norm2(x) # Pre-Norm
x = residual + self.dropout(self.ffn(x))
return x
class TransformerEncoder(nn.Module):
"""多層 Transformer エンコーダ"""
def __init__(self, d_model, n_heads, d_ff, num_layers, dropout=0.1):
super().__init__()
self.layers = nn.ModuleList([
TransformerEncoderLayer(d_model, n_heads, d_ff, dropout)
for _ in range(num_layers)
])
self.final_norm = nn.LayerNorm(d_model)
def forward(self, x, key_padding_mask=None):
for layer in self.layers:
x = layer(x, key_padding_mask)
return self.final_norm(x)
# テスト
d_model, n_heads, d_ff = 128, 4, 512
num_layers, seq_len, batch = 3, 20, 4
encoder = TransformerEncoder(d_model, n_heads, d_ff, num_layers)
x = torch.randn(batch, seq_len, d_model)
# Padding mask: True = その位置を無視
key_padding_mask = torch.zeros(batch, seq_len, dtype=torch.bool)
key_padding_mask[0, 15:] = True
output = encoder(x, key_padding_mask=key_padding_mask)
print(f"入力: {x.shape} → 出力: {output.shape}")
import torch.nn as nn
import math
class TransformerEncoderLayer(nn.Module):
"""単層 Transformer Encoder"""
def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
super().__init__()
# マルチヘッド自己注意
self.self_attn = nn.MultiheadAttention(
d_model, n_heads, dropout=dropout, batch_first=True
)
# フィードフォワードネットワーク（2 層の線形 + 活性化）
self.ffn = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(), # GELU は ReLU より一般的に使われる
nn.Dropout(dropout),
nn.Linear(d_ff, d_model)
)
# 層正規化（Pre-Norm スタイル、学習がより安定）
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x, key_padding_mask=None):
# サブレイヤー 1: 自己注意
residual = x
x = self.norm1(x) # Pre-Norm
attn_out, _ = self.self_attn(x, x, x, key_padding_mask=key_padding_mask)
x = residual + self.dropout(attn_out)
# サブレイヤー 2: フィードフォワードネットワーク
residual = x
x = self.norm2(x) # Pre-Norm
x = residual + self.dropout(self.ffn(x))
return x
class TransformerEncoder(nn.Module):
"""多層 Transformer エンコーダ"""
def __init__(self, d_model, n_heads, d_ff, num_layers, dropout=0.1):
super().__init__()
self.layers = nn.ModuleList([
TransformerEncoderLayer(d_model, n_heads, d_ff, dropout)
for _ in range(num_layers)
])
self.final_norm = nn.LayerNorm(d_model)
def forward(self, x, key_padding_mask=None):
for layer in self.layers:
x = layer(x, key_padding_mask)
return self.final_norm(x)
# テスト
d_model, n_heads, d_ff = 128, 4, 512
num_layers, seq_len, batch = 3, 20, 4
encoder = TransformerEncoder(d_model, n_heads, d_ff, num_layers)
x = torch.randn(batch, seq_len, d_model)
# Padding mask: True = その位置を無視
key_padding_mask = torch.zeros(batch, seq_len, dtype=torch.bool)
key_padding_mask[0, 15:] = True
output = encoder(x, key_padding_mask=key_padding_mask)
print(f"入力: {x.shape} → 出力: {output.shape}")
5.2 アテンションの可視化
注意重みの可視化は Transformer モデルを理解しデバッグするための重要な手段です。ヒートマップにより、モデルが各位置でどの情報に注目しているかを直感的に確認できます。
例
import torch
import matplotlib.pyplot as plt
import numpy as np
def visualize_attention(attention_weights, tokens=None, save_path=None):
"""
マルチヘッド注意重みの可視化
パラメータ:
attention_weights: [n_heads, seq_len, seq_len]
tokens: トークンリスト（オプション）
save_path: 保存パス（オプション）
"""
n_heads = attention_weights.shape[0]
seq_len = attention_weights.shape[1]
fig, axes = plt.subplots(1, n_heads, figsize=(n_heads * 3, 3.5))
if n_heads == 1:
axes = [axes]
for head in range(n_heads):
ax = axes[head]
attn = attention_weights[head].detach().cpu().numpy()
im = ax.imshow(attn, cmap='Blues', aspect='auto', vmin=0, vmax=attn.max())
ax.set_title(f'Head {head + 1}', fontsize=10, fontweight='bold')
ax.set_xlabel('Key')
ax.set_ylabel('Query')
if tokens is not None:
ax.set_xticks(range(len(tokens)))
ax.set_yticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=90, fontsize=8)
ax.set_yticklabels(tokens, fontsize=8)
plt.tight_layout()
if save_path:
plt.savefig(save_path, dpi=150, bbox_inches='tight')
plt.show()
# 学習後の注意重みをシミュレート
n_heads, seq_len = 4, 8
attention_weights = torch.softmax(torch.randn(n_heads, seq_len, seq_len), dim=-1)
tokens = ['The', 'cat', 'sat', 'on', 'the', 'mat', '.', '[PAD]']
visualize_attention(attention_weights, tokens)
import matplotlib.pyplot as plt
import numpy as np
def visualize_attention(attention_weights, tokens=None, save_path=None):
"""
マルチヘッド注意重みの可視化
パラメータ:
attention_weights: [n_heads, seq_len, seq_len]
tokens: トークンリスト（オプション）
save_path: 保存パス（オプション）
"""
n_heads = attention_weights.shape[0]
seq_len = attention_weights.shape[1]
fig, axes = plt.subplots(1, n_heads, figsize=(n_heads * 3, 3.5))
if n_heads == 1:
axes = [axes]
for head in range(n_heads):
ax = axes[head]
attn = attention_weights[head].detach().cpu().numpy()
im = ax.imshow(attn, cmap='Blues', aspect='auto', vmin=0, vmax=attn.max())
ax.set_title(f'Head {head + 1}', fontsize=10, fontweight='bold')
ax.set_xlabel('Key')
ax.set_ylabel('Query')
if tokens is not None:
ax.set_xticks(range(len(tokens)))
ax.set_yticks(range(len(tokens)))
ax.set_xticklabels(tokens, rotation=90, fontsize=8)
ax.set_yticklabels(tokens, fontsize=8)
plt.tight_layout()
if save_path:
plt.savefig(save_path, dpi=150, bbox_inches='tight')
plt.show()
# 学習後の注意重みをシミュレート
n_heads, seq_len = 4, 8
attention_weights = torch.softmax(torch.randn(n_heads, seq_len, seq_len), dim=-1)
tokens = ['The', 'cat', 'sat', 'on', 'the', 'mat', '.', '[PAD]']
visualize_attention(attention_weights, tokens)
6. 完全なアテンション分類器
6.1 アテンションに基づくテキスト分類
以下は完全なテキスト分類モデルです：Transformer エンコーダで特徴を抽出し、注意力プーリング（Attention Pooling）で可変長シーケンスを固定長ベクトルに圧縮し、最後に分類器に入力します。
例
import torch
import torch.nn as nn
import math
class AttentionClassifier(nn.Module):
"""Transformer + 注意力プーリングに基づくテキスト分類モデル"""
def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes,
n_heads=4, num_layers=2, dropout=0.1):
super().__init__()
self.embed_dim = embed_dim
# 単語埋め込み
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
self.pos_encoding = PositionalEncoding(embed_dim, dropout=dropout)
# Transformer エンコーダ
encoder_layer = nn.TransformerEncoderLayer(
d_model=embed_dim,
nhead=n_heads,
dim_feedforward=hidden_dim,
dropout=dropout,
batch_first=True,
activation='gelu'
)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
# 注意力プーリング：各位置の重要度重みを学習
self.attn_pool = nn.Sequential(
nn.Linear(embed_dim, 1) # 各位置でスカラー重みを出力
)
# 分類ヘッド
self.classifier = nn.Sequential(
nn.Linear(embed_dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, num_classes)
)
def forward(self, x):
"""
x: [batch, seq_len] — token ids
戻り値: logits [batch, num_classes], attn_weights [batch, seq_len]
"""
# 単語埋め込み + 位置エンコーディング
mask = (x == 0) # padding mask
x = self.embedding(x) * math.sqrt(self.embed_dim)
x = self.pos_encoding(x)
# Transformer エンコード
x = self.transformer(x, src_key_padding_mask=mask)
# 注意力プーリング：各位置の重みを計算し、重み付き和を取る
scores = self.attn_pool(x).squeeze(-1) # [batch, seq_len]
scores = scores.masked_fill(mask, float('-inf')) # padding を無視
attn_weights = torch.softmax(scores, dim=1) # [batch, seq_len]
pooled = torch.bmm(attn_weights.unsqueeze(1), x).squeeze(1) # [batch, embed_dim]
# 分類
logits = self.classifier(pooled)
return logits, attn_weights
# テスト
model = AttentionClassifier(
vocab_size=10000, embed_dim=128, hidden_dim=256, num_classes=5
)
# 入力のシミュレーション（padding 含む）
x = torch.randint(1, 10000, (16, 50))
x[:, -10:] = 0 # 最後の 10 位置は padding
logits, attn_weights = model(x)
print(f"入力: {x.shape}")
print(f"Logits: {logits.shape}") # [16, 5]
print(f"注意重み: {attn_weights.shape}") # [16, 50]
import torch.nn as nn
import math
class AttentionClassifier(nn.Module):
"""Transformer + 注意力プーリングに基づくテキスト分類モデル"""
def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes,
n_heads=4, num_layers=2, dropout=0.1):
super().__init__()
self.embed_dim = embed_dim
# 単語埋め込み
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
self.pos_encoding = PositionalEncoding(embed_dim, dropout=dropout)
# Transformer エンコーダ
encoder_layer = nn.TransformerEncoderLayer(
d_model=embed_dim,
nhead=n_heads,
dim_feedforward=hidden_dim,
dropout=dropout,
batch_first=True,
activation='gelu'
)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
# 注意力プーリング：各位置の重要度重みを学習
self.attn_pool = nn.Sequential(
nn.Linear(embed_dim, 1) # 各位置でスカラー重みを出力
)
# 分類ヘッド
self.classifier = nn.Sequential(
nn.Linear(embed_dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, num_classes)
)
def forward(self, x):
"""
x: [batch, seq_len] — token ids
戻り値: logits [batch, num_classes], attn_weights [batch, seq_len]
"""
# 単語埋め込み + 位置エンコーディング
mask = (x == 0) # padding mask
x = self.embedding(x) * math.sqrt(self.embed_dim)
x = self.pos_encoding(x)
# Transformer エンコード
x = self.transformer(x, src_key_padding_mask=mask)
# 注意力プーリング：各位置の重みを計算し、重み付き和を取る
scores = self.attn_pool(x).squeeze(-1) # [batch, seq_len]
scores = scores.masked_fill(mask, float('-inf')) # padding を無視
attn_weights = torch.softmax(scores, dim=1) # [batch, seq_len]
pooled = torch.bmm(attn_weights.unsqueeze(1), x).squeeze(1) # [batch, embed_dim]
# 分類
logits = self.classifier(pooled)
return logits, attn_weights
# テスト
model = AttentionClassifier(
vocab_size=10000, embed_dim=128, hidden_dim=256, num_classes=5
)
# 入力のシミュレーション（padding 含む）
x = torch.randint(1, 10000, (16, 50))
x[:, -10:] = 0 # 最後の 10 位置は padding
logits, attn_weights = model(x)
print(f"入力: {x.shape}")
print(f"Logits: {logits.shape}") # [16, 5]
print(f"注意重み: {attn_weights.shape}") # [16, 50]
6.2 よくあるアテンション変種の比較
| タイプ | Q の由来 | K/V の由来 | 主な用途 | 代表的な応用 |
|---|---|---|---|---|
| 自己注意（Self-Attention） | 自身 | 自身 | シーケンス内部の大域的相互作用 | Transformer、BERT、ViT |
| 交差注意（Cross-Attention） | ターゲットシーケンス | ソースシーケンス | シーケンス間の情報融合 | 機械翻訳、画像キャプショニング、Stable Diffusion |
| チャネル注意（Channel Attention） | グローバルプーリング特徴 | 同上 | チャネルの重要度を学習 | SENet、CBAM |
| 空間注意（Spatial Attention） | チャネル圧縮特徴 | 同上 | 空間位置の重要度を学習 | CBAM、Spatial Transformer |
| スパース注意（Sparse Attention） | 一部の位置 | 一部の位置 | 長いシーケンスの計算複雑度を低減 | Longformer、BigBird |
7. ベストプラクティスとよくある質問
7.1 使用のコツ
- マスクの一貫性: padding mask と causal mask の方向を一致させます（True = 無視 vs 1 = 保持）。API によって規約が異なる場合があります。
- ヘッド数の選択: 通常は $$d_{model}$$ の約数に設定し、8 または 16 が一般的です。ヘッドが多すぎると各ヘッドの次元 $$d_k$$ が小さくなり、性能を損なう可能性があります。
- 残差接続: 注意力層の前後には必ず残差接続 + 層正規化を使用します。これは深い Transformer を学習するための必須条件です。
- 学習率ウォームアップ: Transformer の学習初期は勾配の変動が大きいため、warmup + cosine decay スケジューラを使用すると安定性が大幅に向上します。
- Flash Attention: PyTorch 2.0+ は、次を介してサポートしています
nn.functional.scaled_dot_product_attention自動的に Flash Attention を呼び出し、メモリと速度が大幅に最適化されます。
7.2 よくある質問
| 問題 | 原因 | 解決策 |
|---|---|---|
| 注意重みがすべて一様分布になる | 入力が正規化されていない、または学習率が大きすぎる | 入力のスケールを確認し、LayerNorm を使用し、学習率を下げる |
| 学習 loss が低下しない | 位置エンコーディングの欠落または mask の誤り | 位置エンコーディングが追加されていることを確認し、mask の方向を確認する |
| メモリ不足（長いシーケンス） | 注意行列の $$O(n^2)$$ 複雑度 | Flash Attention、勾配チェックポイント、またはスパース注意を使用 |
| 推論時に注意が退化する | 学習/推論時に mask が一致しない | padding mask が学習時と推論時で同じロジックを使用するようにする |
その他の拡張注意機構は現代の深層学習における中核技術です。NLP から CV、音声からタンパク質構造予測まで、ほぼすべての最先端モデルが注意に基づいています。その原理と実装を理解することは、Transformer、ViT、Stable Diffusion、GPT などのモデルを習得するための重要な第一歩です。