Python实现Transformer模型详解 多头注意力机制实战教程

在TensorFlow框架中构建Transformer模型时，MultiHeadAttention层常被开发者误解为一个即插即用的完整解决方案。实际上，它仅负责核心的注意力机制运算，而Transformer模型的完整架构——从输入嵌入、位置编码到最终的输出生成——都需要开发者自行构建与组装。这其中存在多个关键环节，若处理不当，可能导致模型训练困难、性能低下甚至完全失败。

MultiHeadAttention 层功能有限，需自行构建完整输入流程

首先需要明确：tf.keras.layers.MultiHeadAttention层的职责非常单一，即根据输入的query、key、value计算注意力分布并进行加权聚合。至于模型必需的位置编码（Positional Encoding）、残差连接（Residual Connection）、层归一化（Layer Normalization）以及前馈神经网络（Feed-Forward Network），该层均不负责。许多开发者直接将原始序列输入，导致输出维度不匹配或梯度爆炸，根源即在于此。

该层对输入张量的格式有严格要求：query、key、value三个张量的最后两维必须符合[batch_size, sequence_length, num_heads * head_dim]的格式。同时，query与key的特征维度（feature dimension）必须保持一致（尽管它们的序列长度可以不同）。

因此，在TensorFlow中正确搭建Transformer的推荐步骤如下：

• 首先使用tf.keras.layers.Embedding层将输入的词元ID（token IDs）映射为稠密向量表示。
• 随后，手动添加位置编码。建议直接使用正弦余弦函数（tf.sin和tf.cos）生成，以减少不必要的第三方依赖。
• 在初始化MultiHeadAttention层时，必须确保注意力头数num_heads能够整除每个头的维度key_dim，否则会触发ValueError: key_dim must be divisible by num_heads错误。
• 最后，在训练阶段务必设置training=True以启用Dropout机制，而在模型验证或推理阶段则应将其关闭。

掩码（Mask）设置错误会导致注意力机制泄露信息

掩码是确保Transformer模型（尤其是解码器部分）正确工作的核心机制。它主要处理两种场景：一是因果掩码（Causal Mask），用于防止解码器在生成当前词元时“看到”未来的信息；二是填充掩码（Padding Mask），用于忽略序列中无意义的填充位置。

然而，MultiHeadAttention层的attention_mask参数仅接受特定形状的张量，其形状应为[batch_size, 1, target_seq_len, source_seq_len]或可广播为此形状。一个常见误区是直接将形状为[batch_size, seq_len]的一维填充掩码传入，导致掩码广播错位。这会使模型在训练过程中意外地关注到本应被屏蔽的填充位置，严重影响学习效果。

正确的掩码构造方法如下：

• 生成因果掩码，可使用tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)函数。
• 生成填充掩码，需将形状为[batch_size, seq_len]的布尔掩码扩展为[batch_size, 1, 1, seq_len]，再与因果掩码进行逻辑与（AND）操作。
• 在解码器中，编码器-解码器注意力层通常仅需填充掩码（因为编码器输出无顺序依赖），而解码器的自注意力层则必须包含因果掩码。

构建自定义Transformer模块时，注意LayerNorm的轴（axis）设置

标准Transformer在每个子层（自注意力层、前馈网络层）之后都会应用层归一化（LayerNorm），其作用维度通常是特征的最后一个维度。许多开发者在参考代码时会直接复制axis=-1这一参数。但在某些特定情况下，例如批次大小（batch size）或序列长度（sequence length）为1时，这可能导致数值计算不稳定，甚至产生NaN值。

问题的关键在于，tf.keras.layers.LayerNormalization默认会对所有非批次维度进行归一化。如果你的输入张量形状为[batch, seq_len, features]，那么设置axis=-1是正确的。但如果在数据处理过程中使用了tf.transpose等操作改变了维度顺序（例如变为[batch, features, seq_len]），那么axis=-1指向的就不再是特征维，而是序列长度维，这显然是错误的。

在实现过程中，有几个细节需要特别注意：

• 在层的call方法中，使用print(x.shape)来确认输入张量的确切形状，这是最直接的调试方法。
• 进行残差连接时，必须确保原始输入（query）与注意力层的输出形状完全一致，否则tf.add操作会引发Incompatible shapes错误。
• 前馈网络（FFN）通常由两个全连接层构成，中间激活函数推荐使用GELU（tf.nn.gelu）。无论是原始Transformer论文还是后续的T5等模型，都验证了GELU比ReLU在Transformer架构中表现更为稳定。

训练过程中损失（Loss）突然飙升的常见原因与排查

Transformer架构对超参数，尤其是学习率（Learning Rate），极为敏感。MultiHeadAttention层内部的Q、K、V投影矩阵如果使用默认的glorot_uniform初始化方式，在模型参数量较大或批次较小时，极易引发梯度爆炸。一个典型现象是：训练初期损失从10正常下降至3，但在后续某一步骤突然飙升至200以上，随后变为NaN。

遇到此类问题，可按以下顺序进行排查：

• 启用学习率预热（Warmup）：例如在前1000个训练步中，让学习率从0线性增长至预设峰值。峰值学习率的设置也需谨慎，Base规模的模型可尝试1e-4，Large模型则建议从3e-5开始。
• 检查权重初始化方式：将MultiHeadAttention层的kernel_initializer参数改为tf.keras.initializers.VarianceScaling(scale=0.125, mode="fan_avg", distribution="uniform")，这更接近原始Transformer论文的推荐设置。
• 进行过拟合测试：如果模型在验证集上损失不下降，可先关闭所有Dropout，尝试让模型在单个小批次数据上过拟合。如果模型连此任务都无法完成，那么基本可以断定是模型结构本身存在缺陷，而非数据或优化器的问题。

总而言之，在TensorFlow中实现并调试Transformer模型，大部分时间往往消耗在四个关键点上：遗漏位置编码、掩码传递错误、LayerNorm轴设置不当以及未进行学习率预热。其他超参数可以逐步调优，但这四个核心环节，必须在模型构建之初就予以高度重视并正确实现。