How to Train Long-Context Language Models (Effectively)

发表时间: 2024-10 · arXiv:2410.02660 (Princeton NLP, ACL 2025)

作者/机构: Tianyu Gao∗, Alexander Wettig∗, Howard Yen, Danqi Chen, Princeton Language and Intelligence, Princeton University

A1 主要贡献

本文研究了如何通过持续预训练（continued pre-training）和监督微调（supervised fine-tuning, SFT）来有效利用长上下文信息，从而训练出强大的长上下文语言模型。

• 研究目标与核心问题：尽管许多工作致力于用最少的训练来扩展预训练语言模型的上下文长度，但这些方法在简单的“大海捞针”（Needle-in-a-Haystack, NIAH）任务上都表现不佳。要稳健地学习这类任务，必须在数十亿词元（token）的长文档上继续训练。本文旨在系统性地研究这一持续训练和SFT过程中的各个设计决策，以找到一条高效且有效的路径。

• 关键创新与贡献：

  1. 建立可靠的评估协议：文章指出，单纯依赖困惑度（perplexity）或简单的NIAH测试来指导模型开发是不可靠的。作者主张使用一套广泛的、能反映实际应用的长上下文下游任务（如RAG、长文档摘要、多示例上下文学习等）进行评估。一个关键的洞见是，评估应在SFT之后进行，因为这能更清晰地揭示模型的长上下文能力，某些性能提升仅在SFT后才显现。同时，评估也必须检查模型是否保持了原有的短上下文性能。

  2. 全面的数据工程探索：通过系统的消融实验，文章深入研究了训练数据的构成。发现：

     • 长数据源：代码库和长篇书籍是优质的长上下文数据来源。
     • 数据混合：仅使用长数据会损害长上下文性能。将长数据（约60%）与高质量的短上下文数据（约40%）混合至关重要，这既能提升长上下文任务表现，又能保持模型原有的短上下文能力。

  3. 训练策略的创新发现：

     • 超长序列训练：一个令人意外的发现是，以比评估长度更长的序列进行训练，能显著提升模型在较短评估长度上的表现。例如，使用512K的训练长度能提升64K评估长度下的性能。
     • SFT策略：与先前研究主张使用合成长指令数据相反，本文发现仅使用标准的、短上下文的指令数据集（如UltraChat）进行SFT，就足以在长上下文任务上取得优异表现。

  4. 推出ProLong模型：基于以上发现，本文提出了最终的模型ProLong-8B。该模型从Llama-3初始化，仅用40B词元的长上下文数据进行训练（约为Llama-3.1长上下文训练数据量的5%），但在128K长度的多数长上下文任务上超越了Llama-3.1-8B-Instruct。ProLong的有效上下文长度可达512K，是公开可用模型中上下文窗口最长的模型之一。

本文的主要贡献、发现和文章结构在开篇的“Takeaways”部分有清晰的总结。所有代码、数据和模型都已公开。

A3 背景知识/关键Observation/设计原则

强大的语言模型训练，其先决条件是拥有一套能够指导模型开发、同时追踪其在真实世界应用中效用的稳健评估体系。尽管像“大海捞针”（NIAH; 【48, Needle in a haystack - pressure testing LLMs, 2024, Garrett Kamradt】）和RULER（【43, RULER: What’s the real context size of your long-context language models?, 2024, First Conference on Language Modeling, Hsieh et al.】）这类合成基准因其简单和可控性而广受欢迎，但我们更关注能反映实际用途的更广泛任务，例如对整个文档进行推理的能力。下文将描述我们的评估协议，并展示其为何对我们的模型开发至关重要。

2.1 在多样化和现实的任务上进行评估

我们决定使用HELMET作为我们的评估套件，因为它【117, HELMET: How to evaluate long-context language models effectively and thoroughly, 2025, International Conference on Learning Representations (ICLR), Yen et al.】是最全面的长上下文基准之一。为了快速迭代，我们仅使用HELMET任务的一个子集进行模型开发：

• Recall（回忆）：给定一个包含随机键值对的JSON文件，检索某个键对应的值。
• RAG（检索增强生成）：给定检索到的维基百科文档，回答一个问题（数据集包括NQ, HotPotQA, PopQA）。
• Re-ranking（重排序）：从一个打乱的文档列表中生成前10的排序（数据集为MSMARCO）。
• ICL（上下文学习）：从多个上下文示例中学习分类任务，类别数从6到151不等；在5个数据集（TREC coarse/fine, NLU, Banking77, Clinc-150）上取平均值。
• QA（问答）：给定一本完整的书籍，回答一个问题（数据集为NarrativeQA）。
• Summarization（摘要）：对长篇法律文件进行摘要（数据集为Multi-LexSum）。
  总的来说，这些多样化的任务反映了一系列长上下文能力，包括回忆、推理、从上下文中学习以及对噪声输入的鲁棒性。Yen等人【117, HELMET: How to evaluate long-context language models effectively and thoroughly, 2025, International Conference on Learning Representations (ICLR), Yen et al.】还表明，与其他长上下文基准不同，HELMET产生的模型性能趋势与人类感知更为一致。我们使用HELMET中未参与开发的其余任务来评估最终模型的泛化能力，并在附录B.8中报告了在其他长上下文基准（如RULER【43】和∞Bench【121】）上的最终性能。

一个稳健的评估套件的重要性在表1中得以展示。作为我们工作的前身，Fu等人【34, Data engineering for scaling language models to 128k context, 2024, International Conference on Machine Learning (ICML)】在模型开发中仅考虑了NIAH和困惑度；尽管他们的模型在NIAH上取得了完美分数，但在HELMET的3个任务上的评估揭示了其模型的重大缺陷。我们还看到，对于强大的模型（Llama-3.1-8B vs. 70B），NIAH甚至HELMET的Recall任务都变得饱和，而其他任务类别则能继续检测出它们长上下文能力的差异。

Table 1: HELMET提供了一个更全面的长上下文评估。我们在Llama-3-8B上复现了Fu等人 (2024) 的工作并进行了SFT。我们报告的是Llama的instruct版本。

我们在附录A.1中提供了关于HELMET评估的更多细节，包括其对指标的谨慎选择。如无特别说明，我们将每个类别的性能在所有数据集以及32K和64K的评估长度上取平均；对于最终的长上下文分数，我们对所有类别进行宏平均。

为什么不使用困惑度（Perplexity）？除了合成的回忆任务，许多先前的工作依赖困惑度（PPL）来评估语言模型的长上下文扩展【15, Extending context window of large language models via positional interpolation, 2023, Chen et al.】、【34, Data engineering for scaling language models to 128k context, 2024, International Conference on Machine Learning (ICML), Fu et al.】、【67, A controlled study on long context extension and generalization in llms, 2024, Lu et al.】，这通常在PG19书籍数据集【93, Compressive transformers for long-range sequence modelling, 2020, International Conference on Learning Representations, Rae et al.】上测量。我们使用第3.2节的消融实验来展示为什么困惑度不是一个用于开发长上下文模型的指示性指标。该实验研究了长文档比例对性能的影响。我们报告了我们的评估结果和在PG19的64K长度文档的最后32K词元上测量的困惑度。如图1所示，虽然使用更多的长数据持续改善PPL，但很明显，使用100%的长数据会显著损害下游的长上下文性能。

Figure 1: 基于困惑度（PPL）做出的设计决策对于长上下文下游任务并非最优。

2.2 在监督微调后进行评估

监督微调（SFT）是一个额外的训练阶段，它在少量自然语言指令和相应回复上微调模型【84, Training language models to follow instructions with human feedback, 2022, Advances in Neural Information Processing Systems (NeurIPS), Ouyang et al.】；这使得基础语言模型能够以聊天形式处理用户查询，并已成为生产前沿语言模型的标准步骤。在这里，我们考虑在SFT前后评估模型的差异。

在初步实验中，我们使用Fu等人(2024)的数据混合的5B词元子集继续训练Llama-3-8B-Base。该混合基于SlimPajama【104, SlimPajama: A 627B token cleaned and deduplicated version of RedPajama, 2023, Soboleva et al.】，并上采样长文档，使其约占70%的词元，同时保留了原始的领域比例。然后，我们使用UltraChat【30, Enhancing chat language models by scaling high-quality instructional conversations, 2023, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, Ding et al.】对几个中间检查点进行SFT。

我们在图2中展示了SFT前后的基准测试结果。在SFT后进行长上下文评估显示出更清晰的信号：(1) SFT显示，随着训练词元增多，模型在RAG和重排序任务上持续改进，而在SFT前评估时，这种改进不那么明显或不存在。(2) SFT使得在QA和摘要等现实应用上的评估成为可能，这些任务需要指令遵循能力，否则会完全失败。我们还注意到，经过额外的SFT阶段后，两次随机训练运行的方差并没有显著增加。因此，除非另有说明，我们报告的是SFT后的长上下文性能。

Figure 2: 一些长上下文性能的提升仅在SFT后评估时才能观察到。我们报告了两次训练运行的均值和标准差。

我们在第5节深入探讨监督微调，探索不同的训练数据集，以及合成长指令数据的使用。然而，我们发现简单地在UltraChat上进行微调仍然是一个出人意料的有竞争力的选择。

2.3 检查短上下文性能是否得以保持

长上下文能力不应以牺牲短上下文性能为代价，特别是因为短上下文评估涵盖了更广泛的能力，例如世界知识、常识和数学推理。然而，以往的长上下文研究在很大程度上忽略了短上下文评估。我们在Open LLM Leaderboard【10, Open llm leaderboard, 2023, Beeching et al.】的5个任务上报告性能：HellaSwag【119, HellaSwag: Can a machine really finish your sentence?, 2019, Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, Zellers et al.】、MMLU【41, Measuring massive multitask language understanding, 2021, International Conference on Learning Representations, Hendrycks et al.】、ARC-challenge【21, Think you have solved question answering? Try ARC, the AI2 reasoning challenge, 2018, Clark et al.】、WinoGrande【96, Winogrande: An adversarial winograd schema challenge at scale, 2021, Communications of the ACM, Sakaguchi et al.】和GSM8K【22, Training verifiers to solve math word problems, 2021, Cobbe et al.】。我们在SFT之前评估短上下文性能，因为这允许与作为长上下文训练初始化的基础模型进行直接比较。

先前的技术会降低短上下文性能。我们在表2中显示，无论是免训练的位置外推（position extrapolation），还是使用现有的长数据混合（【34, Data engineering for scaling language models to 128k context, 2024, International Conference on Machine Learning (ICML), Fu et al.】）进行微调，都无法保持Llama-3-8B在标准短上下文任务上的强大性能。这促使我们寻找能够保留初始模型强大短上下文性能的数据源。

Table 2: 对Llama-3-8B应用位置外推（PE）（通过改变RoPE频率基数，§B.1）或在长上下文SlimPajama混合数据（Fu et al., 2024; Soboleva et al., 2023）上对其进行微调，都会降低这个顶级预训练语言模型在短上下文任务上的性能。

A2 方法细节

3 长上下文数据策划

训练数据的质量和构成已被发现是语言模型预训练中最重要的因素（【90, The refinedweb dataset for falcon llm: Outperforming curated corpora with web data only, 2023, Advances in Neural Information Processing Systems, Penedo et al.】; 【112, QuRating: Selecting high-quality data for training language models, 2024, Proceedings of the 41st International Conference on Machine Learning, Wettig et al.】; 【61, Datacomp-lm: In search of the next generation of training sets for language models, 2024, Li et al.】），因此也是我们研究的主要焦点。为了做出数据决策，我们进行了消融实验：我们使用最大长度为64K词元的设置，在5B词元上继续训练Llama-3-8B-Base，并根据第2节进行评估。关于我们消融实验设置的更多细节，请参见附录A.4。

我们的目标是提升长上下文任务性能，同时保持原始模型的短上下文性能。我们从数据应该是长短文档混合的直觉出发，分别研究这些选择。在我们的消融实验中，长数据由单个64K词元的文档块组成，而对于短数据，我们通过打包文档来构建批次，直到每个序列达到64K词元。

3.1 代码库和书籍是长上下文数据的良好来源

SlimPajama。我们分析了SlimPajama (SP; 【104, SlimPajama: A 627B token cleaned and deduplicated version of RedPajama, 2023, Soboleva et al.】)中长数据的数量。表3显示，书籍占了长上下文词元的大部分。在检查CommonCrawl (CC)中的长数据时，我们观察到尽管质量参差不齐，它也包含一些类似书籍的内容，未来的工作可以通过数据选择方法来识别这些内容。

Table 3: 按数据源划分的长文本文件（≥64K词元）。

代码库。虽然GitHub中只有少数文件能达到非常长的长度（Singh等人【103, Brevity is the soul of wit: Pruning long files for code generation, 2024】的研究表明这些文件质量也较低），但我们通过将一个仓库中的所有文件连接成单个文档，从The Stack【53, The stack: 3 TB of permissively licensed source code, 2023, Transactions on Machine Learning Research, Kocetkov et al.】构建了一个丰富的长上下文数据源。与Guo等人【40, Deepseek-coder: When the large language model meets programming–the rise of code intelligence, 2024】不同，我们不根据依赖关系对文件进行排序，这应该会增加相关文件之间的距离，减少近期偏见。

数据混合。我们用60%的长上下文数据和40%的ShortMix（§3.3）来训练模型。表4显示，仅使用代码库在压力测试性的回忆任务上表现最佳。同时，书籍对于上下文学习、摘要和重排序任务更有广泛的益处。书籍和代码库的等量混合实现了最佳的整体性能。请注意，由于我们使用了高质量的短数据混合，短上下文任务的性能保持稳定。

Table 4: 不同长数据源的影响。长上下文性能是在32K和64K长度上取平均值。

3.2 仅在长数据上训练会损害长上下文性能

短/长数据之间的比例是影响下游性能的另一个关键因素。先前的工作要么只在长数据上训练（【91, YaRN: Efficient context window extension of large language models, 2024, The Twelfth International Conference on Learning Representations, Peng et al.】），要么添加一些短训练数据（【116, Long-context language modeling with parallel context encoding, 2024, Association for Computational Linguistics (ACL), Yen et al.】; 【34, Data engineering for scaling language models to 128k context, 2024, International Conference on Machine Learning (ICML), Fu et al.】）。然而，我们是第一个系统研究短/长数据比例影响的工作。

图3显示，随着长数据比例的增加，短任务性能单调下降。长上下文任务的趋势因任务而异，并因SFT而变得更加复杂：在像回忆和RAG这样的任务上，SFT前的性能偏好高比例的长数据，而SFT后的性能随着长数据的增多而急剧恶化。我们假设，仅在长数据上专门训练模型，使其成为通用SFT的一个糟糕的初始化——这凸显了在SFT后评估检查点的重要性（§2.2）。虽然一些长上下文任务（如ICL）从更多的长数据中持续受益，或者（如重排序）没有显示出清晰的模式，但最佳的平均性能是在60%长数据和40%短数据的比例下实现的，我们最终的ProLong模型也采用了这一比例。

Figure 3: 短/长数据比例的影响。所有模型都在书籍/代码库长数据和我们的ShortMix上训练了5B词元。增加长数据最初能提升长上下文性能，但随后会造成损害。更多的长数据也持续降低短上下文性能。

3.3 选择高质量的短上下文混合数据至关重要

我们在§2.3中看到，在长上下文微调期间很难保持Llama-3-8B在短上下文任务上的强大性能。我们采用我们最佳的长上下文设置（书籍/代码库数据和60%长/40%短的比例），并研究不同短上下文训练混合的影响。我们实验了SlimPajama【104, SlimPajama: A 627B token cleaned and deduplicated version of RedPajama, 2023, Soboleva et al.】、FineWeb-Edu【90, The refinedweb dataset for falcon llm: Outperforming curated corpora with web data only, 2023, Advances in Neural Information Processing Systems, Penedo et al.】、DCLM-Baseline【61, Datacomp-lm: In search of the next generation of training sets for language models, 2024, Li et al.】以及我们自己的ProLong ShortMix。我们的ShortMix受到了MiniCPM【44, Minicpm: Unveiling the potential of small language models with scalable training strategies, 2024, Hu et al.】和Dolma-1.7【105, Dolma: an open corpus of three trillion tokens for language model pretraining research, 2024, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Soldaini et al.】中“第二阶段训练”的启发，这些工作在预训练末期使用更多知识密集型、与下游任务相关的数据。表5显示了我们ShortMix的构成。

Table 5: 我们的ShortMix。

表6表明，短数据部分对短上下文和长上下文下游性能都有重大影响。我们精心策划的ShortMix在短上下文和长上下文任务上都优于其他短数据源，并且我们的数据领域对于保持Llama-3-8B在数学推理上的性能尤为重要。令人惊讶的是，我们发现仅使用FineWeb-Edu——一个为帮助像MMLU这样的知识密集型任务而策划的数据集——作为短上下文组件表现不佳，我们在ShortMix中将其与更多样化的数据源结合。DCLM-Baseline在除GSM8K外的所有短上下文任务上都表现良好。这很可能可以通过与数学相关的数据集结合来改善，但由于我们在项目结束时才添加DCLM-baseline消融实验，我们把这个探索留给未来的工作。

Table 6: 不同短数据源的影响。长上下文性能是6个类别在32K和64K长度上的平均值。

与先前工作的比较。为了证实我们长上下文数据策划的有效性，我们在§B.6中与Fu等人【34, Data engineering for scaling language models to 128k context, 2024, International Conference on Machine Learning (ICML)】的先前工作进行了直接比较。我们的数据混合在长上下文和短上下文任务上都显著优于Fu等人（2024），这突显了ProLong的有效性。

Figure 4: 我们的ProLong模型在整个训练过程中的性能（32K和64K的平均值）。

4 扩展训练数据的规模和长度

在常规预训练中，增加训练步数众所周知能改善下游任务，但在长上下文持续训练的背景下，相关的分析很少。我们整合了消融实验的经验，得出了ProLong的配方，具体在§6中详细描述。值得注意的是，我们扩大了训练预算，使用了更长的序列（最高512K）和更多的词元（在最大序列长度为64K时训练20B词元，另外在512K时训练20B词元）。当从64K切换到512K上下文长度时，我们重置了学习率调度并增加了RoPE频率基数。本节中，我们分析这些决策的影响。

增加步数有帮助。在图4中，我们绘制了40B词元训练运行中，中间检查点的下游性能。虽然长上下文性能在整个训练过程中有所波动，但我们观察到在回忆、RAG、重排序和摘要任务上存在积极趋势。对于短上下文任务，我们观察到平均性能最初从初始化时有所下降，但逐渐恢复。当从64K切换到512K序列长度时，性能再次下降，但随着额外的训练也得以恢复。

将训练长度增加到评估长度之外有帮助。人们可能认为，我们应该在模型期望支持的最大序列长度上训练长上下文模型。许多工作甚至强调在推理时外推到更长的序列（【92, Train short, test long: Attention with linear biases enables input length extrapolation, 2022, International Conference on Learning Representations (ICLR), Press et al.】; 【114, Efficient streaming language models with attention sinks, 2024, The Twelfth International Conference on Learning Representations, Xiao et al.】; 【113, InfLLM: Unveiling the intrinsic capacity of LLMs for understanding extremely long sequences with training-free memory, 2024, Xiao et al.】; 【116, Long-context language modeling with parallel context encoding, 2024, Association for Computational Linguistics (ACL), Yen et al.】; 【15, Extending context window of large language models via positional interpolation, 2023, Chen et al.】）。相反，我们观察到，在更长的序列长度（512K词元）上训练，能显著改善在较短评估长度（64K词元）上的长上下文性能。

我们通过以下方式证实了这一点：我们从一个在64K长度上训练了20B词元的模型进行初始化，然后(1) 继续在64K长度上训练，或者(2) 切换到512K长度的训练。我们在两个实验中使用相同的超参数和数据混合。我们在64K的评估长度上评估了经过4B训练词元后的检查点。比较表7中的两次运行，我们看到从切换到512K训练长度中获得了持续的增益。

Table 7: 在不同序列长度上训练模型的影响。所有结果都在64K序列长度下评估。我们看到，在超出评估上下文窗口的最大长度下进行训练，会持续改善长上下文性能。

在比评估时使用的序列更长的序列上训练有什么好处？这是我们的假设：假设为了解决某个特定任务（例如，回忆），模型必须在跨越精确距离d个词元的依赖关系的例子上进行训练，即不同长度的依赖关系之间没有泛化。再假设这些依赖关系在序列的任何位置出现的可能性都相等。那么，一个长度为nd的文档将比n个长度为d的文档多出(n-1)(d-1)个距离为d的依赖关系。虽然这些假设在实践中并不成立，但这个简化的模型仍然为我们的经验发现提供了直觉。

5 用于长上下文语言模型的监督微调

在本节中，我们研究如何最好地使长上下文语言模型能够遵循指令。我们专注于在指令数据集上进行监督微调（【84, Training language models to follow instructions with human feedback, 2022, Advances in Neural Information Processing Systems (NeurIPS), Ouyang et al.】），并将强化学习和偏好优化留给未来的工作。

本节中的所有实验都使用ProLong基础模型，该模型在最大序列长度为512K的情况下训练了40B词元。相比之下，开源的指令数据非常短，例如UltraChat【30, Enhancing chat language models by scaling high-quality instructional conversations, 2023, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, Ding et al.】的对话平均只有1.2K词元，最长也只有4.1K词元。为了弥合这一差距，一些工作（【115, Effective long-context scaling of foundation models, 2023, Xiong et al.】; 【32, The Llama 3 herd of models, 2024, Dubey et al.】; 【116, Long-context language modeling with parallel context encoding, 2024, Association for Computational Linguistics (ACL), Yen et al.】）提出以合成方式生成长指令数据。

我们考虑了三个流行的SFT数据集——UltraChat【30】, Tulu-v2【45, Camels in a changing climate: Enhancing lm adaptation with Tulu 2, 2023, Ivison et al.】和ShareGPT——以及三种合成数据来源：对于合成QA，我们提示Llama-3-8B-Instruct根据长文档中的一个随机块生成问答对；对于合成RAG，我们复用这些QA对，但呈现文档中的随机块列表以模拟检索到的段落；对于合成摘要，我们通过递归摘要【109, Recursively summarizing books with human feedback, 2021, Wu et al.】为长篇书籍生成摘要。对于所有合成数据，我们编写了几个模板，并随机抽样以增加多样性。更多细节可以在§A.5中找到。在我们的合成指令数据集中，我们始终使用40%合成QA、30%合成RAG和30%合成摘要的组合。指令调优实验的超参数可以在表9中找到。

短上下文指令数据产生了强大的长上下文结果。我们首先在表22中确定UltraChat优于Tulu-v2和ShareGPT。因此，在研究表8中合成长上下文指令数据的比例时，我们使用了UltraChat。令人惊讶的是，我们发现添加合成数据并没有改善这些超长上下文任务的性能，在我们的设置中，即使添加少至1%的合成数据也会损害性能。我们还验证了即使使用更强大的数据生成器，如Llama-3-70B，这种现象仍然存在（§B.5）。因此，我们最终的ProLong模型的SFT仅使用短上下文的UltraChat数据。

Table 8: 不同比例的合成SFT数据（与UltraChat混合）的效果。我们报告了32K和64K平均性能，标有†的任务除外，这些任务在512K下进行压力测试评估。百分比基于词元数，而非样本数。

为什么我们关于合成数据的结论与先前的工作不同？我们提出以下假设：(1) 先前的工作，如Xiong等人【116, Long-context language modeling with parallel context encoding, 2024, Association for Computational Linguistics (ACL)】和Bai等人【7, Longalign: A recipe for long context alignment of large language models, 2024, Bai et al.】，可能长上下文训练不足，合成数据起到了额外的长上下文训练数据的作用。(2) 与Llama-3.1【32, The Llama 3 herd of models, 2024, Dubey et al.】使用的私有指令数据相比，我们的指令数据集要小得多——有可能当使用一个广泛的短指令数据集时，混入合成长数据可以避免模型在长上下文任务上性能退化。

Table 9: ProLong的训练配方。

6 ProLong模型：配方与结果

6.1 最终配方

我们在表9中总结了ProLong的训练配方。我们的最终模型从Llama-3-8B-Instruct模型开始，在64K序列长度上训练了20B词元。然后，它在512K序列长度上进一步训练了20B词元（ProLong base），这是我们通过序列并行【62, Sequence parallelism: Long sequence training from system perspective, 2023, Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Li et al.】实现的。我们通过在UltraChat上对基础模型进行SFT来获得最终的ProLong模型。我们的消融实验和最终模型在数据混合方面有一个小差异，即我们在最终模型中混入了3%的高质量教科书（【19, Language models as science tutors, 2024, Forty-first International Conference on Machine Learning, Chevalier et al.】），因为类似书籍的数据被证明对长上下文有益（§3.1），而且教科书具有高度的教育意义。这也稍微改变了ShortMix的比例。您可以在附录中找到关于我们数据处理（§A.2）和训练技术栈（§A.3）的更多细节。

接下来，我们详细阐述配方中几个经过仔细消融实验的设计选择。

RoPE频率基数调整。我们发现，改变RoPE【106, Roformer: Enhanced transformer with rotary position embedding, 2021, Su et al.】的频率基数以实现位置外推（【115, Effective long-context scaling of foundation models, 2023, Xiong et al.】; 【33, Dynamically scaled rope further increases performance of long context llama with zero fine-tuning, 2023, emozilla】）能显著提高长上下文性能，即使在大量训练后也是如此。§B.1展示了我们关于最佳RoPE基数的消融实验。虽然原始的Llama模型使用10^5的RoPE基数，但我们在64K设置中使用8 × 10^6，在512K设置中使用1.28 × 10^8。

禁用跨文档注意力。Ding等人【29, Fewer truncations improve language modeling, 2024, Forty-first International Conference on Machine Learning, Ding et al.】表明，在文档边界处屏蔽注意力可以提高模型性能，Llama-3的预训练【32, The Llama 3 herd of models, 2024, Dubey et al.】也使用了这一点。在§B.2中，我们展示了在持续训练中禁用跨文档注意力对短上下文和长上下文性能都有益。禁用跨文档注意力还可以带来更高的训练吞吐量，我们在§A.3中有更详细的描述。

从Llama-3-8B-Instruct开始。虽然我们所有的长上下文训练消融实验都使用Llama-3-8B的基础模型，但我们使用Llama-3-8B-Instruct作为最终ProLong模型的初始化。§B.3显示，虽然Llama-3-8B-Instruct对长上下文性能的提升很小，但它显著增强了短上下文性能。

A4 实验环境

• 数据集

  • 长上下文评估：主要使用HELMET基准的一个子集，覆盖Recall（JSON KV）、RAG（NQ, HotPotQA, PopQA）、Re-ranking（MSMARCO）、ICL（TREC, NLU, Banking77, Clinc-150）、QA（NarrativeQA）和Summarization（Multi-LexSum）等任务。最终评估还使用了RULER、∞Bench和NoCha基准。
  • 短上下文评估：使用Open LLM Leaderboard的5个任务：HellaSwag, MMLU, ARC-challenge, WinoGrande, 和GSM8K。
  • 训练数据源：长上下文数据主要来自The Stack（代码库）和SlimPajama（书籍）。短上下文数据（ShortMix）混合了FineWeb、FineWeb-Edu、OpenWebMath、Dolma（维基百科）、SlimPajama（StackExchange, ArXiv）和高质量教科书。
  • SFT数据：主要使用UltraChat，并与Tulu-v2、ShareGPT以及自生成的合成长指令数据进行了比较。

• 模型架构

  • 基础模型：从Llama-3-8B（8K上下文）和Llama-3-8B-Instruct初始化。
  • 最终模型：ProLong-8B，上下文窗口扩展至512K。
  • 关键参数：采用RoPE位置编码，并对频率基数进行了调整（64K长度下为8×10^6，512K长度下为1.28×10^8）。采用AdamW优化器，学习率等超参数见原文Table 9。

• 硬件配置

  • GPU：在8个Nvidia H100 GPU上使用FSDP（Fully Sharded Data Parallel）进行训练。
  • 互联：在512K序列长度的训练中，使用了基于节点内8个GPU的序列并行。

• 软件配置

  • 实现框架：使用PyTorch和Hugging Face Transformers库。
  • 依赖库：使用mosaic-streaming进行数据加载和混合，FlashAttention-2实现高效注意力，DeepSpeed-Ulysses实现序列并行。
  • 评估工具：长上下文评估使用HELMET框架，短上下文评估使用lm-eval-harness。

A4 实验结果

ProLong模型最终性能

• HELMET基准测试：在128K上下文长度的HELMET评估中（Table 10），ProLong在10B参数规模的模型中表现最佳。值得注意的是，ProLong在除摘要外的所有类别上都优于Llama-3.1-8B-Instruct，而其长上下文训练数据量仅为后者的5%（40B vs 800B tokens）。
• 超长上下文能力：为了展示ProLong的512K上下文能力，在QA和摘要任务上进行了从32K到512K的压力测试（Table 11）。结果显示，随着上下文窗口的增长，ProLong的性能持续提升。
• NoCha基准测试：在NoCha这个小说事实核查数据集上（Table 12），ProLong在10B规模的模型中表现突出。在极长实例（>180K tokens）上取得了最高准确率；在<75K tokens的实例上，ProLong是唯一一个表现优于随机猜测（25%）的模型，显著超越了其他同类模型，包括Llama-3.1。
• 其他基准：在RULER和∞Bench基准上（Table 26），尽管这些基准存在一些评估上的局限性，ProLong在10B规模模型中仍然在∞Bench上取得了最佳性能。
• 短上下文性能保持：ProLong在长上下文训练后，很大程度上保持了Llama-3-8B-Instruct原有的短上下文性能（Table 25），证明了其数据混合策略的有效性。

Table 10: 我们在HELMET (Yen et al., 2025) 上128K上下文长度的主要评估结果。对于所有模型，我们都使用相应的指令版本。ProLong是性能最好的10B规模语言模型。完整结果集可在§C中找到。RULER和∞Bench的结果可在§B.8中找到。

Table 11: ProLong在512K下的表现。

Table 12: NoCha基准测试结果 (Karpinska et al., 2024)。ProLong是唯一在<75K类别中取得高于随机性能的模型，并且持续击败Llama-3.1。与原始NoCha排行榜不同，我们报告了所有测试实例的平均准确率，未根据模型的上下文窗口长度过滤测试样本。

关键消融实验结论

• 评估指标：实验证明，困惑度（PPL）并非一个好的长上下文性能指标。使用100%长数据能持续降低PPL，但却严重损害了下游任务的实际性能（Figure 1）。
• 数据混合比例：系统研究了短/长数据比例的影响。发现随着长数据比例增加，短上下文性能单调下降。对于长上下文任务，最佳平均性能出现在长数据占60%、短数据占40%的混合比例下（Figure 3）。
• 长数据源：比较了不同长数据源的效果。代码库在回忆任务上表现最佳，书籍在ICL、摘要和重排序上更有益。两者的等量混合取得了最佳的整体性能（Table 4）。
• 短数据源：高质量的短数据混合（ProLong ShortMix）对于保持短上下文性能和提升长上下文性能至关重要，优于SlimPajama、FineWeb-Edu等其他来源（Table 6）。
• 训练长度：一个关键发现是，使用比评估长度更长的序列进行训练（例如用512K长度训练）能显著提升在较短评估长度（例如64K）上的性能（Table 7）。
• SFT策略：令人意外的是，仅使用标准的短上下文指令数据（UltraChat）进行SFT，其效果优于或等于混合了合成长指令数据的方案。增加合成数据甚至可能损害性能（Table 8）。

A7 补充细节

7 相关工作

将现有语言模型适配于长上下文。许多工作探索以最少的训练来扩展LM的上下文窗口，方法包括位置外推（【15, Chen et al., 2023】; 【91, Peng et al., 2024】; 【16, Chen et al., 2024】; 【31, Ding et al., 2024b】; 【64, Liu et al., 2024a】; 【122, Zhang et al., 2024b】; 【124, Zhu et al., 2024】; 【123, Zhao et al., 2024】; 【110, Wu et al., 2024】; 【45, Hu et al., 2024b】）或操纵注意力模式（【16, Chen et al., 2024】; 【114, Xiao et al., 2024b】; 【113, Xiao et al., 2024a】; 【11, Bertsch et al., 2023】; 【47, Jin et al., 2024】）。Yoshida等人【118, Adding recurrence to pretrained transformers for improved efficiency and context size, 2020】、Choromanski等人【20, Rethinking attention with performers, 2021, International Conference on Learning Representations】和Chevalier等人【18, Adapting language models to compress contexts, 2023, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing】则探索了将长上下文压缩成更短形式的想法。然而，Fu等人【34, Data engineering for scaling language models to 128k context, 2024, International Conference on Machine Learning (ICML)】和Lu等人【67, A controlled study on long context extension and generalization in llms, 2024】表明，使用全注意力、应用简单的位置外推并在长文档上微调模型可以达到更强的效果。

Llama 3.1【32, The Llama 3 herd of models, 2024, Dubey et al.】和Jamba【60, Jamba: A hybrid transformer-mamba language model, 2024, Lieber et al.】通过在标准预训练和监督微调之间增加一个长上下文持续训练阶段来获得长上下文能力，我们遵循了这种设置。Fu等人【34, Data engineering for scaling language models to 128k context, 2024, International Conference on Machine Learning (ICML)】研究了这种设置下的数据工程，并认为0.5B词元的领域平衡、长度上采样的数据足以获得长上下文回忆能力——我们表明，如果采用更全面的评估，这是不够的。Xiong等人【115, Effective long-context scaling of foundation models, 2023】、Dubey等人【32】、Lieber等人【60】、Xiong等人【116, Long-context language modeling with parallel context encoding, 2024, Association for Computational Linguistics (ACL)】、An等人【4, Make your llm fully utilize the context, 2024】和Bai等人【7, Longalign: A recipe for long context alignment of large language models, 2024, Bai et al.】也在SFT阶段采用了合成生成的长数据；然而，我们发现在我们的设置中，使用标准的、短上下文的指令数据取得了最好的长上下文结果。

高效的长上下文架构。已有许多努力致力于设计更高效的架构，例如，线性注意力/RNN（【38, Mamba: Linear-time sequence modeling with selective state spaces, 2023, Gu & Dao】; 【24, Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, 2024, Dao & Gu】; 【68, Mega: moving average equipped gated attention, 2022, Ma et al.】; 【107, Retentive network: A successor to transformer for large language models, 2023, Sun et al.】; 【91, RWKV: Reinventing RNNs for the transformer era, 2023, Findings of the Association for Computational Linguistics: EMNLP 2023, Peng et al.】; 【117, Gated linear attention transformers with hardware-efficient training, 2024, Forty-first International Conference on Machine Learning, Yang et al.】）以及替代性注意力架构（【95, Long-range language modeling with self-retrieval, 2023, Rubin & Berant】; 【108, You only cache once: Decoder-decoder architectures for language models, 2024, Sun et al.】; 【116, Long-context language modeling with parallel context encoding, 2024, Association for Computational Linguistics (ACL), Yen et al.】）。然而，它们通常需要从头开始训练，并且许多在长上下文回忆方面存在固有限制（【46, Repeat after me: Transformers are better than state space models at copying, 2024, Jelassi et al.】; 【6, Simple linear attention language models balance the recall-throughput tradeoff, 2024, Forty-first International Conference on Machine Learning, Arora et al.】）。最近的工作探索了混合模型（【111, An empirical study of mamba-based language models, 2024, Waleffe et al.】; 【60, Jamba: A hybrid transformer-mamba language model, 2024, Lieber et al.】）或将现有LM蒸馏为混合模型（【112, The mamba in the llama: Distilling and accelerating hybrid models, 2024, Wang et al.】），并显示出有希望的结果。

长上下文评估。已经提出了许多用于长上下文评估的基准（【98, ZeroSCROLLS: A zero-shot benchmark for long text understanding, 2023, Findings of the Association for Computational Linguistics: EMNLP 2023, Shaham et al.】; 【43, RULER: What’s the real context size of your long-context language models?, 2024, First Conference on Language Modeling, Hsieh et al.】; 【55, LongEval: Guidelines for human evaluation of faithfulness in long-form summarization, 2023, Proceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics, Krishna et al.】; 【121, ∞Bench: Extending long context evaluation beyond 100K tokens, 2024, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Zhang et al.】; 【3, L-eval: Instituting standardized evaluation for long context language models, 2024, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), An et al.】; 【8, LongBench: A bilingual, multitask benchmark for long context understanding, 2024, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Bai et al.】）。也有研究长上下文LM特定方面的工作，如位置偏差（【65, Lost in the Middle: How Language Models Use Long Contexts, 2024, Transactions of the Association for Computational Linguistics, Liu et al.】）、上下文学习（【12, In-context learning with long-context models: An in-depth exploration, 2024, Bertsch et al.】; 【63, Long-context llms struggle with long in-context learning, 2024, Li et al.】）和书籍长度的摘要（【52, FABLES: Evaluating faithfulness and content selection in book-length summarization, 2024, First Conference on Language Modeling, Kim et al.】）。在这项工作中，我们遵循Yen等人【117, HELMET: How to evaluate long-context language models effectively and thoroughly, 2025, International Conference on Learning Representations (ICLR)】的方法，因为它具有多样化的应用覆盖和可靠的评估。

A5 结论

本文研究了如何最有效地对一个短上下文预训练语言模型进行持续预训练和监督式微调，使其具备长上下文能力。我们对每个环节进行了详尽的消融实验，许多发现与现有的实践或观念相悖。我们利用所有这些发现，打造了ProLong，一个新的、达到SOTA水平的长上下文语言模型。我们公开了所有的代码、数据和模型，希望我们的发现能推动长上下文语言模型的研究和应用。

局限性

尽管我们旨在对我们训练配方的主要组成部分进行消融实验，但由于资源限制，我们无法穷尽所有方面，例如优化超参数和额外的数据混合。我们也将研究范围限制在10B参数规模和Llama-3模型上，这可能会限制我们发现和配方的普适性。另一个担忧是我们可能对用于模型开发的任务过拟合了——然而，我们并没有直接在这些数据集上进行训练，而且用基准任务来指导预训练语言模型的开发已成为一种常见做法。我们还表明，我们最终的配方和模型在额外的评估数据集上表现良好，例如NoCha、RULER、∞Bench以及HELMET中的保留数据集。

A6 附录

A 实验细节

A.1 评估

表13详细列出了我们用于长上下文评估的HELMET数据集。值得注意的是，出于效率考虑，我们没有使用HELMET的所有数据集，也不希望对HELMET过拟合。我们在此强调HELMET相比以往基准实现的一些评估协议改进：
- 足够的上下文长度和细粒度控制。HELMET可以评估模型在128K甚至更长上下文长度下的表现。该评估协议还允许在不同长度下报告结果，为开发者提供了针对不同长上下文需求的细粒度控制。
- 更好的合成回忆任务。如HELMET所示，大海捞针任务【48, Needle in a haystack - pressure testing LLMs, 2024, Garrett Kamradt】因其简单性而大多饱和——模型只需在一些不相关内容中找到一个“针”。我们转而使用更具挑战性的JSON KV任务，该任务由Liu等人【65, Lost in the Middle: How Language Models Use Long Contexts, 2024, Transactions of the Association for Computational Linguistics】首次提出并被纳入HELMET，要求模型在一个大型JSON文件中为给定键找到对应的值。
- ICL使用类别平衡的示例和抽象标签。为了将模型从示例中学习的能力与其对任务或数据集标签分布的预训练偏见解耦【85, What in-context learning “learns” in-context: Disentangling task recognition and task learning, 2023, Findings of the Association for Computational Linguistics: ACL 2023, Pan et al.】，HELMET为每个类别采样相同数量的示例，并使用数字标签（1, 2, ...）代替自然语言标签（例如，location, description, ...）。
- 长上下文QA和摘要的基于模型的评估。HELMET没有使用ROUGE等传统指标（这些指标已被证明不能很好地反映真实模型性能：【26, Deutsch & Roth, 2021】; 【27, Deutsch et al., 2022】; 【37, Goyal et al., 2023】; 【14, Chang et al., 2024】），而是使用基于模型的评估来比较参考答案和模型输出。对于QA，HELMET使用GPT-4o根据问题和参考答案对模型输出进行0-3分的评分。对于摘要，HELMET采用了与Zhang & Bansal【120, Finding a balanced degree of automation for summary evaluation, 2021, Empirical Methods in Natural Language Processing (EMNLP)】和Gao等人【36, Enabling large language models to generate text with citations, 2023, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing】类似