ALiBiAttention with Linear Biases是一种用线性偏置替代位置编码的注意力机制变体。它在注意力分数上减去一个与位置距离成正比的偏置让模型在短序列上训练之后还能向更长序列外推。这个思路在很多长文本模型的实验里都有出现实用价值是明确的。但它有一个非常容易被忽视的问题当序列变长、距离变大、精度降到 fp16/bf16 时偏置会把部分 logits 推到数值不稳定区softmax 输出的概率分布随之退化注意力对内容不再敏感。说得直白一点就是 attention 在数值层面“失明”了。这篇文章会把 ALiBi 的数值失败机制、复现条件、检测方法和缓解手段完整拆一遍。适合正在做长上下文模型、使用 ALiBi 或类似线性偏置方案、以及排查注意力输出异常的人。1. 先搞清楚 ALiBi 的机制和“失明”到底指什么1.1 ALiBi 不改变词向量只改变注意力分数的“地势”Transformer 的位置信息有两种常见注入方式。一种是把可学习位置向量或正弦位置向量加到词向量上比如绝对位置编码另一种是在 attention 内部引入相对位置关系比如相对位置编码。ALiBi 走的是后一类但实现更简单。它不改变 query、key、value 的输入只在一个地方动手QK 点积得到注意力分数之后先乘上缩放系数1 / sqrt(d)再减掉一个和距离成正比的偏置。用公式展开大致是这样score(i, j) q_i · k_j / sqrt(d) - m * |i - j|其中m是每个注意力头自己的斜率。为什么用“斜率”而不是“向量”因为 ALiBi 的核心假设是两个 token 距离越远当前 query 就越不该关注远处的 key。这个“不该关注”的程度用一个线性衰减项直接近似比学习位置向量更省参数也更容易外推。这里有个重要特性softmax 对一个固定常数的偏移是不敏感的。所有 logits 同时加一个常数softmax 结果完全不变。所以 ALiBi 真正改变的不是 logits 的绝对高低而是不同位置之间的“相对差异”。离得近的 key 相对优势变大离得远的 key 相对劣势变大。这就是为什么它能做长度外推训练时模型只见过短距离的相对差异推理时即使序列边长公式里的距离项依然可算不需要重新训练位置向量。1.2 “失明”的具体表现均匀化、锁死、梯度消失理解了机制之后就明白“失明”不是模型崩了或者输出 NaN而是注意力分布退化。我见过三种典型表现。第一种是注意力分布均匀化。如果 logits 之间的相对差异太小softmax 就会退化成接近均匀分布。此时每个 query 对所有 key 一视同仁attention 等于没看内容。这种情况通常出现在偏置没有起到区分作用、或者计算精度丢失了大量相对信息的时候。第二种是注意力被锁死到极小的局部窗口。长序列里距离项不断增大远处 key 的 logits 会被压到非常负的区域。低精度下这些 logits 经过 exp 之后直接变成零softmax 概率全部集中在相邻几个 token 上。从指标上看注意力熵极低模型只看得见局部看不到长距离依赖。这其实比均匀化更隐蔽因为训练 loss 看起来还在下降但模型实际上已经失去长程建模能力。第三种是梯度消失。低精度下 exp 把概率变成 0反向传播时这些位置的梯度也是 0。如果某个 key 长期拿不到梯度它的 token 表征就学不到信息等于这个位置被“放弃”了。更麻烦的是这种失效通常是渐进的不是某一个距离突然全部失效而是从尾部开始一点一点变稀疏很不容易被人工发现。2. 从注意力分数到 softmax数值风险发生的三个位置2.1 距离线性增长logits 会被压到很负的区域先做一个最直观的计算。假设某个注意力头的斜率m 0.1在 1024 长度的序列里最远距离是 1023偏置项就是-0.1 * 1023 ≈ -102。如果序列拉到 8192最远偏置是-819。这些数值本身还在 fp32 的可表示范围内问题不在“存不下”而在 softmax。softmax 的核心操作是exp(logit)。exp(-20)大约是2e-9已经低于 fp16 的正常表示下限exp(-50)大约是1.9e-22在 fp16 下直接下溢到零。也就是说在 fp16 训练或推理时距离达到几百个 token远处位置的注意力概率就已经是精确的 0。这意味着 ALiBi 的“线性惩罚”在一个精度有限的系统里会退化成“硬截断”超过某个距离之后所有 key 的概率都是零。注意这不是设计里的理想行为设计者希望的是概率平滑衰减但数值实现把它变成了阶跃行为。更要命的是如果 logits 区域整体下溢softmax 分母里的那些exp(logit)也可能变成 0。虽然 softmax 数学上要求分母大于 0但浮点实现里分母可能变成一个极小的数或者干脆是 0。这时候可能出现 NaN也可能因为分母精度不足导致概率值完全失真。2.2 fp16/bf16 下 exp 很快触底不同精度的失效边界差别很大。fp16 的指数范围大约在[-14, 15]以 2 为底的指数正常最小非零数约6e-5次正规数最小约6e-8。对于exp(logit)来说logit 大约小于-16就开始进入次正规区小于-24左右基本就下溢成 0。所以 fp16 下距离几百的 token 就很难拿到有效概率。bf16 的指数范围和 fp32 一样大所以不会那么快触发 underflow但它的尾数只有 8 位相对精度很低。当 logits 都很大或很负时相邻两个概率之间的差值可能被舍入掉。换句话说bf16 更早出现的是“概率被量化糊在一起”表现为注意力分布没有层次感。这不是说低精度一定不能用而是说要用的时候必须知道边界你的序列长度、头部斜率、计算精度三者共同决定哪些位置的注意力概率还能保留有效信息。2.3 和 FlashAttention 的线上 softmax 叠加后更难察觉实现层面还有一个容易被忽略的交叉点FlashAttention 这类工具不会一次性把所有 logits 都算出来而是分块做“线上 softmax”。它依赖每个 block 的局部统计量配合 running max 和 rescaling 来还原全局 softmax。ALiBi 偏置如果是在 kernel 内部逐 block 加的那么每个 block 内部的 running max 可能差异很大。理论上 rescaling 会修正但浮点计算里当一个 block 的局部最大值和另一个 block 的最大值相差极大时小的那一方的 exp 项会下溢或严重丢失精度。尤其是在 causal mask 和 ALiBi 同时存在时很多位置既有-inf掩码又有很大的负偏置某些高性能 kernel 会对这种组合做特殊处理不同版本的实现行为还不完全一样。所以实际项目里经常出现这种情况同一个模型用官方 FlashAttention 版本跑没问题换一个包装库或者换一个 kernel 后端长序列效果突然变差。很多人以为是随机性或者环境差异其实根因就是线上 softmax 对极端偏置的数值处理方式不同。3. 复现和检查怎么判断 ALiBi 是否数值失败3.1 先看 logits 范围和分布检测的第一步不是看 loss是看 logits 的实际数值。写一个最小检查脚本把注意力分数和 ALiBi 偏置相加之后先打印 min、max、mean 和分位数。下面是一个通用结构实际使用时按你的模型接口调整import torch def build_alibi_bias(seq_len, num_heads): # 偏置矩阵形状是 [num_heads, seq_len, seq_len] # 每个 head 有自己的斜率 m这里只是示意不是论文原值 slopes torch.tensor([1.0 / (2 ** (i 1)) for i in range(num_heads)]) rows torch.arange(seq_len, dtypetorch.float32) cols torch.arange(seq_len, dtypetorch.float32) distances (rows[:, None] - cols[None, :]).abs() bias -slopes[:, None, None] * distances[None, :, :] return bias # 假设 scores 形状是 [batch, num_heads, seq_len, seq_len] logits scores bias.to(scores.dtype) print(logits min:, logits.min().item()) print(logits max:, logits.max().item()) print(logits mean:, logits.mean().item()) # 统计每个 head 里 logits 小于 -20 的比例 low_ratio (logits -20).float().mean(dim-1) print(below -20 ratio per head:, low_ratio)如果低比例很高后续的 exp 一定会出问题。这里不用纠结阈值到底是 -20 还是 -24重要的是看分布趋势同一层不同 head 之间最低 logits 是否相差很多个数量级随着序列长度翻倍这个差距是不是也在翻倍。3.2 再用注意力熵值判断是否退化logits 范围能暴露数值风险但判断模型是否“失明”更直接的方法是看注意力熵值。熵值的极端情况很清晰熵接近均匀分布的熵ln(seq_len)说明注意力没有区分度。熵特别低接近 0说明注意力集中在极少数 token 上可能退化成局部窗口。probs torch.softmax(logits, dim-1) # 注意力熵: [batch, num_heads, seq_len] entropy -(probs * torch.log(probs 1e-12)).sum(dim-1) print(entropy min:, entropy.min().item()) print(entropy mean:, entropy.mean().item()) print(entropy max:, entropy.max().item())我一般会对比同一批输入在短序列和长序列下的熵分布。如果短序列时熵值还正常长序列时某一层或者某几个 head 的熵值突然塌到一个极小值基本可以判定是 ALiBi 偏置把 logits 推到低精度失效区了。还有一个很有用的信号不同 head 的熵值如果出现明显的“两极分化”——部分 head 熵很高部分 head 熵极低——通常不是数据语义造成的而是斜率设置和数值精度共同作用的结果。3.3 追踪梯度和输出现象训练场景里还要看梯度。定位办法很朴素记录注意力概率矩阵中精确为 0 的位置比例。如果某个 query 的注意力概率有大量精确 0这些位置在反向传播时梯度也是 0对应 key 的 token 表征长期得不到更新。我更习惯从模型输出反推。现象一般是这样的训练 loss 前期下降正常序列长度超过某个阈值后开始震荡或停滞。推理时短文本效果正常长文本的生成内容突然变得很“散”缺乏对前文的依赖。有分类或检索任务时长序列的 attention pooling 结果只由最近几个 token 决定。这几个现象如果同时出现先不要着急调学习率、换优化器。先按 3.1 和 3.2 的脚本把 logits 和注意力熵打出来确认是不是数值问题。排查顺序永远是先看最底层数值再看模型结构最后才是超参数。4. 常见触发场景长序列、低精度、大斜率4.1 长上下文外推距离段从训练分布漂移ALiBi 设计初衷就是外推但外推恰恰也是它最容易出问题的地方。训练时模型见过的最大距离是训练长度 T。推理时序列长度如果到 2T、4T那些超出训练范围的距离对应的偏置值是训练中从未出现过的极端情况。如果训练时用了 fp16模型可能已经隐式地习惯了“尾部概率全部为零”的状态。到推理时序列边长新的尾部距离产生的概率同样是零看起来好像只是继续了训练时的行为。问题在于这些新位置对应的是新的语义依赖关系模型并没有学会如何处理只是因为数值下溢“恰好”变成零概率给了一个虚假的安全感。我见过一个比较典型的案例同一个模型在 2048 长度上训练推理扩展到 8192前 2048 部分表现稳定后面部分几乎没有模型需要的信息loss 还正常。用熵值检查后才发现超过 2048 的 token 对每个 query 的注意力概率都是 0。这不是“外推成功”而是“幸运地没有报错”。4.2 低精度训练与量化推理低精度是另一个高发场景。fp16 训练时ALiBi 偏置如果被转成 fp16 再和 logits 相加距离项稍微大一点就会触发下溢。bf16 训练时虽然很少出现下溢到 0 的情况但概率分布可能被舍入误差“抹平”注意力熵偏高模型有效分辨率下降。量化推理同理。把模型权重量化到 int8、int4 之后注意力计算通常是先反量化为低精度浮点再做 softmax。如果注意力头本身就差量化之后 ALiBi 的偏置误差会被放大原本还能工作的长距离依赖会明显劣化。遇到这类问题优先排查一个点ALiBi 偏置是在哪个精度下加进去的。很多实现会先把偏置转成和 logits 相同的精度再相加。这种写法对短序列没影响长序列或者大批次时会放大数值风险。能在 fp32 下构造偏置、最后再参与注意力计算的尽量保持 fp32。4.3 头部斜率设置和实现版本差异第三个触发点来自斜率的初始化。ALiBi 的斜率通常是按几何级数递减不同 head 有不同的惩罚强度。如果手写实现时把斜率整体放大或者头数设置导致某些头的斜率偏大数值失败会提前出现。我之前排查过一个案例模型在测试集短文本上表现都正常但某些 head 的注意力熵一直特别低把 logits 打出来之后发现那几个 head 的 logits 最小值和正常 head 差了 3 到 4 个数量级明显是斜率初始化偏大。换成更合理的斜率范围之后长序列外推效果立刻改善。实现版本差异也值得记录。不同库对 ALiBi 的支持方式不一样有的是在 attention 外面用attention_bias参数传进去有的是在 kernel 内部动态生成还有的会把 ALiBi 和相对位置编码组合在一起。同一个模型在不同实现下跑出的长文本效果可能不同重点不是抱怨实现差异而是要在你的环境里固定一个版本做一个长序列回归测试。5. 缓解方案和工程建议5.1 对 ALiBi 偏置做裁剪限制最大惩罚最直接的解法是给距离项一个上限。不是让偏置无限线性变负而是当距离超过某个阈值后偏置不再继续下降。bias(i, j) -m * min(|i - j|, C)这样做的好处很明显保证任何位置的 logits 不会被压到低于某个下限至少在距离维度上不会出现无限下溢。代价是模型对超出 C 的距离不再有“区分能力”但这在很多场景下是可接受的。因为长距离依赖本身应该是“选择性关注”而不是“距离越远越不关注”。阈值 C 怎么选可以观察训练集中任务依赖的最大距离也可以直接看 logits 分布把 logits 的下限控制在精度安全区内。这里给一个通用经验fp16 训练时把偏置绝对值限制在 20 到 30 之间是比较常见的选择具体要结合头部斜率调整。不要一个值套所有模型。5.2 计算精度分开管理bias 在 fp32 算softmax 保持稳定第二个建议是精度分层。ALiBi 偏置本身可以很便宜地在 fp32 下构造并缓存不需要参与梯度计算纯粹是一个与位置相关的常量矩阵。所以在实现里最好做到偏置用 fp32 存logits 计算时再把偏置临时转成计算精度或者干脆在 fp32 下完成 logits 相加再向下转。对于梯度训练更稳妥的做法是让 softmax 的输入保持 fp32。很多框架允许你在注意力内部做一个 fp32 的 softmax 回退。代价是速度和显存会增加但可靠性会提升。下面是一段示意逻辑展示如何把偏置单独保持在 fp32def attention_with_alibi(q, k, v, bias, scale1.0 / 8.0): # q, k, v 可能是低精度张量 logits torch.matmul(q, k.transpose(-2, -1)) * scale # bias 保持 fp32这里显式转换 logits_fp32 logits.float() bias # 在 fp32 下做 softmax再转回原精度 probs torch.softmax(logits_fp32, dim-1).to(q.dtype) out torch.matmul(probs, v) return out, probs如果你的训练框架支持算子融合尽量选择在融合内核里对 softmax 保持高精度的实现。不要为了省一点显存而全程 fp16。5.3 实现选型和编码替换如果目标不是研究 ALiBi 本身而是要用一个稳定的长文本模型可以考虑替换实现而不是硬修数值问题。常见的选择包括换成旋转位置编码 RoPE它没有像 ALiBi 那样的负偏置衰减数值上相对温和。使用支持可变偏置的 attention 封装在接口层面传入可控 bias而不是写死在 kernel 里。采用 FlashAttention 兼容 ALiBi 的版本时先确认它的数值处理方式最好在目标序列长度上做一次单测。如果使用相对位置编码、coordinate attention、deformable attention 这类带偏移或坐标加权的机制也要检查它们的偏移量是否会出现超大负值原理和 ALiBi 类似。这个话题多说一句coordinate attention 的坐标权重、deformable attention 的可学习偏移本质都是在给注意力位置信息“画地形”。任何“地形”一旦出现极端的负向区域就可能在低精度或长序列下遇到类似的数值塌缩。ALiBi 只是其中结构最简单的一个所以拿来当分析样例正合适。5.4 加一个“注意力健康度”检查落地到工程里不应等到模型效果变差才人工排查。建议把注意力健康度检查做成一个可定时触发的脚本在训练和推理的关键节点输出几个指标指标说明正常信号异常信号logits min每层 logits 最小值不会随序列长度急剧下降线性下降且逼近精度下限零概率比例注意力概率精确为 0 的位置占比比例较低或分布均匀尾部大量出现 0注意力熵均值每层、每头熵的平均值随难度合理变化某一头熵长期过小或接近 ln(seq_len)熵的 head 方差同层不同 head 的熵差异保持一定层次两极分化明显我一般会在训练 config 里加一个 hook每隔固定步数跑一次短序列和长序列的对比检查。长序列可以只取一小批数据重点看 logits 的分布偏移。这个检查在故障发生前就能提前预警比训练到一半发现 loss 异常再回头查日志要划算得多。6. 排查顺序和最后一点经验6.1 从现象到根因的排查顺序如果已经出现了长文本效果变差、注意力分布异常、训练不平滑这类问题推荐的排查顺序是先确认最外层现象是 loss 震荡、输出重复、还是注意力熵异常。记下第一个出现问题的序列长度和层数。再查输入数据序列长度分布、截断策略、是否混入了超长样本。接着查数值把 logits min、零概率比例、注意力熵按层输出确认是否和 ALiBi 偏置的极端负值相关。然后查精度偏置在哪个精度下计算、softmax 在哪个精度下执行、kernel 是否做了高精度 fallback。最后查配置每层的 head 数量、斜率初始值、是否用了混合精度、是否同时叠加因果 mask。这几个层级从现象逐步往底层走大多数情况下能在第 3 步就定位到问题。6.2 落地建议ALiBi 这类线性偏置方案不是不可用而是要用在合适的条件和精度管理下。我的建议是小规模验证时用默认配置没问题但一旦要上线长文本、大批次、低精度训练必须把数值检查和精度策略前置。如果只是学习实验最短的路径是先用 fp32 跑一条短序列再用 fp32 跑一条长序列对比注意力熵和 logits 分布然后换 fp16 重跑长序列观察数值差异。这一个对比基本就能让你看出 ALiBi 在这个模型里的失效边界在哪里。踩过几次之后我发现很多注意力“失明”问题不是模型能力不够而是数值实现把模型本来可以表达的长距离依赖在底层抹掉了。先看清楚这一点再决定是改配置、改精度还是换编码方案才不会浪费大量调参时间。