Transformer自注意力机制解析：复杂度、多头注意力与复现排查-北京尧图网络

简介《Attention Is All You Need》NIPS 2017 论文的中文翻译文档面向自然语言处理学习者、Transformer 入门者及需要追溯自注意力机制原始论述的研究人员帮助跨越英文原文阅读门槛系统理解完全基于注意力、免除递归与卷积的序列转导架构。压缩包共 1 个 docx 文件约 856KB内容为论文中文译本便于检索、批注与对照阅读。已有 1194 人学习下载具备一定的参考热度。文档保留摘要、引言、背景、模型架构、编码器与解码器堆栈、缩放点积注意力、多头注意力等核心章节并呈现 WMT 2014 英德 28.4 BLEU、英法 41.0 BLEU 等实验结果适合用于论文精读、课程讲义整理和翻译质量对照也能为复现实验与理解并行化优势提供原文级参考。1. 反直觉的起点一个没有循环的网络凭什么把序列任务做得更好2017 年那篇会议论文给出的结论第一次读会觉得不讲道理把编码器—解码器里的循环层整个删掉只留 self-attention 和逐位置前馈网络翻译质量反而更高。WMT 2014 英德任务上大模型拿到 28.4 BLEU比当时最好的结果含集成高出 2 个点以上英法任务 41.0 BLEU八张 P100 训练 3.5 天基础模型只要 12 小时。更值得琢磨的不是“它能跑”而是“去掉递归之后为什么既更快又更准”——这背后是三笔账并行度、路径长度、每层计算量。搞懂这三笔账你才知道今天所有注意力架构的取舍边界在哪。2. 自注意力的账本复杂度、并行度与路径长度三笔账2.1 递归层被记账的三项隐性成本循环网络处理序列的方式本质上是把“位置”和“时间步”绑死第 t 个隐藏状态 h_t 必须等 h_{t-1} 算完才能算。这个依赖关系在推理时看不出问题但在训练时是致命的——一条长度为 n 的序列前向要串行 n 步反向传播还要沿时间展开再串一遍。原文说得很直白这种固有的顺序性排除了训练样本内的并行化而序列一长显存就限制了跨样本的批处理你连“用 batch 换时间”这条路都被堵死。有人试过用分解技巧和条件计算来缓解计算效率确实提升了但“顺序计算”这个根本限制还在原地。这是第一笔账递归的成本不在 FLOPs 上而在“不能同时算”上。GPU 的算力是按并行吞吐卖的你让它一步一步来等于买了八车道只开一条。第二笔账是路径长度。信号从输入位置 i 传到输出位置 j在循环网络里必须走 |i-j| 步步数越多梯度在反向传播时经过的连乘就越多长距离依赖就越难学。这不是玄学是梯度流文献里反复讨论过的老问题。第三笔账才是每层复杂度而这一笔恰好是自注意力唯一“亏”的地方——后面单独说。2.2 复杂度表怎么读O(n²d) 到底亏在哪、赚在哪原文表 1 把三种层类型并排放自注意力每层复杂度 O(n²·d)循环层 O(n·d²)卷积层 O(k·n·d²)顺序操作数分别是 O(1)、O(n)、O(1)任意两位置间最大路径长度分别是 O(1)、O(n)、O(log_k n)。先看亏损项。自注意力的 n² 来自“每个位置都要和所有位置算一次相似度”序列翻倍计算量翻四倍。这是硬伤也是后来所有长上下文优化稀疏注意力、窗口注意力、线性注意力要还的债。但注意表里的对比条件n 是序列长度d 是表示维度。机器翻译里的实际情况是 n 通常远小于 d——词块表示、字节对编码之后一句话几十个 token而 d_model 是 512。当 n d 时n²·d 反而小于 n·d²自注意力在每层计算量上并不吃亏。再看收益项这一项权重更大。顺序操作数 O(1) 意味着整层可以在一个矩阵乘法里同时算完GPU 利用率直接拉满路径长度 O(1) 意味着不管两个位置隔多远它们之间的信号只经过一次注意力加权梯度不用走长征。原文把这个当成核心论据输入输出任意位置组合之间的路径越短长程相关性越容易学。卷积那条路也值得看一眼。单层卷积核宽 k n 时连不上所有位置对要堆 O(n/k) 层连续核或 O(log_k n) 层膨胀卷积路径长度随之变长而卷积层通常还比循环层贵一倍。可分离卷积能把复杂度压到 O(k·n·d n·d²)但即便 k n它的复杂度也就和“自注意力逐位置前馈”打平——既然打平那选路径更短的那个。层类型每层复杂度顺序操作数最大路径长度自注意力O(n²·d)O(1)O(1)循环层O(n·d²)O(n)O(n)卷积层O(k·n·d²)O(1)O(log_k n)受限自注意力O(r·n·d)O(1)O(n/r)表格最后一行是原文留的后门把注意力限制在以输出位置为中心、大小为 r 的邻域内复杂度降到线性代价是路径长度涨到 O(n/r)。这条思路后来被大量长文本模型继承但在 2017 年它只是“未来工作”里的一句伏笔。2.3 路径长度决定了梯度能不能走回家把路径长度单独拎出来说因为它最容易被当成理论装饰。实际训练里路径长度影响的是“梯度到达某个参数时还剩多少信号”。循环网络里位置 1 的信息要影响位置 50 的输出反向时得穿过 49 次隐藏状态转移的雅可比连乘只要谱半径稍微偏离 1梯度就指数衰减或爆炸这就是当年 LSTM 门控机制存在的理由。自注意力把这条链压成一步输出位置 j 对输入位置 i 的依赖直接体现在注意力权重上反向只有一条路径。代价是这条路径上的权重是 softmax 出来的所有位置的权重和为 1相当于做了一次加权平均——原文自己承认这会“降低有效分辨率”。平均是一种信息压缩位置一多每个位置的贡献就被摊薄。多头注意力就是用来对冲这个副作用的不是让一个头去做更精细的平均而是让八个头各自在不同的表示子空间里做平均再把结果拼起来。这部分留到下一章拆。我一般会这么理解这三笔账并行度是“能不能吃满硬件”路径长度是“能不能学得动”每层复杂度是“跑一轮要多少钱”。循环网络在第一项上输得最惨在第二项上勉强及格第三项看着便宜但省下的钱全被串行等待吃回去了。Transformer 的选择是拿第三项的一部分余量去换前两项的大幅改善。3. 缩放点积与多头softmax 前的那个 √d_k 在防什么3.1 点积的方差随维度线性增长注意力函数的定义很朴素查询 q 和一组键值对 (k, v) 映射到输出输出是 v 的加权和权重由 q 与对应 k 的兼容性函数算出。原文选的是点积公式是 softmax(q·kᵀ / √d_k)·v。这里那个 √d_k 不是拍脑袋加的它防的是一件很具体的事。假设 q 和 k 的各分量独立、均值 0、方差 1那么点积 q·k Σ q_i·k_i 是 d_k 个独立同分布项之和方差就是 d_k标准差是 √d_k。d_k 64 时点积的典型量级在 ±8 附近d_k 512 时量级到 ±22。softmax 对输入尺度极其敏感输入一大最大项就吃掉几乎全部概率输出接近 one-hot其余位置梯度趋近于零。原文的说法是“把 softmax 推入梯度极小的区域”翻成训练现场的话就是模型看着在学实际大部分注意力头已经不再更新。除以 √d_k 之后点积的方差被拉回 1 附近softmax 工作在梯度健康的区间。这就是为什么原文说小 d_k 时加性注意力和点积注意力表现相近大 d_k 时点积会变差——差的就是这个尺度。import math import torch import torch.nn.functional as F def scaled_dot_product_attention(q, k, v, maskNone): d_k q.size(-1) # 关键一除以 sqrt(d_k)把点积方差从 d_k 拉回 1 附近 scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: # 关键二非法位置填 -inf而不是 0 scores scores.masked_fill(mask 0, float(-inf)) attn F.softmax(scores, dim-1) return torch.matmul(attn, v), attn这段是几乎所有框架里scaled_dot_product_attention的骨架两个参数值得盯住。d_k取的是 q 的最后一维也就是每个头的键维度不是 d_model——多头切分之后每个头只有 64 维缩放因子必须跟着头走。mask的填充值是-inf而不是 0原因在下一章讲因果掩码时展开softmax(0) 是一个均匀分布等于让模型去“关注”本不该看的位置只是权重被摊平了而已。顺便说一个常被忽略的细节原文提到点积注意力在实践中比加性注意力快得多、空间效率也高得多因为它能直接调用高度优化的矩阵乘法例程。这不是工程上的小便宜而是设计动机的一部分——一个在理论上等价、但在硬件上跑不动的机制在工程上就是不成立的。3.2 多头不是 ensemble是把表示空间切开单头注意力的输出是“所有位置加权平均”的一个向量这个向量被迫承载该位置上所有类型的关系句法依存、指代消解、局部搭配、语义角色全挤在同一个 512 维空间里互相干扰。原文的观察是用单一注意力头这种平均会抑制掉一部分信息。多头的做法是把查询、键、值分别用 h 组不同的、可学习的线性投影降到 d_k、d_k、d_v 维在这 h 份投影上并行跑注意力各自得到 d_v 维输出拼接后再投影一次得到最终结果。原文的参数是 h 8d_k d_v d_model / h 64。关键在于“总计算成本与全维单头注意力相似”——因为每个头的维度降到了 1/8八份加起来和一份 512 维的点积量级相当。这个设计的真正含义是不同头在不同的表示子空间里做注意力模型因此可以同时关注来自不同位置的不同表征信息。原文附录里检查过注意力分布发现单个头确实学会了执行不同任务不少头还表现出与句法、语义结构相关的行为。这句话的分量在于“可解释性”是作为副作用被提到的不是设计目标——先有机制后有人去观察它学到了什么。多头切分有个容易翻车的边界h 必须整除 d_modeld_k 就是整除的结果。原文做了消融单头比最佳设置差 0.9 BLEU但头太多也会掉点——切得太碎每个头的 d_k 太小兼容性函数的表达力不够。8 头 64 维是这篇论文在 512 维模型下的甜点不是普适常数。3.3 一份注意力三种连线方式同一个多头注意力模块在模型里被接成了三种不同的形态区别只在 q、k、v 从哪来。编码器—解码器注意力层查询来自前一个解码器层键和值来自编码器输出。这是最接近传统序列到序列模型里那种跨注意力的一层作用是让解码器每个位置都能覆盖输入序列的所有位置。注意这里 k 和 v 是同一份编码器输出q 是解码器自己的状态——这个非对称结构是整个翻译任务的“对齐”发生的地方。编码器自注意力层q、k、v 全部来自编码器前一层的输出每个位置可以和前一层所有位置交互。解码器自注意力层结构相同但多了一条限制每个位置只能关注到它自己及之前的位置。这条限制是靠掩码实现的原文说得很清楚——在缩放点积注意力里把对应非法连接的所有值设为 -inf。第三种形态最容易被误解成“和第二种一样”。差别在掩码编码器不需要掩码因为它的输入是完整的源句解码器必须掩码否则训练时模型能直接看到正确答案的后半段loss 会掉得飞快推理时立刻原形毕露。这种“训练指标漂亮、上线全崩”的翻车绝大多数是掩码写错了。4. 位置编码、掩码与训练配方顺序信息和稳定性从哪来4.1 正弦位置编码为什么能被线性表示模型里没有循环也没有卷积那就没有任何东西告诉它“谁在前谁在后”。自注意力本身是置换等变的把输入序列打乱输出只是跟着打乱权重一模一样。顺序信息必须外挂进去。原文的做法是在编码器和解码器栈底部的输入嵌入上加一份同维度的位置编码直接相加而不是拼接——因为两者维度都是 d_model 512相加之后维度不变后续所有层不用改。公式用的是不同频率的正弦余弦import math import torch def positional_encoding(max_len, d_model): pe torch.zeros(max_len, d_model) pos torch.arange(0, max_len).unsqueeze(1).float() # 分母按 10000^(2i/d_model) 增长等价于波长构成几何级数 div torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(pos * div) # 偶数维走 sin pe[:, 1::2] torch.cos(pos * div) # 奇数维走 cos return pepos是位置i是维度下标偶数维用 sin(pos/10000^(2i/d_model))奇数维用 cos 同一个频率。每一对 sin/cos 对应一个正弦波波长从 2π 到 10000·2π 构成几何级数低维变化快高维变化慢相当于给每个位置生成一个二进制式“指纹”同时保留粗细两种粒度。选这个函数而不是可学习位置嵌入原文给的理由是对任意固定偏移 kPE(posk) 可以表示成 PE(pos) 的线性函数模型因此容易学会相对位置关系。这个性质来自三角函数的和角公式——sin(ω(posk)) sin(ωpos)cos(ωk) cos(ωpos)sin(ωk)是一个只和 k 有关的线性变换。换句话说模型不需要为每个绝对位置单独记忆只要学一个“平移算子”就能处理相对距离。原文还做了对照实验换成可学习的位置嵌入结果几乎相同表 3 行 E。仍然选正弦版本是因为它可能让模型外推到比训练时更长的序列——固定函数在训练没见过的位置上依然有定义可学习嵌入在那些位置上是没训练过的随机向量。这个取舍的收益在 2017 年的翻译任务里不明显但它的思路后来被各种长度外推方案反复引用。4.2 因果掩码解码器不能偷看未来掩码这件事机制上一句话能说完实现上错了就是灾难。解码器自注意力层要保证位置 i 的预测只依赖小于 i 的已知输出做法是在 softmax 之前把注意力分数矩阵的上三角不含对角线全部置为 -inf。加上输出嵌入整体右移一位模型在预测第 i 个词时手里只有前 i-1 个词。def causal_mask(size): # 下三角为 True可见上三角为 False屏蔽 return torch.tril(torch.ones(size, size)).bool() # 用法mask 形状 [1, 1, size, size]可广播到 [batch, heads, size, size]torch.tril取下三角行是查询位置列是键位置第 i 行第 j 列为 True 表示位置 i 可以看位置 j。对角线保留因为当前位置自己要参与计算。形状上补两个 1 维是为了广播到 batch 维和头维每个头共用同一份掩码。为什么填充值必须是 -inf 而不是 0因为 softmax 之后要归一化某一行的非法位置如果填 0exp(0)1它们会分走一部分概率质量合法位置的权重被稀释模型等于在看一堆“内容无意义但有权重”的位置。填 -infexp(-inf)0非法位置的概率严格为零。这是踩坑里最常见的一类——代码能跑、loss 能降只是质量莫名其妙差一截。4.3 embedding 缩放与权重共享原文有两处容易被跳过的实现细节。第一处嵌入层权重乘以 √d_model。原因是位置编码的数值量级在 ±1 附近而随机初始化的嵌入向量分量方差是 1/d_model量级偏小不缩放的话位置信号会盖过词义信号模型前期基本在学位置、不在学语义。乘 √d_model 把两者拉到同一量级。第二处编码器嵌入层、解码器嵌入层、解码器输出前的线性变换三者共享同一个权重矩阵。这个做法在当时的语言模型里已有先例逻辑是“输入的词向量空间和输出的打分空间应该是同一个空间”共享既省参数又起到正则作用。注意共享的是那个 d_model × vocab 的矩阵而 pre-softmax 那层是 vocab × d_model实现上取的是转置。4.4 学习率 warmup 与正则化三件套优化器用 Adam学习率不是常数而是一条先线性上升、再按步数平方根倒数下降的曲线warmup 步数 4000def lr_lambda(step, d_model512, warmup4000): step max(step, 1) # 防止第 0 步除零 # 前 warmup 步线性升温之后按 step^-0.5 衰减 return (d_model ** -0.5) * min(step ** -0.5, step * warmup ** -1.5)min的两个分支对应两个阶段step * warmup^-1.5随步数线性增长step^-0.5随步数衰减交点正好在 warmup 步。系数 d_model^-0.5 让曲线的绝对高度跟着模型宽度走宽模型初始学习率更小。为什么要 warmup注意力层的参数在初始化时各位置的分数接近均匀梯度方向噪声很大直接上大学习率容易在最初几百步把参数推到坏区域之后再怎么降都救不回来。这条经验后来变成了训练所有 Transformer 类模型的默认动作。正则化用了三样残差 dropout加在每个子层输出上、加回子层输入并归一化之前嵌入与位置编码之和上的 dropout以及标签平滑。基础模型 dropout 率 0.1大模型英德用 0.3英法用 0.1。标签平滑的取值原文写在训练小节效果是“伤害困惑度但提升准确率和 BLEU”——因为模型学会了不那么确定输出分布不再逼近 one-hot泛化更好。checkpoint 平均也很关键基础模型平均最后 5 个、大模型平均最后 20 个检查点间隔 10 分钟。推理用 beam searchbeam 大小 4长度惩罚 α 0.6最大输出长度设为输入长度 50能提前终止就提前终止。这些超参原文说是开发集上调出来的——翻译成工程语言它们不是模型的一部分是后处理的一部分但足以让 BLEU 差出一个点。4.5 表 3 的消融怎么读原文表 3 是一张信息密度很高的表逐行看能读出设计边界。行 A 变注意力头数和键值维度保持计算量不变单头比最佳设置差 0.9 BLEU头太多也掉点。行 B 单独减小 d_k质量明显受损原文的解读是“确定兼容性并不容易比点积更复杂的兼容性函数可能有益”——这句话给后来的各种注意力变体留了口子。行 C、D 加大模型符合预期地变好同时 dropout 那一列显示没有 dropout 时过拟合非常明显。行 E 换成可学习位置嵌入结果与正弦版本几乎一致。表 3 最有价值的不是哪个配置最好而是哪几个旋钮拧了会疼d_k 太小会疼头太多会疼没有 dropout 会疼。这三个都是实现时最容易“为了省事”而动的地方。5. 复现排查五类最容易翻车的现象与定位路径5.1 loss 在前几千步突然发散或出 NaN现象前几百步 loss 正常下降到两三千步突然抬升几轮之后变 NaN。原因通常是学习率曲线没做 warmup或者 warmup 步数设得太短注意力分数在初始化阶段接近均匀、梯度噪声大大学习率把参数推进了饱和区。另一条常见原因是缩放因子写错——用 d_model 而不是 d_k 去除多头切分后缩放不足softmax 早早饱和。解决先确认缩放因子取的是每个头的键维度再把 warmup 拉到 4000 量级同时打印每一步的梯度范数看它在哪个区间开始异常放大。5.2 解码输出重复、乱序、停不下来现象翻译结果里同一个词反复出现或者词序明显乱或者生成到最大长度才停。原因多半在掩码如果解码器自注意力的掩码用了 0 填充而不是 -inf模型能看到未来位置的一部分概率质量训练时 loss 很低推理时没有“未来”可看行为立刻劣化。另一种可能是输出嵌入没有右移一位模型在学着复制当前输入。解决把掩码矩阵打印出来确认第 i 行只有前 i 列为可见再检查训练时的输入是不是[BOS] y[:-1]这种错位形式。5.3 位置信号像没生效现象模型能翻译但调换词序、处理长句时表现明显变差注意力图上看不出任何位置结构。原因是嵌入没有乘 √d_model词向量量级远小于位置编码位置信息在加法时被淹没或者位置编码加在了错误的位置比如加在归一化之后。解决检查嵌入层输出和位置编码的标准差是否在同一量级正常应该在 1 附近再确认加法发生在 dropout 之前、第一层之前。5.4 长序列显存直接爆现象序列长度到一两千显存占用陡增batch 只能降到 1。原因是自注意力的 n² 复杂度——注意力分数矩阵是 [batch, heads, n, n]n 翻倍这项占用翻四倍。这不是 bug是设计代价。解决推理侧可以分块计算注意力训练侧考虑限制注意力窗口或用稀疏模式如果只是想让 batch 大一点先降 heads 数或序列截断别急着换模型。5.5 复现指标和原文对不上现象训练 loss 曲线看着正常BLEU 就是差一两个点。原因往往不在模型而在评估链路忘了做 checkpoint 平均、beam 大小用的是 1 而不是 4、长度惩罚没设、或者 BLEU 的分词方式和原文不一致。原文那 28.4 和 41.0 都是“大模型最后 20 个检查点平均 beam 4 α0.6”这套组合的结果拆掉任何一项都会掉。解决先把后处理链路对齐再动模型结构否则你永远不知道差的那一点来自哪。6. 设计取舍二次复杂度是这份优雅的账单Transformer 的漂亮之处在于它把“序列建模”这件事统一成了一种操作查、键、值的加权求和。统一带来的是可并行、路径短、结构简单——六个相同的层每层两个子层残差加层归一化没有门控、没有时间步展开。但统一也意味着它必须为所有位置对付出 n² 的代价这个代价在 2017 年的翻译任务里可以忽略n 远小于 512在今天的长上下文场景里就是主矛盾。回头看原文留下的两条伏笔其实已经承认了这笔账。一条是受限自注意力把每个位置限制在半径 r 的邻域内复杂度降到线性路径长度涨到 O(n/r)——用可学习性换计算量。另一条是位置编码的外推性选正弦而非可学习嵌入理由是可能处理比训练更长的序列。这两条都在说同一件事模型的结构假设和任务的长度分布必须匹配不匹配的时候优雅会变成负担。验证这套机制是否真的按预期工作我一般会做两件事。一是把注意力权重矩阵抽出来看编码器自注意力应该呈现出与句法结构相关的稀疏模式编码器—解码器注意力在英德这种语序差异大的语言对上应该出现明显的对角线偏移。二是做小规模消融固定其他条件只改 d_k 或头数看验证集指标的变化方向是否和原文表 3 一致。如果方向都对不上多半是实现细节有偏差而不是论文结论有问题。需要说清的取舍还有一处层归一化放在子层之后post-LN训练时需要 warmup 才稳定这是原配置的代价后来大量实现改成了 pre-LN收敛更稳但表达力略有差异。原文没有讨论这个选择它只是当时能跑通的那一个。读论文时把“设计”和“当时的最佳实践”分开看很多争论就不成立了。我自己的血泪经验是每一次重写注意力模块都会强制走一遍三件事——打印掩码矩阵、核对缩放因子用的维度、确认嵌入和位置编码的量级。这三步加起来不到十分钟但省下的是几天对着 loss 曲线猜谜的时间。这份中文翻译版把公式、表格和消融结果都保留了下来配合原文的符号去读实现代码比只看二手解读要踏实得多需要的话可以取来对照着看。希望帮到你。本文还有配套的精品资源点击获取

Transformer自注意力机制解析：复杂度、多头注意力与复现排查

相关推荐

Objective-C块(Block)详解：从语法到内存管理

链表操作技巧与双指针法实战解析

中心对称数判定：从回文数到字符映射与双指针实现

更多资讯

回形针不只能夹文件：从办公小工具到收纳高手的实用指南

低温参量放大器IMPA测试：构建可复现物理标定链

OpenRig详解：Node.js驱动的Codex协议CLI工具链

代驾系统开发：功能模块与技术架构详解

hindsight系统：用四步复盘法把经验变成决策导航

CNCC2024技术落地观察：能效比、可解释性与人本接口的实战突围

今日更新

PLFM_RADAR：从被动告警到主动感知的微服务健康监控系统

SPI隔离设计实战：MAX14431BAEE+与STM32F407工业级应用

数字隔离器如何根治SPI通信地环路干扰：MAX14431与TM4C129实战

本周精选

面向复杂光照与角度的宠物面部识别鲁棒性增强策略

基于DRV8818与STM32的双极步进电机驱动方案实战解析

DRV8818PWPR+PIC18F45K80双极步进电机控制方案实战

本月热门

面向复杂光照与角度的宠物面部识别鲁棒性增强策略

基于DRV8818与STM32的双极步进电机驱动方案实战解析

DRV8818PWPR+PIC18F45K80双极步进电机控制方案实战

读完文章，想聊聊您的网站？