Transformer核心组件详解:Multi-Head、位置编码与Add  Norm实战笔记 老实说Transformer是我啃过的模型里“反直觉”最多的一个。第二周的学习计划排到自己头上时我以为只要理解了Self-Attention就万事大吉结果真到了复现阶段才发现Multi-Head、位置编码、Add Norm这三个部件任何一个理解不到位写出来的模型要么训练不收敛要么效果和单层线性模型差不多。这篇博文就围绕这三块把我两周里从公式推导、代码实现到踩坑的全部过程整理出来给正在啃Transformer的朋友一份可以直接参考的学习笔记。我不会去复述整个Transformer论文的每个细节而是集中在标题里点名的三个东西Multi-Head Attention到底在“多头”什么语义、位置编码为什么非加不可且要用那种公式、Add Norm又是怎么让网络能一路叠到几十层。每部分我都会从原理讲起再给出一段能直接跑通思路的代码逻辑最后附上我在实际动手中遇到的问题和解决办法。适合谁来读如果你是刚学完Self-Attention、准备进入Transformer完整结构的初学者或者你已经看完了理论但自己实现时总在维度上翻车这篇文章应该能帮你把最后几块拼图对上位。1. 先回答一个绕不开的问题Self-Attention已经有了全局建模能力多头到底是“多”在哪里很多人第一次看Multi-Head时都会冒出同一个疑问既然单头Self-Attention已经能让每个位置“看到”所有其他位置那多头无非就是把注意力多算几遍这有什么本质区别我在第二周开始时也是这么想的直到我把单头注意力的输出拿去可视化才明白问题远没有那么简单。1.1 单头注意力的真实局限一种注意力分配只能表达一种“关系”回顾一下Self-Attention的计算输入序列的每个token经过线性变换得到Query、Key、Value三份向量然后计算softmax(QK^T / sqrt(d_k))得到注意力权重再用这份权重去加权求和Value。整个过程确实让每个token都获得了整个序列的信息但关键在于在一句话里一个词往往同时需要参与多种不同的“关系”。举个例子“小明把球传给了小红他跑得很快”。这里的“他”既需要指代前面的“小明”又需要从语义上和“球”“小红”发生关联还需要留意“跑”这个动作的主语是谁。如果只用单头注意力模型只能学出一种加权策略它得在语法关系、指代关系、语义相似性之间做折中最后往往谁的权重都分到一点但谁都不突出。我在实际测试中发现单头模型的注意力图经常是“弥散”的每个位置几乎都在均匀地关注所有位置没有一个清晰的模式。这倒不是模型学错了而是它根本没有能力在一份注意力分布里同时表达多种关系——所谓的“多义性”被强行压进了一个标量权重里。1.2 多头提供了“多个子空间视图”而不是多次平均Multi-Head的做法是把原始的Q、K、V各自投影到h个低维子空间每个头独立计算注意力最后把结果拼接起来再做一次线性投影。这个设计的核心逻辑是每个头相当于一个独立的“关系探测管道”它可以自由地学会关注某一种特定的关系类型不用和其他关系挤在一起。我在复现时做过一个直观的小实验把8个头在某个中文字符串上的注意力图分别画出来能看到一个头特别关注相邻词之间的修饰关系另一个头明显在追踪代词和指代对象之间的远距离连接还有的头在关注词与词之间的位置偏移。这些模式如果放在单头里是永远不可能同时出现的。所以说多头并不是“把注意力算几遍求平均”它是把“注意力应该关注什么”这个模糊问题拆解成了h个更清晰的子问题每个子问题在一个低维空间里解决最后汇总。这也是Multi-Head在参数总量不大幅增加的情况下能显著提升模型表达能力的原因。1.3 一个值得注意的数学细节为什么缩放因子是sqrt(d_k)这个细节属于看论文时最容易一带而过、实现时却最容易出错的地方。Q和K做点积后结果的方差会随维度d_k增大而增大。如果d_k64点积结果的方差大约是64这意味着很多值会落到±8甚至更大的范围。把这样一组数送进softmax会出现一个后果某些项会直接压过其他所有项softmax输出集中到个别位置上注意力图变成“硬选择”梯度接近零训练基本停滞。除以sqrt(d_k)的本质是把点积结果的方差重新拉回到1附近让softmax有正常的温度。我在自己的实现里试过去掉缩放模型前几步loss确实还在下降但很快就卡住不动了这种情况基本可以断定是梯度问题。这个细节不需要背但要能自己推导一遍两个维度为d_k的随机向量的点积方差就是d_k除以它的平方根方差归位。2. Multi-Head Attention的工程实现维度变化、头数选择与注意力掩码原理搞清楚了真正动手写就进入“差一个维度都跑不起来”的阶段。第二周我在这一块反复翻车的次数最多所以把实现相关的细节全部记在这里。2.1 从输入到输出完整维度流转记录假设输入序列形状为(batch_size, seq_len, d_model)d_model512头数h8。Multi-Head的全过程可以拆成五步生成Q/K/V通过3个独立的线性层把d_model512投影到d_model512。这一步每个tensor形状都是(batch, seq_len, 512)。拆头把512维切成8份每份64维。操作上是先reshape成(batch, seq_len, 8, 64)再转置成(batch, 8, seq_len, 64)让头维度排在batch之后方便并行计算。逐头计算注意力每个头在自己的64维空间里做scaled dot-product attention输出形状依然是(batch, 8, seq_len, 64)。并头转置回(batch, seq_len, 8, 64)然后reshape成(batch, seq_len, 512)把所有头的信息拼接回完整维度。输出投影经过最后一个线性层输出形状仍是(batch, seq_len, 512)。这里最反直觉的一点是多头计算完成后输入输出的维度完全一致模型只能靠这512维向量的“内部结构”来承载不同头带来的多视角信息。所以后面为什么要用残差连接和LayerNorm也就不难理解了——网络每一层都要在这些高维空间里稳定地传递信息。2.2 我的实现风格用一次reshape完成拆头和并头我在复现时没有用for循环逐个头计算而是直接靠reshape和transpose一次性做完。这样做有两个好处一是代码简短二是所有头共享同一个计算过程在深度学习框架里能获得更好的并行效率。# 伪代码以主流框架的张量操作为例 def multi_head_attention(Q, K, V, d_model, h): dk d_model // h # 线性投影 Q linear(Q, d_model, d_model) # (batch, seq, d_model) K linear(K, d_model, d_model) V linear(V, d_model, d_model) # 拆头 (batch, seq, h, dk) - (batch, h, seq, dk) Q Q.reshape(batch, seq, h, dk).transpose(1, 2) K K.reshape(batch, seq, h, dk).transpose(1, 2) V V.reshape(batch, seq, h, dk).transpose(1, 2) # 缩放点积注意力 scores matmul(Q, K.transpose(-2, -1)) / sqrt(dk) # (batch, h, seq, seq) attn softmax(scores, dim-1) out matmul(attn, V) # (batch, h, seq, dk) # 并头 (batch, h, seq, dk) - (batch, seq, d_model) out out.transpose(1, 2).reshape(batch, seq, d_model) return linear(out, d_model, d_model), attn这段代码里有几个容易写错的地方值得特别说明。reshape之后必须紧跟transpose如果顺序反了会把每个头的数据切得乱七八糟transpose的两个轴分别是1和2也就是把“seq”和“h”互换matmul要特别注意维度对齐Q和K^T相乘得到的是(batch, h, seq, seq)的注意力分数矩阵。2.3 头数、维度与初始化的经验选择我见过不少初学者把d_model512和h8认为是唯一组合其实不是。合理的做法是保证d_model % h 0并且d_k d_model // h不要太小。我自己测试下来d_k如果低于32每个头能表达的信息量太有限整体效果会出现明显退化d_k64是各方面都比较均衡的配置。关于头数h8是论文里的经典值12、16也很常见。头数增多意味着每个头的维度变小模型有机会学到更细粒度、更多样的关系模式但同时训练会更慢对小数据集也更容易过拟合。我在一个小规模文本分类任务里对比过8头还是4头差距不到1个点差异远没有“能不能收敛”那么显著所以不必在这里过度纠结先跑通8头再说。权重初始化方面线性层默认的初始化方式在实际训练中基本够用。有一点需要提醒如果自己手动初始化不要把线性层的scale设得太大否则第一批forward的attention分布就会接近one-hot梯度直接消失。我在最早的一个版本里犯过这种错误后来统一改用常见框架的默认初始化问题再没出现过。2.4 掩码处理两个必须分别对待的maskTransformer里有两种maskPadding Mask和Causal Mask它们在实现上有本质区别。Padding Mask用于把无效的padding位遮掉。做法是在softmax之前把对应位置设为一个很大的负数比如-1e9这样softmax算出来的概率接近于0。如果直接在softmax之后把attention矩阵里某些位置置零效果看似一样实际上softmax已经把这些位置的梯度传给了其他位置数值上并不等价。Causal Mask用于自回归场景要求每个位置只能看到它自己以及之前的位置。做法是把注意力分数矩阵的上三角部分全部置为-inf然后才做softmax。我一开始图省事把-inf写成了-1e9结果训练出的模型在生成测试时会逐步出现输出质量下降——数值上-1e9虽然能近似无穷小但在某些情况下残留的微小概率还是会被模型利用到。建议直接用-inf。3. 位置编码Attention天然“没有顺序”必须用数学把上下文顺序找回来这一节是我第二周花了最多时间才真正想通的部分因为位置编码的表面公式并不复杂难的是理解“为什么非用它不可”以及“为什么它能表达位置关系”。3.1 先看问题本质Self-Attention是“无序”的Self-Attention对齐每个位置之间计算相似度但它并不关心两个词之间的距离。把“我打你”和“你打我”分别输入模型如果把词序打乱注意力分数会变得一模一样——因为“我”“打”“你”这三个词的向量两两之间的点积关系没有任何变化。这意味着词序信息对Attention来说完全是不可见的。而自然语言里词序往往决定语义所以必须主动把位置信息“注入”到输入向量里。选择注入方式时有个朴素思路在词向量后面拼接一个位置向量。但Concatenate会显著增加维度训练开销变大。Transformer采用的是直接相加——把位置向量加到词向量上。这个做法的前提是模型可以通过后续的线性层把加和后的向量重新“拆解”出语义部分和位置部分。后来实验也证明加法的表达力完全够用。3.2 正弦位置编码的公式直觉一种能够“近似表达相对位置”的编码论文里用的位置编码公式有两组分别对应偶数和奇数维度PE(pos, 2i) sin(pos / 10000^(2i/d_model)) PE(pos, 2i1) cos(pos / 10000^(2i/d_model))其中pos是词在序列中的位置i是维度下标。这个公式的表面含义是不同维度使用不同频率低维i小频率高高维i大频率低。用一个容易理解的类比低频率维度像时钟的“秒针”对位置差异敏感但能表达的范围小高频率维度像“时针”能表达更远的距离但对局部变化不敏感。这个设计更深的巧思在于它能让模型感知到“相对位置”。利用三角恒等式PE(posk)可以表示为PE(pos)的某种线性组合sin(posk) sin(pos)cos(k) cos(pos)sin(k) cos(posk) cos(pos)cos(k) - sin(pos)sin(k)这就意味着任何位置的编码向量都可以由任意另一个位置加上一个固定的“旋转矩阵”变换而来。模型不需要单独学习“相对位置k的值是多少”只需要学会这个线性变换就能隐式地获得相对位置信息。这也是为什么后来不少新方法都沿着“旋转”这个思路做文章。我建议读者实际做一次可视化把PE按维度画出来你会看到低维图形是密集的锯齿波高维图形是平缓的正弦波。这种“不同维度不同频率”的设计本质上等于用一组正交基函数去编码位置像信号处理里的傅里叶展开理解到这个层面就能明白它为什么有一定的外推能力——即使遇到比训练时更长的序列频率范围仍然是完整的。3.3 为什么后来的模型很多都改用“可学习位置编码”正弦编码不是唯一的方案。我在学习时对比过两种主流做法简单列一下方案优点缺点正弦位置编码无需训练参数理论上能外推到更长序列固定在数学公式里无法针对任务自动调优可学习位置编码参数由任务数据驱动实现极简单需要预设最大长度超出部分没有参数无外推能力相对位置编码显式建模相对距离符合自然语言直觉实现复杂度高与自注意力机制耦合深我在自己的小型任务上比较过正则可学习位置编码发现效果几乎相当。如果做的是固定长度的任务比如句子分类可学习位置编码完全够用且更简单如果涉及长文本生成、长序列推理正弦编码至少给了模型一个“位置表在训练长度外仍是有意义的实数”的机会搭配一些长度外推技巧可能更稳。但从当前大模型的实际做法来看主流方案更偏爱旋转类相对位置编码。学习这些经典方案不是为了追逐最新而是为了理解“为什么发展到旋转”——本质都是为了让位置信息与注意力机制更自然地融合。4. Add Norm残差连接和层归一化是怎么撑起“深”这个字的很多教程把Add Norm一笔带过好像它只是一个工程技巧。但我的体验是没有正确理解这一块把Transformer叠个6层试试训练loss分分钟变成NaN或者震荡不停。Add Norm是整个模型能稳定变深的“地基”。4.1 残差连接解决的不是“梯度消失”而是“网络退化”先分清两个概念。梯度消失是网络层数增加时反向传播的梯度连乘导致梯度指数级变小网络前几层学不到东西。解决梯度消失的经典手段是合理的初始化和归一化。而残差连接解决的是另一个更微妙的问题随着层数增加即使训练集上深层网络的误差也会比浅层网络更大这种现象叫退化。为什么会出现退化因为直接让网络学习一个恒等映射H(x)x是很困难的特别是在激活函数和复杂非线性层的共同作用下。而残差结构把目标改成了学习F(x)H(x)-x如果恒等映射是最优解网络只需把F(x)学成0。学习“输出0”对网络来说远比学习“完整输出x”容易得多。所以残差连接本质上给了网络一个“退路”我不确定这层该做什么时就什么都不做把输入原样传给下一层。我在实现Transformer Block时把Add放在Sublayer之后、Norm之前。这形成了一个循环每一层的输出都是“这层学到的变化”加上“上一层的原始输入”。这种结构天然给了梯度一条连乘不衰减的“高速公路”同时保证信息在深层传递过程中不会丢失。所以理论上Transformer叠多少层都还有继续训练的空间这在没有残差结构的模型里是不可想象的。4.2 为什么Transformer选LayerNorm而不是广泛使用的BatchNorm很多人刚接触时会疑惑图像领域BatchNorm用的好好的为什么Transformer要用LayerNorm我的理解是LayerNorm的计算方式不同于BatchNormBatchNorm在batch维度上统计均值方差而LayerNorm在特征维度上统计。具体到NLP场景有两个核心原因。第一BatchNorm依赖batch中样本的统计量序列任务中batch大小往往不大且样本长短不一这使得全局均值和方差波动剧烈训练不稳定。LayerNorm每个样本独立计算自己的统计量不受其他样本影响这在变长输入和在线推理时显得尤其重要。第二Transformer的输入是词向量加位置向量不同维度代表不同的语义子空间LayerNorm在特征维度上做归一化相当于把每个token的表示拉回到一个标准尺度这对于后续注意力分数的稳定计算非常有帮助。打个比方如果词向量某些维度的数值范围很大点积后的方差就会被少数维度主导模型难以训练LayerNorm把这组维度整体缩放到相近尺度等于给每层之前的信号做了一个“标准化”。LayerNorm的实现公式本身不难LayerNorm(x) gamma * (x - mean) / sqrt(var eps) beta其中gamma和beta是可学习的仿射参数。我记得原论文的LayerNorm实现里有个小的均值计算顺序问题不影响主流程但大家在复现时最好确认一下标准化是在最后一个维度上做的。4.3 “先加后归一”还是“先归一后加”Post-Norm与Pre-Norm的选择这个问题是我在实现时偶然发现的。原论文的做法是每个Sublayer结束后先做Add再在最后统一做LayerNorm这种方式叫Post-Norm。但随着网络层数加深Post-Norm的训练变得不太稳定因为残差路径上的数值会不断累积Norm的统计量也跟着变化。后来很多开源实现改为Pre-Norm先对输入做LayerNorm再进入Sublayer最后再做残差相加。我在自己尝试6层和12层Transformer时明显体会到差异Pre-Norm版本训练曲线平滑很多初始阶段loss下降更稳定对学习率也不那么敏感。因此除非严格复现论文实验否则我建议直接使用Pre-Norm结构。网上对这两种方式的对比讨论已经很多实践上Pre-Norm成为事实标准是有原因的。5. 把三个模块串成完整前向过程一个Transformer Block的维度全记录理论部分讲完了最后把Multi-Head、位置编码、Add Norm合起来用一个具体的例子完整走一遍前向传播看维度是怎么一步步变化的。这一步做完整套Transformer结构的“骨架”就彻底立起来了。5.1 分步推演从输入序列到Block输出设定一个具体场景batch_size2序列长度为seq_len10d_model16为了好算用个小维度头数h4所以每个头维度d_k4。先看一张我从第一行推到最后一行时记录的步骤表步骤操作输出形状输入词嵌入结果 位置编码(2, 10, 16)多头注意力线性投影生成Q/K/V各 (2, 10, 16)拆头reshape transpose(2, 4, 10, 4)缩放点积QK^T / sqrt(dk)(2, 4, 10, 10)softmax后加权attn V(2, 4, 10, 4)并头transpose reshape(2, 10, 16)输出投影线性层(2, 10, 16)第一个残差Add加上未进入Attention的原始输入(2, 10, 16)第一次LayerNorm在最后一维归一化(2, 10, 16)前馈网络线性激活线性(2, 10, 16)第二次残差Add加上进入前馈前的输入(2, 10, 16)第二次LayerNorm在最后一维归一化(2, 10, 16)这张表里有几个值得注意的点。整个Block的输入输出形状完全一致这是Transformer能层层堆叠的前提。注意力分数矩阵的形状是(2, 4, 10, 10)含义是每个batch、每个头、每一行查询位置都要对10个键位置计算相似度。d_k4虽然只有4维但每个头仍然能独立学习自己的模式这说明“多头的力量来自并行子空间而不是来自每个头有多大”这个结论真的成立。5.2 我在完整复现前向过程时踩过的三个低Level错误第一处是维度转置方向。拆头时如果不小心多转置了一个轴或者转置顺序错了结果不是报错而是每个头的输入数据被“串味”。因为维度形状没变程序不会提示错误但准确率会诡异下降。排查这类问题的最好办法是挑一个固定的小tensor手动算一遍对比中间形状和数值。第二处是位置编码在Padding位置的处理。Padding部分虽然被mask遮掉了但位置编码还是会加到Padding的向量上。这个小误差通常不会对训练产生大影响但如果做推理端优化、想在KV Cache里省内存就必须考虑“Padding位不参与计算”的合理性。更常见的做法是保持现有逻辑让mask统一处理。第三处是Pre-Norm与残差的顺序。我一开始套用Post-Norm的模板改成Pre-Norm时在输入和输出之间拉错了分支线导致Norm放在残差之后又被重复使用训练曲线一路乱跳。后来养成一个习惯写Block时先画出数据流谁先谁后标清楚再写代码。这类结构性错误看公式不明显跑起来却非常折磨。5.3 一个可供直接参考的Block级伪代码最后给出一段带有Add Norm的完整Transformer Block伪代码结构里同时包含Multi-Head、前馈网络、两个残差和两个LayerNorm方便大家对照自己的实现检查# 伪代码一个标准Transformer BlockPre-Norm风格 def transformer_block(x, d_model, h, d_ff): # x: (batch, seq, d_model) # 子层1多头注意力 res1 x x layer_norm(x) # Pre-Norm attn_out, attn_weights multi_head_attention(x, x, x, d_model, h) x res1 attn_out # 残差Add # 子层2前馈网络 res2 x x layer_norm(x) # Pre-Norm ff_out linear(x, d_model, d_ff) ff_out gelu(ff_out) ff_out linear(ff_out, d_ff, d_model) x res2 ff_out # 残差Add return x, attn_weights每一个LayerNorm都需要分别拥有自己的可学习参数gamma和beta不能两个层共享同一组参数。我见过有人为了省事共用了参数结果模型表达能力明显下降。另外d_ff一般取d_model的4倍左右这是论文里的经验值后续的变体也都遵循类似的扩大策略。大的前馈层提供了非线性变换的空间是整个Block里参数量最大的部分通常也是过拟合最容易发生的地方之一。结尾第二周实操后的三点体会两周啃下来我个人最大的体会是理解Transformer不能把Multi-Head、位置编码、Add Norm当成三个孤立的知识点背公式而要始终追问“这个设计到底解决了什么问题”。多头解决的是单种注意力无法表达多关系的问题位置编码解决的是Attention天然无序的问题残差和LayerNorm解决的是深层网络难以稳定训练的问题。三条线索汇到同一个目标让模型既能精确捕获长距离依赖又能稳稳地越叠越深。实操层面我也积累了一句话总结代码跑不通先检查形状形状对了但效果差先看mask和缩放因子train loss爆炸再检查初始化、Norm顺序和学习率。第二周实验过程中我确实反复被这些细节拖着走但也正是这些坑让我把Transformer从“看懂论文”真正变成了“能自己写出来”的状态。如果你正在学同样的内容记住别贪快把你的Block亲手写一遍并观察每一步的shape绝对比刷十遍视频课更有用。