精读笔记: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
阅读地图
本文篇幅不长,结构清晰,适合新手完整阅读。精读顺序建议:
- 先看本文第一节("背景:什么是长度外推"),建立直觉
- Abstract — 摘要,2 分钟了解论文贡献
- Introduction — 引言,搞清问题背景和动机
- Section 2(原文)— 为什么旧方法不行
- Section 3(原文)— ALiBi 的具体机制(精读重点)
- 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 位置?
传统正弦位置编码的答案是:基本不行。
段落二:现有方法对比
原文核心句(意译):
作者测试了三类位置编码方法的外推能力:
- 正弦位置编码(Sinusoidal PE):外推能力极差,超出训练长度仅 20~50 个 token 就开始明显退化
- 旋转位置编码(Rotary PE / RoPE):比正弦稍好,能多撑 100~200 个 token,但仍然退化
- T5 相对位置偏置(T5 Bias):外推能力最强(可达训练长度之外 600~800 token),但训练速度慢了将近一半
翻译与讲解:
旋转位置编码(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 设计中非常精妙的地方。我们可以做一个比喻:
想象注意力机制的多个"头"是几个不同性格的同事在阅读同一份文档:
- 斜率大的头(如 m=0.5):性格急躁,对远处的词惩罚很重,主要关注当前词附近的内容——"只看最近几个词"
- 斜率小的头(如 m=0.004):性格耐心,对远处的词惩罚很轻,能关注很远处的上下文——"能看到很远的词"
这样,模型就同时拥有了短程关注能力(局部语法、短语结构)和长程关注能力(主题、实体关联)。不同头各司其职,共同完成对序列的理解。
为什么这样能外推到更长的序列?
新手讲解:
这是本论文最核心的洞见。
传统位置编码的问题:位置编码是一个"绝对位置的查找表"。模型在训练时只见过第 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 上训练,但测试时序列变长:
- 正弦编码在 1024 token 时困惑度从 20 跳到 43(翻倍),3072 token 时直接崩到 179
- ALiBi 不仅没有退化,困惑度反而从 19.73 下降到 18.73,序列越长反而越好
为什么测试时序列变长,困惑度反而降低?这正是"早期词汇诅咒(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 研究的重要基础参考:
- 验证了"位置编码不一定要加在输入上"的思路,启发了后续 RoPE、YaRN 等方法
- 明确提出了"长度外推"作为独立研究问题,推动了该方向的系统性研究
- 工程简洁性使其被多个开源模型(如 BLOOM、MPT 系列)实际采用
- "训练短、测试长"的思想直接影响了后来的 LongLoRA、YaRN 等上下文扩展方法
精读笔记完。建议配合原论文图表阅读,尤其是 Figure 1(各方法外推性能曲线)和 Figure 2/3(ALiBi 偏置矩阵示意图)。