精读笔记:ALiBi — 用线性偏置让 Transformer 突破训练长度限制

论文全名:Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation
arXiv 编号:2108.12409(2021 年 8 月提交)
发表会议:ICLR 2022
作者团队:Ofir Press(华盛顿大学 + Facebook AI Research)、Noah A. Smith(华盛顿大学 + Allen Institute for AI)、Mike Lewis(Facebook AI Research)
原文链接:https://arxiv.org/abs/2108.12409


阅读地图

本文篇幅不长,结构清晰,适合新手完整阅读。精读顺序建议:

  1. 先看本文第一节("背景:什么是长度外推"),建立直觉
  2. Abstract — 摘要,2 分钟了解论文贡献
  3. Introduction — 引言,搞清问题背景和动机
  4. Section 2(原文)— 为什么旧方法不行
  5. Section 3(原文)— ALiBi 的具体机制(精读重点)
  6. Section 4(原文)— 实验验证(看关键数字即可)

背景知识:什么是"长度外推"问题?(新手必读)

通俗解释

想象你在准备一场考试。你的复习资料只有 10 页,所有练习题也只涉及 10 页以内的知识。考试那天,老师突然给了你一份 20 页的考卷。

对于传统 Transformer 来说,这就是"长度外推"问题:模型在短文本上训练,考试时却被要求处理更长的文本

"长度外推"(Length Extrapolation) 的正式定义(来自论文):

模型在验证时,随着输入 token 数量超过训练时所见的 token 数量,仍然能够保持良好性能的能力。


为什么传统位置编码做不到?

Transformer 不像 RNN 那样按顺序处理文字——它同时"看到"所有词。为了告诉模型哪个词在前哪个词在后,研究者发明了位置编码(Positional Encoding):给每个词的向量里加一个"位置标签"。

最常见的是正弦位置编码(Sinusoidal PE),出自 Transformer 原论文。它的问题在于:它只见过训练时的那些位置。模型在训练时如果只见过 1~512 位置的词,一旦推理时出现第 600、第 1000 位置的词,模型就"没见过"这个位置的编码,不知道该怎么处理,性能急剧下滑。

这就像你只学过 1~10 的乘法表,突然让你做 15×7,完全不知所措。


问题的代价

训练 Transformer 处理很长的文本(比如 4096 个 token)非常昂贵,因为注意力机制的计算量是 O(n²),序列长度翻倍,计算量翻四倍。如果能在短序列上训练、在长序列上推理,就能大幅节省成本。ALiBi 就是解决这个问题的方案。


一、Abstract(摘要)

原文核心句

"ALiBi does not add positional embeddings to word embeddings; instead, it biases query-key attention scores with a penalty that is proportional to their distance."

"A 1.3 billion parameter LM trained on 1024 tokens with ALiBi achieves the same perplexity as a sinusoidal model trained on 2048 tokens when both are tested on sequences of 2048 tokens, even though our model trains 11% faster and uses 11% less memory."

翻译

ALiBi 不把位置编码加到词向量里,而是直接在 query-key 注意力分数上施加一个"与距离成正比的惩罚"。

用 ALiBi 训练的 13 亿参数语言模型,在 1024 token 长度上训练,测试时(在 2048 token 序列上)能达到与正弦编码模型(训练长度 2048 token)相同的困惑度——训练速度还快 11%,内存少用 11%。

新手讲解

困惑度(Perplexity,PPL) 是衡量语言模型好坏的指标。困惑度越低,说明模型对文本的"预测"越准确,模型越聪明。你可以把它理解为:模型每次预测下一个词时,平均有多"困惑"——数字越小越好。

摘要的核心信息是三点:
1. ALiBi 不用位置编码,改用"对远处词扣分"的方式
2. 短序列训练 → 长序列推理,效果和直接长序列训练一样好
3. 还省内存、省时间


二、Introduction(引言)核心段落

段落一:问题动机

原文核心句(意译)

自 Vaswani 等人(2017)提出 Transformer 以来,有一个根本问题始终没有得到解答:模型能否在推理时处理比训练时更长的序列?

训练语言模型处理长序列的成本很高。对于 RNN,研究者曾认为短序列上训练的模型能自然泛化到更长序列——Vaswani 等人也猜测 Transformer 可能也有此能力。但实验证明,情况并非如此。

翻译与讲解

"外推"(extrapolation)来自数学:已知曲线在一段范围内的形状,能不能预测范围之外的走势?语言模型的"长度外推"就是类似的问题——训练时见过 1~1024 位置,推理时能不能正确处理 1025~2048 位置?

传统正弦位置编码的答案是:基本不行


段落二:现有方法对比

原文核心句(意译)

作者测试了三类位置编码方法的外推能力:

翻译与讲解

旋转位置编码(Rotary Position Embedding,RoPE) 是 Su 等人提出的改进方法,通过旋转向量来编码相对位置,被 LLaMA、GPT-NeoX 等主流模型采用。

T5 相对位置偏置 是 Google T5 模型使用的方案,通过在注意力矩阵中加"桶化"的相对距离偏置实现位置感知。

作者发现:现有方案要么外推能力差,要么计算慢。能不能既外推得好,又不慢?这就是 ALiBi 的出发点。


段落三:ALiBi 概览

原文核心句(意译)

作者提出 ALiBi(Attention with Linear Biases)——不在词向量里添加位置编码,而是在 query-key 点积分数上直接加一个线性递减的偏置值,以此隐式地传递位置信息。

新手讲解

传统做法:位置信息放在输入里(词向量 + 位置向量 → 送进注意力层)
ALiBi:位置信息放在注意力计算里(词向量直接送进去,但注意力分数里扣掉距离惩罚)

类比:就像课堂点名,传统方法是给每位同学发一个"座位号牌"(位置编码),而 ALiBi 是老师心里有个规则:"越是坐在远处的同学,我叫他回答问题的欲望越低"(距离惩罚),不需要额外的号牌。


三、方法章节:ALiBi 机制(精读核心)

本节对应原文 Section 3:"Attention with Linear Biases (ALiBi)"

3.1 标准注意力 vs ALiBi 注意力

标准注意力计算回顾

原文核心句(意译)

在标准 Transformer 注意力机制中,第 i 个 query 向量 qᵢ 与所有 key 向量组成的矩阵 K 做点积,得到注意力分数,再经过 softmax 归一化:

softmax(qᵢ · K^T)

然后用这个注意力权重对 value 矩阵 V 加权求和,得到输出。

同时,位置编码在输入层被加到词向量上:词向量 + 位置向量 一起送进网络。


ALiBi 的修改

原文核心句(意译,基于原文公式)

ALiBi 把注意力计算改为:

softmax(qᵢ · K^T  +  m · [-(i-1), ..., -2, -1, 0])

其中:
- qᵢ · K^T:原来的 query-key 点积分数(不变)
- m:该注意力头(attention head)专属的斜率(slope),一个固定的正数,在训练前就设好,不参与梯度更新
- [-(i-1), ..., -2, -1, 0]:一个线性递减的偏置向量,代表"当前词 i 距离各个 key 位置的距离"

重要:ALiBi 不在任何网络层添加位置编码。词向量就是纯粹的词向量,没有位置信息混入。


用大白话理解这个公式

设当前词(query)在位置 5,它要"看"其他词(key):
- 看位置 5 的词(自己):偏置 = 0(不扣分)
- 看位置 4 的词:偏置 = -1(扣 1 分)
- 看位置 3 的词:偏置 = -2(扣 2 分)
- 看位置 1 的词:偏置 = -4(扣 4 分)

位置越远,扣的分越多,经过 softmax 后这些位置的注意力权重就越低。

类比:想象你在听朋友讲话(当前时刻 = query),你对他刚说的话(近处 key)注意力最集中,对十分钟前说的话(远处 key)已经有些模糊,注意力权重自然下降。ALiBi 就是把这个"越久远越记不清"的规律,直接硬编码进了注意力计算里。


3.2 每个注意力头用不同斜率 m

原文核心句(意译)

不同注意力头(attention head)使用不同的斜率 m,形成一个等比数列。

对于 8 个注意力头的模型,斜率序列为:

1/2¹, 1/2², 1/2³, ..., 1/2⁸
= 0.5, 0.25, 0.125, 0.0625, 0.03125, 0.015625, 0.0078125, 0.00390625

对于 16 个注意力头的模型,在 8 头序列之间插入几何平均值,从 2^(-0.5) 开始,公比为 2^(-0.5):

1/2^0.5, 1/2^1, 1/2^1.5, ..., 1/2^8

通用公式:对于 n 头模型,斜率序列以 2^(-8/n) 为首项,以该值为公比的等比数列。

这些斜率不是学习得到的,训练前就固定下来。作者在论文中强调,这组斜率不需要针对不同数据集微调,跨领域泛化性很好。


为什么要不同斜率?

新手讲解

这是 ALiBi 设计中非常精妙的地方。我们可以做一个比喻:

想象注意力机制的多个"头"是几个不同性格的同事在阅读同一份文档:

这样,模型就同时拥有了短程关注能力(局部语法、短语结构)和长程关注能力(主题、实体关联)。不同头各司其职,共同完成对序列的理解。


为什么这样能外推到更长的序列?

新手讲解

这是本论文最核心的洞见。

传统位置编码的问题:位置编码是一个"绝对位置的查找表"。模型在训练时只见过第 1~512 位置,推理时遇到第 600 位置,这个位置对应的向量从没出现过,模型不知道怎么处理。

ALiBi 的优雅之处:ALiBi 的偏置是根据"相对距离"计算的,没有绝对位置的概念:
- 当前词到前一个词:距离 = 1,偏置 = -m
- 当前词到前两个词:距离 = 2,偏置 = -2m
- ……无论序列有多长,这个规律都适用

就算序列长度从 512 扩展到 2048,这个"离得越远扣分越多"的规则依然有效,不需要"见过"第 2048 个绝对位置才能运作。

用数学语言说:偏置只依赖相对距离,与绝对位置无关,因此长度外推是自然的


3.3 实现方式(极简)

原文核心句(意译)

ALiBi 极易实现,只需修改几行代码。具体做法是:在 softmax 之前,将已有的"mask 矩阵"(因果语言模型中用于遮蔽未来词的矩阵)加上对应头的线性偏置矩阵即可。

不引入任何新参数,训练速度和内存占用与正弦编码模型几乎相同(内存增加 0~0.7%)。

新手讲解

"mask 矩阵"是自回归语言模型(GPT 类模型)中的一个技巧:为了让模型在训练时不"作弊"偷看后面的词,会用一个下三角矩阵把注意力矩阵的右上部分全部置为负无穷(-∞)。

ALiBi 的实现就是在这个 mask 矩阵里,同时加上线性偏置,两件事一次完成。工程上非常干净。


四、关键实验结果

本节对应原文 Section 4:"Experiments"

4.1 实验设置

数据集:WikiText-103(英文维基百科,较小规模验证实验)和 CC100+RoBERTa(4610 亿字节,大规模实验)

基线方法
- 正弦位置编码(Sinusoidal PE)
- 旋转位置编码(RoPE / Rotary PE)
- T5 相对位置偏置(T5 Bias)
- 可学习位置编码(Learned PE)

评估指标:困惑度(Perplexity,PPL),越低越好


4.2 核心实验:短训练 → 长推理的外推能力

WikiText-103,训练长度 L=512,测试不同长度

评估长度 正弦编码 PPL ALiBi PPL
512(训练长度) 20.05 19.73
1024 43.54(退化) 18.73(提升!)
3072 178.97(崩溃) 18.40(依然平稳)

翻译与讲解

这个表格展示了最惊人的对比。两个模型都只在 512 token 上训练,但测试时序列变长:

为什么测试时序列变长,困惑度反而降低?这正是"早期词汇诅咒(early token curse)"的缓解效果——测试序列越长,每个词平均能获得越充分的上下文,预测准确性提高。ALiBi 能利用这个好处,而正弦编码却因为遇到了"没见过的位置"而崩溃。


WikiText-103,训练长度 L=1024,测试不同长度

评估长度 正弦编码 PPL T5 Bias PPL ALiBi PPL
1024(训练长度) 19.34 18.80 18.66
2048 51.09(退化) 18.34 18.05
3072 96.46(崩溃) 约 17.5 17.92

ALiBi 和 T5 Bias 都展现出很好的外推能力,但 T5 Bias 的训练速度慢得多(见下表)。


4.3 计算效率对比(WikiText-103,L=1024)

方法 训练速度(WPS) 推理速度(WPS) 内存(GB)
正弦编码 26,000 77,800 19.2
旋转编码 17,700 39,400 22.8
T5 Bias 13,000 20,200 20.9
ALiBi 25,800 76,400 19.3

WPS = Words Per Second(每秒处理词数),越高越快

讲解

T5 Bias 虽然外推能力不错,但训练速度只有正弦编码的一半(13k vs 26k WPS),推理速度更是只有四分之一(20k vs 78k WPS)。

ALiBi 几乎和正弦编码速度持平,内存也几乎一样,但外推能力远超正弦编码,和 T5 Bias 相当甚至更好。这是一个"三全其美"的方案:快、省、能外推。


4.4 大规模实验:13 亿参数模型(CC100+RoBERTa)

原文核心句(意译)

一个用 ALiBi 在 L=1024 上训练的 13 亿参数语言模型,在 L=2048 的测试集上,能达到与正弦编码模型(在 L=2048 上训练)相同的困惑度——而前者训练快 11%,内存少用 3.1 GB(19.2 GB vs 约 29.3 GB)。

模型 训练长度 测试困惑度 (L=2048) 内存 训练速度
正弦编码 L=2048 8.83 29.3 GB 基准
ALiBi L=1024 8.92 26.2 GB 快 11%

讲解

这个结果非常有说服力。ALiBi 模型只用了正弦编码模型一半的训练序列长度,但在同等测试条件下性能相当(8.92 vs 8.83,差距不到 0.1)。与此同时还省了 3.1 GB 显存,训练还快了 11%。

对于工业界来说,这意味着可以用便宜一半的训练成本,获得同等质量的长文本处理能力。


4.5 "早期词汇诅咒"(Early Token Curse)

原文核心句(意译)

研究发现,ALiBi 改善困惑度的主要原因,是缓解了"早期词汇诅咒",而非模型真正学到了更长的依赖关系。

所谓早期词汇诅咒:当用不重叠推理(non-overlapping inference,即把长文本切成若干段,每段独立处理)评估模型时,每一段的开头几个词缺乏充分的上下文(因为上一段的内容被截断了),预测质量很差,拉高了整体困惑度。

当测试时序列更长时,每段开头词占总词数的比例更低,困惑度自然下降。ALiBi 能更好地利用这个效果,而正弦编码在长序列上已经崩溃,所以没法利用。

讲解

打个比方:你在读一本书,每次只能翻看一页(1024 词),看完后"强制遗忘",再翻下一页。那么每页的第一句话都会特别难理解(没有上文)。如果换成每次能看两页(2048 词),第一句话难理解的问题还在,但它占总阅读量的比例更小了,整体理解质量就会提升。


五、相关工作与结论(概览)

相关工作(原文 Section 5):作者回顾了各类位置编码方案(绝对位置、相对位置、可学习位置编码等),分析了 T5、XLNet、Transformer-XL 等模型的位置表示策略,指出 ALiBi 的核心贡献是在效率和外推能力之间找到了新的平衡点。(附录部分含消融实验和更多细节,感兴趣可参阅原文附录 A/B/C。)

结论(原文 Section 6):

使用 ALiBi,我们加速了一个 13 亿参数模型的训练,该模型的测试序列长度与 GPT-3(2048 token)相同。我们训练得更短,测试得更长——这正是本文标题的含义。

ALiBi 的核心贡献总结:
1. 去掉位置编码,改为在注意力分数上施加线性距离惩罚
2. 不同注意力头用不同斜率,兼顾短程和长程依赖
3. 训练短、测试长,外推能力来自偏置规律与绝对位置无关
4. 工程上几乎零开销,几行代码改动,速度和内存与基线相当


六、ALiBi 核心原理总结(一图抵千言)

传统注意力(加位置编码):
  词向量 + 位置向量 ──→ [注意力计算] ──→ 输出
                        只见过1~512位置,超过512就不知所措

ALiBi 注意力:
  词向量(纯粹)──→ [注意力计算] + [线性距离惩罚] ──→ 输出
                      · 距离1:-1m
                      · 距离2:-2m
                      · 距离k:-km
                    (无论k多大,规律永远适用)

不同头的斜率 m(以 8 头为例):

头 1:m = 1/2   = 0.500  → 只关注最近几个词
头 2:m = 1/4   = 0.250
头 3:m = 1/8   = 0.125
头 4:m = 1/16  = 0.063
头 5:m = 1/32  = 0.031
头 6:m = 1/64  = 0.016
头 7:m = 1/128 = 0.008
头 8:m = 1/256 = 0.004  → 能关注到很远的词

斜率越大 → 惩罚越猛 → 视野越短;斜率越小 → 惩罚越轻 → 视野越远。


七、关键术语速查表

术语 英文 解释
长度外推 Length Extrapolation 在比训练序列更长的文本上推理,且不明显退化
线性偏置 Linear Bias 与距离成正比的惩罚值,离得越远偏置越大(越负)
斜率 Slope (m) 每个注意力头专属的偏置比例系数,固定不学习
困惑度 Perplexity (PPL) 语言模型预测质量指标,越低越好
注意力头 Attention Head 多头注意力中的单个并行注意力计算单元
早期词汇诅咒 Early Token Curse 分段推理时每段开头缺乏上文导致预测质量差
旋转位置编码 RoPE / Rotary PE 通过旋转操作编码相对位置,LLaMA 等模型使用
T5 相对偏置 T5 Bias Google T5 使用的桶化相对位置偏置,外推强但计算慢
WPS Words Per Second 每秒处理词数,衡量训练/推理速度
非重叠推理 Non-overlapping Inference 将长文本切成不重叠片段分别处理的推理方式

八、为什么 ALiBi 重要?(历史意义)

ALiBi 发表于 2021 年,此后成为长上下文 Transformer 研究的重要基础参考:

  1. 验证了"位置编码不一定要加在输入上"的思路,启发了后续 RoPE、YaRN 等方法
  2. 明确提出了"长度外推"作为独立研究问题,推动了该方向的系统性研究
  3. 工程简洁性使其被多个开源模型(如 BLOOM、MPT 系列)实际采用
  4. "训练短、测试长"的思想直接影响了后来的 LongLoRA、YaRN 等上下文扩展方法

精读笔记完。建议配合原论文图表阅读,尤其是 Figure 1(各方法外推性能曲线)和 Figure 2/3(ALiBi 偏置矩阵示意图)。