怎么在Python TensorFlow中实现Transformer_通过MultiHeadAttention解决

在TensorFlow里用MultiHeadAttention层搭Transformer，很多开发者会误以为它是个“开箱即用”的完整模块。实际上，它只负责核心的注意力计算，而整个Transformer的骨架——从输入嵌入到最终的输出——都需要你亲手搭建。这中间有几个关键环节，如果处理不当，轻则模型不收敛，重则训练直接崩溃。

MultiHeadAttention 层不是万能的，得自己搭好输入结构

首先得明确一点：tf.keras.layers.MultiHeadAttention这个层，它的职责非常纯粹，就是计算注意力权重，然后把value向量加权聚合起来。至于位置编码、残差连接、层归一化（LayerNorm）或者前馈网络（FFN），它一概不管。很多人直接把原始序列扔进去，结果要么输出维度对不上，要么梯度瞬间爆炸，问题就出在这里。

这个层对输入格式有明确要求：query、key、value这三个张量的最后两维，必须是[batch, seq_len, num_heads * head_dim]。而且，query和key的feature维度必须一致（虽然它们的序列长度可以不同）。

所以，正确的搭建顺序应该是这样的：

• 先用tf.keras.layers.Embedding把输入的token ID映射成稠密向量。
• 然后，手动加上位置编码。这里建议直接用正弦余弦函数（tf.sin和tf.cos）构造，避免引入不必要的第三方依赖。
• 在初始化MultiHeadAttention层时，要特别注意num_heads（注意力头数）必须能整除key_dim（每个注意力头的维度），否则会直接抛出ValueError: key_dim must be divisible by num_heads的错误。
• 最后，训练时务必设置training=True来启用dropout，而在验证或推理时，则要记得关掉。

mask 传错位置会导致 attention 看到 padding 或未来 token

掩码（mask）是Transformer，尤其是解码器部分，正确工作的关键。它主要处理两种情况：一是因果掩码（causal mask），防止解码时看到未来的信息；二是填充掩码（padding mask），用来忽略序列中无意义的填充位置。

但MultiHeadAttention的attention_mask参数只接受一个张量，并且其形状必须是[batch, 1, seq_len, seq_len]，或者能广播成这个形状。一个常见的坑是，直接把一维的padding_mask（形状为[batch, seq_len]）传进去，导致掩码广播错位。这样一来，模型在训练时可能就“偷看”到了本应被屏蔽的填充位置，严重影响效果。

正确的构造方法如下：

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

自定义 Transformer block 时，LayerNorm 的 axis 别写成 -1

标准的Transformer在每个子层（自注意力、前馈网络）之后都会接一个LayerNorm，作用在特征的最后一个维度上。很多人在参考代码时，会直接照抄axis=-1这个参数。但在某些特定情况下，比如batch size为1或序列长度为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）。无论是原始论文还是后来的T5等模型，都验证了GELU比ReLU在Transformer中表现更稳定。

训练时 loss 突然飙升，大概率是 learning rate 或初始化问题

Transformer架构对超参数，尤其是学习率，异常敏感。MultiHeadAttention层内部的Q、K、V投影矩阵如果使用默认的glorot_uniform初始化，在模型较大或批次较小时，很容易引发梯度爆炸。一个典型的现象是：训练前10步损失从10顺利降到3，但第11步突然飙升到200以上，紧接着就变成NaN了。

遇到这种情况，可以按以下顺序快速排查：

• 启用学习率预热（warmup）：例如在前1000步，让学习率从0线性增长到峰值。峰值学习率的设置也有讲究，Base规模的模型可以尝试1e-4，Large模型则建议从3e-5开始。
• 检查初始化方式：将MultiHeadAttention层的kernel_initializer改为tf.keras.initializers.VarianceScaling(scale=0.125, mode="fan_a vg", distribution="uniform")，这更接近原始Transformer论文的实现。
• 做过拟合测试：如果模型在验证集上损失不降，先关掉所有的dropout，尝试让模型在单个小批次数据上过拟合。如果连这都做不到，那基本可以断定是模型结构存在bug，而不是数据或优化器的问题。

说到底，在TensorFlow里实现Transformer，大部分调试时间都花在了四个地方：忘了加位置编码、掩码传反了、LayerNorm的轴设错了、学习率没预热。其他参数可以慢慢调优，但这四个点，必须从一开始就牢牢盯紧。