ALiBi位置编码数值失效的根因、复现与修复

在 Transformer 生态中位置编码一直是“看似简单、实则敏感”的模块。很多人用 ALiBi 替换传统位置编码之后会碰到一个很隐蔽的问题模型在训练前期正常跑几个 step 后 loss 出现 NaN或者推理时某几个 head 的 attention 输出变成全零。这种问题往往不是数据问题也不是梯度爆炸而是ALiBi 位置编码在特定数值条件下发生的数值失效。本文将围绕 ALiBi 位置编码的原理、数值失效的根因、复现方式与修复方案展开。1. 从 Attention 到 ALiBi位置编码的演进1.1 为什么 Transformer 需要位置信息Self-Attention 的核心计算是Q和K的点积它天生是“对顺序不敏感”的。无论词在句子中的位置如何交换只要Q、K的向量内容不变注意力分数就不会变。但自然语言是强顺序相关的猫追老鼠和老鼠追猫的语义完全不同。为了让模型感知顺序Transformer 需要在输入中注入位置信息。经典做法有两类绝对位置编码为每个位置学习一个可训练向量或使用正弦/余弦函数生成固定向量然后把位置向量加到输入 embedding 上。相对位置编码在注意力分数计算时根据两个 token 之间的相对距离调整分数。RoPE旋转位置编码和 ALiBi 都属于相对位置编码的范畴。RoPE 通过旋转矩阵把位置信息混合进Q、KALiBi 则直接在注意力分数上增加线性偏置。1.2 ALiBi 的设计动机ALiBiAttention with Linear Biases来自论文Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation核心动机是解决 Transformer 的长度外推问题。它不再为 token 添加位置 embedding而是在注意力分数计算后直接加上一个与 token 距离成正比的负偏置token 之间距离越远偏置的绝对值越大注意力分数被压得越低token 之间距离越近偏置越接近 0注意力分数几乎不受影响。这种设计的优势是不需要额外的可学习位置向量省了 embedding 参数同时某种程度上具备序列长度外推能力。1.3 数值失效问题概述“数值失效”不是 ALiBi 必然导致的问题而是多种因素叠加后的结果。在常见的 ALiBi 实现中注意力分数的计算是scores Q K^T / sqrt(d) scores scores alibi_bias scores scores attention_mask probs softmax(scores)如果某些位置被 mask 设为-inf再与 ALiBi 的负偏置叠加或者scores本身在 fp16 下超过浮点数表示范围softmax 就可能计算出NaN或全零梯度。下面我们会从数值层面拆解这个问题。2. ALiBi 原理与通用实现2.1 ALiBi 的偏置公式ALiBi 的核心公式可以写成alibi_bias[i, j] -m_head * |i - j|其中i、j表示序列中两个 token 的位置m_head是每个 attention head 专有的斜率slopealibi_bias是一个形状为[num_heads, seq_len, seq_len]的矩阵。在因果语言模型中j只能取 i的位置所以通常还会对alibi_bias先加一个上三角 mask再参与 softmax 计算。下面是一个最小实现import torch def build_alibi_bias(seq_len: int, num_heads: int) - torch.Tensor: 构建 ALiBi 相对位置偏置矩阵。 返回 shape: [num_heads, seq_len, seq_len] # 1. 计算每个 head 的斜率 closest_power 2 ** torch.floor(torch.log2(torch.tensor(num_heads, dtypetorch.float32))) base 2 ** (-8 / closest_power) powers torch.arange(1, num_heads 1, dtypetorch.float32) slopes torch.pow(base, powers) # 2. 构建相对距离矩阵 positions torch.arange(seq_len, dtypetorch.float32) rel_dist positions[None, :] - positions[:, None] # [seq_len, seq_len] # 3. 对 rel_dist 取绝对值并乘上负斜率 alibi_bias -slopes[:, None, None] * rel_dist.abs()[None, :, :] return alibi_bias # 示例8 个 head序列长度 64 bias build_alibi_bias(64, 8) print(bias.shape) # torch.Size([8, 64, 64])这个实现是理解 ALiBi 的最小骨架。实际工程中你可能会看到它被塞进xformers、flash-attention或自定义 CUDA kernel 中形式略有不同但核心逻辑一致。2.2 slope 的计算方式每个 head 的斜率不是随机生成的而是按几何级数排列的。如果 head 总数是 8那么斜率取值是2^(-8/8), 2^(-16/8), 2^(-24/8), ..., 2^(-64/8)也就是从2^-1到2^-8。斜率越小对远距离 token 的惩罚越弱斜率越大对远距离 token 的惩罚越强。这相当于让不同的 head 分别关注“近程”和“远程”信息。2.3 与 mask 叠加的两种顺序ALiBi 偏置和 attention mask 的叠加顺序非常关键。常见有两种写法# 顺序 A先加 ALiBi再加 mask scores scores alibi_bias scores scores attention_mask # 顺序 B先加 mask再加 ALiBi scores scores attention_mask scores scores alibi_bias看起来两者只是加法的顺序不同但在数值上可能产生完全不同的结果。我们来看后面的分析。3. 数值失效根因拆解3.1 从 softmax 的数值特性讲起softmax 的公式是softmax(x_i) exp(x_i) / sum_j exp(x_j)计算exp(x)时如果x是较大的正数exp(x)可能溢出为inf如果x是一个很小的负数exp(x)会下溢为 0。标准做法是在 softmax 前减去最大值exp(x_i - max(x))这样可以把分母限制在有限范围内避免直接溢出。但这里有个前提所有参与 softmax 的值都必须是有限数。一旦输入中出现NaN或-inf即使减去最大值也无法恢复NaN - max仍然是NaN-inf - (-inf)是NaN-inf - 有限值仍然是-inf。3.2-inf与NaN的来源在 ALiBi 的数值失效案例中-inf和NaN主要有三个来源。来源一因果 mask 的-inf因果注意力需要把未来位置 masked 掉常见的实现方式是给未来位置的分数加上一个非常大的负数很多代码直接写为-inf。如果 ALiBi bias 是先加然后再加上-inf此时某个位置的 scores 是scores[i, j] qk_score alibi_bias[i, j] - inf-inf和alibi_bias相加仍然是-inf看起来没问题。但如果在某些框架或自定义 kernel 中加法顺序被改写成了scores[i, j] (qk_score - inf) alibi_bias[i, j]此时qk_score - inf得到-inf-inf alibi_bias依然是-inf。问题在于有些算子对-inf的传播处理不一致可能在中间产生NaN。特别是当alibi_bias中也存在inf或极端值时。来源二fp16 下的溢出在混合精度训练中Q K^T的结果可能很大。以 fp16 为例最大表示范围是 65504。如果Q、K的数值较大内积结果很容易超过这个范围并变成inf。在 bf16 下最大表示范围更大但精度更低。两者的风险不同fp16溢出风险更高表现为inf或NaNbf16溢出风险相对较低但精度不够小斜率容易被吞掉。来源三softmax 的 0/0 情况当 scores 中存在-inf时exp(-inf)得到 0。如果某个位置的所有 scores 都是-inf那么 softmax 的分子和分母都是 0结果是NaN。这种情况在加了因果 mask 后不太容易出现因为对角线位置通常是可见的。但在某些实现中如果 mask 错误地把对角线也屏蔽了或者alibi_bias的负偏置过大导致所有可见位置的分数都非常低最终 softmax 结果也会偏向极端。3.3 fp16/bf16 精度下的放大器效应ALiBi 对精度的敏感程度比普通位置编码更高原因是偏置的量级和 attention score 的量级可能相差很大。例如第 6 个 head 的斜率是2^-48非常小。在 fp16 下这个值约等于2^-48 ≈ 3.55e-15它比 fp16 的最小正规格化数6.10e-5要小得多。也就是说fp16 很难精确表示2^-48这个斜率它会被下溢为 0或者被舍入到一个不精确的值。一旦斜率变成 0这个 head 的 ALiBi 偏置就完全失效退化成普通的 attention。这不是“效果差一点”的问题而是模型结构行为发生了变化。3.4 ALiBi 特有的退化场景除了常规的数值溢出问题ALiBi 还有一个特殊场景值得注意推理时序列长度大于训练长度。假设模型训练时最大长度是 2048推理时输入长度变成了 4096。此时alibi_bias的最大绝对值变成-m_head * 4095而训练时最大只有-m_head * 2047。对于斜率较大的 head这个偏置会变得非常大导致长距离 token 的 attention score 被压到非常低甚至接近 0。经过 softmax 后这些 token 的注意力权重完全消失。这就是标题中 “Attention Goes Blind” 的一个直观体现模型在两三千 token 以内工作正常一旦超过训练长度远端上下文就完全“失明”了。4. 复现数值失效问题4.1 复现环境为了验证数值失效问题我们不需要完整的模型只需要一个包含 ALiBi 的注意力模块。复现环境Python 3.10PyTorch 2.x一个支持 fp16 的 GPU或 CPU 模拟 fp16版本不是必须固定的重点是展示数值现象。4.2 一个最小复现脚本下面这个脚本构造了一个小规模的注意力场景对比 fp32 和 fp16 下的输出。import torch import torch.nn.functional as F def build_alibi_bias(seq_len: int, num_heads: int, dtype: torch.dtype torch.float32): closest_power 2 ** torch.floor(torch.log2(torch.tensor(num_heads, dtypetorch.float32))) base 2 ** (-8 / closest_power) powers torch.arange(1, num_heads 1, dtypetorch.float32) slopes torch.pow(base, powers) positions torch.arange(seq_len, dtypetorch.float32) rel_dist positions[None, :] - positions[:, None] alibi_bias -slopes[:, None, None] * rel_dist.abs()[None, :, :] return alibi_bias.to(dtype) def causal_attention(q, k, v, alibi_bias): q/k/v: [batch, heads, seq_len, head_dim] scores torch.matmul(q, k.transpose(-2, -1)) / (k.size(-1) ** 0.5) # 先加 ALiBi scores scores alibi_bias.unsqueeze(0) # 构造因果 mask seq_len q.size(-2) mask torch.triu(torch.ones(seq_len, seq_len, deviceq.device, dtypetorch.bool), diagonal1) scores scores.masked_fill(mask.unsqueeze(0).unsqueeze(0), float(-inf)) probs F.softmax(scores, dim-1) output torch.matmul(probs, v) return output, probs torch.manual_seed(42) batch, heads, seq_len, head_dim 2, 8, 128, 64 # fp32 输入 q32 torch.randn(batch, heads, seq_len, head_dim) k32 torch.randn(batch, heads, seq_len, head_dim) v32 torch.randn(batch, heads, seq_len, head_dim) # fp16 输入 q16 q32.half() k16 k32.half() v16 v32.half() bias32 build_alibi_bias(seq_len, heads, torch.float32) bias16 build_alibi_bias(seq_len, heads, torch.float16) out32, probs32 causal_attention(q32, k32, v32, bias32) out16, probs16 causal_attention(q16, k16, v16, bias16) print(fp32 output contains NaN:, torch.isnan(out32).any().item()) print(fp16 output contains NaN:, torch.isnan(out16).any().item()) print(fp32 probs sum row[0]:, probs32[0, 0, 0].sum().item()) print(fp16 probs sum row[0]:, probs16[0, 0, 0].sum().item()) print(fp32 max prob:, probs32.max().item()) print(fp16 max prob:, probs16.max().item())在这个脚本里我们没有刻意做任何“错误”的操作只是使用常规的 ALiBi 实现并切换到 fp16。你可能会看到以下现象之一fp16 output contains NaN输出为Truefp16 下某一行 softmax 的权重分布极端化fp32 和 fp16 的probs差异很大。4.3 预期现象与输出分析如果q/k/v的数值范围比较大fp16 下scores很容易出现inf。例如在 fp16 中如果某个scores元素是1e4加上 ALiBi 偏置后变成9e3再经过 softmax 的exp结果溢出为inf最后输出变成NaN。我们可以做一个更直接的实验把q/k/v放大一定倍数观察数值变化。# 放大 Q/K/V模拟模型训练早期数值不稳定阶段 scale 3.0 q16_large (q32 * scale).half() k16_large (k32 * scale).half() v16_large (v32 * scale).half() out16_large, probs16_large causal_attention(q16_large, k16_large, v16_large, bias16) print(fp16 (scaled) output contains NaN:, torch.isnan(out16_large).any().item()) print(fp16 (scaled) probs sum row[0]:, probs16_large[0, 0, 0].sum().item())这种情况下几乎必然出现NaN。核心原因是Q K^T的结果与Q/K的模长直接相关当模型训练前期参数初始化数值较大或 layer norm 位置设计不当时scores会突破 fp16 的表示范围。5. 修复与防御方案5.1 统一 mask 与 bias 的加法顺序首先建议在代码中固定一种安全顺序先加 ALiBi 偏置再加 mask。scores scores alibi_bias scores scores.masked_fill(causal_mask, float(-inf))这个顺序能保证-inf只来自于 maskALiBi 偏置不会与-inf混合运算。如果你先加 mask再加 ALiBiscores scores.masked_fill(causal_mask, float(-inf)) scores scores alibi_bias # -inf finite 仍然等于 -inf但某些算子可能产生 NaN在 PyTorch 中-inf finite通常还是-inf但在部分自定义 CUDA kernel 或低精度算子中行为可能不一致。统一顺序可以减少一个不稳定变量。5.2 使用安全的 softmax 实现不要手工写torch.exp(scores) / torch.sum(torch.exp(scores), dim-1)除非你确定输入范围。优先使用 PyTorch 内置的F.softmax因为它内部会做max减法数值上更稳定。如果你使用的是flash-attention需要注意它的 mask 接口。Flash Attention 对 mask 的处理与标准 PyTorch 不同它采用attn_mask参数且不同版本行为有差异。在实际项目中要确认你使用的 flash-attention 版本对 ALiBi 偏置的支持方式避免“以为已经加了实际被忽略”的情况。5.3 使用 float32 累加在混合精度训练中一个常见优化是让Q K^T以 fp16 计算但在累加偏置和 mask 之前把 scores 转换为 fp32。# 伪代码示例 scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(k.size(-1)) scores scores.to(torch.float32) scores scores alibi_bias # alibi_bias 保持 fp32 scores scores.masked_fill(mask, float(-inf)) probs F.softmax(scores, dim-1) output torch.matmul(probs.to(v.dtype), v)这样可以避免scores在累加过程中溢出。代价是注意力计算的显存占用更高。如果是长序列场景可以只在早期训练阶段开启稳定后再切回 fp16。5.4 针对 ALiBi 的偏置裁剪策略针对长序列推理时的 “Attention Goes Blind” 问题可以对 ALiBi 偏置做一个裁剪clipping操作。具体做法是给alibi_bias设置一个下限min_bias不让它无限变大def build_alibi_bias_clipped(seq_len: int, num_heads: int, min_bias: float -8.0) - torch.Tensor: bias build_alibi_bias(seq_len, num_heads) bias torch.clamp(bias, minmin_bias, max0.0) return bias这是对 ALiBi 的一种局部修改会改变模型的注意力分布需要重新验证效果。但在长序列推理中如果目标是“不崩溃”这比直接输出 NaN 或全零权重要好。实际项目中是否使用裁剪需要结合模型效果来评估。如果你的任务对长距离依赖要求很高裁剪可能伤害远端信息的传递如果只是为了避免数值崩溃裁剪是简单有效的方案。6. 常见问题与排查清单在实际使用 ALiBi 的过程中你可能会遇到下面这些情况。我把高频现象、可能原因和解决思路整理如下。问题现象可能原因解决思路训练早期 loss 变成 NaNfp16 下 QK^T 溢出为 inf将 scores 累加分转换为 fp32或减小初始学习率推理时恰好第 N 个 token 之后输出异常ALiBi bias 在长序列下过大对 alibi_bias 做 clip某个 head 的位置信息失效低精度下极小的 slope 被下溢为 0在 fp32 中计算偏置再转换到 fp16使用 flash-attention 后行为不一致不同版本对 ALiBi 的支持方式不同检查 flash-attention 的 mask 接口确认偏置是否生效softmax 输出出现全 0 权重mask 与 ALiBi 叠加后出现-inf行检查 mask 是否正确保留对角线位置fp32 正常fp16 训练 loss 波动大动态 loss scaling 窗口不佳调整 loss scaling 策略或对 scores 做 clip下面这个 checklist 可以作为排查参考先确认scores中是否出现inf或NaN。打印alibi_bias的最小值和最大值确认低精度下 slope 是否下溢。确认 mask 中非-inf的位置是否满足预期。把注意力分数改成 fp32 累加看问题是否消失。如果使用 flash-attention尝试关闭它改用标准 attention 验证。逐步增加序列长度找到输出开始异常的长度阈值。7. 最佳实践与工程建议7.1 偏置矩阵统一在 fp32 下构建ALiBi 的偏置矩阵应该始终用 fp32 计算最后再转换为目标精度。不要把rel_dist.abs()的乘法放在 fp16 或 bf16 下进行否则极小的 slope 精度会被严重破坏。alibi_bias_fp32 build_alibi_bias(seq_len, num_heads, torch.float32) alibi_bias_fp16 alibi_bias_fp32.half()7.2 记录 attention 熵值在长序列场景下建议评估 attention 权重的熵值。如果某个 head 的 entropy 持续接近 0说明它的注意力分布过于尖锐大概率已经退化。def attention_entropy(probs: torch.Tensor) - torch.Tensor: probs: [batch, heads, seq_len, seq_len] 返回 [batch, heads, seq_len] return -torch.sum(probs * torch.log(probs 1e-12), dim-1)在训练或推理过程中定期记录 entropy能提前感知 ALiBi 偏置是否破坏了注意力分布。7.3 监控远距离 token 的注意力权重如果你正在做长文本任务推荐单独统计“远距离 token 的注意力权重总和”。例如当序列长度为 4096 时统计每个 query 对前 1024 个 token 的注意力权重和。如果这个值长期为 0说明远端上下文确实被 ALiBi 屏蔽了需要调整 bias 的 slope 或裁剪策略。7.4 不要盲目修改 ALiBi 的实现ALiBi 的 slope 计算方式在很多开源实现里已经被验证过不要为了“增加稳定性”随意修改公式。如果你必须增加稳定机制建议通过外层操作实现对scores做固定阈值的 clip对最终output做 dropout 或残差连接在 fp32 下完成注意力分数计算。这些方式对模型原结构影响相对较小。7.5 工程上保留一份 FP32 注意力实现在开发阶段最好保留一份纯 fp32 的注意力实现作为基准。遇到数值问题先用它跑一遍判断问题是来自“模型结构”还是“低精度实现”。这能大大缩短排查时间。8. 总结ALiBi 位置编码解决了一个实际问题却也引入了一类容易被忽视的数值风险。它的核心问题在于偏置的大小和 attention score 的量级差异过大在低精度训练和超长序列推理时很容易触发溢出、下溢和 softmax 退化。本文从 ALiBi 原理出发拆解了数值失效的四个来源给出了可复现的最小代码示例并总结了修复方案与工程建议。最核心的三件事是mask 与偏置的叠加顺序要固定、低精度计算前保护 scores 范围、长序列推理时要审视 ALiBi 是否让远端注意力失明。如果你正在使用 ALiBi 或相关相对位置编码方案建议先把本文的复现脚本跑一遍观察你自己模型在 fp16 和长序列下的表现再做针对性优化。