从Transformer到代码：手把手拆解PyTorch中MultiheadAttention的前向传播（附mask原理动画）

在自然语言处理和计算机视觉领域，Transformer架构已经成为革命性的模型框架。而MultiheadAttention作为其核心组件，理解其内部实现机制对于掌握现代深度学习至关重要。本文将带您深入PyTorch的nn.MultiheadAttention实现，通过代码级解析和可视化演示，揭示注意力机制在前向传播中的完整流程。

1. MultiheadAttention架构全景

Transformer中的注意力机制并非单一视角，而是通过"多头"设计实现并行特征提取。想象一下，当人类观察一幅画时，有人会先注意色彩，有人会先关注构图——MultiheadAttention正是模拟这种多角度观察方式。

PyTorch中的实现核心参数包括：

• embed_dim: 输入token的嵌入维度
• num_heads: 注意力头的数量
• dropout: 注意力权重的随机丢弃率
• bias: 是否在投影层添加偏置项

关键设计特点：

python复制# 典型初始化示例
mha = nn.MultiheadAttention(
    embed_dim=512,
    num_heads=8,
    dropout=0.1,
    bias=True
)

多头注意力的核心思想是将高维空间分割为多个子空间：

1. 输入嵌入被分割为num_heads个低维子空间
2. 每个子空间独立计算注意力
3. 结果拼接后通过线性投影恢复原始维度

2. 前向传播的输入输出解析

forward方法是理解MultiheadAttention的关键入口。其标准调用形式为：

python复制attn_output, attn_weights = mha(
    query, 
    key, 
    value,
    key_padding_mask=None,
    need_weights=True,
    attn_mask=None
)

2.1 输入张量结构

三个核心输入张量具有相同形状(L, N, E)：

• L: 序列长度（如句子中的t