38 · ALiBi — Attention with Linear Biases

位置编码 Press et al. · 2022
arXiv:2108.12409

上一代的问题:所有位置编码都会拖慢推理速度或限制外推能力 — Sinusoidal PE 不能外推到训练长度之外;RoPE 有一定外推能力但有 30% 速度损失;T5 Bias 外推好但有 49% 速度损失。所有方法都试图将位置信息注入 Q/K/V 表征之中,增加了计算开销。

这代改了什么:彻底不碰 Q/K/V,直接在 attention score 上加一个静态的、非学习的线性距离惩罚:softmax(QKT/√d + B),其中 Bij = -|i-j|·mh。每个 head 有自己独立的几何斜率 mh = 2-8h/H

效果:零速度代价 (0-0.7%),完美外推 (train 512 测 8K+),零额外参数。被 MPT、BLOOM、GPT-NeoX 采用。

一、核心洞察

位置编码是长度外推能力的唯一决定因素(控制所有其他变量不变时)。外推能力不来自于模型"理解"位置,而来自于一种自然的追新偏置 (recency bias)——邻近的 token 天然应该获得更高的注意力权重。

This insight is both elegant and counterintuitive. 传统观点认为位置编码要为模型提供精确的位置坐标(Sinusoidal 通过三角函数周期映射,RoPE 通过复数旋转编码绝对位置)。但 ALiBi 的实验证明这些复杂的机制并不是外推的必要条件 —— 简单的距离惩罚就足够了。模型不需要知道 token 在哪个位置,只需要知道两个 token 隔了多远。

更深刻地说,ALiBi 揭示了 Transformer 的 attention 机制本质上是局部优先的 (locality-prior)。在 long-range 场景下,模型更依赖附近 token 的上下文,而非远处的精确位置信息。ALiBi 通过对远距离 tokens 施加惩罚来强化这种自然偏置,使得模型在面对更长的序列时,注意力分布能够自动适应。

关键洞见:ALiBi 是所有主流长度外推方案中唯一一个不去尝试"编码位置"的 —— 它认为位置信息不需要被编码到表征里,只需要在注意力计算时对距离做出响应。这是一种从"编码主义"到"行为主义"的范式转变。

二、公式

2.1 核心公式

Attention(Q,K,V) = softmax( QKT / √d + B ) · V
Formula 1: ALiBi Attention

其中 B 是偏置矩阵,每个元素定义为 query token i 和 key token j 之间的线性距离惩罚:

Bij = -|i - j| · mh
Formula 2: Linear Bias

2.2 Head-specific Slopes

每个注意力头使用不同的斜率 mh,构成一个从大到小的几何序列:

mh = 2-8h/H   for h = 1, 2, ..., H
Formula 3: 斜率定义

对于 H=8 头的典型设置,各头斜率分别为:

Headmh解释
h=11/2 = 0.5000最激进局部 — 每远一个 token,分数衰减 0.5
h=21/4 = 0.2500中等局部
h=31/8 = 0.1250中等
h=41/16 = 0.0625中等
h=51/32 = 0.0313中等远程
h=61/64 = 0.0156远程
h=71/128 = 0.0078很远程
h=81/256 = 0.0039最温和远程 — 几乎无惩罚

2.3 为什么选择几何序列

几何序列意味着高编号的 head 斜率呈指数下降。这种设计使得 heads 自动覆盖不同的距离范围:低号 heads 聚焦局部上下文(高距离惩罚),高号 heads 可以关注远距离 token(低距离惩罚)。这是 ALiBi 能够同时兼顾局部精度和远程能力的关键。

公式中的 8 是一个缩放因子 —— 它确保第一个 head 的斜率约为 1/2,最后一个 head 的斜率足够小(接近 0),从而覆盖有意义的距离范围。对于不同数量的 heads,这个公式会自动适应并维持相同的几何分布。

三、工作原理

ALiBi 的工作流程极其简单,不需要修改 Transformer 的任何内部结构:

  1. 标准 Q/K/V:像往常一样计算 query、key、value 矩阵 —— 不加任何位置编码。
  2. 计算裸分数:计算 QKT/√d,得到原始的注意力分数矩阵。
  3. 加偏置矩阵:为每个 head 生成一个固定的线性偏置矩阵 B,其中 Bij = -|i-j| · mh。这个矩阵在整个推理过程中保持不变,一次生成可复用。
  4. Softmax + 加权求和:在 softmax 之前将偏置矩阵加到注意力分数上,然后正常进行加权求和。

Head 1 对距离惩罚最激进 (m=0.5)—— 接近严格局部注意力。Head 8 惩罚最温和 (m=1/256)—— 可以关注到远处 token。每个 head 相当于一个不同"测距传感器",共同覆盖从近到远的全部范围。

实现细节:偏置矩阵 B 在推理时不需要每次重新计算。由于 B 只依赖于序列中 token 的相对距离 i-j,而非具体数值,可以在构建计算图时预先分配并缓存。这进一步降低了 ALiBi 的运算开销,使其几乎为零成本。

四、核心实验结果

4.1 外推能力

ALiBi 最令人印象深刻的结果是其外推能力:

训练长度测试长度SinusoidalALiBi优势
51251217.9617.99持平
512102418.4518.07ALiBi 优
512204819.6318.26ALiBi 优
512307223.9218.40ALiBi 大幅优

注意最后一行的对比:ALiBi 在 512 token 上训练,在 3072 token 上测试,perplexity 为 18.40,甚至优于 Sinusoidal PE 在 3072 token 上从头训练的 18.67。用 6 倍短的训练序列,外推后性能超过训练过的 Sinusoidal!

测试长度Sinusoidal (训 3072)ALiBi (训 512)ALiBi (训 3072)
307218.6718.4016.12

4.2 训练速度对比

方法Words/sec相对 Sinusoidal原因
Sinusoidal28,500基准标准 Transformer
ALiBi28,300-0.7%仅加一次偏置矩阵,几乎零开销
RoPE20,000-30%每个 token 都需要做复数旋转操作
T5 Bias14,400-49%需要查询-计算共享偏置参数矩阵

ALiBi 的速度优势非常明显:几乎没有计算开销,而 RoPE 有 30% 的速度损失,T5 Bias 高达 49%。这意味着如果你把模型中的 RoPE 替换为 ALiBi,可以用同样的算力训练更深或更宽的模型,或者在同样的时间预算内处理更多数据。

五、Early Token Curse 现象

一个重要发现:ALiBi 的外推能力并不来自于更好的长距离建模,而是来自缓解了"早期 token 诅咒" (early token curse)。

在标准的因果语言模型中,序列开头的 token 看到的上下文很少,其表征质量很差。验证序列越长,这些开头 token 占的比例就越小,整体 perplexity 自然下降。ALiBi 的 sliding window 评估实验发现,perplexity 随着序列变长保持平稳而非持续下降,说明其"外推"不是神奇的远程建模能力,而主要是统计分布偏移的结果。

sliding window 实验:当用固定窗口(如 512 token)滑动评估长序列时,ALiBi 没有表现出对 Sinusoidal 的优势。这说明 ALiBi 的优势主要存在于 causal 模式的完整序列评估中 —— 其本质是减轻了前缀 token 因缺乏左侧上下文而产生的表征退化。

这个发现对于理解"外推能力"具有深远意义:

六、ALiBi vs RoPE 哲学对比

维度ALiBiRoPE
机制对注意力分数加线性偏置对 Q/K 向量做复数旋转
参数数零参数零参数
速度代价0-0.7%~30%
外推能力强 (train 512 测 8K+)中等 (适度外推)
数学哲学"惩罚距离" — 让模型自动适应远近"编码位置" — 让表征携带位置信息
对模型影响完全不影响 Q/K/V 表征改变了 Q/K 的内积空间结构
旋转不变性有 (RoPE 保持相对位置的旋转对称性)
工程复杂度极低 (写几行代码)中等 (需要复数乘法)

ALiBi 和 RoPE 代表了两种完全不同的位置编码哲学:

有趣的是,尽管 RoPE 在理论上更优雅(保持了内积的旋转对称性,且能自然处理相对位置),但 ALiBi 在实际外推任务上表现更好。这说明在 Transformer 的上下文学习中,"能关注到多远"比"如何优雅地编码位置"更重要。

七、工程影响

7.1 被哪些模型采用

7.2 工程上的优势

7.3 局限性

八、上下游关联

上一篇:Sinusoidal Position Encoding (Vaswani et al., 2017) — 通过固定频率三角函数编码绝对位置,但无法外推。Learned Position Embedding (Devlin et al., 2019) — 可学习的位置嵌入,更灵活但仍无法外推。T5 Bias (Raffel et al., 2020) — 在注意力分数上加可学习的偏置,但 49% 速度代价。

下一篇:YaRN (Peng et al., 2023) — 基于 RoPE 的扩展方案,通过缩小旋转频率来扩展上下文。NTK-aware scaling (Bloc97 et al., 2023) — 通过插值和温度缩放扩展 RoPE 上下文。LongRoPE (Chen et al., 2024) — 通过非均匀频率插值扩展上下文 32x+。

并行工作:Kerple (Chi et al., 2022) — 将 ALiBi 的线性偏置推广为可学习的核函数偏置,提供更多灵活性。Sandwich (Chi et al., 2022) — 将 ALiBi 和可学习偏置组合使用。

后续改进:ALiBi 的线性偏置后来被证明在非常大的上下文窗口中(如 128K+)性能会衰减,因此后来的 XPos (Sun et al., 2023) 和 ReRoPE (Su, 2023) 等方案尝试结合 ALiBi 的速度优势和 RoPE 的旋转对称性。

九、算法与实现

9.1 PyTorch 伪代码

import math
import torch

def get_alibi_slopes(n_heads: int) -> torch.Tensor:
    """生成 head-specific 斜率"""
    def get_slope(h):
        # m_h = 2^{-8h/H}
        return 2 ** (-8 * h / n_heads)
    return torch.tensor([get_slope(h) for h in range(1, n_heads + 1)])

def build_alibi_bias(seq_len: int, n_heads: int, device: str) -> torch.Tensor:
    """构建 ALiBi 偏置矩阵 [n_heads, seq_len, seq_len]"""
    # 相对距离矩阵 [seq_len, seq_len]
    pos = torch.arange(seq_len, device=device)
    distance = pos[:, None] - pos[None, :]  # i - j for all i, j
    abs_distance = distance.abs()            # |i - j|

    # 斜率 [n_heads, 1, 1]
    slopes = get_alibi_slopes(n_heads)
    slopes = slopes.view(-1, 1, 1).to(device)

    # 偏置 [n_heads, seq_len, seq_len]
    bias = -abs_distance[None, :, :] * slopes
    return bias

class ALiBiAttention(nn.Module):
    def forward(self, Q, K, V, alibi_bias):
        # Q, K, V: [batch, heads, seq, dim]
        scores = torch.matmul(Q, K.transpose(-2, -1))
        scores = scores / math.sqrt(Q.size(-1))

        # 加 ALiBi 偏置 (广播到 batch 维度)
        scores = scores + alibi_bias

        attn = torch.softmax(scores, dim=-1)
        return torch.matmul(attn, V)

# 使用示例
alibi_bias = build_alibi_bias(seq_len=512, n_heads=8, device="cuda")
# 整个推理过程只需构建一次,后面复用

9.2 关键实现要点

十、个人思考

最让人惊叹的是 ALiBi 的简洁性。在位置编码变得越来越复杂的趋势下(从 Sinusoidal 到相对位置编码到 RoPE 到 T5 偏置),ALiBi 用最少的公式和零额外参数达到了最好的外推效果。这提醒我们在 AI 研究中:复杂性并不等同于有效性。

Speed matters。30% 的速度差异在数据中心级别的部署中意味着数百万美元的电费和硬件成本。ALiBi 的 0.7% 开销使其成为工业部署的首选方案。这反映了一个趋势:实用主义的创新往往能在工程上取得更广泛的 adoption。

"Early token curse"的发现是这篇论文中最被低估的洞察。它提醒我们评估外推能力时要格外谨慎 —— 很多看似神奇的外推效果可能只是统计伪影。这也指出了未来工作的方向:真正需要的是提升长距离的语义建模能力,而不仅仅是靠缓解前缀诅咒来降低 perplexity。

哲学之争:ALiBi vs RoPE 的对比本质上是两种工程哲学的对决。RoPE 追求数学优雅和对称性保持,ALiBi 追求极简高效和实用外推。从 adoption 来看,两者都获得了广泛使用 —— RoPE 在 LLaMA 系列中成为标准,ALiBi 在 MPT、BLOOM 中占据主导。这说明问题不能一概而论,不同的架构和场景可能适合不同的方案。

未来的方向可能不是找到一个"万能"的位置编码,而是设计出能够根据任务自适应选择偏置机制的方案。近年来出现的动态位置编码 (如 XL, 2024) 和混合位置编码 (如将 ALiBi 和 RoPE 组合) 正是在这个方向上的探索。