Unlimiformer：支持无限长度输入的长距离Transformer（NeurIPS 2023）

这是论文的官方实现：

Amanda Bertsch、Uri Alon、Graham Neubig 和 Matthew R. Gormley： Unlimiformer：支持无限长度输入的长距离Transformer（即将发表于 NeurIPS 2023）

Unlimiformer 是一种在不改变注意力机制数学定义的情况下，通过检索式注意力增强预训练编码器-解码器模型的方法。 这使得任何预训练的编码器-解码器模型都能处理无限长度的输入！ 另请参阅我们的 推文。

Unlimiformer 可用于提升已训练模型的性能。为了获得最佳效果，可以使用 Unlimiformer 训练方法对模型进行训练。

如果您对这项工作有任何疑问，请在 GitHub 问题 中提出，或发送邮件至 abertsch@cs.cmu.edu, ualon@cs.cmu.edu。

2023年10月 - Unlimiformer 将在 NeurIPS 2023 上亮相！

2023年8月 - Unlimiformer 现在支持 Llama-2（及其所有衍生模型）！

例如，要使用 Llama-2 处理极长的输入，比如一本 整本书 的内容，可以使用以下命令：

python src/run_generation.py --model_type llama --model_name_or_path meta-llama/Llama-2-13b-chat-hf \
    --prefix "<s>[INST] <<SYS>>\n 您是一位乐于助人的助手。请根据完整的指令或问题提供详细的回答。\n<</SYS>>\n\n 请总结以下书籍：" \
    --prompt example_inputs/harry_potter_full.txt \
    --suffix " [/INST]" --test_unlimiformer --fp16 --length 200 --layer_begin 16 \
    --index_devices 1 --datastore_device 1 

• 最终的提示将是 --prefix、--prompt 和 --suffix 这三个标志内容的拼接。
• --prompt 标志既可以是文本文件的路径（如 example_inputs/harry_potter_full.txt），也可以直接指定提示字符串。
• 必须使用 --test_unlimiformer 标志才能启用 Unlimiformer。
• --length 标志决定了期望的输出长度。
• --layer_begin 标志决定了从哪一层开始应用 Unlimiformer。例如，如果设置为 --layer_begin 20，模型的前 20 层将像往常一样仅对提示中的最后 context_window_size 个标记执行标准注意力操作，而第 21 层及以后的层则会关注整个长输入。根据我们的初步实验，--layer_begin 的值应超过模型总层数的一半，且调整该值会显著影响输出质量。
• --datastore_device N 和 --index_devices N1 N2 N3 ... 标志指定了在哪些 GPU 上存储 Unlimiformer 的数据存储和索引（基础模型将存储在 GPU #0 上）。
• 添加 --stream_output 标志可以使生成的标记在生成时逐个显示。

开始使用

一般说明

将 src 目录下的文件复制到您的源代码文件夹中。

您需要为 usage.py 中列出的 Unlimiformer 特定参数设置值——您可以将这些参数添加到平时处理超参数的地方。要使用该模型，必须设置 test_unlimiformer=True。对于数据存储的使用，模型必须处于评估模式（例如，在推理之前调用 model.eval()）。

inference-example.py 提供了一个使用默认参数运行序列到 Unlimiformer 模型的最小示例。

run.py 是一个整合了 Unlimiformer 的完整训练设置示例，改编自 SLED。完整的命令行如下所示。

复现论文中的实验——命令行

要对 GovReport 数据集上的 BART-base 模型进行标准微调和评估（作为示例），可以使用以下命令：

python src/run.py \
    src/configs/training/base_training_args.json \
    src/configs/data/gov_report.json \
    --output_dir output_train_bart_base_local/ \
    --learning_rate 1e-5 \
    --model_name_or_path facebook/bart-base \
    --max_source_length 1024 \
    --eval_max_source_length 1024 --do_eval=True \
    --eval_steps 1000 --save_steps 1000 \
    --per_device_eval_batch_size 1 --per_device_train_batch_size 2 \
    --extra_metrics bertscore

• 要在 训练 时使用 Unlimiformer（论文中称为“检索训练”），请使用：--unlimiformer_training --max_source_length 16384
  • 在这种情况下，您可能还希望在 测试/验证时使用 Unlimiformer，并同时使用：--test_unlimiformer --eval_max_source_length 999999
• 或者，要在 训练 时使用计算成本更低的“随机编码”方法，可以使用 --random_unlimiformer_training --max_source_length 16384
• 若要在“检索训练”和“随机编码训练”之间交替使用，可以同时使用两个标志：--unlimiformer_training --random_unlimiformer_training --max_source_length 16384

有关其他标志和选项，请参阅 usage.py。

推荐设置

使用 Unlimiformer 进行评估

在评估时，我们建议使用每个设置的默认值。

使用 Unlimiformer 进行训练

对于低成本的方法，我们建议按常规方式训练，并在早停时使用 Unlimiformer。为此，设置 knn=True 并将其他所有值保持默认。

为了获得最佳性能，有三种高成本的训练设置。哪种设置最好取决于具体的数据集。

1. 设置 random_unlimiformer_training=True：这是论文中的“随机编码训练”设置。
2. 设置 unlimiformer_training=True：这是论文中的“检索训练”设置。
3. 同时设置 random_unlimiformer_training=True 和 unlimiformer_training=True：这是论文中的“交替训练”设置。

有关相对训练成本的更详细分解，请参阅论文中的表 5。

处理超大输入的技巧

对于训练

• 您可能需要在训练时截断输入，例如将其限制为 8k 或 16k 个标记。而在评估时则可以使用完整的输入。
• 您还可以尝试将输入分割成 16k 标记的块，并将每一块作为单独的示例进行训练。

用于评估（包括早停）

• 如果您经常遇到 CUDA 显存不足的问题，请设置 use_datastore=True，以使用 Faiss 数据存储来保存隐藏状态。
• 如果问题仍然存在，请将 gpu_datastore=False 或 gpu_index=False，但请注意，这会降低性能。

训练好的模型

论文中提到的以下模型已在 Hugging Face 上发布。请注意，您必须将 Unlimiformer 特有的文件添加到您的仓库，并通过 test_unlimiformer=True 来加载这些模型。如果您从 Hugging Face 下载这些模型，默认情况下它们可能不会使用 Unlimiformer！

表3：低成本训练方法

数据集	方法	Hugging Face 链接
GovReport	基线：BART-base	abertsch/bart-base-govreport
GovReport	BART-base + Unlimiformer 早停	abertsch/unlimiformer-bart-govreport-earlyk
SummScreen	基线：BART-base	abertsch/bart-base-summscreen
SummScreen	BART-base + Unlimiformer 早停	abertsch/unlimiformer-bart-summscreen-earlyk

表4：长距离训练方法

数据集	方法	Hugging Face 链接
GovReport	BART + Unlimiformer（交替训练）	abertsch/unlimiformer-bart-govreport-alternating
SummScreen	BART + Unlimiformer（检索训练）	abertsch/unlimiformer-bart-summscreen-retrieval

表5：BookSum

数据集	方法	Hugging Face 链接
BookSum	基线：BART-base	abertsch/bart-base-booksum
BookSum	BART-base + Unlimiformer 早停	abertsch/unlimiformer-bart-booksum-earlyk
BookSum	BART-base + Unlimiformer（随机编码训练）	abertsch/unlimiformer-bart-booksum-random-encoding
BookSum	BART-base + Unlimiformer（交替训练）	abertsch/unlimiformer-bart-booksum-alternating

结果

引用

如果您使用了我们的方法或模型，请引用我们的论文：

@article{bertsch2023unlimiformer,
  title={Unlimiformer: Long-Range Transformers with Unlimited Length Input},
  author={Bertsch, Amanda and Alon, Uri and Neubig, Graham and Gormley, Matthew R},
  journal={arXiv preprint arXiv:2305.01625},
  year={2023}
}

Unlimiformer 快速上手指南

Unlimiformer 是一种无需改变注意力机制数学定义，即可通过检索增强预训练编码器 - 解码器模型的方法。它允许任何预训练的 Encoder-Decoder 模型处理无限长度的输入（如整本书），特别适用于长文本摘要和生成任务。

环境准备

系统要求

• 操作系统: Linux (推荐) 或 macOS
• GPU: 支持 CUDA 的 NVIDIA GPU（建议显存充足，或使用多卡分布存储索引）
• Python: 3.8+

前置依赖

确保已安装 PyTorch 和 Transformers 库。建议使用 pip 安装基础依赖：

pip install torch transformers faiss-gpu sentencepiece protobuf

注意：本项目依赖 Faiss 进行高效检索。如果无法访问官方源，可尝试使用国内镜像加速安装：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple torch transformers faiss-gpu sentencepiece protobuf

安装步骤

Unlimiformer 没有独立的 PyPI 包，需直接克隆源码并使用其中的模块。

1. 克隆仓库

   git clone https://github.com/abertsch72/unlimiformer.git
   cd unlimiformer
   

2. 配置代码 将 src 目录下的文件复制到你的项目源代码文件夹中，或者直接在当前目录下运行脚本。

   核心逻辑在于需要在你的推理或训练代码中引入 usage.py 定义的参数，并设置 test_unlimiformer=True。

基本使用

场景一：推理（处理超长输入）

以下示例展示如何使用 Unlimiformer 让 Llama-2 模型处理整本书的内容进行摘要。

关键参数说明：

• --test_unlimiformer: 必须添加，用于启用 Unlimiformer。
• --layer_begin: 指定从哪一层开始应用检索注意力。通常设为总层数的一半以上（例如 16 或 20）。
• --prompt: 可以是文本文件路径（如整本书），也可以是字符串。
• --datastore_device / --index_devices: 指定存储 datastore 和索引的 GPU 设备 ID，避免显存溢出。

运行命令：

python src/run_generation.py --model_type llama --model_name_or_path meta-llama/Llama-2-13b-chat-hf \
    --prefix "<s>[INST] <<SYS>>\n You are a helpful assistant. Answer with detailed responses according to the entire instruction or question. \n<</SYS>>\n\n Summarize the following book: " \
    --prompt example_inputs/harry_potter_full.txt \
    --suffix " [/INST]" --test_unlimiformer --fp16 --length 200 --layer_begin 16 \
    --index_devices 1 --datastore_device 1 

若希望实时看到生成结果，可添加 --stream_output 标志。

场景二：最小化代码集成示例

如果你想在自定义 Python 脚本中集成，请参考 src/inference-example.py 的逻辑。核心步骤如下：

1. 加载预训练模型。
2. 调用 model.eval() 进入评估模式（使用 datastore 必须此步骤）。
3. 在生成参数中设置 test_unlimiformer=True 及相关 Unlimiformer 参数。

# 伪代码示例，具体实现请参考 src/inference-example.py
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
# 引入 unlimiformer 相关逻辑 (需将 src 文件加入路径)
from usage import add_unlimiformer_args, setup_unlimiformer 

# 1. 加载模型
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")
tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")

# 2. 设置为评估模式 (必须)
model.eval()

# 3. 准备输入 (超长文本)
input_text = "..." # 你的长文本

# 4. 执行生成 (需在内部逻辑中开启 test_unlimiformer=True)
# 具体调用方式取决于你如何整合 src/run.py 或 inference-example.py 的逻辑
# 关键在于传递 unlimiformer 特定的 args

场景三：微调训练（可选）

若需获得最佳效果，建议使用 Unlimiformer 进行训练（Retrieval training）。

基础微调命令（BART-base on GovReport）：

python src/run.py \
    src/configs/training/base_training_args.json \
    src/configs/data/gov_report.json \
    --output_dir output_train_bart_base_local/ \
    --learning_rate 1e-5 \
    --model_name_or_path facebook/bart-base \
    --max_source_length 1024 \
    --eval_max_source_length 1024 --do_eval=True \
    --eval_steps 1000 --save_steps 1000 \
    --per_device_eval_batch_size 1 --per_device_train_batch_size 2 \
    --extra_metrics bertscore

启用 Unlimiformer 训练： 添加 --unlimiformer_training 并增大 --max_source_length（如 16384）。

python src/run.py ... \
    --unlimiformer_training --max_source_length 16384 \
    --test_unlimiformer --eval_max_source_length 999999

提示：对于显存受限的情况，训练时可截断输入（如 8k 或 16k），仅在评估时使用完整长输入。若评估时显存不足，可设置 use_datastore=True 使用 Faiss 存储隐藏状态。