ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

精读Transformer原文:从零理解Attention Is All You Need的架构与工程细节

精读Transformer原文:从零理解Attention Is All You Need的架构与工程细节 这篇论文我从研二第一次读到现在前前后后翻了几十遍。2017年《Attention Is All You Need》刚挂出来的时候我还在用LSTM跑机器翻译当时觉得Attention只是RNN的一个辅助模块压根没想到它能把循环结构整个端掉。后来Transformer横扫NLP又席卷CV、语音、多模态甚至连推荐系统都在用我每次回看这篇原文都能读出以前忽略的细节。这次重新精读我想换个角度不按论文原本的章节顺序去讲而是站在“一个想从零复现并真正理解Transformer的人”的视角把这篇论文拆成六大模块它解决了什么问题、整体架构怎么搭、自注意力机制为什么有效、位置编码到底在做什么、训练里哪些细节决定成败、以及多年后重读有哪些容易踩的误解。把这六块啃透你就掌握了Transformer最核心的骨架以后再看BERT、GPT、ViT、扩散模型里的各种注意力变体都会觉得理所当然。这篇文章适合三类人刚入门深度学习、想认真啃下Transformer原文的学生或工程师已经会调包用BERT/大模型、但没系统读过原论文的从业者以及像我一样想“温故而知新”的算法老兵。我尽量把每个公式背后的直觉都讲明白同时把论文里那些容易被一笔带过的工程细节拿出来单独说。1. 在Transformer出现之前序列建模到底难在哪1.1 前Transformer时代的三种主流方案现在大家谈论Transformer觉得天经地义但回到2017年序列建模的主流方案是递归神经网络RNN家族尤其是LSTM和GRU。这类模型的核心思路是“逐步递推”把输入序列拆成一个个时间步每个时间步维护一个隐藏状态h_t这个状态既携带历史信息又参与当前步的输出。以机器翻译为例当时的经典框架是Encoder-Decoder也就是Seq2Seq。Encoder把源语言句子从左到右读一遍把整个句子的语义压缩进最后一个隐藏状态Decoder再从这个固定向量开始逐词生成目标语言。问题很直观句子越长前面的信息在传递过程中流失越严重。LSTM和GRU通过门控机制缓解了梯度消失但本质上仍然是“接力棒式的信息传递”每多传一步都会打折扣长距离依赖依然是大问题。还有一种思路是用卷积神经网络CNN处理序列。CNN的好处是能并行计算因为不同位置的卷积核是同时滑动的。但CNN有个天生短板局部感受野。一个卷积核只能看到窗口内的词要捕捉长距离依赖就得堆很多层卷积来扩大感受野。论文里专门对比了这一点指出要让CNN覆盖到所有输入位置需要堆叠O(n/k)层n是序列长度k是卷积核大小这个代价在长句子上相当可观。1.2 循环结构的三大痛点RNN系列的痛点我总结下来有三条这也是Transformer能一战封神的原因。第一个是顺序计算的硬约束。RNN每一步的隐藏状态依赖于前一步的结果所以无法并行。训练时GPU利用率非常难看尤其句子长的时候基本是在用时间换内存。当年跑一个翻译模型要等一周很大一部分时间都耗在串行计算上。第二个是长距离依赖的信息衰减。即便LSTM把梯度路径拉直了一些但在超长序列上靠门控“记住”早期信息仍然很吃力。论文里引用了相关研究明确指出RNN在长距离依赖建模上的瓶颈是个结构性问题不是简单堆参数就能解决的。第三个是编码瓶颈。标准Seq2Seq用一个固定维度的向量承载整个源句子的信息句子稍微复杂点这个向量就成了信息瓶颈Decoder生成时能参考的上下文极其有限。后来的注意力机制最开始就是用来解决这个问题的让Decoder在每一步生成时直接回看Encoder的所有隐藏状态而不是只靠最后一个。事实证明效果拔群但那时Attention还只是搭在RNN外面的“外挂”不是主角。1.3 论文给出的核心答案《Attention Is All You Need》的核心主张一句话就能概括既然Attention这么有用那就干脆全用Attention把循环结构整个扔掉。作者设计了一个完全基于注意力机制的架构叫Transformer实现了两个RNN做不到的目标一是所有位置的Token可以完全并行计算二是任意两个Token之间的依赖路径长度恒为1不受序列长度的拖累。注意“恒为1”这个表述很关键。在RNN里第i个词要影响第j个词j远大于i信息必须沿着隐藏状态一步步传过去路径长度是|j-i|在Transformer里第i个词直接和所有其他词做Attention交互一步到位。论文把这个路径长度上升到了“学习长距离依赖”的核心衡量指标这也是后来很多做长文本模型的论文都会沿用的一套分析框架。2. Transformer整体架构一眼看懂全局2.1 Encoder-Decoder框架的沿袭与改造Transformer并没有发明全新的“编码-解码”范式而是沿用了机器翻译里成熟的Encoder-Decoder结构。整体分两半左边的Encoder负责把输入序列编码成一组向量表示右边的Decoder负责基于这些表示逐词生成输出序列。这个框架的好处在于任务边界清晰——编码器只负责理解解码器只负责生成。但内部的模块和RNN时代有了本质区别。Encoder和Decoder都不再是循环神经网络而是由若干个相同的层堆叠而成。论文里用了一个很形象的词叫“块”block每个块内部包含多头注意力、前馈网络、残差连接和层归一化。Transformer没有用“时间步”这个概念取而代之的是“位置”和“Token”这让整个计算图变得像一张真正的图结构每一条边都是可以并行计算的。2.2 六大关键组件先建一个全局清单我在第一次看论文时最大的困惑是“注意力、残差、归一化、前馈网络这么多东西优先级是什么”。后来我总结了一张全局清单用这个清单对照架构图读效率会高很多。多头自注意力Multi-Head Self-Attention让每个Token观察其他所有Token并聚合信息。前馈网络Feed-Forward NetworkFFN对每个Token独立地做非线性特征变换。残差连接Residual Connection防止深层网络梯度消失让信息可以无损跨层传递。层归一化Layer Normalization稳定训练过程的数值。掩码多头注意力Masked Multi-Head AttentionDecoder特有的保证生成第t个词时看不到t之后的词。位置编码Positional Encoding给注意力机制补充顺序信息。我建议你画一张横向流程图从左到右把Input Embedding、位置编码、多头注意力、残差、LayerNorm、FFN、输出映射依次排开再把Decoder里的Masked版本也画上去整篇论文的物理结构就八九不离十了。2.3 输入输出是怎么变成向量的Transformer本身并不理解文字所有输入必须先转换成向量。论文里的做法是先做Token Embedding也就是把每个词或子词映射成一个d_model维的向量论文里d_model默认512。这个Embedding矩阵是可学习的训练一开始是随机初始化随着训练不断更新。但光有词向量还不够。注意力机制本身对Token的“顺序”完全不敏感它只看Token两两之间的相似度所以如果我们直接把词向量丢进去“我爱你”和“你爱我”在Transformer眼里没有区别。为了解决这个问题论文在输入Embedding上叠加了一个位置编码Positional Encoding用不同频率的正弦和余弦函数生成每个位置独有的向量然后再和Token向量相加。这一步在论文里属于“最不显眼但绝不能少”的设计后面我会单独拿出一节来聊。Decoder侧的输出类似它每生成一个Token就把已经生成的所有Token作为输入经过多头注意力处理最后接一个线性层加Softmax输出下一个词的概率分布。这个“自回归”式的生成方式和老式Seq2Seq没有本质区别区别在于内部计算用的全是注意力。3. 自注意力机制与多头设计Transformer最核心的引擎3.1 从“字典查找”的视角理解Q、K、V自注意力机制Self-Attention是整个Transformer的心脏很多人在理解它时卡在Q、K、V这三个字母上。其实把它们想象成一个“档案馆里查资料”的过程就好懂了。Query查询像你手里拿着的问题Key键是档案馆里每份档案的标签Value值是档案的实际内容。你拿着Query去和每个Key做相似度匹配匹配的分数越高说明这份档案和你的问题越相关于是你就花更多注意力去读取这份档案的Value。在序列模型里每个Token都会作为Query去查询序列里的所有其他Token同时它自己也是一个Key和Value被别的Token查询。也就是说在一次自注意力里序列中每个Token都既是查询者又是被查询的档案。具体到计算对于第i个Token它要计算一个输出向量这个向量的计算方式是先用它的Query和序列里每个Token的Key做点积得到一组未归一化的匹配分数再用Softmax把这些分数变成总和为1的权重最后把每个Token的Value按对应权重加权求和。整个过程可以压缩成一个公式Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V。理解这行字就理解了Transformer的40%。3.2 为什么要除以sqrt(d_k)小数字大玄机公式里有一个容易被忽略但极其关键的设计——缩放因子1/sqrt(d_k)。论文里给的解释是当d_k比较大时点积结果的方差会变大导致Softmax函数的输入落在梯度极小的饱和区训练会变得很慢。这是一个非常工程化的解释我最初读的时候没太在意后来真去跑了实验才发现不除以这个缩放因子模型在训练初期几乎不收敛Loss下降慢得像蜗牛。从数学上看假设Q和K的每个分量都是均值0、方差1的独立随机变量它们的点积相当于d_k个这样的乘积相加方差会变成d_k。除以sqrt(d_k)后方差又回到1Softmax的输入分布保持稳定。这个“保持方差恒定”的思路和很多初始化方法如Xavier初始化异曲同工。你可以在代码里试一下同样的多头注意力层去掉缩放因子和不去掉前几百步的Loss差异非常明显。3.3 为什么一个注意力头不够要“多头”论文里用的不是单个注意力而是多头注意力Multi-Head Attention。做法是把d_model维的Query、Key、Value经过不同的线性映射投影成h组低维向量每组独立计算注意力最后拼接起来再过一次线性层。论文默认h8所以每个头的维度是d_k d_v d_model / h 64。为什么非要拆成多个头我的理解是这样一种注意力只能学到一种“关系模式”而语言里的关系太复杂了。有的头可能擅长捕捉语法依赖比如动词和它的主语有的头可能负责捕捉指代关系比如代词指向的名词有的头则在前缀和后续内容之间建立桥梁。多头相当于请了一群各有所长的“小专家”每个专家负责一种观察视角最后汇总起来得到更全面的表示。论文里还专门做了可视化实验展示了不同注意力头关注的词对效果非常直观。我建议初学者都可以复现一下这个可视化你会发现有些头专注句法关系有些头专注远程依赖这种“各司其职”的分工机制比单头注意力强大的多。后来很多研究也都证实多头之间确实存在功能分化不是单纯的重复计算。3.4 Decoder里的掩码注意力防止看到未来Decoder里的注意力有两个特殊之处。第一个是它接收的Query来自Decoder当前生成的部分而Key和Value来自Encoder的最终输出这是一种“交叉注意力”Cross-Attention它的作用是让解码器每一步生成时都能从源语言句子中获取相关信息。第二个特殊之处在Decoder的“自注意力”部分它需要加一个掩码Mask。掩码的逻辑很简单当Decoder预测第t个词时它不应该看到第t个位置之后的任何信息否则就成了作弊。具体做法是在计算QK^T之后、Softmax之前把未来位置的分数设为一个极大的负数比如-1e9这样Softmax之后它们对应的权重就趋近于零。这个技术在训练时特别重要因为训练阶段是并行输入的我们一次性给Decoder整个目标句子的嵌入如果没有掩码每个位置都能看到后面的词预测任务毫无意义。加上掩码后每个位置只能看到自己和之前的位置训练效率才能拉满。有个实践层面的小坑在实现Transformer训练时Mask的形状和位置一不注意就会出错。我的建议是先把“Batch x Head x Seq x Seq”这个分数矩阵画出来标出哪些位置要mask掉再写代码非常省事。4. 位置编码给无序的注意力补上顺序感4.1 自注意力天生没有“顺序概念”前面提到自注意力计算的是“两两之间的关系”它天然是对排列不敏感的。把“我打你”三个词随便换顺序每个词和其他词的交互矩阵完全一样。这对于翻译任务来说就是灾难“我打你”和“你打我”意思完全相反。RNN天然有顺序感是因为它按顺序往前进CNN有顺序感是因为卷积核从左往右滑动而Transformer从头到脚都是注意力必须额外引入位置信息。论文选择了在输入侧直接加入一个位置编码向量和Token Embedding相加后一起进入网络。需要注意的是这个加法的前提是位置编码的维度和Token Embedding维度完全一致都是d_model512。这样每个输入Token就同时携带了“我是谁”和“我在哪”两重信息。4.2 三角函数编码的数学直觉论文里的位置编码公式是这样的对位置pos编码维度2i取sin(pos / 10000^(2i/d_model))对维度2i1取cos(pos / 10000^(2i/d_model))。用不同频率的三角函数来生成位置编码有两个好处。第一是值域有界所有值都落在[-1, 1]之间相对稳定不会像可学习的位置向量那样出现极端值。第二也是更漂亮的数学性质第posk个位置的位置编码可以被第pos个位置的位置编码的一个线性变换表示。这意味着模型可以很容易地学会“相对位置”关系——它不需要牢牢记死每一个绝对位置只需要通过线性映射就能捕捉到位置之间偏移量。这个性质对长句子的泛化特别重要因为位置编码一旦与训练时的长度绑定太死对未见过的新长度就很不友好。论文也提到他们也试过可学习的位置编码实验效果和三角函数版本几乎一样。于是在后续的BERT等预训练模型里大家直接用可学习位置Embedding省心且效果相当。两种方案我有都亲测过可学习版本在小规模数据上更容易收敛三角函数版本在序列长度外推上略占优势实际选型看业务需要。4.3 一个“常学常新”的细节重读这篇论文时我注意到一个每次看都有新感受的细节位置编码只被加在Encoder和Decoder最底层的输入处而不是每一层都加。这样做的隐含假设是底层的位置信息会随着层数的加深被逐层传递和增强不需要每层重复注入。后来很多变体比如Transformer-XL里的相对位置编码打破了这一点但在原始Transformer里这个“只加一次”的简单选择已经足够。另一个细节是位置编码和Token向量“相加”而不是“拼接”。相加意味着两部分信息共享同一批维度这会压缩一些表达能力但它保持了模型参数规模不变计算代价也更低。论文作者显然做过权衡与其增加维度搞拼接不如让网络自己学会如何去分离位置信息和语义信息。事实上训练完成后这两个信息确实可以被注意力头分别利用这也是多头注意力“各司其职”的一个侧面证据。5. 训练细节与工程干货论文里容易被跳过的重点5.1 学习率预热机制Learning Rate Warmup论文中有一个很容易被忽略的训练细节动态学习率。具体公式是 lr d_model^(-0.5) * min(step^(-0.5), step * warmup_steps^(-1.5))其中warmup_steps默认4000步。这个学习率策略分两个阶段前4000步学习率线性从0增长到峰值4000步之后学习率按step^(-0.5)的比例逐步衰减。我在实际训练Transformer时几乎无脑沿用这套调度非常稳健。为什么要预热直觉上是因为Transformer在第一步时模型各部分还没有“热身”参数处于一个比较随机的状态此时如果学习率过大很容易把模型推向一个糟糕的局部区域。用一个较小的学习率先走几步等模型有了基本的梯度方向再逐步加大步长能显著减少训练初期的震荡。这个细节在论文的实验部分并不显眼但对复现结果的影响极大。我用base配置训练英语到德语翻译模型时去掉预热直接上大学习率训练Loss在前几百步就发散了不得不重启。后来老老实实加载了预热的调度器才复现出接近论文的BLEU分数。5.2 正则化组合拳Residual Dropout与Label Smoothing论文里用了两种正则化手段一个是Dropout一个是标签平滑Label Smoothing。Dropout应用在两个位置一是在每个子层的输出、进入残差连接和LayerNorm之前对数据做随机置零二是在Embedding和位置编码相加之后、进入第一个注意力层前也做一次Dropout。论文默认的dropout_rate是0.1数值不大但训练时没有这层扰动模型过拟合的风险明显上升尤其在训练数据量不够大的时候。Label Smoothing更有趣论文里把平滑系数设为epsilon_ls 0.1。传统训练使用one-hot标签目标词概率为1其他词为0而标签平滑把一部分概率“匀”给其他词让模型不要对训练数据过于自信。论文还专门指出标签平滑会让困惑度Perplexity稍微变差但BLEU分数却会提高。原因是过分自信的模型在生成时容易输出过于确定但未必正确的词而平滑后的模型保留了更多的“选择余地”泛化更强。这个反直觉的经验不跑实验真体会不到。5.3 关键超参数对照表论文里给出了base和big两套配置很多后续工作都是在base上做增量改进。我把关键参数列成表格方便对照实现。参数量base约6500万big约2.13亿d_model向量维度base 512big 1024头数hbase 8big 16前馈网络维度d_ffbase 2048big 4096Dropout比率0.1两套相同标签平滑系数0.1两套相同批大小约25000个源Token和25000个目标Token组成一个batch训练轮数base约10万步big约30万步对应具体GPU时间big模型在机器翻译任务上相比base有稳定提升但训练时间和显存消耗也明显上涨。复现时建议从base起步先把每个模块的正确性验证无误再考虑加大模型。5.4 实验结论给我们的三点启示论文报告了两个机器翻译任务的实验结果英德翻译BLEU达到28.4英法翻译BLEU达到41.8两个分数都超越当时最先进的模型同时训练成本只有对手的一小部分。论文明确写了base模型在8张P100 GPU上训练3.5天big模型训练3.5天换算成训练成本大约是当时最强SOTA模型的几十分之一。我从实验结果里读出的核心信息有三点第一去掉循环结构不会让模型变笨反而让信息流动更高效第二训练成本的大幅下降主要来自并行化这在算力有限的团队里简直就是救命稻草第三这些指标在当年是SOTA在今天当然已经被大模型远远超越但论文里提出的“并行化全局依赖结构简单化”的工程思想仍然影响着现在的模型设计。6. 重读多年后我踩过的最常见误解6.1 常见误解一自注意力就是“找相似词”我以前在讲课时很多人第一反应以为自注意力就是“计算两个词之间的相似度”。这个理解太片面了。自注意力输出的是一个加权求和后的向量它真正做的是“从整个序列中提取与当前Token相关的上下文信息”而不是单纯给两个词打一个相似度分数。相似度计算只是第一步更关键的是后续的聚合和变换。类比一下一个研究员不只是判断哪本参考资料和自己课题相关更要把参考书里的精华摘出来综合成一份新的调研报告。6.2 常见误解二多头注意力各头是等价的有人以为多头注意力就是把一个注意力计算重复8遍然后把结果拼起来而已。实际上多个头在训练过程中会自发分化各自关注不同类型的依赖关系。论文里的可视化已经直观展示了这一点后来也有研究用探针分析验证了这种分工。理解这一点对设计后续的稀疏注意力、混合注意力帮助很大。6.3 常见误解三位置编码只是“给词加个位置标签”三角函数位置编码并非简单的独热编码或整数计数它背后承载了相对位置信息可被线性表示的性质。如果你只把它看成“位置ID”就不会理解为什么它能对长序列有一定泛化能力也无法理解后来RoPE、ALiBi等位置编码变体的设计动机。6.4 复现时最容易踩的三个工程坑我在复现过程中遇到过不少问题挑三个有代表性的分享一下。第一个坑是Mask的维度匹配错误。多头注意力里Mask的形状应该是(Batch, Head, Seq, Seq)但在实际代码里经常出现少一个Head维度的情况广播机制又不会报错导致模型行为和预期完全不同。排查这类问题的一个技巧是单独拿一个小张量跑一遍Forward打印中间结果一步步核对shape。第二个坑是Dropout与训练/评估模式切换。Transformer里Dropout很密集如果在评估时忘了调用model.eval()切换模式Dropout仍然生效生成结果会带着随机性复现实验时甚至会出现多次结果不一致的状况。这个问题老手也会栽。第三个坑是LayerNorm与BatchNorm的混用。Transformer用的是LayerNorm是沿着特征维度d_model做归一化和CNN里常用的BatchNorm完全不同。有些框架默认的归一化接口容易搞混一旦用错训练稳定性和收敛速度都会受到明显影响。6.5 论文留下的“开放式作业”每次重读这篇论文我都会问自己几个问题为什么Transformer能扩展到大模型注意力计算O(n²)的复杂度瓶颈有没有更优解法位置编码还有哪些更好的设计这些问题未必在原文里有明确答案但它们都是由这篇论文引出的思考方向。从GPT系列到Llama从Swin Transformer到ViT从FlashAttention到线性注意力后辈们的每一步创新几乎都可以追溯到这篇论文所开创的范式。对我来说《Attention Is All You Need》已经不单纯是一篇学术论文更像是一座里程碑提醒我“简洁的设计能走多远”。每次重读我都会提醒自己不要只是重复论文里的结论要去追问那些“为什么”把公式变成自己脑海里能顺畅运行的系统。最后再分享一个我自己的小习惯每次看到一篇和注意力机制相关的新论文我都会回到这篇原始论文去对比看看新方法是在哪里做了改动又是为什么改动。这样一来新论文的很多创新点都会变得更清晰理解深度也远高于“刷论文标题”。这篇Transformer原文常读常新每一次细读都值得。
返回列表