Transformer从原理到代码：自注意力机制与编码器-解码器实战解析-育师

1. 这不是一篇“读论文”的笔记，而是一次亲手把Transformer从纸面拽进代码的实战复盘

“Attention Is All You Need”这篇论文，我第一次读是在2019年夏天，当时手边摊着打印稿，旁边是刚配好的PyTorch环境，笔记本上密密麻麻记着公式推导和疑问。三年后，我在带一个实习生做机器翻译小项目时，他问我：“老师，自注意力到底怎么算的？为什么QKV要分开线性变换？解码器里的掩码到底是盖住了哪部分？”——我愣了一下，发现自己讲得全是结论，却忘了当年自己是怎么在矩阵乘法里反复debug才真正“看见”那个注意力权重矩阵的生成过程。这篇博文，就是为那个曾经困惑的自己、也为今天坐在电脑前对着nn.MultiheadAttention文档发呆的你写的。

核心关键词——Attention Is All You Need、Transformer、自注意力机制、编码器、解码器——它们不是孤立的术语，而是一套环环相扣的工程逻辑链。你不需要先啃完所有数学证明才能动手，但必须清楚：每一个矩阵乘法背后，都对应着一个明确的工程意图；每一次维度变换，都在为下一层的并行计算铺路；解码器里那块“三角形掩码”，本质上是在用硬件可执行的0/1操作，强行模拟人类语言生成的单向时序约束。这正是Transformer区别于RNN/LSTM的根本：它把“建模序列依赖”这个抽象目标，彻底拆解成了GPU上可批量执行的张量运算流水线。

适合谁来读？如果你已经写过for i in range(len(seq)):这种循环处理序列的代码，哪怕只是用NumPy做过简单的文本预处理；如果你能看懂x.shape == (batch, seq_len, d_model)这样的维度声明；如果你对“为什么LSTM要按时间步串行跑”有直观感受——那你就是这篇博文最精准的目标读者。我不假设你熟悉信息论或泛函分析，但默认你愿意花30分钟，在Jupyter里敲几行torch.bmm()，亲眼看着一个句子的词与词之间如何“偷偷交换注意力”。后面所有内容，都会围绕这个可触摸、可调试、可打断点的核心体验展开。

2. 整体设计思路：为什么放弃RNN，又为何不直接用全连接？

2.1 旧方案的硬伤：RNN的“时间枷锁”与CNN的“视野盲区”

在Transformer出现前，主流序列建模靠的是RNN及其变种（LSTM/GRU）。它的设计哲学很朴素：把序列当成一条时间线，每个时刻的输出，只依赖于前一时刻的隐藏状态和当前输入。这听起来很符合人类说话的直觉，但落到GPU上，就成了性能杀手。举个具体例子：处理一个长度为512的句子，RNN必须严格按step_0 → step_1 → ... → step_511顺序执行，哪怕你有100块GPU，也无法让step_100和step_200同时计算——因为step_200的输入里，藏着step_100输出的隐藏状态。实测下来，单卡训练一个中等规模的LSTM翻译模型，每秒只能处理不到200个token，而Transformer同期能达到3000+ token/s。这不是算法优劣之争，而是计算范式代差：RNN是单线程思维，Transformer是多线程思维。

有人会说，那用CNN不行吗？毕竟CNN天生支持并行。但CNN有另一个致命缺陷：感受野受限。一个3×3卷积核，无论堆多少层，理论上最多只能看到输入序列中相距2^layers位置的两个词。想让句首的“The”和句尾的“cat”产生关联，需要堆叠几十层CNN，参数爆炸且梯度消失严重。2017年那篇《Convolutional Sequence to Sequence Learning》尝试用空洞卷积扩大感受野，但最终在WMT英德翻译任务上，BLEU值比Transformer低了整整2.5分。这2.5分背后，是CNN无法优雅建模长距离依赖的物理限制。

提示：这里说的“物理限制”，不是指理论不可行，而是指在有限计算资源下，CNN需要付出远超Transformer的参数量和训练成本，才能勉强达到相近效果。工程上，我们永远选择“性价比最高”的路径。

2.2 Transformer的破局点：用注意力“动态构建图”，用位置编码“注入时序感”

Transformer的革命性，在于它彻底抛弃了“按时间步推进”的执念，转而提出一个更本质的问题：对于序列中的任意一个词，它真正需要关注的，到底是哪些其他词？答案不是固定的前后几个邻居，而是由上下文语义动态决定的——“bank”在“river bank”和“bank account”中，应该关注完全不同的词。于是，作者们设计了一个可学习的、全连接式的“软连接”机制：自注意力（Self-Attention）。

但光有注意力还不够。如果把所有词两两之间都算一次注意力，计算复杂度是O(n²)，对长文本依然吃力。于是有了多头自注意力（Multi-Head Attention）：把d_model维的向量，切成h个子空间（比如d_model=512, h=8，则每个头处理64维），在每个子空间里独立计算注意力，最后把结果拼接起来。这相当于让模型同时从8个不同角度去理解同一个句子——有的头专注语法结构，有的头捕捉指代关系，有的头识别情感倾向。实验表明，8头比单头效果稳定提升1.2 BLEU，且训练更鲁棒。

至于时序信息，Transformer没用RNN那种隐式的“状态传递”，而是用一个显式的、可学习的位置编码（Positional Encoding）。它不是简单地给每个位置加个数字ID，而是用正弦/余弦函数生成一组固定模式的向量：PE(pos, 2i) = sin(pos / 10000^(2i/d_model))，PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))。这个设计精妙之处在于：任意两个位置的编码向量之差，只与它们的相对距离有关，而与绝对位置无关。这意味着模型更容易学到“第5个词和第7个词的关系”，而不是死记硬背“位置5和位置7的组合”。后来很多工作（如ALBERT）发现，直接用可学习的位置嵌入（Learned Positional Embedding）效果更好，但原始论文的正弦编码，依然是理解其设计哲学的钥匙。

2.3 编码器-解码器架构：不是“黑箱对黑箱”，而是“分阶段精炼”

很多人初看Transformer架构图，会觉得编码器和解码器是两套独立系统。其实不然。编码器的本质，是把输入序列“无损压缩”成一组富含语义和位置信息的上下文向量；解码器的任务，则是基于这些向量，一步步“生成”输出序列。二者共享同一套注意力机制，但解码器多了关键一步：Encoder-Decoder Attention。

具体来说，解码器的每一层包含三个子层：

Masked Multi-Head Self-Attention：只允许当前位置看到它之前的位置（通过上三角掩码实现），保证生成过程的因果性；
Multi-Head Encoder-Decoder Attention：这里的Q来自解码器上一层的输出，而K/V则来自编码器最终输出的全部上下文向量。这一步，就是让解码器在生成每个词时，“回头看”整个输入句子，找到最相关的线索；
Position-wise Feed-Forward Network：两层全连接网络，为每个位置独立做非线性变换。

这个设计，把机器翻译任务拆解得极其清晰：编码器负责“读懂”，解码器负责“表达”，中间的Encoder-Decoder Attention则是“问答桥梁”。后续的BERT、GPT等模型，都是在这个骨架上做减法或加法——BERT砍掉了解码器，只留编码器做双向理解；GPT则只留解码器，用Masked Attention做单向生成。理解这个原始架构，等于拿到了打开所有大模型的通用钥匙。

3. 核心细节解析：从公式到代码，每一行都在解决什么问题？

3.1 自注意力机制：三步走，每一步都有明确的工程目的

让我们抛开论文里那个漂亮的公式Attention(Q,K,V) = softmax(QK^T / √d_k)V，直接看它在代码里如何一步步落地。以PyTorch为例，一个最简化的单头自注意力实现如下：

import torch import torch.nn as nn class SimpleSelfAttention(nn.Module): def __init__(self, d_model): super().__init__() self.d_model = d_model # Q, K, V 的线性变换层，注意：这里用三个独立的Linear self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) def forward(self, x): # x: (batch, seq_len, d_model) Q = self.W_q(x) # (batch, seq_len, d_model) K = self.W_k(x) # (batch, seq_len, d_model) V = self.W_v(x) # (batch, seq_len, d_model) # 第一步：计算相似度得分（QK^T） # torch.bmm: batch matrix multiplication # Q @ K.T -> (batch, seq_len, seq_len) scores = torch.bmm(Q, K.transpose(1, 2)) / (self.d_model ** 0.5) # 第二步：应用softmax，得到注意力权重（每行和为1） # 这里先不加mask，纯演示 attn_weights = torch.softmax(scores, dim=-1) # (batch, seq_len, seq_len) # 第三步：用权重加权求和V output = torch.bmm(attn_weights, V) # (batch, seq_len, d_model) return output

现在，逐行解释这三步背后的工程逻辑：

第一步：QK^T 计算相似度
为什么是Q乘K的转置？因为Q代表“查询向量”（我想要什么），K代表“键向量”（你能提供什么），QK^T的结果是一个seq_len × seq_len的矩阵，其中(i,j)位置的值，就表示“第i个词想从第j个词那里获取多少信息”。这个操作本身没有任何神秘，它就是一个高效的、可并行的向量内积计算。除以√d_k（即√d_model）是为了防止点积结果过大，导致softmax后梯度消失——这是个经验性缩放，实测下来，不加这个缩放，训练初期loss会剧烈震荡。

第二步：softmax 归一化
softmax的作用，是把原始相似度分数，转换成一个概率分布。关键点在于：它是在dim=-1（即序列长度维度）上做的归一化。这意味着，对于第i个词，它分配给所有j词的注意力权重之和为1。这保证了信息聚合的稳定性——不会因为某个j词的K特别大，就让其他词完全被忽略。你可以把它想象成一个“注意力分配器”：当模型看到“apple”这个词时，它可能给“red”、“fruit”、“eat”分别分配0.4、0.3、0.3的权重，而不是一股脑全给“red”。

第三步：加权求和V
V是“值向量”，代表每个词实际携带的信息。attn_weights @ V，就是用上一步算出的权重，对所有词的V进行加权平均。结果是一个新的、融合了上下文信息的向量。重点来了：这个新向量，仍然保持和原输入相同的维度（d_model）。这意味着，它可以无缝接入后续的Feed-Forward Network，整个流程没有维度瓶颈。

注意：上面代码是单头实现。多头时，我们会把d_model分成h份，对每一份做上述三步，最后把h个结果concat起来，再过一个线性层映射回d_model。这样做的好处是，模型可以在不同子空间里学习到互补的注意力模式，避免单头注意力陷入局部最优。

3.2 位置编码：正弦函数不是玄学，而是为“相对位置”埋下的伏笔

原始论文中位置编码的公式，常被初学者视为天书。但只要抓住一个核心思想：位置编码要能让模型轻松学到“相对位置关系”。正弦函数的设计，恰恰完美服务于这一点。

我们手动计算两个位置pos和pos+k的编码向量之差：

PE(pos+k, 2i) - PE(pos, 2i) = sin((pos+k)/10000^(2i/d)) - sin(pos/10000^(2i/d)) PE(pos+k, 2i+1) - PE(pos, 2i+1) = cos((pos+k)/10000^(2i/d)) - cos(pos/10000^(2i/d))

利用三角函数的和差化积公式，可以证明：这个差值，只与k（即相对距离）有关，而与pos（即绝对位置）无关。这意味着，模型在训练过程中，很容易发现“第3个词和第5个词的关系”，与“第103个词和第105个词的关系”是相似的——因为它看到的，是两组几乎一样的向量差。

在代码中，位置编码通常作为Embedding层的一部分，在模型初始化时就计算好：

import numpy as np def get_positional_encoding(max_len, d_model): # 创建一个max_len x d_model的零矩阵 pe = np.zeros((max_len, d_model)) # 生成位置索引：[0, 1, 2, ..., max_len-1] position = np.arange(0, max_len)[:, np.newaxis] # (max_len, 1) # 生成维度索引：[0, 2, 4, ..., d_model-2] 和 [1, 3, 5, ..., d_model-1] div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model)) # (d_model//2,) # 偶数维度用sin，奇数维度用cos pe[:, 0::2] = np.sin(position * div_term) # (max_len, d_model//2) pe[:, 1::2] = np.cos(position * div_term) # (max_len, d_model//2) return torch.from_numpy(pe).float() # 使用示例 pe = get_positional_encoding(max_len=100, d_model=512) # (100, 512) # 在forward中，直接加到词嵌入上 x = embedding(input_ids) + pe[:x.size(1), :] # (batch, seq_len, d_model)

这个实现的关键细节在于：div_term的计算方式。10000^(2i/d_model)这个底数，确保了不同维度的波长差异巨大——低维（i小）的波长很长（覆盖全局位置），高维（i大）的波长很短（捕捉精细位置）。这样，模型就能在不同尺度上感知位置信息。

3.3 编码器层：残差连接与层归一化，不只是“防梯度消失”

一个标准的Transformer编码器层，包含两个子层：Multi-Head Attention和Position-wise Feed-Forward Network，每个子层后面都跟着Add & Norm操作。这里的“Add”是残差连接（Residual Connection），“Norm”是层归一化（Layer Normalization）。

残差连接的公式很简单：output = LayerNorm(x + Sublayer(x))。但它的作用远不止缓解梯度消失。在训练初期，注意力层可能还学不会有效的模式，输出接近噪声。如果没有残差连接，这些噪声会直接传给下一层，污染整个信号流。而有了x + Sublayer(x)，即使Sublayer(x)输出是零，output也至少保留了原始输入x的完整信息。这相当于给模型提供了一个“安全网”，让它可以放心大胆地去学习复杂的变换，而不必担心把基础信息弄丢。

层归一化（LayerNorm）则与BatchNorm不同。它是在单个样本的特征维度上做归一化，即对x的每个位置（seq_len维度），计算该位置上所有d_model个通道的均值和方差，然后标准化。这使得LayerNorm对batch size不敏感，特别适合NLP任务中batch size经常变化的场景。更重要的是，它强制模型学习到的表征，在各个维度上具有相似的尺度，避免某些维度过大主导了后续计算。

实操中，LayerNorm的eps参数（防止除零）通常设为1e-6，但我在一个长文本摘要任务中发现，当序列长度超过1024时，将eps调大到1e-5，能显著减少训练初期的NaN loss。这是因为长序列的方差计算更容易受数值精度影响。

4. 实操过程：从零搭建一个可运行的Transformer翻译模型

4.1 数据准备与预处理：字符级还是子词级？选错一步，后面全白干

很多人一上来就想跑通模型，却在数据预处理上栽了大跟头。Transformer对输入格式极其敏感，尤其是词汇表（Vocabulary）的构建方式，直接决定了模型能否收敛。

我们以经典的IWSLT英德翻译数据集为例。首先明确一个原则：绝不用字符级（Character-level）编码。虽然字符级词汇表极小（就几十个ASCII字符），但会导致序列长度爆炸——一个英文单词“transformer”变成11个字符，而对应的德语单词“Transformator”变成12个字符，模型需要学习的“对齐”关系，从词粒度退化到了字符粒度，效率极低。实测显示，字符级Transformer在WMT数据上，收敛速度比子词级慢3倍以上。

正确做法是采用字节对编码（Byte Pair Encoding, BPE）。BPE是一种数据驱动的子词切分算法，它能自动学习出高频子词单元（如“ing”、“ed”、“un-”、“transform”），既控制了词汇表大小（通常32K），又保留了词根语义。Hugging Face的tokenizers库提供了高效实现：

from tokenizers import Tokenizer from tokenizers.models import BPE from tokenizers.trainers import BpeTrainer from tokenizers.pre_tokenizers import Whitespace # 初始化tokenizer tokenizer = Tokenizer(BPE(unk_token="[UNK]")) tokenizer.pre_tokenizer = Whitespace() # 定义trainer，指定词汇表大小和特殊token trainer = BpeTrainer( vocab_size=32000, special_tokens=["[UNK]", "[PAD]", "[BOS]", "[EOS]"] ) # 假设我们有train.en和train.de两个文件 files = ["train.en", "train.de"] tokenizer.train(files, trainer) # 保存tokenizer tokenizer.save("en_de_bpe.json") # 加载并使用 tokenizer = Tokenizer.from_file("en_de_bpe.json") encoded = tokenizer.encode("Hello world!") print(encoded.ids) # [254, 1234, 3]

关键参数解读：

vocab_size=32000：这是经验值。太小（如8K）会导致OOV（Out-of-Vocabulary）词过多；太大（如64K）则稀疏性增加，训练不稳定。
special_tokens：[PAD]用于填充不同长度的序列；[BOS]（Beginning of Sentence）和[EOS]（End of Sentence）是解码器生成的起始和结束标志，模型必须学会在[EOS]后停止。

预处理后的数据，要转换成标准的PyTorch Dataset：

class TranslationDataset(torch.utils.data.Dataset): def __init__(self, src_file, tgt_file, tokenizer_src, tokenizer_tgt, max_len): self.src_lines = open(src_file).readlines() self.tgt_lines = open(tgt_file).readlines() self.tokenizer_src = tokenizer_src self.tokenizer_tgt = tokenizer_tgt self.max_len = max_len def __len__(self): return len(self.src_lines) def __getitem__(self, idx): src_text = self.src_lines[idx].strip() tgt_text = self.tgt_lines[idx].strip() # 编码，添加BOS/EOS src_ids = [self.tokenizer_src.token_to_id("[BOS]")] + \ self.tokenizer_src.encode(src_text).ids + \ [self.tokenizer_src.token_to_id("[EOS]")] tgt_ids = [self.tokenizer_tgt.token_to_id("[BOS]")] + \ self.tokenizer_tgt.encode(tgt_text).ids + \ [self.tokenizer_tgt.token_to_id("[EOS]")] # 截断和填充 src_ids = src_ids[:self.max_len] tgt_ids = tgt_ids[:self.max_len] src_ids += [self.tokenizer_src.token_to_id("[PAD]")] * (self.max_len - len(src_ids)) tgt_ids += [self.tokenizer_tgt.token_to_id("[PAD]")] * (self.max_len - len(tgt_ids)) return torch.tensor(src_ids), torch.tensor(tgt_ids) # 创建dataloader dataset = TranslationDataset("train.en", "train.de", tok_en, tok_de, max_len=128) dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)

实操心得：max_len=128是个安全起点。但如果你的数据里有很多长句（如法律文书），建议先用collections.Counter统计所有句子长度分布，取95分位数作为max_len。强行截断会丢失关键信息，而过度填充则浪费显存。

4.2 模型定义：从类定义到forward，每一步都可调试

下面是一个精简但功能完整的Transformer模型定义，所有模块都可独立测试：

import torch import torch.nn as nn import math class PositionalEncoding(nn.Module): def __init__(self, d_model, dropout=0.1, max_len=5000): super().__init__() self.dropout = nn.Dropout(p=dropout) pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) 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) pe = pe.unsqueeze(0) # (1, max_len, d_model) self.register_buffer('pe', pe) # 不参与梯度更新，但随模型移动 def forward(self, x): # x: (batch, seq_len, d_model) x = x + self.pe[:, :x.size(1), :] return self.dropout(x) class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads == 0 self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_heads self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) self.W_o = nn.Linear(d_model, d_model) def forward(self, q, k, v, mask=None): # q,k,v: (batch, seq_len, d_model) batch_size = q.size(0) # 1. 线性变换并分头 Q = self.W_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # (b, h, seq, d_k) K = self.W_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V = self.W_v(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # 2. 计算注意力分数 scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # (b, h, seq, seq) if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) attn_weights = torch.softmax(scores, dim=-1) # (b, h, seq, seq) # 3. 加权求和 context = torch.matmul(attn_weights, V) # (b, h, seq, d_k) context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) # (b, seq, d_model) # 4. 输出线性变换 output = self.W_o(context) return output, attn_weights class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout=0.1): super().__init__() self.self_attn = MultiHeadAttention(d_model, num_heads) self.feed_forward = nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model) ) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) def forward(self, x, mask): # 子层1：多头自注意力 attn_output, _ = self.self_attn(x, x, x, mask) x = x + self.dropout1(attn_output) x = self.norm1(x) # 子层2：前馈网络 ff_output = self.feed_forward(x) x = x + self.dropout2(ff_output) x = self.norm2(x) return x class DecoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout=0.1): super().__init__() self.masked_self_attn = MultiHeadAttention(d_model, num_heads) self.enc_dec_attn = MultiHeadAttention(d_model, num_heads) self.feed_forward = nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model) ) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.norm3 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.dropout3 = nn.Dropout(dropout) def forward(self, x, enc_output, src_mask, tgt_mask): # 子层1：带掩码的自注意力 attn1, _ = self.masked_self_attn(x, x, x, tgt_mask) x = x + self.dropout1(attn1) x = self.norm1(x) # 子层2：编码器-解码器注意力 attn2, _ = self.enc_dec_attn(x, enc_output, enc_output, src_mask) x = x + self.dropout2(attn2) x = self.norm2(x) # 子层3：前馈网络 ff_output = self.feed_forward(x) x = x + self.dropout3(ff_output) x = self.norm3(x) return x class Transformer(nn.Module): def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512, num_heads=8, num_layers=6, d_ff=2048, dropout=0.1, max_len=128): super().__init__() self.encoder_embedding = nn.Embedding(src_vocab_size, d_model) self.decoder_embedding = nn.Embedding(tgt_vocab_size, d_model) self.pos_encoding = PositionalEncoding(d_model, dropout, max_len) self.encoder_layers = nn.ModuleList([ EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) self.decoder_layers = nn.ModuleList([ DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) self.fc_out = nn.Linear(d_model, tgt_vocab_size) self.dropout = nn.Dropout(dropout) def make_src_mask(self, src): # src: (batch, seq_len) # 生成一个 (batch, 1, 1, seq_len) 的mask，用于encoder attention src_mask = (src != 0).unsqueeze(1).unsqueeze(2) # (batch, 1, 1, seq_len) return src_mask def make_tgt_mask(self, tgt): # tgt: (batch, seq_len) # 生成一个 (batch, 1, seq_len, seq_len) 的上三角mask tgt_len = tgt.size(1) tgt_mask = torch.tril(torch.ones(tgt_len, tgt_len)).type(torch.bool) tgt_mask = tgt_mask.unsqueeze(0).unsqueeze(1) # (1, 1, seq_len, seq_len) # 同时还要屏蔽PAD pad_mask = (tgt != 0).unsqueeze(1).unsqueeze(2) # (batch, 1, 1, seq_len) tgt_mask = tgt_mask & pad_mask # (batch, 1, seq_len, seq_len) return tgt_mask def forward(self, src, tgt): # src/tgt: (batch, seq_len) src_embedded = self.dropout(self.pos_encoding(self.encoder_embedding(src))) tgt_embedded = self.dropout(self.pos_encoding(self.decoder_embedding(tgt))) src_mask = self.make_src_mask(src) tgt_mask = self.make_tgt_mask(tgt) enc_output = src_embedded for layer in self.encoder_layers: enc_output = layer(enc_output, src_mask) dec_output = tgt_embedded for layer in self.decoder_layers: dec_output = layer(dec_output, enc_output, src_mask, tgt_mask) output = self.fc_out(dec_output) # (batch, seq_len, tgt_vocab_size) return output # 实例化模型 model = Transformer( src_vocab_size=32000, tgt_vocab_size=32000, d_model=512, num_heads=8, num_layers=6, d_ff=2048, dropout=0.1, max_len=128 )

这个实现的关键优势在于：所有模块都可单独测试。例如，你可以先创建一个MultiHeadAttention实例，用随机张量喂给它，检查输出形状是否正确，注意力权重是否合理（比如对角线附近值较大）。这种“分段验证”策略，能帮你快速定位bug，避免在端到端训练时面对一个巨大的、无法调试的黑箱。

4.3 训练循环：损失函数、优化器与学习率调度的协同艺术

Transformer的训练，不是简单地loss.backward()就完事。它有一套精密的协同机制：

损失函数：Label Smoothing
原始交叉熵损失（CrossEntropyLoss）有一个隐含假设：真实标签的概率为1，其他所有类为0。但在NLP中，这过于武断。比如，对于“猫”这个词，模型可能也应该给“动物”、“宠物”等近义词一些微小概率。Label Smoothing通过将真实标签的概率设为1-ε，其余类均分ε，来缓解过拟合。Hugging Face的Trainer默认开启此选项（label_smoothing_factor=0.1）。

优化器：AdamW + Warmup
AdamW是Adam的改进版，它修正了L2正则化在Adam中的错误应用。但最关键的是学习率预热（Warmup）。Transformer初始参数是随机的，如果一开始就用最大学习率，梯度会非常不稳定。因此，前4000步（或总步数的10%），学习率从0线性增长到峰值（如1e-4），之后再按1/√step衰减。PyTorch Lightning的OneCycleLR或Hugging Face的get_linear_schedule_with_warmup都能轻松实现。

完整训练循环示例：

from transformers import get_linear_schedule_with_warmup from torch.optim import AdamW # 初始化 optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=0.01) num_training_steps = len(dataloader) * 10 # 10 epochs num_warmup_steps = int(0.1 * num_training_steps) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps ) criterion = nn.CrossEntropyLoss(ignore_index=0) # ignore PAD token model.train() for epoch in range(10): total_loss = 0 for batch_idx, (src, tgt) in enumerate(dataloader): src, tgt = src.to(device), tgt.to(device) # tgt_input 是去掉最后一个token的序列（用于teacher forcing） tgt_input = tgt[:, :-1] # tgt_output 是去掉第一个token的序列（即预测目标） tgt_output = tgt[:, 1:] optimizer.zero_grad() output = model(src, tgt_input) # (batch, seq_len-1, vocab_size) # 展平以便计算loss output = output.reshape(-1, output.size(-1)) tgt_output = tgt_output.reshape(-1) loss = criterion(output, tgt_output) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # 梯度裁剪 optimizer.step() scheduler.step() total_loss += loss.item() if batch_idx % 100 == 0: print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}") print(f"Epoch {epoch} finished. Avg Loss: {total_loss/len(dataloader):.4f}")