训练短，测试长：带有线性偏置的注意力机制（ALiBi）实现输入长度外推

本仓库包含我们2022年ICLR论文《训练短，测试长》中使用的ALiBi代码及模型。本文档将说明如何在WikiText-103数据集上运行我们的实验。

论文：这里

视频：这里

+ -----------
+ 新内容：
+ -----------

1. 最新发表的一篇论文（arXiv:2310.13017）表明，位置插值技巧同样适用于ALiBi（2023年10月）。
2. 如果您希望将ALiBi应用于双向Transformer模型（例如编码器），可以参考此处提到的方法。
3. 下方的FAQ部分解答了关于ALiBi与TransformerXL比较的问题，以及为什么我认为ALiBi效果较好。

带有线性偏置的注意力机制（ALiBi）非常简单！我们没有在Transformer堆栈底部添加位置嵌入，而是如上图所示，为每个注意力分数添加一个线性偏置。超参数“m”是按头设置的，且不进行学习——它在训练开始时就被固定下来。我们提供了一个函数，可以根据模型中的头数自动生成这些m值。

使用ALiBi，模型可以在1024个标记上进行训练，随后无需任何微调即可在2048个甚至更多标记上进行推理。此外，在资源有限的语言建模场景中，即使不进行外推，ALiBi也能提升性能。

其实现过程非常简单：

0. 从模型中移除位置嵌入：https://github.com/ofirpress/attention_with_linear_biases/blob/master/fairseq/models/transformer.py#L941
1. 设置相对偏置矩阵，见：https://github.com/ofirpress/attention_with_linear_biases/blob/master/fairseq/models/transformer.py#L742
2. 将偏置矩阵添加到掩码中，该掩码会在每次计算注意力分数时被应用：https://github.com/ofirpress/attention_with_linear_biases/blob/master/fairseq/models/transformer.py#L1011
3. （这一步在其他框架中可能不需要。）将掩码计算移到层循环之前，以使Transformer运行速度稍快一些：https://github.com/ofirpress/attention_with_linear_biases/blob/master/fairseq/models/transformer.py#L949

就这些！

FAQ：

您认为ALiBi为何有效？

我也不太确定，但我的想法如下：我认为使用位置嵌入（无论是可学习的、正弦的还是旋转式的）并不是最优的选择。它们可能会导致语言模型过度依赖这些位置嵌入，而未能真正理解位置的概念。关于这一点，我在这篇推文中有更详细的讨论。

ALiBi与TransformerXL的位置嵌入方法相比如何？

好问题！尽管尚未经过充分验证，但TransformerXL的位置编码方法也可能支持外推。在我们之前的论文（Shortformer，表5）中，我们已经证明该方法会使注意力机制比原始版本慢两倍以上，并且需要将注意力子层的内存消耗增加一倍。而ALiBi的运行速度与原始注意力机制相当，额外内存开销最多仅为100MB。

如何将ALiBi应用于双向模型，比如仅编码器模型（如BERT）或编码器-解码器模型（如T5）？

请参阅此链接。

如果我想对更长的序列进行外推，是否可以直接对使用正弦/可学习位置嵌入的模型应用滑动窗口掩码？

不行，那样不起作用（我试过）。对于可学习的位置嵌入，一旦超过训练上下文长度，模型就无法为超出部分的标记生成位置嵌入（例如，如果只在1000个标记上训练，那么第1001个标记就没有位置嵌入）。而对于正弦位置嵌入，当输入的位置嵌入与训练期间观察到的不同时，模型会变得不稳定。我在这段视频讲座中对此有进一步的说明。

ALiBi中的偏置是线性增长的。您尝试过让其以其他方式增长吗，比如指数增长？

是的，在开发ALiBi的过程中，我也尝试过使用指数增长的偏置函数，但效果不如线性增长的好。我还试验了其他形式的增长函数，最终发现线性增长的效果最佳，因此我们在ALiBi的最终版本中选择了线性偏置。

---

引用：

@inproceedings{
alibi,
title={Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation},
author={Ofir Press and Noah Smith and Mike Lewis},
booktitle={International Conference on Learning Representations},
year={2022},
url={https://openreview.net/forum?id=R8sQPpGCv0}
}

WikiText-103

要求与安装

本仓库是Fairseq仓库的一个分支，因此具有相同的依赖要求。

安装完依赖后，您可以通过以下命令安装本仓库：

pip install --editable .

数据准备

要下载并预处理数据，请执行以下命令：

cd examples/language_model/
bash prepare-wikitext-103.sh
cd ../.

TEXT=examples/language_model/wikitext-103
fairseq-preprocess \
    --only-source \
    --trainpref $TEXT/wiki.train.tokens \
    --validpref $TEXT/wiki.valid.tokens \
    --testpref $TEXT/wiki.test.tokens \
    --destdir data-bin/wikitext-103 \
    --workers 20

训练与推理

要使用线性偏置注意力机制（ALiBi）训练一个语言模型，输入序列为512个标记时，请运行以下命令：

python train.py --task language_modeling data-bin/wikitext-103 --save-dir wt103/ --arch transformer_lm_wiki103 --max-update 286000 --lr 1.0 --t-mult 2 --lr-period-updates 270000 --lr-scheduler cosine --lr-shrink 0.75 --warmup-updates 16000 --warmup-init-lr 1e-07 --stop-min-lr 1e-09 --optimizer nag --min-lr 0.0001 --clip-norm 0.1 --criterion adaptive_loss --max-tokens 9216 --update-freq 1 --tokens-per-sample 512 --seed 1 --sample-break-mode none --skip-invalid-size-inputs-valid-test --ddp-backend=legacy_ddp --fp16 --required-batch-size-multiple 1

对于长度大于512（最多2048）个标记的输入序列，只需更改--tokens-per-sample参数即可。

若要训练输入长度为3072的模型，则需将--update-freq参数改为3，并将--max-tokens参数降低至3072（同时将--tokens-per-sample也设置为3072）。

如果在训练过程中出现内存不足的情况： 可以将--max-tokens设置为之前值的50%，并将--update-freq设置为之前值的两倍。这样会得到一种在数学上等价于原始命令但所需内存更少的批处理计算方式。如果仍然无法解决，可进一步将--max-tokens设置为之前值的25%，并将--update-freq设置为之前值的四倍，依此类推……

已保存的检查点

如果您希望下载我们在WikiText-103数据集上训练好的模型，它们可以在这里找到：

输入长度	模型	日志
64	⬇️	⬇️
128	⬇️	⬇️
256	⬇️	⬇️
512	⬇️	⬇️
1024	⬇️	⬇️
1536	⬇️	⬇️
2048	⬇️	⬇️
3072	⬇️	⬇️

如果您想按照下面的说明进行操作，请将下载的文件重命名为checkpoint_best.pt。

推理

对于验证集上的非重叠评估，运行以下命令：

l=1024; fairseq-eval-lm data-bin/wikitext-103/     --path wt103/checkpoint_best.pt  --sample-break-mode none --gen-subset valid   --max-sentences 1 --model-overrides "{'max_tokens':$l, 'tokens_per_sample':$l, 'max_target_positions':$l}"  --tokens-per-sample $l --max-tokens $l  --max-target-positions $l  --context-window 0

其中l被设置为验证过程中输入子序列的长度（如上例中l=1024）。

对于验证集上的滑动窗口评估，运行以下命令：

l=1024; fairseq-eval-lm data-bin/wikitext-103/     --path wt103/checkpoint_best.pt  --sample-break-mode none --gen-subset valid   --max-sentences 1 --model-overrides "{'max_tokens':$l, 'tokens_per_sample':$l, 'max_target_positions':$l}"  --tokens-per-sample $l --max-tokens $l  --max-target-positions $l  --context-window $((l-1)) 

（我们只需将--context-window参数修改为l-1，这意味着每次前向传播时，评估窗口都会向后滑动1个标记）。

ALiBi 快速上手指南

环境准备

• 系统要求：Linux/macOS + Python 3.8+
• 前置依赖：

  # 推荐使用国内镜像加速安装
  pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://pypi.tuna.tsinghua.edu.cn/simple
  pip install fairseq --no-cache-dir  # 原生fairseq依赖
  

安装步骤

# 克隆仓库（建议使用国内GitHub镜像加速）
git clone https://hub.fastgit.org/ofirpress/attention_with_linear_biases.git
cd attention_with_linear_biases

# 安装依赖（使用清华源加速）
pip install --editable . -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

数据准备（WikiText-103）

# 下载并预处理数据集
cd examples/language_model/
bash prepare-wikitext-103.sh
cd ../..

# 二值化处理（使用20个进程加速）
TEXT=examples/language_model/wikitext-103
fairseq-preprocess \
    --only-source \
    --trainpref $TEXT/wiki.train.tokens \
    --validpref $TEXT/wiki.valid.tokens \
    --testpref $TEXT/wiki.test.tokens \
    --destdir data-bin/wikitext-103 \
    --workers 20

模型训练

# 训练基础命令（512长度序列）
python train.py --task language_modeling data-bin/wikitext-103 \
    --save-dir wt103/ \
    --arch transformer_lm_wiki103 \
    --tokens-per-sample 512 \  # 调整此参数控制输入长度
    --max-tokens 9216 \
    --update-freq 1 \
    --fp16  # 自动混合精度训练

# 内存不足时调整方案：
# 将max-tokens减半，update-freq翻倍（保持等效batch size）
--max-tokens 4608 --update-freq 2

模型推理

# 非重叠验证（输入长度1024）
l=1024
fairseq-eval-lm data-bin/wikitext-103/ \
    --path wt103/checkpoint_best.pt \
    --tokens-per-sample $l \
    --max-tokens $l \
    --context-window 0  # 非重叠模式

# 滑动窗口验证（输入长度1024）
fairseq-eval-lm data-bin/wikitext-103/ \
    --path wt103/checkpoint_best.pt \
    --tokens-per-sample $l \
    --context-window $((l-1))  # 滑动窗口模式

模型下载

预训练模型（直接下载）：

输入长度	模型文件
1024	下载链接
2048	下载链接

下载后重命名为 checkpoint_best.pt 可直接使用推理命令

---

💡 提示：训练长序列时，建议通过 --tokens-per-sample 逐步增加长度，配合 --update-freq 参数动态调整显存占用。