应直接使用tf.keras.layers.multiheadattention而非手写，因其自动处理head拆分、scale因子、mask应用位置及维度对齐；positionalencoding需用add_weight创建trainable=false的tf.variable；decoder两层attention须区分causal mask与encoder padding mask；encoder输出须显式转tf.float32以避免混合精度训练失败。
直接用 TensorFlow 原生 API 从零搭一个能训、能跑的 Transformer，不是不行，但容易卡在 MultiHeadAttention 的输入 shape、LayerNormalization 的 axis 设置、以及位置编码的广播对齐上——尤其当你想复现原始论文（Vaswani et al., 2017）的细节时。
为什么别手写 MultiHeadAttention 层
TensorFlow 已内置 tf.keras.layers.MultiHeadAttention，它默认支持 query/key/value 分离输入、带 mask 的因果注意力（用于解码器），且自动处理 head 拆分与拼接。手写容易出错的点包括：
-
q @ k^T后的 scale 因子漏除sqrt(d_k)，导致 softmax 梯度饱和 - mask 应用位置错误：应在
softmax前加-1e9 * (1 - mask)，而不是乘法掩码 - batch 维度和 seq_len 维度混淆，尤其在 decoder self-attention 中，
causal_mask必须是上三角
推荐直接用官方层，并显式传入 num_heads 和 key_dim（注意：key_dim * num_heads == embed_dim，否则会报 ValueError: The last dimension of query should be equal to key_dim * num_heads）。
PositionalEncoding 必须用 tf.Variable 而非 tf.constant
原始论文用正弦/余弦函数生成固定位置编码，但若你用 tf.constant 初始化，在模型保存/加载或分布式训练中会丢失该张量；更严重的是，它无法参与 eager execution 下的梯度追踪（虽然位置编码本身不更新，但作为 layer 的一部分，需被正确 trace）。正确做法是：
- 在
__init__中创建self.pos_encoding = self.add_weight(..., trainable=False) - shape 设为
(1, max_seq_len, embed_dim)，开头加1是为了后续能广播到任意 batch size - 初始化值用
np.sin/cos计算后转tf.convert_to_tensor，再赋给 weight
否则你会遇到 AttributeError: 'Tensor' object has no attribute 'numpy' 或训练时 loss 不下降——因为位置信息根本没进模型。
Decoder 的两层 attention 必须区分 mask 类型
Transformer decoder 包含两个 MultiHeadAttention 子层：第一个是 causal self-attention（仅看左侧 token），第二个是 encoder-decoder attention（key/value 来自 encoder 输出）。它们的 mask 完全不同：
- 第一层：用
tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)生成下三角 mask，再转 bool - 第二层：mask 来自 encoder 的 padding mask（shape
(batch, 1, 1, seq_len_enc)），不能复用 decoder 自身的 causal mask - 若第二层也喂了 causal mask，模型会“假装看不见” encoder 的全部输出，loss 爆涨且预测全为 padding token
典型错误是把同一个 look_ahead_mask 传给两层——检查 call() 中两处 attention_mask= 参数是否指向不同变量。
最易被忽略的是：encoder 的输出必须和 decoder 的输入保持相同 dtype（比如都为 tf.float32），否则 MultiHeadAttention 在混合精度训练（mixed_float16）下会静默失败，loss 为 NaN——连 warning 都不抛。务必在 encoder 输出后加 tf.cast(..., tf.float32) 再传入 decoder。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
