扩展LLM上下文长度

如何为Transformer模型编码位置信息，一直是LLM架构中的关键组成部分之一。

近来，我们以及社区中的其他研究者都对这样一个问题颇感兴趣：能否将LLM的上下文长度扩展到更长？

我们针对Llama模型的不同上下文长度扩展方案进行了大量实验。Llama模型在预训练时采用RoPE（旋转位置嵌入）编码，其上下文长度限制为2048。在此，我们分享部分实验结果以及相应的训练和评估脚本，希望能对社区有所帮助。对于表现最佳的两个模型——分别以缩放因子4和16进行线性缩放并结合指令微调——我们也公开了模型权重，供有兴趣的研究者使用或进一步测试。我们认为，缩放因子16的模型在实际任务中能够较好地处理高达16k的上下文长度，甚至可能在约20–24k的上下文长度上也有不错的表现。

缩放因子16的模型

技术论文

概述

我们开展了多种实验，试图扩展模型的上下文长度。首先，我们直接以零样本方式使用基础Llama模型。正如预期，该模型在2048个token的上下文长度内表现良好，但超过此长度后性能迅速下降。

接下来，我们尝试了微调方法，在RedPajama数据集上以4096个token的上下文长度对模型进行训练。这使得模型在4096个token的上下文长度内性能有所提升，然而再往上调时效果便不再明显。

另一种扩展上下文长度的方法是修改RoPE编码。为此，我们尝试了多种思路：

• 线性缩放，如kaiokendev.github.io所述。
• 对RoPE的傅里叶基底按幂次缩放，使低频成分比高频成分被拉伸得更多。
• 对傅里叶基底进行截断。我们的想法是让模型只关注那些足够快、在训练过程中至少能完成一个完整周期的频率；而更慢的频率则置为0（相当于完全不进行旋转，即在所有上下文长度上都同等重要）。
• 随机化位置向量。

特别地，我们将上述方法与在RedPajama数据集上的微调以及基于Vicuna数据集的指令微调相结合，这一组合取得了最为显著的效果。

最后，我们实现了并尝试了xPos论文中提出的方法。该方法通过引入随距离衰减的幅度惩罚项，使得傅里叶基底中的高频成分在长距离下影响较小，而低频成分则影响较大（详见我们的博客文章中的相似度热图）。

重点结果

我们最重要的发现之一是：不同的评估方法或任务会导致对上述方法的排名有所不同。这一点将在下文中详细说明。

尽管如此，我们仍得出以下几点普遍观察：

• 线性插值或缩放似乎是提升模型上下文长度最稳健的方法。
• 使用线性缩放因子N并不一定意味着模型的上下文长度会相应地增加N倍。例如，我们的缩放因子16实验通常在上下文长度达到16000左右时性能便开始下降，而非32000（约2048×16）。我们已规划了一些改进这一现象的方案，将在后续工作中实现。
• 截断和随机化虽然在困惑度指标上表现优异，但在检索任务中的表现却相对较差。
• 基于Vicuna数据集的指令微调能够显著提升基础模型在自身可处理长度范围内的检索准确率，但对于超出其能力范围的长度，则无法“修复”模型的不足。

评估任务

在评估中，我们使用了两个不同的数据集：

• LMSys 数据集（“lines”任务），用于在上下文中定位子字符串；
• 我们的开源开放书问答数据集 WikiQA，该数据集基于其他开源的基础问答数据集构建。

此外，我们还观察了训练集和验证集在以下过程中的对数损失：

对于 LMSys 任务，我们生成了新的、更长的测试用例，上下文长度最长可达约 25000 个 token，超过了原始数据集中 16000 个 token 的上下文测试用例。

WikiQA 任务是根据维基百科文档中提供的信息回答问题。我们在 Google Natural Questions 中的短答案格式数据基础上构建了我们的问答任务。该任务以文档和问题的形式呈现。我们确保问题的答案是一个简短的答案，要么是一个单独的词，要么是从文档中直接摘录的一小句话。通过这种结构化的任务设计，我们可以精确地确定大语言模型本应在上下文中“查找”答案的具体位置，从而通过将答案放置在不同位置来有效评估扩展后的各种上下文长度。

我们选取了大型维基百科文档，并对其进行截断，得到同一文档的多个版本，其大小介于 2000 到 16000 个 token 之间。对于每个文档大小，我们还准备了多个版本，将问题和答案文本放置在文档的不同位置——例如，是在文档的前 10%、主体部分还是最后 10%。通过提供同一文档的多个版本，我们能够在不同模型规模之间以及在同一模型的不同上下文位置上进行全面而公平的评估，因为我们本质上始终在询问相同的信息。

基于维基百科的数据集可能存在一个问题：模型可能会从其预训练语料库中正确回答问题，而不是从给定的上下文中获取答案。为了解决这个问题，我们创建了一个“修改版”数据集。该数据集仅包含答案为数字的问题。在此，我们将答案以及文档中所有出现该答案的地方都替换为一个不同的数字。这样可以确保如果大语言模型仅仅依赖其预训练语料库进行回忆，它就会给出错误的答案。具体的修改方式如下：

• 如果答案是年份（通常在 1000 到 2100 年之间），我们会将其改为原值上下浮动 10 年内的一个随机值。我们将年份视为特殊情况，以避免因打乱时间顺序而导致对文档的解读变得荒谬。
• 如果答案是其他数字，则将其替换为具有相同位数的另一个随机数字。

我们将原始的问答任务称为 自由形式问答 (FFQA)，而修改后的任务则称为 修改版数字问答 (AltQA)。

我们通过测量“存在准确率”，即模型生成的答案中是否包含作为子字符串的答案，来评估这两个版本的问答任务中每个示例的成功情况。要对我们的模型在 WikiQA 上运行推理并计算指标，请参阅位于 此处 的 run_inference_WikiQA.py 和 compute_metrics_WikiQA.ipynb。

我们已在 HuggingFace 上发布了这些数据集，以便其他人也能用于开展自己的长上下文实验。

• 扩展版 LMSys 数据集
• WikiQA 自由形式问答数据集
• WikiQA 修改版数字问答数据集

结果

LMSys 评估

关于以下结果的一个总体说明是：作者认为，在这项任务上微小的准确率差异并不能很好地反映模型排名的质量。因此，在解释结果时，我们主要关注整体趋势。

另外，作为基准，标准的 Llama-13b 模型只有在上下文长度不超过 2048 个 token 时才具有非零准确率（经过 Vicuna 指令微调的版本也是如此）。

不同缩放方法的比较

上图比较了不同的缩放方法。“Scale”表示使用指定的缩放值进行线性插值。我们看到，缩放系数为 16 的线性插值是唯一一种在上下文长度超过 9000 个 token 时仍能获得非零准确率的方法。然而，这种方法似乎会牺牲一些较短上下文长度下的准确率。

缩放指数为 0.5 的方法在较短上下文长度下表现尤为出色，但随着上下文长度的增加，其准确率下降得非常快。

值得注意的是，缩放系数为 16 的方法并没有像预期那样广泛适用。直观来看，人们可能会认为，既然缩放系数为 4 的方法在上下文长度达到 8192 个 token 时仍有非零准确率（这也在情理之中，因为原始上下文长度为 2048 个 token，而 8192 等于 2048 乘以 4；超过这个长度后，模型会遇到前所未有的键与查询之间的相对距离），那么缩放系数为 16 的方法应该能在上下文长度达到 2048 乘以 16，即 32768 个 token 时仍然保持非零准确率。

IFT（指令微调）的影响

上图展示了通过 LoRA 使用 Vicuna 指令集进行微调所带来的 IFT 效果。我们看到，IFT 确实能够小幅但不可忽视地提升准确率。然而，这并不足以改变准确率曲线的整体形状，也无法扩大模型在该任务上能够达到非零准确率的上下文长度范围。

在不同于训练时的缩放系数下进行零样本评估

上图展示了在评估时尝试使用与模型训练时不同的缩放系数（用于线性插值）的各种实验。绿色曲线表示对一个基础模型（在 2048 个 token 的上下文上训练）应用某个缩放系数的情况。这种方法确实将非零准确率的范围从 2048 个 token 扩展到了 4096 个 token，但整个范围内的准确率都很低。不过，总体而言，一旦模型已经使用大于 0 的缩放系数进行训练，它似乎可以在评估时以零样本的方式很好地适应更大的缩放系数——从而大幅提高其能够处理的连贯上下文长度范围（例如，对比上方两张图中“训练=4，评估=8”在 16000 个 token 的上下文长度下仍为非零，而在任何超过 8000 个 token 的情况下均为零）。然而，这样做也会导致准确率下降，尤其是在“训练=16，评估=32”的情况下。

“训练=16，评估=12”的运行是我们所见过的非零准确率持续时间最长的案例。它在约 20000 个 token 的上下文长度下仍保持非零分数。

WikiQA 评估

在下表中，两个模型均以 scale=4 的尺度进行评估。然而，“无缩放”模型并未在 scale > 1 的情况下进行微调（即未经历任何训练）。而 scale=4 的模型则确实在该扩展尺度下接受了微调。

存在准确性：

上下文长度	在 FFQA 上使用 scale=4 进行 IFT	在 FFQA 上无缩放进行 IFT	在 AltQA 上使用 scale=4 进行 IFT	在 AltQA 上无缩放进行 IFT
2048	0.3233	0.2217	0.7281	0.2982
4096	0.3783	0.2467	0.7018	0.2829
8192	0.4434	0.2406	0.6582	0.2401
16384	0.3933	0.0	0.5363	0.0

注：对于 16k 的上下文长度，我们在推理时使用了 8 倍的缩放因子。这使得原始的 2k 上下文能够扩展到 2*8=16k。值得注意的是，尽管缩放后的模型是以 4 倍的缩放因子进行训练的，但它在推理时仍能零样本地插值到 16k（即 8 倍的缩放），且性能损失不大。然而，这一点在未缩放的模型中并不成立，从 16k 数据点上准确率降至 0 即可看出。这表明我们的缩放和上下文长度插值方法确实有效。

输入上下文长度统计

如前所述，我们对文档进行了截断和修改，以生成不同版本的 WikiQA 数据集。每个版本旨在充分测试模型在特定上下文长度下的表现，具体长度由版本名称指示。

FFQA

	平均上下文长度	最大上下文长度
ffqa_2k.json	1936.71	3228
ffqa_4k.json	3805.06	5793
ffqa_8k.json	7598.98	9963
ffqa_16k.json	15000.54	16178

AltQA

	平均上下文长度	最大上下文长度
altqa_2k.json	1953.73	2698
altqa_4k.json	3737.39	5172
altqa_8k.json	7481.37	9619
altqa_16k.json	15013.44	16173

性能对不断增加的上下文长度具有鲁棒性

如上所示，我们在 WikiQA 任务上采用的插值嵌入微调技术，似乎能够生成对输入上下文长度不断增加具有鲁棒性的优秀模型。我们在该任务的两个版本上都展示了这一点。由于我们是以 4 倍的上下文长度进行微调的，因此预计准确率不会在输入大小达到 4*2048=8192 之前下降。即便超过这个限制，我们仍然可以看到一些合理的性能。这似乎是 RoPE 嵌入的周期性特征所致，它使得某些特性可以外推到超出缩放上下文所设定的限制之外的位置。

缩放上下文的影响

我们对比了使用和不使用缩放上下文进行指令微调的模型，结果表明，使用缩放上下文进行 IFT 能显著提升模型性能。需要注意的是，对于这两个模型，我们在评估时仍然使用了缩放上下文（scale=4）。有趣的是，即使是在零样本情况下，缩放后的 RoPE 嵌入也能给出非平凡的准确率。然而，显式地对嵌入进行微调确实带来了可观的收益。我们看到，在 FFQA 上几乎实现了 2 倍的提升，而在 AltQA 上则实现了 2.5 倍的提升，且这一优势在所有由缩放上下文因子插值出的位置上都得以体现。

信息的位置

损失曲线

我们在概述中描述的所有实验中都训练了模型。但并非所有实验都显示出进行全面评估的前景。 部分实验因损失曲线表现不佳而被放弃。在某些情况下，我们还发现实验结果并不总是与训练过程中观察到的损失一致。

以下图片展示了我们运行的部分实验的曲线：

例如，XPOS 的损失从未收敛到与其他运行中观察到的损失水平。起初，我们怀疑是 fp16 的精度不足以处理 XPOS 系数。于是我们将实现调整为使用 fp32 来进行核心注意力点积计算。这确实改善了收敛性，但仍不足以使损失与其它模型持平。我们的假设是，XPOS 与基础位置嵌入差异过大，难以在其基础上进行微调。这有些令人意外，因为 XPOS 实际上只是带有相对差异函数缩放因子的 RoPE。我们曾启动过一项尚未完成的实验，即从缩放因子 1.0 开始，逐步通过多次迭代过渡到 XPOS 函数。

Long-Context 快速上手指南

本指南旨在帮助开发者快速使用 AbacusAI 开源的 Long-Context 项目，将 Llama 系列模型的上下文窗口扩展至 16k 甚至更长。该项目通过线性缩放（Linear Scaling）结合指令微调（IFT），显著提升了模型在长文本任务中的表现。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 20.04+) 或 macOS。
• Python 版本: Python 3.8 或更高版本。
• GPU 要求: 建议使用支持 CUDA 的 NVIDIA GPU。
  • 运行 13B 模型进行推理：建议显存 ≥ 24GB (如 RTX 3090/4090, A10)。
  • 进行微调训练：建议多卡环境或高显存专业卡 (如 A100/H100)。
• 前置依赖:
  • PyTorch (与您的 CUDA 版本匹配)
  • Transformers (Hugging Face)
  • Accelerate
  • PEFT (用于 LoRA 微调)
  • Datasets

安装基础依赖命令：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install transformers accelerate peft datasets sentencepiece protobuf

提示：国内用户可使用清华或阿里镜像源加速安装： pip install -i https://pypi.tuna.tsinghua.edu.cn/simple <包名>

安装步骤

本项目主要提供训练脚本、评估脚本以及预训练好的模型权重。您可以直接克隆仓库并安装必要的 Python 包。

1. 克隆仓库

git clone https://github.com/abacusai/Long-Context.git
cd Long-Context

2. 安装项目特定依赖（如有 requirements.txt）

如果仓库根目录包含 requirements.txt，请执行：

pip install -r requirements.txt

3. 验证安装

确保能够导入关键库并无报错：

python -c "import transformers; import peft; print('Environment ready')"

基本使用

本项目最核心的用法是直接加载已发布的 Scale 16 模型权重，该模型在 16k 上下文长度下表现最佳。

1. 加载预训练模型 (推理示例)

以下是最简单的使用示例，展示如何加载 Giraffe-v1-delta-13b-scaled-16 模型并进行长文本推理。

注意：该模型发布为 Delta 权重（基于 Llama-13b 的增量），加载时通常需要合并基础模型，或者直接使用支持加载 delta 权重的最新 Transformers 版本。以下代码假设您已拥有基础 Llama-13b 模型或直接通过 HuggingFace 加载整合后的模型（如果可用）。

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# 模型标识符
model_name = "abacusai/Giraffe-v1-delta-13b-scaled-16"

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-13b") # 需指定基础 Llama 模型或使用整合版

# 加载模型
# 注意：如果是 Delta 权重，可能需要先合并权重到本地，或使用特定的加载逻辑
# 此处演示标准加载流程，具体取决于 HF 仓库的当前配置
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)

# 准备长文本输入 (模拟 10k+ tokens 的场景)
long_context = "..." * 5000 + "请问这篇文章的核心观点是什么？"

inputs = tokenizer(long_context, return_tensors="pt").to(model.device)

# 生成回答
outputs = model.generate(
    **inputs,
    max_new_tokens=256,
    do_sample=True,
    temperature=0.7,
    pad_token_id=tokenizer.eos_token_id
)

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)

2. 使用评估脚本测试长上下文能力

项目提供了专门的脚本来评估模型在长上下文任务（如 WikiQA）上的表现。

运行 WikiQA 推理与指标计算：

# 进入评估目录
cd python/eval/longeval

# 运行推理脚本 (需根据实际参数调整 model_path 和 scale)
python run_inference_WikiQA.py \
    --model_path abacusai/Giraffe-v1-delta-13b-scaled-16 \
    --scale 4 \
    --output_dir ./results

# 计算评估指标
jupyter nbconvert --to script compute_metrics_WikiQA.ipynb
python compute_metrics_WikiQA.py

3. 数据集获取

为了复现实验或进行自定义测试，您可以直接从 HuggingFace 下载作者整理的长上下文数据集：

• Extended LMSys Dataset: abacusai/LongChat-Lines
• WikiQA (Free Form): abacusai/WikiQA-Free_Form_QA
• WikiQA (Altered Numeric): abacusai/WikiQA-Altered_Numeric_QA

下载示例 (使用 datasets 库):

from datasets import load_dataset

# 加载长文本问答数据集
dataset = load_dataset("abacusai/WikiQA-Free_Form_QA")
print(dataset["train"][0])

通过以上步骤，您即可快速上手 Long-Context 工具，体验在 16k+ 上下文长度下的 LLM 应用能力。