Transformer异常层排查：5个模型的数值稳定性实战

1. 测试起因与测评对象：五个 Transformer 的选型逻辑

这几天我泡在一个相当折磨的问题里：手头的一个序列建模任务在换到 Transformer 之后，训练 loss 怎么调都不对劲，曲线看着在下降，但验证集指标就是上不去。顺手把每一层的激活值打出来扫了一眼，发现某个编码器层的输出分布已经乱到和输入完全不是一个数量级了。这一下勾起了我的好奇心——这个所谓的"异常层"到底异常在哪里，是整个结构通病，还是某个实现独有的毛病？

为了把这个事搞清楚，我干脆花了几天时间，把市面上最常被拿来当作"参考实现"的 5 个 Transformer 都拉出来跑了一遍。不是去比谁精度高，而是专门盯着每一层的激活值、梯度范数、权重奇异值这些"数值体检指标"看。先说结论：异常层的存在比我想象的普遍得多，而且几乎每个模型身上的异常表现都不一样。

我选的 5 个测试对象，对应了我平时工作中最常碰到的场景：

编号	模型	用途	实现来源	层数/规模
T1	经典 Transformer（Encoder-Decoder）	翻译/序列到序列复现	手写 PyTorch 实现	6 层 encoder + 6 层 decoder
T2	BERT-base	预训练语言模型微调	HuggingFace transformers	12 层 encoder
T3	Vision Transformer（ViT-Small）	图像分类	开源仓库复刻	12 层 encoder
T4	Swin Transformer	图像分类/检测骨干	官方源码移植	4 阶段，共 12 层
T5	用于时序预测的 Transformer	长序列未来值预测	自研实现（标准 encoder-decoder）	4 层 encoder + 4 层 decoder

选这五个不是随机的。T1 覆盖了"从论文手写实现"的经典路径，T2 覆盖了"开箱即用、拿来微调"的主流路径，T3 和 T4 代表视觉 Transformer 的两条路线（全局 attention 和窗口 attention），T5 则是我实际业务里最关心的一类——Transformer 用在回归型时序任务上，数值稳定性的问题往往比分类任务暴露得更明显。这几个方向合在一起，基本能代表大家平时接触到的绝大部分 Transformer 应用场景。

1.1 我定义的"异常"到底是什么

很多人可能第一反应是：你要找的是不是某个层发生了 NaN？确实 NaN 是异常的一种极端形态，但我的这次测试更多关注的是那些"还没崩，但已经很不健康"的现象。

我给自己定了三条硬指标：

• 激活值分布漂移：某层输出的均值和标准差，和它前面的层相比出现数量级上的偏移，比如 std 从 1 左右突然变成 8 甚至几十。
• 梯度范数失衡：用 backward hook 统计每一层的梯度，发现某些层的梯度范数比同模型其他层高出 2 到 3 个数量级，或者低到几乎消失。
• 注意力熵异常：对 attention 权重矩阵计算熵值，正常的 softmax 输出应该有一定的分散度，如果熵值低到接近 0，说明 attention head 已经退化成了 one-hot 模式。

后面所有排查都会围绕这三条展开。我可以先把测试环境和统一评测脚本固定下来：Python 3.10、PyTorch 2.1、单卡 A100、固定随机种子。每个模型用各自领域里最常规的输入（文本用 token id，图像用 224x224 的 patch）喂进去，统计前向和反向过程中每一层的数值表现。

2. 检测手段与评判基线：先用 hook 把每一层扒开

检测"异常层"最直接的办法不是靠眼睛看代码，而是用 PyTorch 提供的 hook 机制，在前向传播和反向传播时把每一层的输入输出截获下来。我把这个流程封装成了一个非常简单的工具函数，后面所有模型都复用了这套脚本。

import torch import torch.nn as nn def attach_inspector(model, tag=""): def make_hook(name): def forward_hook(module, input, output): if not isinstance(output, torch.Tensor): return t = output.detach().float() # 打印输出张量的均值和标准差，这是最直观的漂移信号 print(f"[{tag}][{name}] out mean={t.mean():.4f} std={t.std():.4f} " f"min={t.min():.4f} max={t.max():.4f}") return forward_hook def make_bw_hook(name): def backward_hook(module, grad_in, grad_out): g = grad_out[0].detach().float() print(f"[{tag}][{name}] grad std={g.std():.4f} " f"grad_norm={g.norm():.4f}") return backward_hook handlers = [] for name, module in model.named_modules(): if isinstance(module, (nn.Linear, nn.LayerNorm, nn.MultiheadAttention)): handlers.append(module.register_forward_hook(make_hook(name))) handlers.append(module.register_full_backward_hook(make_bw_hook(name))) return handlers

这里有个容易被忽略的点是：注册 hook 的时候必须把模块类型过滤得足够精准。我一开始图省事，把所有 SubModule 都挂上了 hook，结果输出日志被 Embedding 和 Dropout 的噪音刷屏，根本没法看。后来把范围缩小到 Linear、LayerNorm、MultiheadAttention 这三类，日志立刻清爽了很多。对于手写 Transformer，MultiheadAttention 在 PyTorch 里默认是不注册为 leaf 模块的，所以实际跑起来要看具体情况再调整。

2.1 先建立对照组：正常的层长什么样

在判断哪一层"异常"之前，必须先知道正常情况下的数值长什么样，不然很容易误判。我用一个完全随机的输入做了一次"基线对照"：一个 6 层、hidden size 512 的 encoder，在标准初始化下，前几层的输出标准差大约在 1 到 1.5 之间，梯度范数在 1 到 20 之间缓慢变化。

这个对照实验其实泄露了一个关键信息：正常的 Transformer 里，各层输出标准差也会随着层数的加深缓慢增大，但只要保持在同一个量级内，模型训练就是稳定的。反过来，一旦某层的 std 突然从 1 跳到 6、8 甚至几十，就说明内部信号的"传导放大"出问题了，这才是异常层需要被揪出来的位置。

2.2 配合权重奇异值分析

只看激活和梯度还不够。为了判断异常层到底是"接错了信号"还是"自身的权重矩阵已经病变"，我还会对目标层的权重矩阵做 SVD，统计最大奇异值和最小奇异值的比值（条件数）。

def check_weight_condition(model, layer_name): for n, m in model.named_modules(): if n == layer_name and isinstance(m, nn.Linear): w = m.weight.data.float() s = torch.linalg.svdvals(w) print(f"[{layer_name}] singular values: max={s.max():.4f} " f"min={s.min():.4f} cond={s.max() / (s.min() + 1e-8):.4f}")

条件数如果超过几百，权重矩阵就有"病态"倾向了——输入稍微偏移一点，输出就会剧烈抖动。这是判断异常层根因的重要证据链。

3. 逐个跑一遍：五个模型的第一轮排查记录

3.1 手写经典 Transformer：尾层编码器的激活值爆炸

第一个出问题的是 T1，那个从论文手写实现的经典 Transformer。我用了一个小规模英德翻译任务来跑，输入 batch size 32、序列长度 64，hidden size 512。

hook 日志打出来的现象非常扎眼：encoder 第 1 到第 3 层的输出标准差稳定在 1.1 到 1.3，第 4 层开始长到 2.8，第 5 层直接跳到 7.6，第 6 层更是到了 13.4。同一时刻 decoder 层还没有这么夸张，所以问题集中在 encoder 的后半段。

这其实暴露了 Post-LN（Post-LayerNorm）这种结构的典型问题。残差连接把每一层的输出全部往一条主线上累加，越往后这条主线的"信号总量"越大，而我的实现里 LayerNorm 是放在残差相加之后、进入下一层之前，理论上它应该把累加后的信号拉回来；但问题就出在 LayerNorm 的 scale 行为上——它做的是逐特征归一化，不是逐 token 归一化，在序列维度上它的归一化力度其实很弱，如果输入张量中存在个别极端大的 token，这个 token 就会把整个层的输入推得很高，随后线性变换把它进一步放大。

这种现象我在复现经典的 Attention Is All You Need 时已经踩过很多次了，但这次的幅度确实夸张。进一步用 weight condition check 看了一眼第 5 层 FFN 的第一个 Linear 层，条件数高达 312，几乎可以断定这个层的权重矩阵已经处在"病态边缘"。

3.2 HuggingFace BERT：梯度范数峰值出现在 embedding 和 LayerNorm

T2 我直接加载了 bert-base-uncased 的预训练权重，在 GLUE 的 MRPC 任务上做微调。因为它是成熟的预训练模型，理论上各层分布都应该是健康的，但我统计梯度的时候发现了一个很有意思的失衡现象：

• embedding 层的梯度范数高达 47.3
• LayerNorm 各层的 gamma 梯度范数在 18 到 25 之间
• 而第 5 到第 8 层的 Transformer 块内部 Linear 层梯度范数只有 0.03 到 0.09

也就是说，模型参数的大部分更新压力落在了 embedding 和 LayerNorm 上，真正的中层特征提取模块反而几乎没有收到有效梯度。这算不算"异常层"？从训练健康度的角度看算——它会导致微调时底层词嵌入被剧烈修改，而中层学到的语言知识被跳过。

我去翻了一下 BERT 的初始化方案，发现它把 embedding 的初始化标准差设成了 0.02，而典型的 Transformer 初始化用的是 1/sqrt(d_model)（大约是 0.03）。这个差异本身不大，真正的放大器是 LayerNorm 的 affine 参数和 embedding 梯度在整个词表上的累加效应——当序列里有大量 token 映射到同一个词嵌入时，这个 embedding 行的梯度会被乘以出现次数，形成梯度尖峰。

3.3 ViT：最后一层 LayerNorm 输出跑出数千倍方差

T3 用的是 ViT-Small 结构，在 CIFAR-10 上训练，hidden size 384、12 层。按照视觉 Transformer 的常规设置，输入 patch embedding 之后马上接一个 LayerNorm。我训练到第 3 个 epoch 的时候，hook 日志抓到了一个让我差点以为代码算错了的现象：

• 第 11 层输出的标准差在 0.9 左右
• 紧接着第 12 层的 LayerNorm 输出，标准差突然跳到 35.7
• 同时这个 LayerNorm 输出的 max 值到达了 141 这个离谱的数字

而且这个现象不是每个 batch 都出现，大概每 20 个 batch 会出现一次，表现得非常"间歇性"。也就是说这一层平时看着正常，一旦某个 batch 里包含了几张纹理特别极端的图片（比如有强烈高频边缘的裁剪），经过 11 层 attention 的叠加放大后，进入最后那个 LayerNorm 的信号里会出现几个异常大的分量。

这里有个很具体的工程细节：PyTorch 的 LayerNorm 默认 eps 是 1e-6，在输入方差本身就很大的情况下，这个 eps 几乎不起作用。我看到那个 LayerNorm 层的输入方差是正常 batch 的 300 倍，此时 eps=1e-6 根本压不住数值，于是 LayerNorm 输出的标准差直接"爆表"。这不是网络设计错了，而是数值边界没有给极端输入留好缓冲。

3.4 Swin Transformer：shifted window 的 softmax 过度尖锐

T4 的问题和前面几个都不一样，不是激活值爆炸，而是注意力权重退化。Swin Transformer 的窗口 attention 在做 shifted window 时，会有一部分窗口处在图像边缘，实际拿到的 token 数量远小于完整的 7x7=49 个。比如左上角第一个窗口，在 shift 之后可能只剩 4x4 个 token。

我统计了每个窗口的平均 attention 熵值，完整窗口的熵大约在 3.8 左右（64 个 token 左右时），而边缘那些小于等于 16 个 token 的窗口，熵值掉到了 2.0 以下。换句话说，这些窗口的 softmax 几乎变成了 one-hot——某个 token 的注意力权重接近 1，其他 token 全部被压到接近 0。

这个问题之所以算"层异常"，是因为它会导致梯度不均匀：窗口内部的注意力头退化后，该区域的梯度贡献近乎为零，网络在训练后期几乎学不到边缘窗口的局部特征。Swin 的作者其实也意识到了，这是相对位置编码在 token 数量过少时噪声增大的副作用。我的排查结果和官方 issue 里的一些报告高度吻合。

3.5 时序预测 Transformer：自回归解码时的 NaN 断层

最后是 T5，我最关心的业务场景——用标准 Encoder-Decoder 结构做带外生变量的多步时序预测。输入是 96 个时间步，输出预测 48 步。

前 30 步生成完全正常，但从第 31 步开始，decoder 的 masked self-attention 里出现 NaN，而且 NaN 只出现在 decoder 的第 2 层和第 3 层，第 1 层和第 4 层都还活着。

我开了torch.autograd.detect_anomaly()去反向追踪，定位到的源头是 attention softmax 层的输入出现了float('inf')。进一步追查，发现问题的起点在 decoder 的第 1 层 FFN 输出：当预测序列推进到第 30 步之后，有些位置的 FFN 输出跑出了 1e38 量级的数值，经过残差连接累加，到第 2 层的 attention logits 时已经大到 softmax 的指数计算直接上溢到inf。

时序预测的 Transformer 有一个和文本/图像任务都不一样的特征：它的输入数值范围是连续的、无界的，不像 token ID 那样有固定上界。只要训练数据里存在个别极大值点，模型就很容易在自回归过程中把误差不断累积放大，直到某一层彻底兜不住，数值崩溃。

4. 异常层的共同根因：不是某一层坏了，而是全局信号放大

五个模型跑完，我对着日志看了很久，慢慢理出了一个共性：所谓的"异常层"，并不是那一层的代码写错了或者结构特殊，而是整个模型内部的信号放大机制在那一层集中暴露了。换句话说，异常层只是"爆点"，真正的病灶在信号通过残差连接和归一化结构的整个传导链条里。

4.1 残差流上的信号累积效应

Transformer 的残差连接是一条贯穿所有层的主干道。每一层的输出都会加回到这条主干道上，深度学习社区把这条主干道称作"residual stream"。正常的 LayerNorm 会控制住每一层子模块的输出幅度，但控制不住主干道上累积起来的总能量。只要有一层输出稍微偏大，后面层接收到的输入就整体变大，这会直接放大所有后续层。

用 T1 的数据来算一笔账：假设每一层正常输出信号的增量是 1.1 倍，那么经过 6 层累积，残差流上的信号幅度大约是 1.1^6 ≈ 1.8 倍，这还在健康范围。但如果其中某一层的输出突然放大到正常值的 3 倍，残差流上就是 1 + 3 = 4 倍的额外能量，再经过下一层 FFN 的线性变换（权重矩阵条件数 300 以上），输出自然暴涨到十几倍。这不是神秘现象，就是最朴素的乘法放大了。

4.2 Pre-LN 和 Post-LN 的不同命运

我这次测的 5 个模型里，队形其实不一样：T1、T3 用的是 Post-LN（每个子层后再加 LayerNorm），T2、T4、T5 用的是 Pre-LN（每个子层前加 LayerNorm）。两种结构对异常层的"容忍度"完全不同。

Post-LN 的 LayerNorm 加在残差相加之后，理论上最平衡，但它在深层容易遇到梯度消失——本文开头说的 T1 第 5 层激活值爆炸就是典型的 Post-LN 深层问题。Pre-LN 的 LayerNorm 加在子层之前，输入被控制得更稳定，所以 T2、T4、T5 都没有出现激活值数量级的爆炸，但它们把不稳定性转移到了别的地方——比如 T5 的 FFN 输出累积导致 softmax 溢出，以及 T2 的梯度失衡。

这个对比告诉我们一件重要的事：当你看到一个 Transformer 的某一层异常，别急着改那一层，先看它到底用的是 Pre-LN 还是 Post-LN，这两类结构最常见的异常表现形式完全不同。

4.3 Softmax 温度的作用被低估了

T4 的窗口 attention 熵值过低，本质上是注意力缩放因子（也就是 softmax 的温度）和 token 数量不匹配。Transformer 的标准做法是点积除以 sqrt(d_model)，这个设计在序列长度较大、token 数量充足时很好用，但 Swin 的 shifted window 在边缘区域只能拿到十几个 token，此时每个 token 的注意力竞争空间只有一个小集合，除以 sqrt(96)≈9.8 之后，logits 的绝对差值仍然很大，softmax 的归一化会把最大权重推到 0.9 以上。

同样的问题在 T5 里也出现过，但不是因为 token 少，而是因为 logits 的绝对数值本身就大了。

4.4 初始化标准差的概率尾巴

还有一个藏在背后的共犯是初始化。五个测试对象的初始化标准差各不相同，但问题都出在"概率分布的尾巴"上。

初始化意味着权重是从某个随机分布里采样的，这个分布只要存在极端的尾部样本，一个 1e4 次方量级的 swin 窗口工况下，即使理论上概率极低，在深层网络几十上百次矩阵乘法中被放大后，也会变成实际观测到的异常值。LayerNorm 虽然在绝大多数时刻能把信号拉回到正常范围，但只能"事后归一化"，不能"事先预防"——输入里一旦有了足够的极端分量，LayerNorm 反而会把极端值放大到整个向量的比例中。

5. 修复思路与验证结果：把异常压回正常范围

排查不是目的，解决才是。针对上面五个模型的异常形态，我分别做了不同的修复调整，并在修复后用同一套 hook 工具做了前后对比。这里先说一个总的原则：修复的目标不是把某层输出调得零零整整，而是让整个模型的数值水平回到安全区间，同时不破坏原本的学习能力。

5.1 手写 Transformer：分层初始化 + 梯度裁剪

T1 的处理比较直接。既然问题出在第 4 到第 6 层的 FFN 权重条件数过高，我就重新用更小的标准差初始化了 FFN 的权重：

def init_ffn(layer, hidden_size): # 原实现: std = 1 / sqrt(hidden_size) ≈ 0.0442 # 改为: std = 1 / sqrt(hidden_size * 4) ≈ 0.0221 nn.init.normal_(layer.weight, mean=0.0, std=0.0221) nn.init.constant_(layer.bias, 0.0)

这个改动等于是给 FFN 的线性变换加了一个小的衰减因子，让残差流上的信号增量更小。同时我给训练循环加了梯度裁剪，阈值为 1.0。修复后第 5 层的输出标准差从 7.6 降到了 2.1，第 6 层从 13.4 降到了 2.6，整个模型恢复到一个健康区间。翻译任务的 BLEU 值也在同等训练步数下提升了约 1.2 个点。

5.2 ViT：调整 LayerNorm 的 eps 值

T3 的那个间歇性异常批让我意识到，eps 在 LayerNorm 里的作用远比平时想象的重要。我直接把所有 LayerNorm 的 eps 从 1e-6 调到了 1e-2，并且把位置编码的初始化幅度从默认的 1.0 缩小到 0.4。

这里解释一下原理：LayerNorm 的 eps 本质上是一个"地板值"，它对特征的方差做了个下界保护。正常情况下 1e-6 够了，但遇到方差暴涨的 batch 时，1e-2 可以把归一化输出限制在更小的波动范围内。位置编码缩小是为了避免 patch embedding 的数值叠加过于激进。

修复后那个间歇性异常 batch 没有再出现过，第 12 层 LayerNorm 输出的标准差稳定在 1.0 左右，CIFAR-10 测试集的 top-1 准确率从 86.1% 提升到了 86.7%。

5.3 Swin：降低窗口注意力的缩放系数

T4 的修复我尝试了两个方案。第一个方案是检查边缘窗口，把 token 数少于设定阈值的窗口中注意力缩放系数从 sqrt(d) 改为 sqrt(d) * 0.6，等于是把温度调低，让 logits 分布更平缓。第二个方案更通用——直接给 attention softmax 之前加一个温度学习参数，让网络自己学。

实测下来第二个方案效果更好，因为不需要手动判断哪些窗口是"边缘窗口"。我在每个 attention head 里加了一个可学习的温度标量，初始值 1.0：

class LearnableScaleAttention(nn.Module): def __init__(self, dim, num_heads): super().__init__() self.num_heads = num_heads self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1)) def forward(self, q, k, mask=None): attn = q @ k.transpose(-2, -1) * (self.temperature * (q.shape[-1] ** -0.5)) if mask is not None: attn = attn.masked_fill(mask == 0, float("-inf")) attn = attn.softmax(dim=-1) return attn

加入温度参数后，边缘窗口的平均 entropy 从 2.0 回升到了 2.9，同时整体模型的 Top-1 准确率没有下降。这个方案我觉得非常适合迁移到其他有"注意力退化"问题的 Transformer 模型上。

5.4 修复前后的整体数值对比

下面这个表可以直观看到五个模型修复前后的关键指标变化：

模型	关键指标	修复前	修复后
T1 经典 Transformer	第 5 层 encoder 输出 std	7.6	2.1
T2 BERT	embedding 梯度范数	47.3	11.5
T3 ViT-Small	第 12 层 LayerNorm 输出 std	35.7（间歇）	1.0
T4 Swin	边缘窗口注意力熵值	2.0	2.9
T5 时序 Transformer	decoder 第 31 步 NaN 出现概率	100% 崩溃	0%（50 步内）

这里需要说明的是，T2 的修复方式是调整微调的 LayerNorm eps 和优化器学习率分层（embedding 层用更小的学习率 1e-5，Transformer 层用 1e-4），修复后梯度比值从大约 500 倍降到 10 倍左右，这个幅度对 BERT 微调来说已经比较健康了。T5 则综合使用了 FFN 输出的梯度裁剪和 decoder 内部的 logits clamp，确保 attention 的输入不会跑到 1e30 这个量级。

6. 我实际踩坑后的几条硬经验

这次为期几天、横跨 5 个 Transformer 的排查，说到底给我最大的收获不是某个具体参数怎么调，而是建立了一套"数值体检"的心智模型。这里我把最有价值的几条经验写下来，都是普通文档里不会直接告诉你的东西。

第一，别只看 loss 曲线。many episodes 卡 loss 其实满满都是问题——loss 下降正常，不代表数值健康。T1 在出现激活值爆炸的同时，loss 还在稳步下降，因为网络找到了另一个方向来补偿不稳定的层，导致验证指标上不去。把每一层的激活分布打出来，往往比看 loss 更快发现问题。

第二，异常层几乎总是出现在残差主干道的中后段，而不是最浅或最深的层。最浅的层通常是 embed 或 patchify，它们直接吃输入，分布和输入强相关；最深的层离输出层近，梯度信号在这里衰减很快，也不太容易爆。最容易出问题的是模型前 1/3 到 2/3 位置的那些层——它们既累加了前面积累的主干信号，又距离输出足够远，梯度不容易帮它们"纠偏"。

第三，给每个 Transformer 建一个"数值体检"脚本，这回事值得做。我自己这次就是临时搭了一套基于 hook 的检查流程，后来发现这套东西完全可以沉淀成一个固定工具。每次训练启动后的前 100 步里跑一下，如果发现任何一层满足"std 超过前一层 3 倍"或"梯度范数超过同模型其他层中位数 20 倍"这两个条件，立刻自动告警。这个阈值组合我测试过，基本不会误报，又能及时把问题暴露出来。

第四，修复时不要用一个方案套所有模型。我之前一度想总结出"万能修复法则"，结果发现 Post-LN 和 Pre-LN 的侧重点完全不同，视觉和时序模型的数值边界也完全不同。最好的做法是：先用统一工具定位异常层，再结合该模型的归一化位置、初始化方式、任务数据分布来判断该往哪个方向修。调参很有用，但比调参更重要的是先搞清楚你要修的那个"异常"到底是从哪条路径传导过来的。

最后再补充一个小技巧。给 T5 做修复的时候，我发现在model.eval()推理阶段也会偶发 NaN，这让我意识到异常层的检查不能只在训练阶段做。推理时如果遇到某个 batch 出现极端值，模型同样会崩溃。所以后来我专门给 T5 的推理流程加了一个"数值哨兵"：对 decoder 每一层的输出做一次有限值检查，发现非有限值就立即用上一个合法 step 的输出做回退。这个机制虽然保守，但在线上服务里比任何参数修复都管用。

这 5 个 Transformer 的排查做下来，我对"异常层"的理解彻底变了：它不是一个神秘的 bug，而是一个信号流走向极端时留下的记号。真正有价值的不是"找到异常层就完了"，而是顺着异常层往回追，把整条传导链路检查一遍。这套思路放到任何模型、任何任务上都通用，值得你在自己的项目里试一试。