Transformer 这个架构从 2017 年那篇著名的论文开始几乎重塑了整个自然语言处理领域的玩法。现在你随便打开一个翻译工具、对话助手或者文本摘要功能底层大概率都跑着 Transformer 的某个变体。但很多人第一次翻开论文或者看源码的时候会被里面多头注意力、位置编码、残差连接这些概念绕得头晕。我当初也是这么过来的所以这篇文章打算把 Encoder 这一块彻底拆开揉碎从设计动机到代码实现再到实际调试中踩过的坑全部讲清楚。不管你是刚入门 NLP 的学生还是想从 RNN 迁移到 Transformer 的工程师看完应该都能自己动手搭一个能跑的 Encoder 出来。1. 为什么 Encoder 值得单独拿出来讲1.1 Encoder 在整个架构里的角色定位Transformer 最初是为机器翻译任务设计的整体是一个 Encoder-Decoder 结构。Encoder 负责把输入序列比如一句中文编码成一堆向量表示Decoder 再基于这些向量逐个生成目标序列比如对应的英文。你可以把 Encoder 想象成一个信息压缩和提炼的工厂原始文本进来经过层层加工最后输出一组富含语义和上下文信息的特征向量。很多人学习的时候容易把注意力全放在 Decoder 上觉得生成才是重点。但实际工作中你会发现大量任务其实只需要 Encoder。比如文本分类、情感分析、命名实体识别、句子相似度计算这些任务根本不需要 Decoder直接拿 Encoder 的输出接一个简单的分类头就能跑。BERT 就是典型的 Encoder-only 架构它在很多理解类任务上表现极其出色。所以把 Encoder 吃透等于拿到了进入现代 NLP 大门的钥匙。Encoder 的核心目标可以概括为一句话让序列中每个位置的向量都融合了整个序列的信息并且这个融合过程是可并行计算的。这跟 RNN 有本质区别RNN 必须从左到右一步步来天然无法并行而 Encoder 通过自注意力机制一次性看到所有位置训练效率高了好几个数量级。1.2 从 RNN 到 Transformer 的必然性在 Transformer 出现之前处理序列数据的主流方案是 RNN 及其变体 LSTM、GRU。它们的工作方式很像人读句子一个字一个字看过去维护一个隐藏状态每读一个新词就更新一次状态。这种方式很直观但有两个致命问题。第一个问题是串行计算。你没法同时处理第 1 个词和第 10 个词必须等前面算完才能算后面。这在 GPU 上非常吃亏因为 GPU 最擅长的就是大规模并行运算。序列一长训练时间就爆炸。第二个问题是长距离依赖。虽然 LSTM 通过门控机制缓解了梯度消失但当序列长度超过一定范围比如 500 个词以上前面信息传到后面还是会严重衰减。你让模型记住开头提到的主语到结尾再用来判断指代关系它经常记不住。Transformer 的解决方案很直接既然串行慢那就全部并行既然距离远会丢信息那就让任意两个位置直接建立连接。自注意力机制允许每个位置直接跟序列中所有其他位置计算关联度路径长度永远是 1不存在信息衰减的问题。这个设计思路的转变是 Transformer 能成为主流的核心原因。1.3 本文要拆解的核心模块清单为了让你有个清晰的预期我先把这篇要讲的内容列一下。一个标准的 Encoder 层包含以下几个关键组件多头自注意力机制Multi-Head Self-Attention让每个位置关注所有位置多头则是从不同子空间捕捉不同类型的关联。前馈神经网络Feed-Forward Network对每个位置的向量独立做非线性变换增强表达能力。残差连接Residual Connection缓解深层网络训练困难的问题。层归一化Layer Normalization稳定训练过程加速收敛。位置编码Positional Encoding因为自注意力本身不感知顺序必须额外注入位置信息。这些模块不是随便堆在一起的每个都有明确的设计动机。接下来我会逐个拆解把“为什么这么设计”讲透然后再给出可以直接运行的代码。2. 核心模块拆解与设计动机2.1 自注意力机制到底在算什么自注意力的核心思想用一句话说就是对于序列中的每个词计算它跟其他所有词的关联程度然后根据关联程度对其他词的信息进行加权求和得到这个词的新表示。具体怎么算呢每个词的输入向量会分别乘以三个不同的权重矩阵得到 Query查询、Key键、Value值三个向量。你可以用一个图书馆检索来类比Query 是你想找的主题Key 是每本书的标签Value 是书的内容。你用 Query 去跟每个 Key 做匹配匹配度高的书就多拿一些内容出来。计算过程分四步Query 和 Key 做点积得到注意力分数。点积越大说明两个向量越相似关联越强。分数除以根号 d_k其中 d_k 是 Key 向量的维度。这一步叫缩放目的是防止点积结果过大导致 softmax 梯度消失。对缩放后的分数做 softmax归一化成概率分布。用这些概率对 Value 向量加权求和得到输出。用公式表示就是Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V为什么要除以根号 d_k假设 Q 和 K 的每个分量都是独立同分布的均值 0 方差 1那么点积结果的方差就是 d_k。当 d_k 很大时点积结果会很大softmax 输出会变得非常尖锐几乎变成 one-hot 向量梯度就消失了。除以根号 d_k 可以把方差拉回 1 左右保持梯度的健康。2.2 为什么要用多头注意力如果只做一次注意力计算模型只能学到一种关联模式。但语言中的关系是多种多样的有语法上的主谓关系有语义上的同义关系有位置上的相邻关系。单一注意力头很难同时捕捉所有这些模式。多头注意力的做法是把 Q、K、V 分别投影到 h 个不同的低维子空间在每个子空间里独立做注意力计算最后把 h 个结果拼接起来再投影一次。这样每个头可以专注于不同类型的关联。比如有的头可能学到了关注相邻词有的头学到了关注句法依存关系。实际实现中假设模型维度 d_model 512头数 h 8那么每个头的维度就是 512 / 8 64。注意这里不是把 512 维切成 8 份分别算而是用 8 组不同的投影矩阵把输入映射到 8 个 64 维空间。计算完成后拼接回 512 维再过一层线性变换。这里有个实操细节多头注意力的参数量和单头注意力维度 512其实差不多因为每个头的维度缩小了。但多头带来的表达能力提升是显著的这也是为什么现在几乎所有模型都用多头。2.3 前馈网络为什么不可少自注意力机制本质上是线性的加权求和虽然 softmax 引入了非线性但整体表达能力还是有限的。前馈网络的作用就是给模型注入更强的非线性变换能力。前馈网络的结构很简单两层线性变换中间夹一个 ReLU 激活函数。第一层把维度从 d_model 扩展到 d_ff通常是 4 倍比如 512 变 2048第二层再压回 d_model。用公式表示FFN(x) max(0, xW1 b1)W2 b2为什么中间要放大 4 倍这是原论文的经验选择更大的中间维度能提供更丰富的特征变换空间。你可以理解为先把信息摊开做复杂的非线性组合再压缩回原来的维度。这个比例在后来的实践中被广泛验证大部分模型都沿用 4 倍这个设置。需要注意的是前馈网络是逐位置独立应用的。也就是说序列中每个位置的向量各自过一遍同样的 FFN位置之间不交互。交互已经在注意力层做完了FFN 只负责对每个位置的特征做深度加工。2.4 残差连接与层归一化的配合深层网络训练有个经典难题梯度消失和退化。理论上网络越深表达能力越强但实际上层数一多训练反而变差。残差连接是解决这个问题的关键。残差连接的做法很简单把某一层的输入直接加到输出上。用公式表示就是output sublayer(x) x。这样梯度在反向传播时有一条高速公路可以直接传回去不用经过所有中间层。即使某些层的变换效果不好最差也就是学到一个恒等映射不会让整体变差。层归一化则是对每个样本的特征维度做归一化把均值拉到 0方差拉到 1然后再做可学习的缩放和平移。它跟批归一化的区别在于批归一化是在 batch 维度上统计层归一化是在特征维度上统计。NLP 任务中序列长度不一batch 统计不稳定所以层归一化更合适。Transformer 里用的是 Post-LN 结构也就是LayerNorm(x sublayer(x))先残差再加归一化。不过后来的很多改进版本用 Pre-LN也就是x sublayer(LayerNorm(x))训练更稳定不需要 warmup 也能收敛。这个选择在实际项目中值得注意。2.5 位置编码的注入方式自注意力有个特点它对输入顺序完全不敏感。你把句子里的词打乱注意力计算的结果是一样的因为它是所有位置的全连接。但语言是有顺序的主语在谓语前面和后面意思完全不同。所以必须额外告诉模型每个词的位置。原论文用的是正弦位置编码用不同频率的正弦和余弦函数生成位置向量直接加到词嵌入上。公式如下PE(pos, 2i) sin(pos / 10000^(2i/d_model)) PE(pos, 2i1) cos(pos / 10000^(2i/d_model))其中 pos 是位置i 是维度索引。这种编码的好处是不需要学习参数而且理论上可以外推到训练时没见过的序列长度。不同维度对应不同频率低维度变化快高维度变化慢组合起来能唯一标识每个位置。后来也出现了可学习的位置编码比如 BERT 就是用一个可训练的 embedding 表。两种方式各有优劣正弦编码外推性好可学习编码在训练长度内拟合能力更强。实际选哪个看你的任务需求如果推理时序列长度可能超过训练长度正弦编码更稳妥。3. 从零实现一个 Encoder 层3.1 环境准备与依赖说明我这里的实现基于 PyTorch版本建议 1.10 以上太老的版本有些 API 不太一样。如果你用 TensorFlow 或者 JAX逻辑完全一致只是 API 调用方式不同。import torch import torch.nn as nn import torch.nn.functional as F import math不需要额外安装什么特殊库PyTorch 自带的模块就够了。如果你想跑 GPU确保 CUDA 环境配好后面我会说明哪些地方需要注意设备一致性。3.2 多头注意力的完整实现先定义多头注意力模块。这里我把 Q、K、V 的投影合并成一个线性层这样计算效率更高代码也更简洁。class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads, dropout0.1): super().__init__() assert d_model % num_heads 0, d_model 必须能被 num_heads 整除 self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads self.W_qkv nn.Linear(d_model, 3 * d_model) self.W_o nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): batch_size, seq_len, _ x.shape qkv self.W_qkv(x) q, k, v qkv.chunk(3, dim-1) q q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) k k.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) v v.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn_weights F.softmax(scores, dim-1) attn_weights self.dropout(attn_weights) context torch.matmul(attn_weights, v) context context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) output self.W_o(context) return output这段代码有几个关键点需要说明。第一W_qkv一次性输出 3 倍维度然后用chunk(3)切成 Q、K、V比定义三个独立线性层效率高。第二view加transpose的操作是把(batch, seq_len, d_model)变成(batch, num_heads, seq_len, d_k)这样每个头可以独立计算。第三mask 的处理用masked_fill把不需要关注的位置设成负无穷softmax 之后这些位置的概率就是 0。注意transpose之后张量在内存中不连续直接view会报错必须先调.contiguous()。这个坑我踩过好几次尤其是在自定义层里忘了加训练时直接崩。3.3 前馈网络的实现细节前馈网络相对简单就是两层线性加激活。class PositionwiseFeedForward(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))))这里的d_ff通常设为4 * d_model。激活函数原论文用的是 ReLU后来也有用 GELU 的效果在某些任务上更好。如果你想跟 BERT 保持一致把F.relu换成F.gelu就行。3.4 组装完整的 Encoder Layer把注意力、前馈网络、残差连接、层归一化拼起来就得到一个完整的 Encoder 层。class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, num_heads, dropout) self.ffn PositionwiseFeedForward(d_model, d_ff, dropout) 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, maskNone): attn_output self.self_attn(x, mask) x self.norm1(x self.dropout1(attn_output)) ffn_output self.ffn(x) x self.norm2(x self.dropout2(ffn_output)) return x这是标准的 Post-LN 结构。如果你想要 Pre-LN把norm1和norm2的位置换一下x x self.dropout1(self.self_attn(self.norm1(x), mask)) x x self.dropout2(self.ffn(self.norm2(x)))Pre-LN 在深层模型中训练更稳定我个人的项目里现在基本都用 Pre-LN省去了 warmup 的麻烦。3.5 位置编码的生成与注入位置编码的实现如下class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000, dropout0.1): super().__init__() self.dropout nn.Dropout(dropout) pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.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) self.register_buffer(pe, pe) def forward(self, x): x x self.pe[:, :x.size(1), :] return self.dropout(x)register_buffer的作用是把pe注册成模型的一部分但不作为可训练参数。这样保存和加载模型时它会自动带上而且会跟着模型一起搬到 GPU 上。如果你用普通属性存搬到 GPU 时容易忘记手动转移导致设备不一致报错。3.6 堆叠多层 Encoder 并跑通前向传播最后把所有东西组装起来class TransformerEncoder(nn.Module): def __init__(self, vocab_size, d_model512, num_heads8, d_ff2048, num_layers6, max_len5000, dropout0.1): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) self.pos_encoding PositionalEncoding(d_model, max_len, dropout) self.layers nn.ModuleList([ EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) self.d_model d_model def forward(self, x, maskNone): x self.embedding(x) * math.sqrt(self.d_model) x self.pos_encoding(x) for layer in self.layers: x layer(x, mask) return x注意self.embedding(x) * math.sqrt(self.d_model)这一步。词嵌入的数值范围通常比较小乘以根号 d_model 是为了放大嵌入的尺度让位置编码的相对影响不至于过大。这个细节原论文里有提但很多人实现时会漏掉。跑一个简单的前向传播验证一下model TransformerEncoder(vocab_size10000, d_model512, num_heads8, num_layers6) x torch.randint(0, 10000, (2, 20)) output model(x) print(output.shape) # torch.Size([2, 20, 512])输出形状是(batch_size, seq_len, d_model)符合预期。4. 训练调试中的常见问题与排查4.1 注意力权重全变成均匀分布怎么办这是新手最常遇到的现象训练一段时间后打印注意力权重发现每一行都差不多接近均匀分布。这意味着模型没有学到任何有意义的关联模式。可能的原因有几个。第一学习率太大导致参数更新过猛注意力分数被推到了 softmax 的饱和区。可以尝试降低学习率或者加 warmup。第二初始化不当Q 和 K 的投影矩阵初始值太大或太小。PyTorch 默认的线性层初始化是均匀分布范围跟输入维度有关一般没问题但如果你自己改了初始化就要检查。第三mask 设置错误把不该遮的位置遮住了导致有效信息太少。排查方法先打印几层注意力权重的均值和方差看看是否有区分度。如果方差极小基本可以确定是上面某个原因。我一般会先用一个很小的学习率比如 1e-5跑几百步确认注意力能分化出来再逐步调大。4.2 梯度爆炸与梯度消失的应对策略Transformer 虽然比 RNN 好训练但深层版本仍然会遇到梯度问题。梯度爆炸表现为 loss 突然变成 NaN梯度消失表现为 loss 长时间不下降。梯度爆炸的常规处理是梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)这行代码放在loss.backward()之后、optimizer.step()之前。max_norm设 1.0 是常见选择也可以试 0.5 或 5.0看具体任务。梯度消失则更多跟学习率调度有关。原论文用了 warmup 策略前几千步学习率线性增加之后按步数的平方根倒数衰减。这个策略对 Post-LN 结构几乎是必须的因为 Post-LN 在训练初期梯度方差很大直接上大学习率容易崩。如果你用 Pre-LNwarmup 可以省掉但加了也没坏处。4.3 位置编码外推失败的场景正弦位置编码理论上可以外推到任意长度但实际用的时候你会发现当推理序列长度远超训练长度时效果还是会下降。原因是模型在训练时只见过特定范围的位置模式超出这个范围后虽然编码本身有定义但模型没学过怎么解读这些新模式。如果你确实需要处理变长序列有几个方案。一是训练时就用足够长的序列覆盖推理时的最大长度。二是用相对位置编码比如 T5 里的相对位置偏置它编码的是位置之间的相对距离而不是绝对位置外推性更好。三是用旋转位置编码RoPE通过旋转矩阵注入位置信息在长序列任务上表现很好现在很多大模型都在用。4.4 显存不够时的优化手段Transformer 的显存占用主要来自注意力矩阵复杂度是 O(n^2)n 是序列长度。序列长度翻倍显存占用翻四倍。如果你显存吃紧可以尝试以下手段。第一减小 batch size这是最直接的。第二用梯度累积把小 batch 的梯度累加几次再更新等效于大 batch。第三用混合精度训练torch.cuda.amp可以显著降低显存占用通常能省 30% 到 50%。第四如果序列特别长考虑用稀疏注意力或者分块注意力比如 Longformer 的做法把全局注意力限制在局部窗口内。我自己的经验是混合精度加梯度累积能解决大部分显存问题实在不行再考虑改注意力结构。4.5 常见问题速查表问题现象可能原因排查方法解决方案Loss 变 NaN梯度爆炸、学习率过大打印梯度范数梯度裁剪、降低学习率、加 warmup注意力权重均匀学习率大、初始化差、mask 错误打印权重方差调小学习率、检查 mask、重设初始化训练 loss 不降梯度消失、结构错误检查各层输出范围换 Pre-LN、检查残差连接显存溢出序列过长、batch 过大查看显存占用曲线混合精度、梯度累积、减小 batch推理长度外推差位置编码外推性不足对比训练/推理长度相对位置编码、RoPE训练慢未用 GPU、串行操作多检查设备、profile确认数据在 GPU、减少 CPU-GPU 同步5. 几个容易被忽略的实操心得5.1 初始化对训练稳定性的影响PyTorch 默认的线性层初始化是 Kaiming 均匀初始化对 Transformer 来说基本够用。但如果你发现训练初期 loss 波动很大可以试试 Xavier 初始化把所有权重矩阵的方差控制在合理范围。具体做法是遍历所有nn.Linear模块用nn.init.xavier_uniform_重新初始化权重偏置置零。词嵌入的初始化也值得注意。默认是正态分布标准差 1。但有些实现会用更小的标准差比如 0.02目的是让嵌入的初始值不要太大避免训练初期注意力分数过于极端。BERT 用的就是 0.02你可以参考。5.2 Dropout 放置位置的讲究Dropout 在 Transformer 里有好几个放置位置注意力权重上、残差连接前、前馈网络中间。原论文在每个子层的输出上都加了 dropout注意力权重上也加了。实际调参时dropout 率通常设 0.1任务数据量小的时候可以调大到 0.2 或 0.3数据量大就调小甚至关掉。有个细节位置编码之后也加了 dropout。这一步很多人会忘但它对防止过拟合有帮助尤其是小数据集上。5.3 学习率 warmup 的必要性再讨论前面提过 Post-LN 需要 warmup这里再展开说一下为什么。Post-LN 结构里层归一化在残差之后训练初期各层的输出方差会逐层累积放大导致深层梯度非常大。Warmup 让学习率从很小的值慢慢升上去给模型时间把各层输出调整到合理范围。Warmup 的步数一般设总训练步数的 5% 到 10%或者固定几千步。原论文设了 4000 步。如果你用 Pre-LN理论上可以不用 warmup但实践中加一点 warmup 通常能让初期更稳代价也不大。5.4 如何验证 Encoder 真的学到了东西训练完之后怎么确认 Encoder 不是个摆设我一般做两个检查。第一看注意力权重是否有明显的模式比如某些头关注相邻位置某些头关注句法相关的位置。第二把 Encoder 输出的向量拿去做相似度计算看语义相近的句子是否真的距离更近。如果注意力权重看起来随机或者相似度计算毫无区分度那大概率是训练出了问题需要回头检查数据、学习率、初始化这些环节。这个验证步骤很多人会跳过直接看下游任务的指标但指标好不一定代表 Encoder 学好了可能只是分类头拟合能力强。5.5 从 Encoder 到完整模型的扩展路径把 Encoder 跑通之后你可以往几个方向扩展。一是加 Decoder 做序列到序列的任务比如翻译、摘要。二是接分类头做文本分类。三是做对比学习用 Encoder 输出做句子表示训练相似度模型。四是在 Encoder 基础上加池化层做句子级别的 embedding。不管哪个方向Encoder 部分都是通用的。把这一块调稳了后面的扩展就是搭积木的事。我个人的习惯是先把 Encoder 单独训练一个语言模型任务确认它能正常收敛再迁移到具体任务上微调这样成功率会高很多。最后分享一个小技巧如果你在调试注意力机制时觉得矩阵运算太抽象可以把一个小序列比如 5 个词的注意力权重打印成表格手动算一遍 softmax 和加权求和跟代码输出对一下。我当初就是这么把整个流程彻底搞明白的比看十遍论文都管用。