MultiHeadAttention层需手动构建输入结构，包括嵌入、位置编码、残差连接、LayerNorm和前馈网络；其输入要求query/key/value维度兼容，num_heads须整除key_dim，mask需正确构造为[batch,1,seq_len,seq_len]，LayerNorm axis须匹配实际shape，残差连接需shape一致，FFN推荐GELU，学习率需warmup且初始化需谨慎。
MultiHeadAttention 层不是万能的，得自己搭好输入结构
TensorFlow 的 tf.keras.layers.MultiHeadAttention 只负责计算注意力权重和加权聚合，不处理位置编码、残差连接、LayerNorm 或前馈网络。很多人直接把原始序列喂进去，结果输出 shape 不对或梯度爆炸——因为它默认要求 query、key、value 的最后两维分别是 [batch, seq_len, num_heads * head_dim]，且三者维度需兼容（query 和 key 的 seq_len 可不同，但 feature 维必须一致）。
实操建议：
- 先用
tf.keras.layers.Embedding把 token ID 映射为向量，再手动加正弦位置编码（tf.sin/tf.cos构造，别依赖第三方库） -
MultiHeadAttention的num_heads必须整除key_dim（即每个 head 的维度），否则报错ValueError: key_dim must be divisible by num_heads - 训练时务必设
training=True，否则 dropout 不生效，验证时记得关掉
mask 传错位置会导致 attention 看到 padding 或未来 token
Transformer 解码器需要两种 mask：causal_mask（防止看未来）和 padding_mask（屏蔽填充值）。但 MultiHeadAttention 的 attention_mask 参数只接受一个张量，且 shape 必须是 [batch, 1, seq_len, seq_len] 或广播兼容形状。常见错误是把一维 padding_mask（如 [batch, seq_len]）直接传进去，导致 mask 广播错位，模型在训练时“偷偷”看到 padding 位置。
正确做法：
- 生成 causal mask：用
tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0) - 生成 padding mask：先将
[batch, seq_len]的布尔 mask 扩展为[batch, 1, 1, seq_len]，再与 causal mask 相乘（逻辑与） - 解码器中，encoder-decoder attention 的
attention_mask只需 padding mask（因 encoder 输出已无序），而 decoder self-attention 必须带 causal mask
自定义 Transformer block 时，LayerNorm 的 axis 别写成 -1
标准 Transformer 每个子层后接 LayerNorm，作用于最后一个维度。但很多人复制 Keras 示例时把 axis=-1 写死，结果在 batch size=1 或序列长度为 1 时出 NaN——因为 tf.keras.layers.LayerNormalization 默认对所有非 batch 维归一化，若输入是 [batch, seq_len, features]，应显式指定 axis=-1；但如果用了 tf.transpose 改变顺序（比如 [batch, features, seq_len]），axis=-1 就错了。
关键细节：
- 确认输入 tensor 的 shape，用
print(x.shape)在 call 前调试 - 残差连接必须保证
query和 attention 输出 shape 完全一致，否则tf.add报错Incompatible shapes - 前馈网络（FFN）推荐用两层 Dense + GELU（
tf.nn.gelu），别用 ReLU——原论文和 T5 都验证过 GELU 更稳
训练时 loss 突然飙升，大概率是 learning rate 或初始化问题
Transformer 对学习率极其敏感。MultiHeadAttention 中的 Q/K/V 投影矩阵若用默认 glorot_uniform 初始化，在大模型或小 batch 下容易梯度爆炸。典型现象是前 10 步 loss 从 10 降到 3，第 11 步跳到 200+，然后 nan。
快速排查项：
- 用学习率预热（warmup）：前 1000 步线性从 0 升到峰值，峰值推荐
1e-4（base size）或3e-5（large） - 检查
MultiHeadAttention的kernel_initializer，换成tf.keras.initializers.VarianceScaling(scale=0.125, mode="fan_avg", distribution="uniform")（模仿原始实现） - 验证集上 loss 不降？先关掉所有 dropout，确认模型能过拟合单个 batch——否则是结构 bug，不是数据问题
位置编码没加、mask 传反、LayerNorm axis 错、学习率没 warmup——这四个点占了 80% 的调试时间。别的可以慢慢调，这几个得优先盯住。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
