如果你调过深层 Transformer一定见过这种场面训练刚开始时 loss 正常下降几十步后突然开始震荡甚至变成 NaN不管你小心翼翼地把学习率调多小都没用。我第一次做文本生成模型时遇到这种事第一反应就是换优化器参数、加梯度裁剪、换初始化折腾了一整周才意识到问题根本不在那些地方而是藏在结构里两个从来不被单独强调的组件残差连接和层归一化。这是 Transformer 系列的第 5 篇我打算把这两个组件彻底摊开讲清楚。它们各自解决什么问题、工作机制到底是怎么回事、放在哪个位置会有完全不同的效果、实战里有哪些坑。对于正在从零搭建 Transformer或者正在为训练不收敛头疼的同学这篇应该能帮你省掉不少冤枉时间。1. 训练发散的第一现场这两个组件分别守的是哪道防线1.1 梯度连乘是深层网络绕不过去的坎要让深层网络能训练最核心的问题从来不是“表达能力强不强”而是“梯度还到不到得了浅层”。反向传播的数学形式决定了这一点损失函数对第 1 层输入的梯度需要把每一层雅可比矩阵乘起来。∂L/∂h₁ ∂L/∂h_L · ∏(∂h_{l1}/∂h_l)每一层的 ∂h_{l1}/∂h_l 由权重矩阵和激活函数导数组成。这些矩阵的谱范数只要普遍小于 1层数一深梯度就会指数级衰减反过来普遍大于 1梯度就会指数级爆炸。衰减的后果很隐蔽——loss 还在正常下降但你如果打印各层梯度范数会发现前面几层的参数基本不更新网络实际上是在“浅层冻结”的状态下训练的。爆炸的后果就比较刺激了loss 直接冲到 NaN。这两种现象不是独立存在的同一个网络的不同层可能同时出现前面梯度几乎为零后面梯度大得离谱。这种“前冻后炸”的状态我后来在好几个训练事故里都见过。1.2 Transformer 把这道题做难了Transformer 的一个 block 里有两条非线性子层多头注意力和前馈网络。每个子层内部都有矩阵乘法、非线性激活和 dropout叠加十几层之后直接沿着这条非线性路径反传梯度连乘的项数非常多。比早期卷积网络更麻烦的是Attention 里的 softmax 还会引入动态的归一化因子让梯度的尺度随序列内容变化很难用一个固定的初始化方案把所有层的谱性质都调好。这时候残差连接和层归一化就登场了。我的理解是它们不是同一个问题的两种解法而是在防守两道完全不同的防线残差连接负责“修路”让梯度即使绕过所有非线性模块也有一条干净的恒等通道直达浅层。层归一化负责“修车”让每层输入的数值范围和分布都落在稳定区间减小层与层之间的动态范围差异。这也是为什么单独讨论它们任何一方都说不全。没有残差LN 只能缓解数值尺度问题挡不住梯度连乘的衰减没有 LN残差通道再干净激活值的范数也会在一层层加法里越积越大。两者是配合关系不是替代关系。2. 残差连接一个加法为什么能保住梯度不衰减残差连接写出来就一行y x F(x)。任何一个子层学习的不再是从 x 到目标输出的完整映射而是目标输出相对于 x 的“增量”。这个改动看似只是把函数包了一层但对反向传播的数学性质是根本性的。2.1 加法在反向传播里做了什么求导一下就清楚了∂y/∂x I ∂F/∂x反向传播时梯度传到这一层天然分成两条路径之和经过 F 的那条以及恒等映射这条。恒等映射的雅可比是单位矩阵 I不放大不衰减梯度可以像穿过空气一样穿过所有深层的复杂变换直达网络浅层。这正是残差连接和“直接把网络做深”的本质区别。没有残差时梯度想从第 12 层流到第 1 层要经过 11 个雅可比矩阵连乘有残差时梯度可以选择不走那些雅可比矩阵而是走加法里凭空多出来的这条捷径。你完全可以把它类比成城市交通网络加深像是在主路上增加更多红绿灯和收费站残差连接则是在旁边修了一条没有信号灯的快速路。交通拥堵时车辆可以选择走快速路整座城市不至于瘫痪。2.2 为什么是加法而不是拼接或者乘法构筑残差连接时有人自然会想直接拼接 x 和 F(x) 不是信息量更大吗或者用乘法门控一下不更灵活吗我在实际对比中体会到的区别是这样的拼接的问题是维度翻倍后面必须加一层投影把维度压回来而且梯度在拼接点上要经过一个可学习的变换才能继续向 x 分支传播快速路不够“干净”。乘法的理论更麻烦y x ⊙ F(x) 对 x 的导数还包含 F(x) 本身如果某处 F(x) 恰好接近 0梯度就会被压死信息通道形同堵塞。加法给出的恒等雅可比是天然存在的不需要网络学出来任何初始化状态下这条梯度通道都是通的。而且加法在实现上几乎零成本参数没有增加计算量只多一个 tensor 相加。2.3 实现残差连接时最容易被忽略的细节残差连接对外部维度是有要求的F(x) 的输出维度和 x 必须严格一致。Transformer 的多头注意力和前馈网络都是把 d_model 作为输入输出维度天然满足条件所以很多教程里根本不需要投影。如果你的自定义结构里子层把维度改了就得额外加一个线性投影 W_s 把 x 映射到 F(x) 的维度这个投影本质上是给残差通道加了一道可学习的变换梯度路径就不再是完全恒等的了。还有一个细节经常被忽略dropout 放在哪里。常见做法是子层内部先做线性变换、激活、再 dropout然后把 dropout 后的结果加到主分支上。有些实现图省事把 dropout 放在残差相加之后相当于把主分支的原始信息也随机丢了一部分训练初期梯度会更不稳定这个顺序值得检查一下。3. LayerNorm样本内部做归一化比 BatchNorm 更适合序列建模层归一化的公式看起来也不复杂μ (1/d) Σ xᵢ σ² (1/d) Σ (xᵢ - μ)² y γ · (x - μ) / √(σ² ε) β但这里有个很容易混淆的点均值和方差到底在哪些维度上统计不同框架的默认行为不一样很多训练事故就是从这里开始的。3.1 LayerNorm 到底对什么做归一化Transformer 的主流实现里LayerNorm 是对每个 token 的特征向量做归一化统计维度只有 d_model不跨 batch也不跨序列长度。你可以把每个 token 在高维空间里的坐标想象成一个点LN 做的事情就是把这一整组坐标平移一下、缩放一下让它们以原点为中心、方差为 1。这和第二篇里讲的词向量归一化有点像但层级不同。词向量归一化是静态的输入处理而 LayerNorm 是每一层都做参数各自独立。有些框架里的 LayerNorm 默认会把最后一维当作归一化维度这和你想要的恰好一致但如果你用的是某些自定义实现就要特别确认统计范围是不是包含了序列长度。一旦把序列维度也纳入统计同一个 batch 里不同长度的句子会互相影响padding 位置也会污染均值方差模型行为就会变得非常怪。3.2 为什么 Transformer 不用 BatchNormBatchNorm 曾经在卷积网络里几乎是标配但搬到 Transformer 上问题很多。最直接的问题是它统计的是整个 batch 的均值和方差而 NLP 场景里 batch 内每个序列长度都不一样为了保证 batch 计算需要 paddingpadding 一堆假 token 进统计量里这个均值方差就不代表真实分布了。更本质的一点是BatchNorm 的行为在训练和推理阶段不一致。训练时用当前 batch 的统计量推理时用训练阶段保存的全局滑动平均。Transformer 经常要在推理阶段处理变长的、单条的输入这种不一致在序列任务上尤其伤。LayerNorm 就没有这些问题它只依赖当前样本当前 token 自己的特征分布输入长度变了、batch 大小变了统计结果都不受影响训练和推理行为完全一致。这也是它在序列建模里胜出的核心原因。3.3 γ 和 β 不是装饰归一化后的恢复能力如果 LN 只是减均值除方差那一层就只能输出均值 0 方差 1 的分布表达能力被锁死了。γ 和 β 的作用就是解锁这个限制网络可以学习把归一化后的标准分布重新拉回任意想要的均值和方差。一般初始化是 γ 全 1、β 全 0意思是初始状态下“只做归一化不改变分布形状”。这个初始化很关键。有人图省事把 γ 初始化成全 0效果上就相当于把这个子层一开始给屏蔽了训练行为会完全不同后面我讲踩坑案例时会再提。关于“残差已经让信息穿过去了为什么还需要 LN”这个问题我的理解是残差连接一个接一个相加激活值的范数会随着层数增长持续累积到了深层可能已经超出激活函数的合适区间。LN 就是在每个子层入口把动态范围重新框定。你把它想成高速公路每开一段就设一个收费站检查一下车辆数据有没有超限超了就规整回正常范围。4. 亚层拓扑的两种摆法Post-LN 与 Pre-LN残差和 LN 的先后顺序不是谁拍脑袋定的它直接决定了梯度通道干不干净。这里有两种经典拓扑实际工程里都有人用但行为差异很大。4.1 Post-LN 的结构与代价经典论文图里画的是 SubLayer 之后 AddAdd 完之后 Normx_new LayerNorm(x SubLayer(x))这种摆法叫 Post-LN因为归一化放在残差加法之后。每一层的整体输出都被 LN 归一化过一次所有进入下一层的数据分布都很规整。但代价是梯度要从深层传回浅层时每经过一个 Post-LN 都要穿过一次 LN 的雅可比矩阵。LN 本身不是恒等映射它的导数不是单位阵梯度经过时会被缩放。层数多了以后这一连串缩放同样可能带来不稳定。Post-LN 最著名的工程问题是需要非常精心地设置学习率预热。我刚接触 Transformer 时不明白为什么要 warmup后来才理解训练初期参数随机各层输出的尺度差异很大LN 在每层的缩放系数会剧烈变化如果一开始就用大学习率很容易把参数推出稳定区域。预热的作用就是让网络先在一个小学习率下把各层的尺度“校准”到比较一致再逐步加大更新。4.2 Pre-LN当前大模型事实上的标准Pre-LN 把 LN 挪到子层之前x_new x SubLayer(LayerNorm(x))这里最关键的变化是残差的恒等分支上不再有任何变换梯度可以走一条完全恒等的快速路直通浅层。实际训练中 Pre-LN 的稳定性确实比 Post-LN 高出一截很多实现可以完全不依赖 warmup学习率也可以拉得更大。我自己的经验是从零搭网络做实验优先用 Pre-LN 能省掉大量调参时间。Flow 训练阶段的梯度范数量级也更平稳。代价是什么呢理论上残差通道上的信息一直没被归一化深层堆叠后主分支上可能会积累一些和任务无关的“直流分量”某些深层任务上最终效果会比精心调好的 Post-LN 略差一点。但对大多数实际项目来说稳定收敛带来的收益远大于这一点潜在差距。两种结构对比如下对比项Post-LNPre-LN归一化位置残差加法之后子层入口之前恒等梯度分支会被 LN 缩放完全干净训练稳定性较敏感更稳warmup 需求通常必需可省深层理论效果可能略优可能略欠工程实用倾向复现经典实验自建模型首选5. 数值、实现和升级FP16 下的 LayerNorm、RMSNorm以及那些容易忽略的细节理论上讲清楚了落到代码里还有一堆数值层面的坑尤其是在混合精度训练普及之后。5.1 混合精度下的方差计算陷阱LayerNorm 的计算里最危险的一步是先算均值再算方差。当输入数值较大且均值不为零时(x - μ)² 的平方运算在 FP16 下很容易出问题如果 x 本身接近 FP16 的最大表示范围 65504平方直接变成无穷。很多训练跑着跑着 loss 变 NaN第一现场就在这一行。实践中我有两个习惯。第一LayerNorm 的统计计算尽量放在 FP32 精度下完成哪怕输入是 FP16也先转成 FP32 算完均值方差再转回去如果用的是支持 BF16 的硬件直接全程 BF16 通常也没问题因为 BF16 的动态范围很大但精度低需要观察是否有明显精度损失。第二动手写 LayerNorm 之前先检查输入张量的绝对最大值如果数值已经逼近 FP16 上限问题往往不在 LN 本身而是前面某一步已经出现溢出要往前排查。5.2 RMSNorm省掉均值后节省了什么LN 的体积大模型时代开始流行 RMSNorm公式是y x / √(mean(x²) ε) · γ它干脆连均值都不算了只做缩放不做中心化。节省的不只是几次开方求和的运算量而是在超大规模模型里每一层少一个均值计算都能省下可观的内存和算力。关于 RMSNorm 一个常见的误解是它“约等于”LN可以无脑替换。实际上 RMSNorm 不做中心化输入分布的均值会原样通过。如果网络前面有残差加法激活的均值会慢慢偏离 0RMSNorm 就相当于在学习一个动态的缩放系数去补偿这个偏移。在从头训练的大规模模型中RMSNorm 的表现通常能和 LN 打平甚至更好但如果你想把一个用 LN 训练好的参数直接迁移到 RMSNorm 结构里均值偏移会以误差的形式残留出来不是简单换层就能兼容的。5.3 参数初始化和加载时的几个判断点LN 的参数是 per-feature 的一个 d_model 维的 γ 和一个 d_model 维的 β。初始化全 1 和全 0 是常识但我见过不少自定义实现里把两个都初始化为 0导致前向输出变成恒定零向量损失一开始就卡死特别隐蔽。另外在把权重从一个实现迁移到另一个实现时要核对两边的 LN 结构语义同样是 LNA 框架可能把 ε 默认成 1e-5B 框架默认 1e-12小数值在深层传播中会被放大体现在输出上就是完全不同的行为。还有的复现实现把 LN 放在残差分支内部加载权重时张量形状对得上但计算图语义完全错位。我的检查方法很土但有效拿一个固定输入在两边各跑一次前向逐层对比输出张量的均值和方差偏差超过 1e-3 就说明语义不对。6. 踩坑实录三个真实训练事故的排查过程6.1 案例一loss 曲线反复震荡最终冲上 NaN有一次我要训练一个跨语言文本分类模型8 层 Transformer结构参照经典论文。训练到 5000 步左右 loss 开始震荡curve 像锯齿一样再过几百步直接 NaN。最初我怀疑学习率调度换了两种方案没变化又怀疑初始化也无果。后来打印每层梯度范数才发现前 3 层梯度几乎为零倒数第 2 层梯度大得吓人典型的“前冻后炸”状态。再看结构写的是 x LN(x dropout(sub(x)))是标准的 Post-LN。因为这套代码里没有实现 warmup从第 1 步起就用比较大的学习率更新LN 的缩放系数在训练初期剧烈变化直接把深层梯度推向极端。修复方式是改成 Pre-LN 拓扑并且把学习率预热补上。改完之后同样的训练步数loss 一路平滑下降。这个事故给我的教训是Post-LN 并不是不能训练但它的稳定性严重依赖 warmup 和整体学习率策略不要拿 Post-LN 的代码配一个“现代大学习率方案”那基本是给训练事故上保险。6.2 案例二把 LayerNorm 换成 RMSNorm 后指标突然垮掉另一个项目是给视觉任务用的 Transformer一开始用 LN想在大模型化改造里换成 RMSNorm 省算力。改完从头训练指标确实接近一切正常。问题出在我要把一个已训练好的参数迁移到新结构上测试本来以为只是把 LN 层替换成 RMSNorm 层加载权重后模型输出完全错乱任务指标跌到接近随机。排查时我对比了两层的前向数值LN 会对输入做中心化把均值减掉这个偏移信息在网络里是有意义的RMSNorm 不做这一步权重的语义直接错位。差值不是简单的常量因为中心化后还会被 γ 缩放。最后只能放弃直接迁移改成重新训练花费了不少算力。这个案例给我一个很深的印象结构升级不是换个层名那么轻巧的事每一层的统计语义改变都会沿着残差路径一路传导到整个网络。6.3 案例三混合精度训练爆 NaN罪魁在方差计算的平方溢出长文本生成任务开了混合精度后 loss 下降正常但每隔一段时间会突然跳到 NaN。关掉混合精度就一切正常所以问题一定出在精度相关的环节。我逐层打印激活值发现某一层 LayerNorm 的输入最大值已经到了 60000 出头这个数在 FP16 下平方之后必炸。而这一层的输入因为它前面一个残差加法把两个大范数张量叠在了一起又没有 LN 及时拉住就一路涨上去了。处理方式有两个一是把 LayerNorm 的计算保持在 FP32二是在子层输出加一个临时的缩放因子把数值压回安全区间做完 attention 或 FFN 后再缩放回去。最终我同时做了这两件事并给优化器配了动态损失缩放NaN 问题彻底消失训练速度也没有明显损失。这里我也想多说一句LN 能解决的是分布尺度问题但它不是越权。如果某个残差分支上多个大范数张量一直叠加LN 只在自己那一个位置拉回尺度拉回之前的加法瞬间仍可能出现溢出。所以看 NaN 不能只看 LN 那一行代码要从数据流上游一路排查到加法发生的地方。最后再分享一个我长期使用的小习惯任何一个新搭的 Transformer 结构我会先用极小的模型、极少的步数用同一份 batch 反复前向反向几十次观察 loss 是否在下降、梯度范数是否量级稳定。这个过程能快速暴露残差连接和 LN 的大部分问题——比如拓扑顺序错了、维度统计错了、初始化错了。等小模型跑顺了再放大到正式规模和混合精度也就很少会遇到那些让人熬夜的“神秘 NaN”了。