面向将语言模型上下文扩展至 128K 的数据工程

我们研究了将语言模型上下文长度扩展至 128K 的持续预训练方案，重点在于数据工程。我们假设长上下文建模，特别是\textit{利用任意输入位置信息的能力}，是一种主要通过大规模预训练获得的能力，并且通过对适当数据混合进行轻量级持续预训练，可以将此能力轻松扩展到远长于训练期间所见（例如从 4K 到 128K）的上下文。我们调查了持续预训练数据的\textit{数量}和\textit{质量}：(1) 就数量而言，我们表明 5 亿到 50 亿个 token 足以使模型能够在 128K 上下文内的任何位置检索信息；(2) 就质量而言，我们的结果同样强调\textit{领域平衡}和\textit{长度上采样}。具体而言，我们发现 naive 地对某些领域（如书籍）的较长数据进行上采样（这是现有工作的常见做法）会导致次优性能，而平衡的领域混合至关重要。我们证明，在此类数据的 10 亿至 50 亿 token 上对完整模型进行持续预训练，是将语言模型上下文长度扩展至 128K 的一种有效且负担得起的策略。我们的方案优于强大的开源长上下文模型，并缩小了与 GPT-4 128K 等前沿模型的差距。