ALiBi 位置编码的数值失效:注意力失明现象 When Attention Goes Blind: Numerical Failure in ALiBi Positional Encodings
发现 ALiBi 线性偏置在长上下文中数值下溢,使注意力头变盲。
前置知识
ALiBi(带线性偏置的注意力)
Press et al. 2022 提出的相对位置编码,不引入额外参数,而是在 attention logits 上叠加一个与 query-key token 距离成正比的负偏置 $B = -m \cdot D$,其中 $m$ 是每个注意力头专属的 slope(斜率)。距离越远偏置越负,softmax 后权重越小。不同 head 按几何级数排列 slope,兼顾局部与长程依赖。
本文核心就是分析这个线性偏置在数值精度下的失效,看懂 $B = -m \cdot D$ 是理解全文的前提。
Softmax 数值下溢
softmax 把 logits 转成概率,分子分母都是指数 $e^{x_i}$。在有限浮点精度下,当指数过小时 $e^{x}$ 会小于可表示的最小正数,被截断为 0,称为下溢(underflow)。本文给出阈值 $\tau_{u,\mathrm{bf16}} \approx -92.18$、$\tau_{u,\mathrm{fp32}} \approx -103.27$:当某项 logit 加偏置后低于该值,对应注意力权重直接变 0。
这正是 ALiBi 失效的物理根源——线性增长的偏置迟早会把指数推过下溢阈值。
注意力头 slope 与失明距离
ALiBi 为每个 head 分配 slope $m_h$(默认 $2^{-8/(h-1)}$ 几何级数)。slope 越陡偏置随距离增长越快,head 越早失明。论文定义失明距离 $\delta_h = \min\{d \in \mathbb{N} \mid \varepsilon - m_h d \le \tau_u\}$,即该 head 对距离超过 $\delta_h$ 的 token 完全看不到。最陡 head 最先部分失明($\delta_1$),最平 head 最后完全失明($\delta_H$)。
slope 是论文实验的核心旋钮——通过调 slope 控制失明发生的早晚,从而隔离数值失效的影响。
长度外推(length extrapolation)
训练用较短上下文(如 2,048),推理/测试时希望模型在更长上下文(4,096、8,192)仍能工作,称为长度外推。ALiBi 当初的核心卖点就是 train short, test long。但本文发现数值下溢恰恰破坏了这种外推能力——距离越长失明越严重,外推在最需要的地方失效。
理解 ALiBi 的设计初衷,才能理解下溢为什么是反噬——它正好在最需要外推的长距离上让注意力失效。
Passkey / Needle-in-a-Haystack 检索探针
评估长上下文检索能力的两类探针。Passkey(密钥检索)在填充文本中先放一个密钥再追问模型;NIHS(大海捞针)把密钥嵌进真实文档,难度更高。本文用词组替代数字以适配小模型,并报告随 token 距离变化的准确率曲线下面积(AUC),按上下文内/外分别统计。
论文最重要的发现正是通过这两个探针暴露的——标准 benchmark 几乎看不出问题,但检索探针显示出巨大差异。
研究动机
ALiBi 自 2022 年提出以来,一直以廉价、无参数、可外推被宣传——它用一个与 token 距离线性增长的负偏置 $B = -m \cdot D$ 调制 attention logits,理论上让远距离 token 权重平滑衰减。然而作者发现,这个线性无界增长在有限浮点精度下会引发一个此前被忽视的失效模式:偏置随距离持续变负,最终把 softmax 里的指数 $e^{A_{i,j}+B_{i,j}}$ 推到浮点格式的下溢阈值之下(bf16 约 $-92.18$,fp32 约 $-103.27$),对应注意力权重被直接截断为 0。结果是该 head 对距离超过失明距离 $\delta_h$ 的所有 token 完全看不见,即部分失明。在 16-head、bf16 的典型配置下,token 距离达到 2048 时已有约 36.6% 的注意力矩阵项越过下溢线。这个失效直接反噬 ALiBi 最核心的承诺——长度外推:距离越长失明越严重,外推恰恰在最需要的地方失效。
本文的目标是本文的目标是系统刻画这一数值失效,并回答四个递进问题:(1)从解析上证明 ALiBi 偏置必然导致 softmax 下溢,并推导每个 head 的失明距离 $\delta_h$ 与 slope $m_h$ 的关系;(2)在真实预训练模型(BLOOM、Falcon-RW、MPT)上实证该问题确实发生,并测量其规模;(3)通过受控的小模型预训练实验,把注意力失明和上下文外退化两种效应解耦,搞清楚失明到底在多大程度上伤害模型;(4)提出并评测四类训练期缓解策略(clamping、robust slopes、log-scaled distances、soft capping),看能否在 passkey/NIHS 检索上恢复性能,并据此给出可操作的模型训练建议。
与已有工作不同的是,此前关于位置编码失效的研究几乎都集中在 RoPE 上(如 Wang et al. 2025 发现 bf16 下 RoPE 编码被破坏、Barbero et al. 2025 发现 Gemma 在低旋转频率上丢失语义信号),而 ALiBi 作为最便宜的替代方案长期被默认是数值稳健的。Chi et al. 2023 虽观察到 ALiBi 陡斜率 head 表现得像滑窗,但没把现象追溯到浮点下溢这一根因。本文的独特切入点是:把 head 行为异常这个模糊现象,精确归因到一个可计算、可验证的浮点下溢阈值,并进一步量化它在检索任务上的实际影响——填补了 ALiBi 到底安不安全这一被忽视的空白。
核心方法
直觉上 ALiBi 像给每个注意力头配一副度数不同的远视眼镜,陡 slope 看近、平 slope 看远;但偏置随距离线性增长无上限,指数迟早下溢为 0,head 便彻底看不见远处。技术路线:把 ALiBi 表述为 logits $A=QK^\top/\sqrt{d_k}$ 加偏置 $B=-m\cdot D$($D_{i,j}=(i-j)\mathbf{1}[j<i]$)再 softmax;证明当 $A_{i,j}+B_{i,j}\le\tau_u$ 时该项下溢为 0,由此定义失明距离 $\delta_h(\varepsilon)=\min\{d\mid\varepsilon-m_h d\le\tau_u\}$,以 $\varepsilon=0$ 的 $\delta_h$ 标记起点。最后提出四类缓解策略:偏置 clamping($B^*=\max(B,c_{\mathrm{clamp}})$)、重设 slope 使失明距离铺满上下文、对数压缩距离($D^*=\log(D+1)$)、对 logits 做 soft capping($A^*=z\tanh(A/z)$)。
核心创新不是提出一个新位置编码,而是揭示并量化一个被忽视多年的既有缺陷:ALiBi 的线性偏置在有限精度下必然下溢,且这种下溢是有结构、可预测的——每个 head 的失明距离 $\delta_h$ 完全由其 slope $m_h$ 和浮点阈值 $\tau_u$ 决定,可闭式算出。这与已有的线性偏置类工作(KERPLE、FIRE、Sandwich)形成本质区别:后者都在改进偏置的形状以利于外推,却没人质疑过线性增长会不会撑爆浮点数。更关键的是,论文把这一纯数值现象与模型行为挂钩:失明不仅造成前向传播里浪费算力(被下溢清零的项仍参与计算),更可能在训练反向传播中切断梯度——因为只有非零权重才带梯度,下溢后的 head 永远学不到如何利用窗口外的信息。这一训练期危害假说是此前完全没有被讨论过的角度。
方法步骤详情
论文分四步。第一步(§3 解析):把 $B=-m\cdot D$ 代入 softmax $\mathrm{softmax}(x)_i=e^{x_i}/\sum_l e^{x_l}$,证明随距离增长指数必然越过阈值 $\tau_u$,推导失明距离 $\delta_h$,并指出失明带来位置盲区与权重重分配。第二步(§4.1):在 BLOOM-560M、Falcon-RW-7B、MPT-7B 上用 FineWeb-Edu 分层测试集测困惑度并统计实际下溢比例。第三步(§4.2):用 148M 的 Llama 架构、20B FineWeb-Edu token、每配置三个种子训练,对比 Steep/Safe/Wide/ALiBi/RoPE 五种 slope,在 0–4096 距离测 passkey/NIHS AUC 并跑 CS/QA/LG 零样本基准。第四步(§4.3):同 setup 下单独及组合评测 clamping($c_{\mathrm{clamp}}=-87$)、robust slopes($\delta_H=4096$)、log 距离、soft capping($z=50$)。
技术新颖性
新颖性体现在三个层面。分析层面:首次给出 ALiBi 失明的闭式刻画 $\delta_h = \min\{d \mid \varepsilon - m_h d \le \tau_u\}$,把模糊的 head 表现异常量化成可计算的几何量;并指出 MPT 早就在实现里夹了偏置($\delta_H \approx 2000$),说明部分人隐约察觉过问题却未公开报告。实验层面:通过设计 Steep/Safe/Wide 三组极端 slope 配置,巧妙地把注意力失明和超出训练上下文两种退化解耦——此前没人这样做过对照。缓解层面:把 clamping、log 距离、soft capping 等散见于其他位置编码(KERPLE 的 log、Gemma 的 soft capping)的技术,首次系统地搬到 ALiBi 上做消融,并发现 log 缩放能把最陡 head 的 $\delta_1$ 从 124 推到 $4.37 \times 10^{53}$,几乎永久消除下溢。
实验结果
预训练模型检验(Figure 2、Table 2)证实失效存在:三模型困惑度越过 2048 训练窗口后骤升,下溢比例与解析预测高度吻合;检索探针差异显著,BLOOM 上下文外近乎崩溃(NIHS out 0.06),而 MPT 因实现带 clamping 达 NIHS out 0.16。受控训练(Table 3a)最反直觉:刻意失明的 Steep($\delta_1=128$、$\delta_H=512$)检索最差(PK out 0.09、NIHS out 0.02)却未崩溃;几乎消除失明的 Safe passkey 最强(in 0.98、out 0.17),NIHS 反远逊 ALiBi(in 0.48 vs 0.68、out 0.02 vs 0.20)——平 slope head 兼作检索头,弱位置信号本身有用。缓解实验(Table 3b)无单一赢家:log 距离把上下文外 passkey 从 0.08 拉到 0.77,与 clamping 组合(C+L)达 0.79,但含 L 配置 NIHS 全线下滑;标准 benchmark 各组仅差 1.6–3.4 个百分点,缺陷几近隐形。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Passkey 检索(上下文外) | AUC | C+L = 0.79 | ALiBi 默认 = 0.08 | +0.71(约 10 倍) |
| Passkey 检索(上下文内) | AUC | E+S = 0.99 | ALiBi 默认 = 0.93 | +0.06 |
| NIHS 大海捞针(上下文内) | AUC | ALiBi 默认 = 0.68 | RoPE = 0.28 | ALiBi 领先 +0.40 |
| NIHS 大海捞针(上下文外) | AUC | ALiBi 默认 = 0.20 | 所有缓解策略均 ≤ 0.17 | 默认 ALiBi 仍最强 |
| Decoder benchmark(语言类 LG) | 平均分 | ALiBi = 54.0 | Steep = 52.4 | +1.6 pp(差异很小) |
局限与改进
作者明确承认两点:实验只在 148M 参数的小模型上做,虽然下溢本身是解析成立的,但缓解策略的效果未必能外推到几十亿、几百亿参数的模型;同时架构和语料固定(Llama + FineWeb-Edu),没有探索不同架构或语料的交互作用,把最有潜力的配置做大规模超参搜索留给了未来。我额外观察到三点局限:一是 MPT 这种实现里夹了 bias 的工程做法暗示工业界早有应对,但论文未深入调研各家推理框架的真实行为;二是 perplexity 与下溢比例在 Figure 2 里并未呈现强相关,作者自己也承认无法判断困惑度尖峰究竟来自下溢还是单纯越界,削弱了下溢导致困惑度爆炸的因果链;三是每个配置只有三个随机种子,且若干组合(C、E、C+S)方差极大(如 CS 标准差达 2.4–2.6),说明若干结论在统计上并不稳健。
独立分析的弱点
第一,缓解策略与 NIHS 任务的内在矛盾未解决。Log 缩放虽然救了 passkey,却系统性地压低 NIHS(in 0.41、out 0.03),原因是它把所有距离压扁、削弱了区分正确 token 所需的位置判别力。改进方向:可考虑混合编码——近距离保留线性 ALiBi 偏置以保判别力,远距离切换到 log 或夹断,做成分段函数式的相对偏置。第二,只在 148M 验证,规模泛化存疑。大模型 head 数更多、slope 更平,失明距离可能整体右移,缓解策略的相对收益会改变。改进方向:至少在一个 1B–3B 模型上复现关键配置(如 C+L、E+S)。第三,未量化训练期梯度危害。作者提出下溢切断梯度假说但未验证。改进方向:直接监测失明 head 的梯度范数与窗口外 attention 的学习动态,用梯度分析坐实或证伪该假说。第四,标准 benchmark 看不出问题这一隐形性本身是隐患——改进方向是提出一个把下溢比例纳入报告的长上下文评测协议。
未来方向
作者明确点名的方向包括:用更大模型和更系统的超参搜索验证最有潜力的组合(如 C+L、E+S);实证下溢在反向传播中切断梯度这一训练期危害假说;探索 nGPT 等架构级归一化能否替代 soft capping 来限制 logit 范围。基于成果可延伸的方向我认为有:把分段线性/对数的混合偏置做成可学习参数,让模型自己学距离-偏置曲线;将失明距离 $\delta_h$ 作为评测指标纳入长上下文模型榜单;把分析框架推广到 TISA 等其他线性加性偏置编码;以及研究推理框架(FlashAttention/FlexAttention)层面对失明 token 的稀疏化优化,把现在算了再丢弃的浪费变成显式滑窗加速。
复现评估
复现门槛中等偏高。有利因素:作者承诺发表后开源代码;训练用 FineWeb-Edu、评测用 lm-evaluation-harness 全是公开数据/工具;模型架构完全照搬 SmolLM 配方,超参在附录 Table 7 全部列出;预训练模型检验部分(BLOOM/Falcon/MPT 的困惑度加下溢测量)几乎不耗算力,任何人都能快速复现。不利因素:148M 模型训练用了约 9,000 GPU·小时的单卡 H100(94 GiB),主实验全部在 bf16 + FlexAttention 下跑,且作者强调 FlashAttention 当时不支持 clamping/log/soft capping,所以必须用他们基于 FlexAttention 的实现——这意味着要复现缓解实验,得先拿到这套定制注意力 kernel。结论:解析部分和预训练模型检验极易复现,受控训练实验需要数千 GPU·小时和定制代码,普通研究者难以完整复刻,但可借助开源代码在中小规模验证关键结论。
论文图表
图分两部分。(a) 画出 ALiBi 各 head 的偏置随 token 距离线性增长,并标出 bf16 下溢阈值(虚线约 $-92.18$),偏置越过该线即开始失明。(b) 从注意力矩阵视角展示随 token 距离增大,下溢(变盲)的注意力权重比例如何上升,并在 2048 距离处达到约 36.6%。
这是全文的图腾——一眼讲清下溢如何把线性偏置变成盲区,是理解整个失效机制的起点。