PyTorch 复现 Transformer听起来像是又一个“照着论文敲代码”的练手项目但真正走完 TR3 周这一程我才发现它远比想象中更考验对细节的掌控。如果你也想从零开始搭建一个完整的 Transformer 翻译模型而不是直接调用现成库里的封装类这篇内容大概能帮你少走不少弯路。我会把这一周里踩过的坑、验证过的想法、还有最终跑通的代码结构一起整理出来通篇都是实操向的适合已经会 PyTorch 基础操作、但还没完整手写过 Attention 机制的读者。1. 为什么我坚持要徒手复现 Transformer而不是直接 pip install1.1 复现的价值不在“能用”而在“知其所以然”做 NLP 相关工作的同人应该都有这种体验transformers库一行from transformers import TransformerEncoder就能拿到一个表现还不错的预训练模型谁还会去关心论文里Multi-Head Attention到底是怎么把Q、K、V拆开又合并的呢但这次项目 TR3 周给我的最大感受是复现的价值不在于你最后跑出多好的 BLEU而在于你被迫去面对那些“文档替你回答”的问题。比如为什么Scaled Dot-Product Attention要除以根号d_k不除行不行Decoder的 Mask 到底要 mask 哪些位置padding 的 token 又要怎么隔离Embedding和Positional Encoding是 concat 还是 add训练时Decoder输入是完整的目标句子那推理时为什么必须一个词一个词地往左移位如果你只写过调用库的代码这些问题永远只会停留在“好像有印象”的水平。而当你真的复现一遍逐个模块去查、去推、去打印中间张量形状你就会建立起属于自己的“为什么”而不是“别人告诉我这么做”。1.2 复现对象的选型为什么拿“序列到序列”翻译任务做实验床Transformer 原论文是拿机器翻译任务做验证的。我也照这个思路在本地模拟了一个中文到英文的小型平行语料实验环境用虚构的简化数据来跑通流程。为什么选翻译任务而不是文本分类因为翻译任务天然需要完整的 Encoder-Decoder 架构能强迫你实现编码器的双向上下文建模虽然标准 Transformer 每个位置都”看到“全局解码器的因果/自回归约束不能偷看未来词交叉注意力机制Decoder 如何从 Encoder 的输出中提取信息如果你只做分类任务那大概只需要一个 EncoderDecoder 部分的知识就完全浪费了。所以我的建议是哪怕你最终的目标不是翻译第一版复现也尽量把 Encoder 和 Decoder 都写全。后面再裁剪起来很容易但一开始就缺一半理解会不完整。1.3 环境与工程准备我的实验环境如下你可以直接参考项目配置Python3.10PyTorch2.1.xGPU单卡显存 8GB 以上建议数据规模数千对平行语料足够验证逻辑不追求精度词表大小源语言 / 目标语言各约 8000-15000主要依赖torch,numpy,jieba中文分词英文端可用spacy或空格分词提示如果你的显存比较紧张可以把d_model从论文默认的 512 降到 256num_heads保持 8batch_size调小 16-32。重要的是流程和代码逻辑正确参数量级是次要的。2. 三个 Mask 的来源与实现细节这是复现的第一道鬼门关2.1 为什么需要 Padding Mask让模型学会“忽略空气”序列数据几乎不可能每条都一样长。我们通常会把一个 batch 内的句子 pad 到相同的长度。但问题是padding 位置是一个“假 token”它不该参与注意力计算否则模型会把注意力分配到无意义的位置上。在 Attention 分数计算完之后我们会生成一个src_key_padding_mask形状是(batch, seq_len)值为True的位置代表“该位置的 key 无效”。在传入scaled_dot_product_attention时这个 mask 会被转换成(batch, 1, 1, seq_len)的形状加到注意力分数矩阵上在这些无效位置上加一个极大的负数值我习惯用-1e9也有不少人直接用float(-inf)再经过 softmax这些位置的注意力权重就会趋近于 0。代码片段如下def masked_fill_scores(scores, key_padding_mask, value-1e9): # key_padding_mask: (batch, seq_len_k), Trueignore if key_padding_mask is not None: scores scores.masked_fill( key_padding_mask.unsqueeze(1).unsqueeze(2), value ) return scores这里大家要记住一个原则key_padding_mask影响的是key 侧的无效位置所以它的维度要跟seq_len_k对齐而不是seq_len_q。2.2 因果 MaskDecoder 必须学会“不偷看未来”Decoder 在预测第t个词的时候只能使用 t 时刻之前包括 t 时刻的信息不能使用 t1、t2 这些未来位置的 token。这个约束叫因果性causality。在训练阶段我们一次性把完整的目标句子从句子起始符开始喂给 Decoder。如果不用 mask 遮挡未来位置模型就可以“作弊”——直接从后面的位置抄答案。因此我们需要一个上三角为 True 的 mask形状是(seq_len_q, seq_len_k)def generate_causal_mask(seq_len): mask torch.triu(torch.ones(seq_len, seq_len), diagonal1).bool() return mask # True 表示该位置需要被遮挡稍微强调一下torch.triu(..., diagonal1)出来的矩阵下三角包括对角线都是 False上三角是 True。它加在注意力分数上之后softmax 后每个 query 只能看到自己当前位置及之前的 key。2.3 为什么是除以根号 d_k一个容易被跳过的数学细节原论文中用的是scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)为什么除以sqrt(d_k)这得从 softmax 的数值稳定性说起。假设q和k的每个分量都是均值为 0、方差为 1 的随机变量那么它们点积的方差就会是d_k因为方差是各维度方差之和。当d_k偏大时比如 512点积的数值范围会变得非常大这会导致 softmax 进入梯度饱和区即经过 softmax 后某个位置的概率极其接近 1其余位置几乎为 0梯度随之变得极小。除以sqrt(d_k)相当于把点积的方差重新拉回到 1 附近让 softmax 保持在梯度比较充足的工作区间。我试过不除这个因子直接训练模型确实能收敛但速度和稳定性明显下降特别是在d_model512的时候前几百步 loss 曲线会非常颠簸。3. 模块搭建从嵌入层到多头注意力再到编码器与解码器3.1 Embedding 和 Positional Encoding为什么是相加而不是拼接每个 token 在输入模型之前首先要查一个可训练的 embedding 表得到形状为(seq_len, d_model)的向量序列。但 Transformer 没有 RNN 那种天然的顺序感为了让模型感知 token 的相对位置关系需要额外注入位置信息。原论文使用了一组正弦位置编码def get_sinusoid_encoding(seq_len, d_model): pe torch.zeros(seq_len, d_model) position torch.arange(0, seq_len).unsqueeze(1).float() div_term torch.exp( torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model) ) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) return pe.unsqueeze(0) # (1, seq_len, d_model)然后在forward里直接相加x self.embedding(tokens) * math.sqrt(self.d_model) x x self.pe[:, :x.size(1)]这里有个初学者容易忽略的细节论文里对 embedding 做了sqrt(d_model)缩放。原因是 embedding 的权重方差通常是 1 左右而位置编码的取值范围只有 [-1, 1]两者直接相加时embedding 的影响力会被稀释。乘上sqrt(d_model)之后两者的量级更匹配模型能更容易捕捉 token 本身的语义特征。至于为什么是 add 而不是 concat我的理解是concat 会把特征维度直接翻倍模型参数和计算量随之增加而实验证明加法已经足够让模型区分不同位置的 token。3.2 多头注意力的核心逻辑分头、缩放、加权、合并这是复现过程中最容易“看着会了一写就错”的部分。很多人包括我自己第一版代码在把Q、K、V拆成多个头的时候维度对不上或者 softmax 作用的维度选错了。我建议你把“多头注意力”在脑中拆成四个阶段来写线性投影输入x经过W_Q / W_K / W_V三个权重矩阵得到Q、K、V。这一步的形状变化是(batch, seq_len, d_model) - (batch, seq_len, d_model)。拆分头把最后一维d_model分成num_heads * head_dim。我习惯用view和transpose组合def split_heads(x, num_heads): batch_size, seq_len, d_model x.size() head_dim d_model // num_heads x x.view(batch_size, seq_len, num_heads, head_dim) x x.transpose(1, 2) # (batch, num_heads, seq_len, head_dim) return x这里为什么要transpose(1, 2)因为我们要对每个头独立计算注意力希望把num_heads变成 batch 维度的“平行世界”让矩阵乘法同时作用于所有头。缩放点积注意力scores Q K^T / sqrt(head_dim)然后对scores在最后一维key 的长度方向做 softmax。注意 softmax 的维度是seq_len_k不是seq_len_q。加权求和attn_output scores V然后把所有头重新拼接回d_model维度再经过一个输出投影矩阵W_O。整个多头注意力块在代码里的样子大致是这样class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads, dropout0.1): super().__init__() assert d_model % num_heads 0 self.d_model d_model self.num_heads num_heads self.head_dim d_model // num_heads self.wq nn.Linear(d_model, d_model) self.wk nn.Linear(d_model, d_model) self.wv nn.Linear(d_model, d_model) self.wo nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) self.scale math.sqrt(self.head_dim) def forward(self, q, k, v, maskNone): batch_size q.size(0) q self.split_heads(self.wq(q)) # (batch, heads, q_len, head_dim) k self.split_heads(self.wk(k)) v self.split_heads(self.wv(v)) scores torch.matmul(q, k.transpose(-2, -1)) / self.scale if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn torch.softmax(scores, dim-1) attn self.dropout(attn) out torch.matmul(attn, v) out out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) return self.wo(out)这里要多说一句很多人的 bug 出现在 mask 的形状上。如果你的mask是(batch, seq_len_q, seq_len_k)想直接用到scores上需要先unsqueeze(1)变成(batch, 1, seq_len_q, seq_len_k)才能广播。如果没做这一步运行时不会报错但结果会静默出错模型完全学不到结构。这种题最可怕。3.3 残差连接与 LayerNorm 的顺序Post-LN 还是 Pre-LN论文原版用的是 Post-LN结构是x → Attention → Dropout → 残差相加 → LayerNorm。这个结构写起来很顺class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout): super().__init__() self.self_attn MultiHeadAttention(d_model, num_heads, dropout) self.feed_forward FeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, mask): x x self.dropout(self.self_attn(x, x, x, mask)) x self.norm1(x) x x self.dropout(self.feed_forward(x)) x self.norm2(x) return x但有经验的人可能知道现在很多新模型改用 Pre-LN即x_norm LayerNorm(x)然后x x sublayer(x_norm)。Pre-LN 的好处是梯度传播路径更干净深层堆叠时更容易训练但收敛后的精度有时略低于 Post-LN。我的建议是复现时先用 Post-LN 贴合原论文如果发现深层训练不稳定再改成 Pre-LN 对比一下这也是复现过程中很好的一种“实验素养”。3.4 Feed-Forward 和激活函数FFN 部分比较直接两个线性层中间夹一个 ReLU原论文或 GELU现在更常用。内部维度d_ff一般是d_model的 4 倍。要注意的是输入输出维度保持一致中间维度只在层内存在。class FeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(dropout) def forward(self, x): return self.linear2(self.dropout(F.relu(self.linear1(x))))如果你想把代码扩展成 BERT 风格可以把 ReLU 换成 GELU。但对于复现原版 TransformerReLU 就够了。3.5 把 Layer 串成 Encoder 和 DecoderEncoder 就是多层EncoderLayer的堆叠。Decoder 稍微复杂一点它包含三个子层第一层是Masked Self-Attention也就是 Self-Attention 加上了因果 mask同时也要处理 padding mask第二层是Cross-AttentionQ来自 Decoder 上一层的输出K和V来自 Encoder 的输出第三层是 FFN。在 Cross-Attention 中我们要把 Encoder 输出的memory作为k和v同时 Encoder 输出对应的key_padding_mask也要传进来避免 Decoder 注意到 Encoder 侧的 padding 位置。class DecoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout): super().__init__() self.masked_attn MultiHeadAttention(d_model, num_heads, dropout) self.cross_attn MultiHeadAttention(d_model, num_heads, dropout) self.ffn FeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.norm3 nn.LayerNorm(d_model) def forward(self, x, memory, src_mask, tgt_mask): x x self.dropout(self.masked_attn(x, x, x, tgt_mask)) x self.norm1(x) x x self.dropout(self.cross_attn(x, memory, memory, src_mask)) x self.norm2(x) x x self.dropout(self.ffn(x)) x self.norm3(x) return x注意这里tgt_mask是一个组合 mask即causal_mask与tgt_key_padding_mask的逻辑或关系。你可以直接在生成时做一个联合处理或者在 forward 里把两个 mask 先合并再送入 Attention。4. 训练循环里最容易出错的三个细节损失函数、标签平滑与学习率预热4.1 损失函数为什么用 CrossEntropy但要 ignore_index输出层是一个nn.Linear(d_model, vocab_size)把 Decoder 最后一层的表征映射到词表大小的 logits。Loss 是标准的交叉熵。这里的关键是目标序列里被 padding 的部分不应该计算 loss。PyTorch 的nn.CrossEntropyLoss提供了ignore_index参数。我建议把 padding 的 index 设为0词表里专门用一个位置给pad或者显式传ignore_indexpad_idx。另外这里的 label 不是完整的目标句子而是右移一位的目标句子。也就是说Decoder 第t步的输入是target[:, :-1]监督标签是target[:, 1:]。这样才能做到“用已知的前 t-1 个词预测第 t 个词”。logits model(src, tgt[:, :-1]) # (batch, tgt_len-1, vocab_size) loss criterion( logits.contiguous().view(-1, vocab_size), tgt[:, 1:].contiguous().view(-1) )4.2 标签平滑不是玄学是让模型“谦虚一点”原论文里用了label_smoothing0.1。它的思路很简单真实标签 one-hot 是[0, 1, 0, ...]现在把它变成[0.1/vocab_size, 0.90.1/vocab_size, 0.1/vocab_size, ...]相当于给模型一点容错空间让它不要过度相信训练集中的硬标签。我实测过这种写法在小规模语料上能比较明显缓解过拟合BLEU 会稍微高一点虽然整体绝对值不大。实现可以用nn.CrossEntropyLoss(label_smoothing0.1)不需要自己手动改标签。4.3 学习率预热Transformer 训练稳定的下半场关键原论文使用的优化器是 Adam但 beta 参数和默认值略有不同。它还设计了所谓的 Noam 学习率策略def lr_lambda(step, d_model512, warmup_steps4000): if step 0: step 1 return (d_model ** -0.5) * min(step ** -0.5, step * (warmup_steps ** -1.5))这个学习率曲线前 4000 步是线性上升的之后缓慢下降。原理上warmup 是为了让模型在最开始少受随机初始化带来的梯度过大的影响先稳定住参数更新方向再逐渐放大步长。如果不做 warmup很多情况下模型初期 loss 震荡很剧烈特别是深层的 Transformer。不过我也要说明一点现在很多开源实现直接使用torch.optim.lr_scheduler.OneCycleLR或者CosineAnnealingLR也能达到类似效果。复现阶段我建议你还是按原论文的 Noam 策略来这样后续做实验对照时能有一个统一的基准。4.4 完整训练循环的骨架放一个精简版的训练循环骨架供参考model Transformer(...).to(device) criterion nn.CrossEntropyLoss(ignore_indexpad_idx, label_smoothing0.1) optimizer torch.optim.Adam(model.parameters(), betas(0.9, 0.98), eps1e-9) scheduler LambdaLR(optimizer, lr_lambda) model.train() for epoch in range(epochs): for batch in data_loader: src, tgt batch optimizer.zero_grad() logits model(src, tgt[:, :-1]) loss criterion(logits.reshape(-1, vocab_size), tgt[:, 1:].reshape(-1)) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step()说一下梯度裁剪。虽然 Noam 学习率能避免很多发散问题但翻译任务里偶尔还是会出现“loss 突然变得很大”的情况。梯度裁剪就是为了防止单步更新过大导致参数直接崩掉。我建议 max_norm 设在1.0到5.0之间。5. 越过理论之后五个复现路上最容易翻车的实操细节5.1 位置编码必须加到 Encoder 和 Decoder 的输入而不是只在某一端我第一版偷懒只在 Encoder 上加位置编码Decoder 没加结果训练 loss 能降但生成的句子完全没有语序经常出现词语全部乱序的“炒单词”现象。后来才意识到Decoder 的输入照样需要位置信息来告诉它“目标词和源词”之间的顺序关系。补上之后效果正常了很多。5.2 训练和推理阶段 Decoder 的用法完全不同别把输入形状搞混训练时 Decoder 是一次性输入整个目标句子右移后用 causal mask 掩盖未来信息推理时没有目标句子只能用自回归方式先把[起始符]喂进去得到第一个词再把这个词拼到输入末尾继续得到下一个词。这个切换过程我最开始也没处理好把训练用的「教师强制」代码直接搬到推理阶段结果生成结果乱七八糟还好后来做了个简单的 greedy decode 才跑通。如果你打算做 beam search更要小心处理结束符和最大生成长度这两个边界条件。5.3 打印中间张量的形状是一个好习惯复现过程中我会在每个关键模块的 forward 里临时打印张量形状确认没有维度问题。特别是多头注意力的scores形状以及 Cross-Attention 里memory的 seq_len 和 Decoder 输入的 seq_len 不一致的情况。这是个笨办法但在复现初期非常管用。5.4 尝试过小的数据集时别急着上超大规模词表这里我自己吃过大亏。一开始直接用了预训练词表 3 万 的配置去跑几千条模拟数据词表里大量词根本没法学到靠谱的 embedding反而拖累整体训练。后面把词表缩小到几千每个词至少出现若干次模型才真正“学得动”。如果你想快速验证代码正确性控制词表大小比控制模型层数更能加速迭代。5.5 很早就要考虑推理路径别等训练收敛之后才想起生成也就是说训练脚本不要写成一个只能算 loss 的纯训练模式。我建议你从第一天开始就写一个简单的 greedy inference 函数每训练若干步拿一条验证集样本生成一次肉眼看输出是不是逐渐变得像话。这个做法能让你在早期就发现“模型在背诵数据”这类潜在问题而不是等到全部训练完才发现生成结果不可用。写在最后我的一点个人体会TR3 周的项目做下来我最强烈的感受是Transformer 值得手动复现一次但不要指望一次就能跑得很顺。我的第一版调试了整整一天最后发现只是一个 mask 拼接时unsqueeze维度放错位置导致所有注意力都失效了这种错误如果不打印中间形状根本不可能发现。后面我把测试集生成结果打印出来才一点点找到合理的方向。如果你准备开始这个项目我的建议是不要急着追求 BLEU 分数先把代码跑通把训练 loss 降下来再慢慢调参、加 trick、做对照实验。一个能跑通的复现比一份看起来很精美但没有亲自调试过的代码价值高得多。祝你也享受到这种“一个模块一个模块啃下来”的踏实感觉。