资讯动态

Transformer代码拆解:从注意力机制到最小实现全解析

发布时间:2026/10/17 17:51:20 来源:尧图企业网站定制
1. 整体设计先把Transformer代码拆成三块来看说到Transformer很多人第一反应就是那篇经典论文和一堆复杂的公式。但真到了看代码的时候容易一头扎进某个文件的细节里出不来。我的建议是先把整个项目在脑内切成三块数据准备、模型结构、训练推理。不管代码库多大核心流程跑不出这三个部分。拿我自己最早接触Transformer源码的经历来说当时拿到的是一个最小实现版本全部代码加起来不到一千行。第一件事不是逐行读而是先看目录结构找到模型定义文件、数据处理文件、训练脚本分别在哪。理清楚这三块之后再去逐段读逻辑心里才会有底。数据准备要解决的是原始文本如何变成模型能吃的张量。具体来说就是分词、建词表、转成索引序列、做padding和mask。模型结构则是Transformer的核心包括多头注意力、前馈网络、层归一化、位置编码、Encoder和Decoder的堆叠。训练推理环节则是把模型跑起来的加速器部分包括损失函数、优化器、学习率调度、以及推理时的自回归解码。理解这三块之间的关系就像看一条流水线原始语料进来经过数据模块变成标准化的数值张量张量进入模型模块前向传播算出预测结果预测结果和真实标签一起交给训练模块算出损失、回传梯度、更新参数。推理时则把训练好的模型权重加载进来输入一句话逐步生成输出。如果一上来就陷入某一行的实现细节比如注意力分数到底怎么算的其实很容易迷失方向。先把全貌建立起来后续看任何一行代码都会知道它在这个流水线里的位置理解成本直线下降。2. 注意力机制代码逐行拆解终于把Q、K、V搞明白了注意力机制是Transformer的心脏。很多初学者看到Q、K、V三个字母就头疼不理解为什么要搞三个矩阵出来。先从直觉上说Q的意思是你想查什么K是内容索引V是真正的内容本身。类比一下你去图书馆找书Q是你脑子里的查询需求K是书架上的标签索引V是书里的实际内容。系统先通过Q和K计算匹配度也就是注意力分数再按分数去加权提取V中的信息。代码实现中Q、K、V通常是三个线性变换得到的。初始化时会有nn.Linear层分别映射出query、key、value三个张量。初学者容易困惑的一点是输入不是已经有一个embedding向量了吗为什么还要再做线性变换因为embedding向量本身是静态的而注意力机制需要让每个token根据当前任务动态地决定“该关注谁”。这三个线性变换是可学习的训练过程中会不断调整让模型找到最合适的映射方式。注意力分数的计算是理解整个机制的关键一步。代码通常写成torch.matmul(Q, K.transpose(-2, -1))除以一个缩放因子。为什么要除以这个缩放因子因为当向量维度变大时Q和K的点积结果方差会变大导致softmax的输出分布过于尖锐梯度很容易消失。除以sqrt(d_k)相当于把方差拉回一个合理范围这是论文里明确给出的做法实际操作中非常管用。softmax这一步把分数转换成0到1之间的权重分布。注意这里有一步mask操作有些位置会被设置成负无穷经过softmax之后权重就变成0。这个mask在处理padding位置时是必须的一个batch里的序列长度不同短序列会补零到和长序列一样长这些填充位置不应该参与注意力计算所以要在softmax之前把它们屏蔽掉。最后一步是把权重和V相乘再求和得到每个token的输出表示。代码上就是torch.matmul(attn_weights, V)。到这里单个注意力头的完整计算就结束了。我建议学到这里时自己动手把矩阵的shape变化全部写出来比如输入是[batch, seq_len, d_model]经过线性变换还是[batch, seq_len, d_model]拆成多头后变成[batch, heads, seq_len, d_head]算完注意力再拼回去。shape理清楚了基本就吃透了大半。2.1 多头注意力为什么要切好几份再并行算多头注意力的直觉是一个注意力头只能学到一种关注模式比如有的头倾向于关注语法关系有的头倾向于关注语义相似性。多个头并行计算然后把结果拼接起来模型就能同时捕捉多种不同维度的信息。代码实现里常见的做法不是建多个独立的注意力层而是把d_model切成heads份每一份作为头维度d_head d_model / heads。比如d_model是512heads是8那么每个头负责64维。切分之后Q、K、V分别变成[batch, heads, seq_len, d_head]这个维度变换靠reshape和transpose两个操作完成。容易出错的地方在于reshape的顺序不能想当然。PyTorch里通常先view成[batch, seq_len, heads, d_head]然后transpose(1, 2)变成[batch, heads, seq_len, d_head]。如果反过来先把seq_len放前面再view就会把不同头的维度混到一起结果完全错乱。这个坑我踩过调试的时候张量形状打印出来怎么看怎么对但结果就是不对劲最后发现是view和transpose的顺序反了。多头计算完之后需要把多个头的结果合并回d_model做法是transpose(1, 2)变回[batch, seq_len, heads, d_head]然后view成[batch, seq_len, d_model]最后再过一层线性输出投影。到这里多头注意力才算完整走完。2.2 掩码的实现细节padding mask和look-ahead mask的区别掩码是Transformer代码里最容易懵的地方。不同位置用的掩码还不一样至少分两种一种是padding mask一种是look-ahead mask。padding mask用于屏蔽掉序列中无意义的填充tokenlook-ahead mask用于解码器因为解码器在做self-attention的时候每个位置只能看到它前面的token不能看到未来的token否则等于提前泄露了答案。从代码实现上说padding mask通常是一个布尔张量shape是[batch, seq_len]标明哪些位置是有效token哪些是填充token。在送入注意力之前会转换成[batch, 1, 1, seq_len]的形状方便和注意力分数矩阵做广播。look-ahead mask则是一个上三角矩阵用torch.triu生成对角线上方的位置为True这些位置在softmax之前会被替换成极小值。我见过很多人在写mask的时候把两种掩码搞混或者只做了其中一种。最简单的方法是单独写一个函数生成掩码然后在模型里明确标注当前用的是哪一类。调试阶段把mask打印出来看一眼基本就能发现问题。还有个容易被忽略的细节有些实现会在softmax之后再做一次mask清零确保填充位置的注意力权重确实是0。这两种方式效果类似但前者更常见实现也更高效。3. Encoder和Decoder的代码骨架从位置编码到堆叠结构Transformer的整体骨架由Encoder和Decoder两大部分组成。Encoder负责把输入序列编码成一组上下文相关的向量表示Decoder则基于这些表示逐词生成目标序列。无论是哪一侧都是由若干个相同的层堆叠而成每层内部的结构完全一致只是参数独立。先看位置编码。RNN天然有顺序感因为token是一个挨一个进入模型的。Transformer没有这个特性它的Self-Attention把一句话当成一个集合来处理顺序信息完全丢失。所以必须把位置信息显式加进去。原论文用的是正弦余弦位置编码代码里通常先算好一个固定的矩阵[seq_len, d_model]偶数位置用sin奇数位置用cos然后加到输入embedding上。为什么用sin和cos而不是直接学一个位置向量原论文的解释是正弦余弦函数可以让模型容易学习到相对位置关系而且对于训练中没见过的更长序列也有一定的外推能力。实操中用可学习位置编码的也很多效果差距不算大。代码实现上位置编码在注册为buffer不参与梯度更新。如果拿到的预训练模型是用固定位置编码做的自己在代码里换成了可学习的加载权重时维度不匹配直接报错这种坑提前知道就能绕开。再看层结构。Encoder的每一层由两个子层组成第一个是多头自注意力第二个是前馈网络。每个子层外面都有残差连接和层归一化。代码里通常写成x x sublayer(x)然后再过LayerNorm。层归一化和BatchNorm的区别在于BatchNorm是跨样本做的归一化而LayerNorm对每个样本自己的所有特征维度做归一化这在变长序列的NLP任务里更稳定。前馈网络由两个线性层夹一个激活函数组成常见的是ReLU或者GELU中间的隐层维度通常是d_model的4倍。比如输入输出都是512维中间层就是2048维。这个结构很朴素但参数量占了模型的大头。为什么需要前馈网络因为多头注意力本质上是线性变换的组合如果两个注意力层直接叠在一起模型整体上还是线性的。前馈网络引入非线性变换让模型有了更强的表达能力。Decoder的结构比Encoder复杂一些。每一层有三个子层先是带look-ahead mask的自注意力再是交叉注意力最后是前馈网络。交叉注意力是Decoder和Encoder之间的桥梁Q来自Decoder前一层的输出K和V来自Encoder的输出。这样每个解码位置都能从输入序列中“查询”相关信息。代码上只是把注意力函数的参数来源换了一下复用同一个注意力类即可。3.1 从代码角度看残差连接为什么重要残差连接的作用是缓解深层网络中的梯度退化问题。Transformer堆了那么多层如果每个子层的变换都很复杂反向传播时梯度在长链路中容易消失。残差连接让梯度多了一条“直通”路径相当于每个子层只需要去学习相对于输入的“增量”部分。代码里就是很简单的add操作没有多余的门控机制。但这里有个细节值得注意残差连接和层归一化的顺序。目前主流实现是Pre-LN即先做层归一化再做子层计算最后相加原论文用的则是Post-LN即先子层计算再归一化再相加。实际训练中Pre-LN更稳定尤其在深层模型里收敛更快。代码里如果看到结构略有不同大概率是这个顺序的差异不算bug。我们在自己搭建模型时强烈建议先按Pre-LN来写因为Post-LN对学习率特别敏感调起来很痛苦。等模型跑通了再回头对比也行至少省下心来去处理真正会报错的地方。3.2 编解码器交互中的张量形状与广播交叉注意力是新手最容易丢维度的地方。Decoder的self-attention输出Q形状是[batch, tgt_len, d_model]Encoder的输出提供K和V形状是[batch, src_len, d_model]。两者长度不一样是很正常的因为源语言和目标语言的序列长度本来就可能不同。广播机制会把两个序列长度同时纳入计算注意力分数矩阵的形状是[batch, heads, tgt_len, src_len]。有人会问self-attention和交叉注意力在代码上是不是同一个函数通常是的。区别只是在传参时谁是Q、谁是K、谁是V。如果代码里封装了一个attention函数内部严格按“Q、K、V三个参数”来写那在Encoder里QKV在Decoder的交叉注意力里Q来自上一层解码结果K和V来自Encoder输出。这样统一设计代码简洁且不易出错。层归一化在Decoder中也扮演重要角色。它让数值分布稳定尤其在训练初期模型还没学会合理的注意力分配输出分布可能非常极端归一化能把极端值拉回到相对稳定的区间训练更容易跑下去。很多问题排查时怀疑模型结构写错了最后发现只是归一化位置不对或者漏了残差加法这类问题一定要趁早用shape调试法确认清楚。4. 从零实现一个最小可用的Transformer数据批次、训练与推理理论说再多不如自己动手撸一遍训练流程。我的建议是先用一个小数据集比如做一个英文到法语的简单翻译任务或者做一个文本分类任务把整个pipeline打通。这一步的目的不是训练出一个SOTA模型而是理解每个环节怎么衔接。先看数据准备。一个batch的数据要经过这几步原始句子列表按词表转成索引序列然后做padding补到batch内最大长度。对翻译任务还需要在源序列开头加起始符、结尾加结束符。代码里常用一个collate函数来处理这些逻辑它接收一个batch的样本输出模型所需的输入张量和目标张量。比较关键的还有mask的生成。只要batch中存在padding就一定要生成对应的padding mask只要做Decoder训练就一定要生成look-ahead mask。我习惯把mask的生成逻辑单独抽成一个函数并且把形状打印出来核对。mask写错是训练效果奇差的常见原因不报错但损失降不下去令人头大。然后是模型构建。我这里定义一个最小Transformer参数设置可以参考d_model256heads8d_ff1024Encoder和Decoder各3层。这样参数量不过几百万在CPU上也能跑动调试成本很低。构建时要注意位置编码要和词表embedding维度一致否则加不上去输出层需要把Decoder的输出映射到目标词表大小的logits用的通常是Linear层。损失函数用交叉熵但要注意目标序列通常也做了padding需要把填充位置的损失给屏蔽掉。做法是使用nn.CrossEntropyLoss时传入ignore_index参数对应填充token的索引。忽略这个细节会导致模型花费大量力气去“预测”填充token训练效果明显变差。优化器用Adam学习率建议先按论文里的动态调度来warmup阶段线性上升之后按步数倒数衰减。如果觉得麻烦先用一个固定的3e-4也会有效果只是收敛略慢。训练循环的典型代码如下optimizer.zero_grad() src, tgt batch src_mask, tgt_mask create_masks(src, tgt) logits model(src, tgt[:, :-1], src_mask, tgt_mask) loss criterion(logits.reshape(-1, logits.size(-1)), tgt[:, 1:].reshape(-1)) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step()这里一个常见位置是logits和目标序列的对应关系。训练时Decoder的输入是目标序列去掉最后一个token模型预测的目标则是去掉第一个token因为解码器每个位置负责预测下一个词。不少人在这一步对齐错误导致模型学到一些奇怪规律但loss数值上看着还行。推理阶段流程不同。训练时可以并行喂整个目标序列推理时必须一个一个词生成先输入起始符得到下一个词的分布取概率最大的词作为输出再把它拼回输入继续预测直到遇到结束符或达到最大长度。这就是自回归解码。简单实现可以在for循环里完成也可以用torch.no_grad()包住加速。如果要做Beam Search就需要额外维护候选序列的分数列表代码复杂度上一个台阶新手建议先用贪心解码跑通。关于训练终止条件我实践的判断标准是验证集loss不再下降就可以停了。过多训练轮次在小数据集上很容易过拟合没必要追求过低的训练loss。模型参数保存的时候记得带上词表、配置、以及当前步数方便断点续训和复现。5. 排查问题与调参心得那些实践里被坑过的地方Transformer跑起来的第一次很顺利但后续调优时踩的坑远比想象中多。这里把自己实操中遇到的经典问题整理一下可以当作排查手册用。张量形状不匹配是最高频的错误。注意力矩阵、多头拼接、mask广播任何一个环节对不上就会在调用torch.matmul时报错。我的建议是在编写模型时先在关键位置打印shape尤其是Q、K、V计算完之后以及多头合并之后。很多问题根本不用等到训练前向一次就能发现。自己在代码里留debug开关比如环境变量控制是否输出中间形状对排查很有帮助。mask写错是另一类非常隐蔽的问题。有些错误不报错只是训练效率或者精度低。现象往往是模型loss降得还行但生成结果中包含大量填充token或者模型只关注了序列头部的词对后面的内容完全忽略。检查方法很简单把某一层的注意力权重打印出来看看padding位置是否有非零权重。如果有说明mask没生效立即检查mask的生成和传递链路。模型不收敛也常见。新手喜欢一开始就把模型堆到很大然后发现loss根本不动。我建议先缩小模型尺寸比如d_model128、层数2层跑一个小数据集确保能过拟合。能过拟合说明模型学习能力正常再逐步扩大。如果连小数据都学不动问题大概率在数据预处理或者loss计算上。另一个容易忽视的问题是学习率。Transformer对学习率比较敏感设定的学习率过大训练前期loss直接发散到NaN过小则收敛速度极慢。论文里warmup decay的调度方案不是玄学是有实际作用的。自己调参时可以把初始峰值设为1e-3warmup步数占总步数的5%左右如果出现NaN就先调低到3e-4。很多人为了省事固定学习率效果总是差一截不是没有原因的。梯度裁剪是我强烈建议加的。就算模型结构没问题数据梯度极大时也会触发NaNclip_grad_norm等于给训练上了一道保险几乎零成本。还有一个容易被新手忽略的点Dropout的位置和比例。Transformer在注意力输出、前馈网络、embedding加位置编码之后都可能加Dropout作用域很广。训练和推理行为不同推理时Dropout自动失效这一点PyTorch已经处理好了不用自己操心。如果发现训练集loss很低但验证集表现差很大概率是Dropout比例太小尝试从0.1调整到0.3。自定义学习率调度的代码如下可以封装为一个函数方便实验def lr_lambda(step): warmup_steps 4000 return min(step**-0.5, step * warmup_steps**-1.5)关于复现性固定随机种子对调试很有用。在main函数开头设置torch.manual_seed和random.seed能让两次运行的初始权重一致排错时不会因为随机性分不清是代码问题还是运气问题。最后说说我在动手实践中的体会。Transformer看似结构复杂但只要把它拆成“数据准备、模型堆叠、训练推理”三个脉络每个脉络逐点击破完整实现一个可用模型并不难。真正难的是理解每个设计选择背后的动机知道为什么这里要缩放、那里要mask、残差为什么要连、位置信息为什么要加进embedding。代码只是这些设计选择的具体表达理解了动机代码就是自然而然的事情。如果有精力下一步可以尝试自己动手改模型比如把位置编码换成可学习的或者把注意力替换成线性注意力变体观察效果差异。这种探索会比单纯读代码有意思得多也能让你对Transformer的理解更上一层楼。我个人觉得Transformer就像一套乐高基础零件就是那几个怎么拼、怎么改、怎么搭出自己的结构才是这趟学习的乐趣所在。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑