ALiBi Positional Bias/zh

Article
Topic area	transformers
Prerequisites	Transformer, Self-Attention, Positional Encoding

概述

线性偏置注意力（ALiBi）是一种向 Transformer 注意力中注入位置信息的方法，它通过在注意力分数上添加一个固定的、与距离相关的偏置，而不是将位置嵌入到 token 表示中。该方法由 Press、Smith 和 Lewis 于 2021 年提出。ALiBi 用一个静态的、按头独立的偏置矩阵替代了正弦位置编码和可学习位置嵌入，其元素随查询位置与键位置之间的距离线性增长。每个注意力头对远距离 token 的关注会受到不同程度的惩罚，惩罚速率由每个头固定的斜率控制。

ALiBi 的核心特性是能够外推到远超训练时所见的序列长度。标准的位置编码在评估长度超过训练上下文的输入时性能会急剧下降，而采用 ALiBi 在长度 1024 上训练的模型，可以在长度 2048、4096 甚至更长的序列上进行评估，仅出现轻微的困惑度下降，且无需任何重新训练或插值。这一特性使 ALiBi 成为早期长上下文语言建模的主力，并成为 BLOOM、MPT 和 Replit Code 中首选的位置方法。

动机

原始的 Transformer 使用正弦位置编码，将其与输入 token 的嵌入相加。可学习的绝对位置嵌入（如 BERT 和 GPT-2 中所用）将固定的正弦曲线替换为可训练的向量。这两种方法都存在一个根本性的局限：超出训练范围的位置要么不在嵌入表中（可学习方法），要么落在模型从未需要解释的正弦曲线区域（正弦方法）。经验上，超出范围位置的困惑度会急剧上升。

Press、Smith 和 Lewis 观察到，Transformer 的位置信息最终是通过注意力分数进入网络的。与其让位置信息通过输入嵌入传递，为什么不直接在它被使用的地方进行编码呢？ALiBi 完全移除了位置嵌入，转而通过查询与键之间相对距离的函数来对注意力 logits 施加偏置。由于偏置是关于距离的简单线性函数，没有任何可学习参数，因此在前所未见的距离上对其求值是良定义的，并且表现平滑。

公式

标准的缩放点积注意力对于查询 $ Q $、键 $ K $ 和值 $ V $ 计算如下：

$ {\displaystyle \operatorname{Attention}(Q, K, V) = \operatorname{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V.} $

ALiBi 在softmax内部、按行归一化之前为每个头添加一个偏置矩阵 $ B^{(h)} $，从而修改上述公式为：

$ {\displaystyle \operatorname{Attention}^{(h)}(Q, K, V) = \operatorname{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}} + B^{(h)}\right) V.} $

在因果语言建模设置中，对于头 $ h $，查询位置 $ i $ 与键位置 $ j $（其中 $ j \le i $）之间的偏置为：

$ {\displaystyle B^{(h)}_{ij} = -m_h \cdot (i - j),} $

其中 $ m_h > 0 $ 是每个头固定的（非可学习的）斜率，$ (i - j) \ge 0 $ 是因果距离。位置 $ j > i $ 按惯例被屏蔽为 $ -\infty $。对于双向模型，偏置是对称的：$ B^{(h)}_{ij} = -m_h \cdot |i - j| $。

其关键特性是：不向输入添加任何位置嵌入；偏置仅依赖于距离 $ i - j $，而非绝对位置；且斜率 $ m_h $ 不可学习。因此，每个头都实现了对邻近 token 呈指数衰减的偏好，衰减速率由其斜率控制。

斜率方案

ALiBi 按几何方式在各注意力头之间分配斜率。对于一个具有 $ n $ 个头的注意力层，斜率为：

$ {\displaystyle m_h = 2^{-8 h / n}, \quad h = 1, 2, \ldots, n.} $

对于 $ n = 8 $ 个头，该方案给出的斜率为 $ 1/2, 1/4, 1/8, \ldots, 1/256 $。对于 $ n = 16 $，斜率以几何方式在相同范围内插值，每个斜率都是其在 8 头方案中相邻斜率的几何平均值。斜率较大的头几乎只关注近邻 token；斜率较小的头则偏置近乎平坦，可以关注遥远的上下文。该方案由作者经验性地选定，并在不同模型规模上保持不变；尽管简单，至今仍未被显著改进。

关键的是，这些斜率是常数，而不是参数。它们不通过 Backpropagation 更新，不出现在优化器状态中，也不为模型增加任何参数。因此，就参数量而言，ALiBi 比可学习位置嵌入或 Rotary Position Embedding 都更省。

长度外推

原始 ALiBi 论文最显著的成果是长度外推。一个具有 13 亿参数的 Transformer，在 1024 token 上下文上训练，在 3072 token 上下文上评估时，保持的困惑度与直接在 3072 token 上训练时所能达到的相比仅有微小差距。而正弦编码和可学习位置嵌入在相同的分布外长度下评估时，困惑度会上升一个数量级或更多。

其机制非常直观。偏置 $ -m_h (i - j) $ 是关于距离的线性函数，对于任意非负的 $ i - j $ 都有良定义，包括大于训练中任何所见值的距离。模型从不需要解释它从未遇到过的位置输入；它只需对一个其形状与始终见过的完全相同的线性偏置应用softmax。在远距离上，注意力分布会自然地集中到邻近的 token 上，比例由每个头的斜率决定，而模型在训练过程中学会利用这一点。

在实际中，外推并非免费。质量会随长度逐渐下降，而非常长的评估（超过训练长度数倍时）会丢失远距离 token 的信息，因为所有头的偏置都会变得很大且为负。尽管如此，这种平缓的退化使 ALiBi 成为训练时上下文长度受限但推理时上下文可能增长的场景下的一个可靠选择。

与其他位置方法的比较

正弦位置编码 将位置的固定函数加到输入嵌入上。它们是绝对的（每个位置都有一个独立的编码），无参数，原则上对任何位置都有定义；但实际上无法外推，因为训练好的网络没有学会解释超出训练范围的高频分量。

可学习的绝对位置嵌入为每个位置分配一个可训练的向量。它们完全无法外推；超出表的位置根本不存在。

相对位置编码（T5 偏置，Shaw 等人）根据查询与键之间分桶的相对距离，向注意力分数添加一个可学习的标量偏置。ALiBi 与之密切相关，但使用固定的线性偏置而非可学习的分桶偏置，这既消除了参数，又使外推到未见过的距离成为可能。

旋转位置嵌入（RoPE）按与位置成正比的角度对 $ Q $ 和 $ K $ 中的维度对进行旋转。在现代大型语言模型（LLaMA、Mistral、Qwen、Gemma）中，RoPE 现在比 ALiBi 更流行，因为它在分布内通常能给出略好的质量，并通过 NTK 感知或 YaRN 缩放支持基于插值的上下文扩展。开箱即用时 RoPE 的外推不如 ALiBi 干净，但结合位置插值技巧后，它在长上下文微调方面已超过 ALiBi。

变体与扩展

已经提出了若干改进。

对称 ALiBi 对双向编码器（如 BERT 类模型）使用 $ B_{ij} = -m_h |i - j| $。原始论文聚焦于因果模型，但对称变体已被用于仅编码器和编码器-解码器场景。

可学习斜率让 $ m_h $ 可训练，而非固定不变。这通常不会改善质量，并部分破坏了外推保证，因为学到的斜率可能朝着依赖训练分布的方向漂移。

Sandwich 和 KERPLE 用其他关于距离的单调函数（对数函数、核化函数）替代线性偏置，在保留外推性质的同时，有时还能给出分布内更好的困惑度。这些方法在概念上是 ALiBi 的衍生物，并占据相同的架构位置。

动态 NTK ALiBi 在推理时调整斜率方案，以进一步扩展有效上下文，其思路类似于针对 RoPE 的 NTK 感知缩放。

实现细节

ALiBi 不需要修改输入流水线，不需要位置嵌入层，也不需要修改 $ Q $ 或 $ K $。实现通常会按每个序列长度预先计算一次偏置矩阵 $ B^{(h)} $，并将其加到注意力 logits 上。在 FlashAttention 及类似的融合内核中，偏置是根据行列索引在线计算的，只需要一个大小为 $ n $ 的小型按头斜率缓冲区。

对于使用键值缓存的推理，偏置仅依赖于已缓存键相对于当前查询的绝对位置，而这些位置是已知的。因此，上下文长度不断增长的流式推理无需任何特殊处理，只需将偏置扩展到新的位置即可。

ALiBi 与 Layer Normalization、解码器与编码器架构以及标准的注意力掩码都兼容。由于没有位置信息混入残差流，因此输入嵌入表也比使用可学习位置嵌入的模型略小。

局限性

ALiBi 存在若干已知的弱点。首先，尽管外推过程平缓，但在位置精度至关重要的基准上（例如精确匹配的片段检索、字符级任务、某些代码补全场景），其分布内质量通常略低于经过良好调优的 RoPE。其次，斜率方案是固定的，不针对任务或层级进行自适应调整；部分注意力头实际上几乎未被使用，因为其斜率过于陡峭，无法关注紧邻环境之外的内容。第三，线性距离偏置是一种强烈的偏向局部性的归纳偏置；需要关注单个远距离 token 的任务（例如从长上下文开头检索某个键）会随着上下文增长而变得更困难，因为偏置会主动抑制远处的键。

这些权衡解释了为何在近期的大型语言模型中，ALiBi 在一定程度上被 RoPE 加位置插值方法取代，尽管它仍是长度外推中一种干净、无参数的基线方法，并且仍在 2023 年之前训练的生产系统中被广泛使用。