资讯详情

从零实现Transformer Encoder:自注意力机制与PyTorch实战解析

📅 2026/10/13 1:12:25 | 华诺云谱 👁 阅读
从零实现Transformer Encoder:自注意力机制与PyTorch实战解析
1. 从零理解Transformer Encoder到底在做什么如果你最近在折腾NLP相关的项目大概率绕不开Transformer这个词。不管是做文本分类、情感分析还是想自己搭一个小的翻译模型Transformer的Encoder部分几乎是最先要啃下来的硬骨头。我刚开始接触的时候看论文里的那张架构图第一反应是“这箭头也太多了吧”但真正把数据流一步步拆开跑通之后发现它的设计逻辑其实非常干净——本质上就是把一句话里的每个词通过自注意力和前馈网络反复提炼成包含上下文信息的向量表示。这篇文章主要面向两类人一类是刚入门深度学习、想搞明白Transformer内部到底怎么运转的初学者另一类是用过BERT之类的预训练模型但对其底层Encoder结构只有模糊认知、想自己动手实现一遍的开发者。我会从整体设计思路讲起然后逐层拆解Self-Attention、Multi-Head Attention、位置编码、残差连接和Layer Normalization这些核心组件最后给出一个可以直接跑的PyTorch实现并分享我在调试过程中踩过的坑。需要提前说明的是这篇文章不会涉及任何与网络访问工具相关的内容所有代码和实验都在本地环境完成。另外文中提到的所有项目名称、数据集名称均为虚构代称仅用于说明技术原理。2. Encoder的整体架构设计思路拆解2.1 为什么Encoder要采用堆叠结构Transformer的Encoder并不是单独一层而是由N个完全相同的层堆叠而成。原始论文里N6也就是把同样的结构重复了6次。这个设计思路其实和CNN里堆叠卷积层是类似的逻辑——浅层捕捉局部、简单的特征深层捕捉全局、抽象的特征。具体到Encoder里每一层包含两个子层一个是多头自注意力机制另一个是前馈神经网络。每个子层外面都套了残差连接和Layer Normalization。数据从底层输入经过每一层的处理后输出的向量维度保持不变原始论文里是512维这样才能一层层往上叠。我一开始有个疑问为什么要保持维度不变后来想明白了如果维度变了残差连接就没法直接相加了。残差连接要求输入和输出的形状一致所以整个Encoder里所有子层的输出维度都必须和输入维度对齐。这也是为什么在实现的时候不管是Attention还是FFN最后都会把维度投影回原始的d_model。注意堆叠层数不是越多越好。层数增加会带来参数量和计算量的线性增长同时梯度消失的风险也会上升。实际项目中如果数据量不大4层甚至2层也能work不必死磕原始论文的6层。2.2 自注意力机制解决了什么问题在Transformer之前处理序列数据主要靠RNN和LSTM。RNN的问题是必须按顺序一步步计算前一个时间步算完才能算下一个没法并行。而且当序列很长的时候早期的信息传到后面会严重衰减这就是所谓的长期依赖问题。自注意力机制的核心思想是对于序列中的每一个词都让它和序列中所有其他词计算一个相关性分数然后根据这些分数对所有词的信息进行加权求和。这样一来每个词都能直接“看到”序列中任意位置的词不管距离多远路径长度都是1。同时因为每个词的计算不依赖前一个词的结果整个矩阵运算可以并行完成。举个例子假设有一句话“小猫追着尾巴跑”当模型处理“尾巴”这个词的时候自注意力机制会让它和“小猫”“追”“跑”都计算相关性。很明显“小猫”和“尾巴”的相关性会比较高所以“尾巴”的表示里会融入更多“小猫”的信息。这就是上下文感知的来源。2.3 多头注意力为什么要“多头”如果只做一次自注意力模型只能学到一种关注模式。但语言里的关系是多种多样的——有的词关注语法关系有的词关注语义关系有的词关注位置关系。多头注意力就是让模型在不同的子空间里并行地做多次自注意力每次用不同的Q、K、V投影矩阵最后把结果拼接起来再投影一次。原始论文里用了8个头每个头的维度是512/864。这样总计算量和单头512维差不多但模型能捕捉到更丰富的关系模式。我实测下来头数太少比如1个或2个效果确实会打折扣但头数太多比如16个以上在小数据集上反而容易过拟合而且计算效率会下降。2.4 位置编码为什么必须存在自注意力机制本身是位置无关的——把输入序列打乱顺序算出来的注意力分数是一样的。但语言是有顺序的“猫追狗”和“狗追猫”意思完全相反。所以必须额外注入位置信息。原始论文用的是正弦余弦位置编码公式如下import math import torch import torch.nn as nn class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000): super().__init__() 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): return x self.pe[:, :x.size(1), :]这个编码的好处是每个位置的编码是唯一的而且任意两个位置之间的编码差异可以通过线性变换表示这让模型能比较容易地学到相对位置关系。后来也有用可学习位置编码的方案比如BERT就是直接训练一个位置嵌入矩阵效果也不错但需要更多数据来训练。3. 核心组件细节解析与实操要点3.1 Self-Attention的矩阵计算过程自注意力的计算可以概括为四个步骤生成Q、K、V矩阵计算注意力分数缩放softmax归一化加权求和。用公式表示就是Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V这里的d_k是每个头的维度。为什么要除以sqrt(d_k)因为当d_k比较大的时候Q和K的点积结果会很大softmax之后会变得非常尖锐梯度会趋近于0训练会变得困难。除以sqrt(d_k)可以把方差拉回到1附近保证梯度稳定。我用一个具体的例子来说明。假设输入序列长度是4d_model是8头数是2那么每个头的维度就是4。Q、K、V的形状都是(4, 4)。QK^T得到(4, 4)的注意力分数矩阵每一行表示某个位置对所有位置的关注程度。除以sqrt(4)2之后再做softmax就得到了归一化的注意力权重。import torch import torch.nn.functional as F def scaled_dot_product_attention(q, k, v, maskNone): d_k q.size(-1) scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn F.softmax(scores, dim-1) output torch.matmul(attn, v) return output, attn提示mask的使用很关键。在Encoder里如果处理的是padding后的序列需要用mask把padding位置遮掉否则这些无意义的位置也会参与注意力计算影响效果。mask的形状通常是(batch_size, 1, 1, seq_len)通过广播机制作用到分数矩阵上。3.2 Multi-Head Attention的实现细节多头注意力的实现思路是先把输入通过线性层投影成Q、K、V然后把维度拆成(num_heads, head_dim)并行计算注意力最后把多个头的结果拼接起来再过一次线性层。这里有个实现上的细节容易踩坑拆分维度的时候不能简单地用view或者reshape因为这样会把不同头的数据混在一起。正确的做法是用view(batch_size, seq_len, num_heads, head_dim)然后transpose或者用einops库的rearrange操作。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.head_dim 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, x, maskNone): batch_size, seq_len, _ x.size() q self.w_q(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k self.w_k(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) v self.w_v(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) attn_output, _ scaled_dot_product_attention(q, k, v, mask) attn_output attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) output self.w_o(attn_output) return output我实测下来这个实现在batch_size32、seq_len128、d_model512、num_heads8的配置下单次前向传播在普通GPU上大约需要几毫秒完全能满足中小规模任务的需求。3.3 前馈网络的设计考量前馈网络其实就是一个两层的全连接网络中间夹一个ReLU激活函数。原始论文里第一层把维度从512扩展到2048第二层再压回512。这个扩展比例是4倍后来成了默认配置。为什么要有这个FFN一种解释是自注意力层主要负责整合上下文信息而FFN负责对每个位置的信息做非线性变换和特征提取。两者分工不同缺一不可。我在实验里试过去掉FFN模型效果直接掉了一大截验证了它的必要性。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) self.relu nn.ReLU() def forward(self, x): return self.linear2(self.dropout(self.relu(self.linear1(x))))注意d_ff的选择需要根据任务和计算资源来定。4倍是默认值但如果你的d_model本身就很大比如10244倍就是4096参数量会非常可观。这时候可以考虑用2倍或者用GLU变体来减少参数量。3.4 残差连接与Layer Normalization的配合每个子层外面都包了一层残差连接和LayerNorm具体是output LayerNorm(x Sublayer(x))残差连接的作用是缓解梯度消失让梯度能直接流过网络。LayerNorm的作用是对每个样本的特征维度做归一化稳定训练过程。这里要注意的是LayerNorm和BatchNorm不同它是在特征维度上做归一化而不是在batch维度上所以不受batch_size的影响适合变长序列的场景。原始论文用的是Post-LN也就是先做残差相加再做LayerNorm。后来很多工作发现Pre-LN先做LayerNorm再进子层训练更稳定不需要warmup也能收敛。我在实际项目里更倾向于Pre-LN尤其是层数比较多的时候。4. 完整Encoder实现与实操过程4.1 搭建单层Encoder把前面讲的组件组装起来就得到了一个完整的Encoder层。代码结构很清晰多头注意力子层、前馈网络子层每个子层外面套残差和LayerNorm。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) self.feed_forward 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)) ff_output self.feed_forward(x) x self.norm2(x self.dropout2(ff_output)) return x这里用的是Post-LN结构和原始论文一致。如果你想改成Pre-LN只需要把norm的位置换一下def forward(self, x, maskNone): x x self.dropout1(self.self_attn(self.norm1(x), mask)) x x self.dropout2(self.feed_forward(self.norm2(x))) return x两种写法我都跑过Pre-LN在层数较多时确实更稳loss曲线更平滑但最终效果差异不大。建议新手先用Post-LN和论文对齐方便对照理解。4.2 堆叠多层Encoder并接入嵌入层完整的Encoder模型需要包含词嵌入层、位置编码层以及N个EncoderLayer的堆叠。class TransformerEncoder(nn.Module): def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, max_len5000, dropout0.1): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) self.pos_encoding PositionalEncoding(d_model, max_len) self.layers nn.ModuleList([ EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) self.dropout nn.Dropout(dropout) self.d_model d_model def forward(self, x, maskNone): x self.embedding(x) * math.sqrt(self.d_model) x self.pos_encoding(x) x self.dropout(x) for layer in self.layers: x layer(x, mask) return x注意嵌入层后面乘了一个sqrt(d_model)这是原始论文里的做法目的是让嵌入向量的尺度和位置编码的尺度匹配避免位置编码被淹没。4.3 参数选择与计算过程假设我们要做一个中文文本分类任务词表大小30000最大序列长度128d_model设为256头数设为8FFN维度设为1024层数设为4。来算一下参数量词嵌入30000 × 256 7,680,000每层多头注意力4 × (256 × 256) 262,144每层FFN256 × 1024 1024 × 256 524,288每层LayerNorm2 × 2 × 256 1,024单层合计约787,4564层合计约3,149,824总计约10,829,824也就是大约1080万参数这个规模在单卡上训练完全没问题batch_size可以设到64甚至128。如果显存不够可以适当减小d_model或者层数。提示d_model必须是num_heads的整数倍否则拆分多头的时候会报错。常见的组合有256/8、512/8、768/12等。选的时候优先保证head_dim在32到128之间太小了表达能力不够太大了计算效率低。4.4 训练过程中的关键配置训练Transformer Encoder的时候有几个配置项特别关键学习率调度原始论文用了warmup策略先线性增加学习率再按步数的平方根倒数衰减。这个策略对Transformer特别重要因为一开始梯度比较大直接上大学习率容易发散。class WarmupScheduler: def __init__(self, optimizer, d_model, warmup_steps4000): self.optimizer optimizer self.d_model d_model self.warmup_steps warmup_steps self.step_num 0 def step(self): self.step_num 1 lr self.d_model ** (-0.5) * min(self.step_num ** (-0.5), self.step_num * self.warmup_steps ** (-1.5)) for param_group in self.optimizer.param_groups: param_group[lr] lrDropoutEncoder里多处用了Dropout包括注意力权重、FFN中间层、嵌入层后面。dropout率一般设0.1数据量小的时候可以适当调大到0.2或0.3。梯度裁剪Transformer训练时梯度偶尔会爆炸加上梯度裁剪max_norm1.0能有效防止loss变成nan。Mask处理如果batch里有不同长度的序列padding位置的mask一定要正确设置。我见过不少人忘了加mask结果模型在padding位置也计算注意力效果差了很多。5. 常见问题与排查技巧实录5.1 Loss不下降或者变成NaN这是新手最常遇到的问题。根据我的经验排查顺序应该是这样的问题现象可能原因排查方法解决方案Loss一开始就NaN学习率太大打印每步的loss和梯度范数降低学习率加warmupLoss缓慢下降但不收敛缺少warmup检查学习率曲线加上warmup策略Loss震荡严重batch_size太小增大batch_size试试增大batch或调小学习率Loss下降后突然NaN梯度爆炸监控梯度范数加梯度裁剪Loss一直很高mask设置错误检查mask形状和值修正mask逻辑我印象最深的一次是loss在第几百步突然变成nan查了半天发现是位置编码的max_len设小了序列长度超了之后索引越界导致出现了inf。所以位置编码的max_len一定要设得比实际最大序列长度大一些。5.2 注意力权重看起来不合理有时候把注意力权重可视化出来发现所有位置的权重都差不多或者都集中在某一个位置。这种情况通常有几个原因一是没有除以sqrt(d_k)导致softmax后的分布太尖锐或者太平坦。二是初始化有问题Q和K的投影矩阵初始化方差太大或太小。三是训练不充分模型还没学到有意义的注意力模式。我的做法是先检查缩放因子是否正确然后用Xavier初始化或者Kaiming初始化重新初始化投影矩阵最后多训几个epoch再看。正常情况下训练充分的模型注意力权重应该呈现出一定的结构比如相邻位置的权重较高或者某些关键词的权重特别突出。5.3 训练速度慢的优化思路Transformer Encoder虽然可以并行但如果实现不当速度也会很慢。几个优化点用nn.MultiheadAttention替代手写实现PyTorch官方的版本做了很多底层优化把mask提前算好不要在每次forward的时候重新计算用混合精度训练torch.cuda.amp速度能提升30%到50%如果序列长度固定可以提前把位置编码算好存下来不用每次重新算batch_size尽量设大一些GPU利用率会更高我实测下来在同样的硬件上优化前后的训练速度能差2到3倍。尤其是混合精度训练几乎不损失精度但速度提升非常明显。5.4 过拟合了怎么办Transformer的参数量通常比较大在小数据集上很容易过拟合。我常用的几招增大dropout率从0.1调到0.3甚至0.5加权重衰减一般设1e-4到1e-5减少层数或者d_model把模型变小用早停策略验证集loss不下降就停数据增强比如同义词替换、随机删除等注意不要一上来就上大模型。我见过不少人拿几万条数据去训12层、768维的Encoder结果训练集loss降到很低验证集却一塌糊涂。先从2层、256维开始效果不够再加。5.5 位置编码的常见误区位置编码这块有几个容易搞错的地方。一是max_len设得太小序列一长就报错。二是把位置编码和嵌入相加之后忘了做dropout。三是用了可学习位置编码但没把它加到优化器参数里。还有一个细节如果序列长度变化很大比如有的样本长度10有的长度500那么位置编码的max_len要按最大长度来设。但这样会浪费一些计算因为短序列也会算到最大长度的位置编码。一种优化是按batch内的最大长度动态截取位置编码能省一点计算。6. 从Encoder出发的扩展思路把Encoder搞明白之后其实很多变体都可以自己推导出来。比如BERT就是只用了Transformer的Encoder部分通过掩码语言模型和下一句预测任务来预训练。再比如一些文本分类模型直接在Encoder输出上加一个池化层和全连接层就能用。如果你想继续深入我建议下一步可以看看Decoder部分理解一下掩码自注意力和交叉注意力是怎么工作的。然后可以尝试把Encoder和Decoder拼起来做一个简单的序列到序列模型比如英译中或者文本摘要。再往后可以研究各种预训练任务的设计以及如何把预训练模型迁移到下游任务上。我个人在实际操作中的体会是Transformer Encoder看起来组件多但每个组件的设计都有很明确的动机。把每个“为什么”想清楚了实现起来就是顺理成章的事。最怕的就是照着代码抄一遍跑通了但不知道每一步在干什么那样遇到问题就无从下手。建议至少自己从零实现一遍然后逐层打印中间结果的形状和数值观察数据是怎么流动的。这个过程花不了多少时间但对理解整个架构的帮助非常大。
📝

华诺云谱内容团队

资深建站顾问 · 行业研究员

10年+企业数字化服务经验,专注智能建站、SEO优化与品牌营销,持续输出建站技巧、行业洞察与营销干货,已帮助5000+企业实现数字化增长。

你可能需要的服务

订阅华诺云谱资讯周报

每周一封,精选建站技巧、SEO与营销干货,直达邮箱。已有 8,000+ 企业主订阅,助你少走弯路。

↑