PyTorch torch.nn.MultiheadAttention 関数

PyTorch torch.nn リファレンスマニュアル

torch.nn.MultiheadAttentionは PyTorch のマルチヘッド注意機構モジュールです。

これは Transformer アーキテクチャの中心的なコンポーネントであり、モデルが異なる位置にある異なる表現部分空間の情報に同時に注意を向けることを可能にします。

関数定義

torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, kdim=None, vdim=None, batch_first=True)

パラメータ説明：

  • embed_dim(int): 入力埋め込み次元。
  • num_heads(int): 注意ヘッドの数。
  • dropout(float): dropout 確率。デフォルトは 0。
  • kdim(int): キーベクトルの次元。デフォルトは None（embed_dim と同じ）。
  • vdim(int): バリューベクトルの次元。デフォルトは None（embed_dim と同じ）。
  • batch_first(bool): True の場合、入出力の最初の次元は batch になります。デフォルトは True。

使用例

例 1: 基本的な使い方

実例

import torch
import torch.nn as nn

# マルチヘッド注意を作成：512次元、8ヘッド
mha = nn.MultiheadAttention(embed_dim=512, num_heads=8)

# 入力：batch=4、系列長=100、次元=512
query = torch.randn(4, 100, 512)
key = torch.randn(4, 100, 512)
value = torch.randn(4, 100, 512)

# 順伝播
output, attn_weight = mha(query, key, value)

print("Query 形状:", query.shape)
print("Output 形状:", output.shape)
print("注意重み形状:", attn_weight.shape)

例2：自己注意

実例

import torch
import torch.nn as nn

mha = nn.MultiheadAttention(embed_dim=256, num_heads=4)

# 同じ入力を Q、K、V として使用（自己注意）
x = torch.randn(2, 50, 256)

# self-attention: q=k=v=x
output, weights = mha(x, x, x)

print("入力形状:", x.shape)
print("出力形状:", output.shape)
print("注意重み形状:", weights.shape)

例 3: マスク付き注意

実例

import torch
import torch.nn as nn

mha = nn.MultiheadAttention(embed_dim=128, num_heads=4)

# 入力
x = torch.randn(1, 20, 128)

# 上三角マスクを作成（デコーダー用）
mask = torch.triu(torch.ones(20, 20), diagonal=1).bool()

output, _ = mha(x, x, x, attn_mask=mask)

print("入力形状:", x.shape)
print("Output 形状:", output.shape)
print("Mask 形状:", mask.shape)

例 4: 完全な Transformer エンコーダー層

実例

import torch
import torch.nn as nn

class TransformerLayer(nn.Module):
    def __init__(self, d_model, nhead):
        super(TransformerLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.ReLU(),
            nn.Linear(d_model * 4, d_model)
        )

    def forward(self, x):
        # 自己注意 + 残差
        attn_out, _ = self.self_attn(x, x, x)
        x = self.norm1(x + attn_out)

        # FFN + 残差
        ffn_out = self.ffn(x)
        x = self.norm2(x + ffn_out)

        return x

# テスト
layer = TransformerLayer(d_model=512, nhead=8)
x = torch.randn(4, 100, 512)
output = layer(x)

print("入力形状:", x.shape)
print("出力形状:", output.shape)

例 5: 注意重みの確認

実例

import torch
import torch.nn as nn
import numpy as np

mha = nn.MultiheadAttention(embed_dim=64, num_heads=2, batch_first=True)

# 短い動画シーケンス
x = torch.randn(1, 5, 64)

_, attn = mha(x, x, x)
attn = attn.squeeze(0)  # batch を除去

print("最初のヘッドの注意重み（最初の3つの位置）:")
print(attn[0, :3, :3].tolist())

print("n可視化 - 位置0からすべての位置への注意:")
print(np.array2string(attn[0, 0].numpy(), precision=2))

よくある質問

Q1: num_heads はどのように選びますか？

embed_dim は num_heads で割り切れる必要があります。一般的な値：8、12、16。

Q2: なぜ Q、K、V の3つの行列が必要ですか？

モデルが異なる射影を学習できるようにし、表現能力を高めます。

Q3: key_padding_mask とは何ですか？

padding 位置をマスクし、注意が padding に対して計算されるのを防ぎます。

使用シーン

  • Transformer: エンコーダーとデコーダー
  • 自己注意モデル: BERT、GPT
  • 系列モデリング: RNN の代替

ヒント：batch_first=True の場合、入力形状は (batch, seq, embed_dim) です。

PyTorch torch.nn リファレンスマニュアル

その他の拡張