上一代的问题:所有位置编码都会拖慢推理速度或限制外推能力 — 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 施加惩罚来强化这种自然偏置,使得模型在面对更长的序列时,注意力分布能够自动适应。
其中 B 是偏置矩阵,每个元素定义为 query token i 和 key token j 之间的线性距离惩罚:
每个注意力头使用不同的斜率 mh,构成一个从大到小的几何序列:
对于 H=8 头的典型设置,各头斜率分别为:
| Head | mh | 解释 |
|---|---|---|
| h=1 | 1/2 = 0.5000 | 最激进局部 — 每远一个 token,分数衰减 0.5 |
| h=2 | 1/4 = 0.2500 | 中等局部 |
| h=3 | 1/8 = 0.1250 | 中等 |
| h=4 | 1/16 = 0.0625 | 中等 |
| h=5 | 1/32 = 0.0313 | 中等远程 |
| h=6 | 1/64 = 0.0156 | 远程 |
| h=7 | 1/128 = 0.0078 | 很远程 |
| h=8 | 1/256 = 0.0039 | 最温和远程 — 几乎无惩罚 |
几何序列意味着高编号的 head 斜率呈指数下降。这种设计使得 heads 自动覆盖不同的距离范围:低号 heads 聚焦局部上下文(高距离惩罚),高号 heads 可以关注远距离 token(低距离惩罚)。这是 ALiBi 能够同时兼顾局部精度和远程能力的关键。
公式中的 8 是一个缩放因子 —— 它确保第一个 head 的斜率约为 1/2,最后一个 head 的斜率足够小(接近 0),从而覆盖有意义的距离范围。对于不同数量的 heads,这个公式会自动适应并维持相同的几何分布。
ALiBi 的工作流程极其简单,不需要修改 Transformer 的任何内部结构:
Head 1 对距离惩罚最激进 (m=0.5)—— 接近严格局部注意力。Head 8 惩罚最温和 (m=1/256)—— 可以关注到远处 token。每个 head 相当于一个不同"测距传感器",共同覆盖从近到远的全部范围。
ALiBi 最令人印象深刻的结果是其外推能力:
| 训练长度 | 测试长度 | Sinusoidal | ALiBi | 优势 |
|---|---|---|---|---|
| 512 | 512 | 17.96 | 17.99 | 持平 |
| 512 | 1024 | 18.45 | 18.07 | ALiBi 优 |
| 512 | 2048 | 19.63 | 18.26 | ALiBi 优 |
| 512 | 3072 | 23.92 | 18.40 | ALiBi 大幅优 |
注意最后一行的对比:ALiBi 在 512 token 上训练,在 3072 token 上测试,perplexity 为 18.40,甚至优于 Sinusoidal PE 在 3072 token 上从头训练的 18.67。用 6 倍短的训练序列,外推后性能超过训练过的 Sinusoidal!
| 测试长度 | Sinusoidal (训 3072) | ALiBi (训 512) | ALiBi (训 3072) |
|---|---|---|---|
| 3072 | 18.67 | 18.40 | 16.12 |
| 方法 | Words/sec | 相对 Sinusoidal | 原因 |
|---|---|---|---|
| Sinusoidal | 28,500 | 基准 | 标准 Transformer |
| ALiBi | 28,300 | -0.7% | 仅加一次偏置矩阵,几乎零开销 |
| RoPE | 20,000 | -30% | 每个 token 都需要做复数旋转操作 |
| T5 Bias | 14,400 | -49% | 需要查询-计算共享偏置参数矩阵 |
ALiBi 的速度优势非常明显:几乎没有计算开销,而 RoPE 有 30% 的速度损失,T5 Bias 高达 49%。这意味着如果你把模型中的 RoPE 替换为 ALiBi,可以用同样的算力训练更深或更宽的模型,或者在同样的时间预算内处理更多数据。
一个重要发现:ALiBi 的外推能力并不来自于更好的长距离建模,而是来自缓解了"早期 token 诅咒" (early token curse)。
在标准的因果语言模型中,序列开头的 token 看到的上下文很少,其表征质量很差。验证序列越长,这些开头 token 占的比例就越小,整体 perplexity 自然下降。ALiBi 的 sliding window 评估实验发现,perplexity 随着序列变长保持平稳而非持续下降,说明其"外推"不是神奇的远程建模能力,而主要是统计分布偏移的结果。
这个发现对于理解"外推能力"具有深远意义:
| 维度 | ALiBi | RoPE |
|---|---|---|
| 机制 | 对注意力分数加线性偏置 | 对 Q/K 向量做复数旋转 |
| 参数数 | 零参数 | 零参数 |
| 速度代价 | 0-0.7% | ~30% |
| 外推能力 | 强 (train 512 测 8K+) | 中等 (适度外推) |
| 数学哲学 | "惩罚距离" — 让模型自动适应远近 | "编码位置" — 让表征携带位置信息 |
| 对模型影响 | 完全不影响 Q/K/V 表征 | 改变了 Q/K 的内积空间结构 |
| 旋转不变性 | 无 | 有 (RoPE 保持相对位置的旋转对称性) |
| 工程复杂度 | 极低 (写几行代码) | 中等 (需要复数乘法) |
ALiBi 和 RoPE 代表了两种完全不同的位置编码哲学:
有趣的是,尽管 RoPE 在理论上更优雅(保持了内积的旋转对称性,且能自然处理相对位置),但 ALiBi 在实际外推任务上表现更好。这说明在 Transformer 的上下文学习中,"能关注到多远"比"如何优雅地编码位置"更重要。
上一篇: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 的旋转对称性。
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")
# 整个推理过程只需构建一次,后面复用
seq_len 和 n_heads,与 batch 无关,可以预计算并缓存最让人惊叹的是 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 组合) 正是在这个方向上的探索。