Transformer：自注意力、多头机制、位置与掩码

问题与目标

RNN 按时间步顺序处理序列，长距离信息需要经过多次状态传递。Transformer 改用自注意力，让每个位置直接聚合其他位置的信息，并用并行矩阵运算处理整个序列。

本篇拆开缩放点积注意力、多头机制、位置表示、掩码、残差、层归一化和前馈网络，并使用 PyTorch 检查张量形状。重点不是默写结构图，而是理解每个模块解决什么问题。

编码器允许输入位置彼此可见；解码器的自注意力使用因果掩码，只能读取当前位置及其左侧内容。编码器—解码器注意力再让解码器读取源序列表示。

核心概念

缩放点积注意力

给定 Query、Key 和 Value：

Attention(Q, K, V) = softmax(QKᵀ / √d_k + mask) V

QKᵀ 得到位置之间的相关性。除以 √d_k 是为了防止维度增大时点积过大，使 Softmax 过早饱和。Mask 在不允许访问的位置加入极小值，Softmax 后权重接近零。

假设批次为 B、序列长度为 T、隐藏维度为 D，常见形状是：

输入 X:          [B, T, D]
Q、K、V:         [B, H, T, D/H]
注意力分数:      [B, H, T, T]
合并多头后输出:  [B, T, D]

多头不是简单复制

每个头拥有独立的投影参数，可以学习不同的关联模式。各头输出拼接后再经过线性变换。头数必须与隐藏维度兼容，通常要求 D % H == 0。

为什么需要位置

纯自注意力只根据内容计算关系，交换 Token 顺序会得到相同的集合关系。模型需要额外的位置表示来区分“节点检查服务”和“服务检查节点”。位置可以使用固定函数、可学习向量或旋转位置编码等方式注入。

残差、归一化与前馈网络

• 残差连接保留输入并改善深层网络的梯度传播。
• LayerNorm 在单个样本的特征维度上归一化，适合变长序列。
• 前馈网络对每个位置独立应用同一组非线性变换，扩展特征后再压回隐藏维度。

不同模型可能采用 Pre-Norm 或 Post-Norm。阅读实现时应确认归一化发生在子层之前还是之后。

可运行实现

下面构造一层多头自注意力，同时传入 Padding Mask 和 Causal Mask。需要已经安装 PyTorch。

import torch
from torch import nn


torch.manual_seed(7)

batch_size = 2
seq_len = 5
hidden_size = 8
num_heads = 2

x = torch.randn(batch_size, seq_len, hidden_size)
attention = nn.MultiheadAttention(
    embed_dim=hidden_size,
    num_heads=num_heads,
    batch_first=True,
)

# True 表示该 Key 位置应被忽略。
padding_mask = torch.tensor([
    [False, False, False, True, True],
    [False, False, False, False, True],
])

# 上三角为 True，表示当前位置不能读取未来 Token。
causal_mask = torch.triu(
    torch.ones(seq_len, seq_len, dtype=torch.bool),
    diagonal=1,
)

output, weights = attention(
    query=x,
    key=x,
    value=x,
    key_padding_mask=padding_mask,
    attn_mask=causal_mask,
    need_weights=True,
    average_attn_weights=False,
)

print("input:  ", tuple(x.shape))
print("output: ", tuple(output.shape))
print("weights:", tuple(weights.shape))
print("sample head 0:\n", weights[0, 0].detach().round(decimals=3))

输出形状应为：

input:   (2, 5, 8)
output:  (2, 5, 8)
weights: (2, 2, 5, 5)

第一维是批次，第二维是注意力头，最后两个维度分别对应 Query 位置和 Key 位置。因果掩码使每一行右侧未来位置的权重为零，Padding Mask 则屏蔽补齐 Token。

常见问题与排查

Mask 的布尔含义写反

不同 API 对 True、1 和 0 的解释不完全相同。有的表示保留，有的表示屏蔽。先用一个很小的矩阵打印权重，确认被屏蔽位置是否接近零。

维度无法被头数整除

多头注意力通常把隐藏维度均分到各头。若 hidden_size=10、num_heads=3，无法等分。调整隐藏维度或头数，不要靠 reshape 强行绕过。

输出出现 NaN

检查某一行是否所有位置都被 Mask，Softmax 对全负无穷输入可能产生 NaN；同时检查输入和参数中是否已经出现非有限值。

把 Attention 当作整个 Transformer

完整层还包括投影、残差、归一化、前馈网络和 Dropout。只有注意力矩阵并不能代表完整模型行为。

小结

Transformer 的关键不是单独一个注意力公式，而是注意力、位置、掩码、残差、归一化和前馈网络共同组成的可堆叠结构。调试时从张量形状和 Mask 开始，比只看最终生成文本更容易定位错误。