符尧：仅靠数据工程我能将LLM的上下文检索能力拓展到128K

符尧：仅靠数据工程我能将LLM的上下文检索能力拓展到128K

大家好，我是HxShine

今天要分享的是一篇由Yao Fu、Rameswar Panda、Xinyao Niu、Xiang Yue、Hannaneh Hajishirzi、Yoon Kim和Hao Peng共同撰写的文章。Title: Data Engineering for Scaling Language Models to 128K Context。该文章**主要通过研究语言模型的数据工程问题，来将模型上下文长度到128K。**其方法也非常简单，作者在前面研究的工作上，详细研究了训练数据的分布情况，通过对训练数据的数量、分布、采样方法、训练长度等方面进行适当的调整，即可将大模型的上下文检索能力成功拓展到128K。可以看出作者在大模型数据工程侧做了非常细致的思考和分析，或许大模型数据工程和算法同等重要，期待LLM更多的能力也可以通过数据工程来拓展和挖掘！

👇关注公众号NLP PaperWeekly，对话框输入“128K”，即可领取上述论文👇

• Title: Data Engineering for Scaling Language Models to 128K Context
• URL: https://arxiv.org/abs/2402.10171
• 代码： https://github.com/FranxYao/Long-Context-Data-Engineering
• Authors: Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, Hao Peng

一、概述

1 Motivation

• 原来的工作例如[1]认为，长期上下文建模是一种新的能力，必须通过大规模训练来“注入”。这导致在基础模型上（例如Llama系列）进行上下文长度拓展，需要大量的数据继续训练才能获得这个能力。

• 作者假设在大规模预训练中通常已经获得了利用任意输入位置信息的长上文建模能力，且这种能力可以通过轻量级的持续预训练与适当的数据混合轻松扩展到训练期间未见过的更长上下文（例如，从4K扩展到128K）。

• LLM的上下文长度非常重要，作者尝试从数据工程 的角度，通过继续预训练来扩展LLM上下文的长度。 主要在数据工程侧探索两个点：1）需要多大的数据量能够使模型检索任何位置的信息。2）数据质量以及数据分布上有什么要求？

2 Methods

本文通过数据工程实验，探索post-training情况下，采用何种数据策略可以提升LLM上下文的长度，主要从post-train的训练上下文长度、数据的均衡性、是否使用上采样这几个方面来考察，最终实验选择的数据工程策略和其他开源版本的数据工程策略对比如下图所示：

作者发现这些差异对于扩展语言模型的上下文长度至128K至关重要，具体差异表现为：

1. 持续预训练数据的长度： 作者使用80K长度的数据进行持续预训练，而Together的工作使用的是32K长度，这限制了模型在超过32K长度的泛化能力。
2. 数据混合： 作者使用SlimPajama数据集，它数据分布比较平衡，与YaRN Mistral使用的仅包含书籍的PG19数据集形成对比。
3. 长度上采样： 作者的方法中包含了对长序列数据的上采样，而LongLoRA没有进行长度上采样。

SlimPajama数据集说明：

• SlimPajama（Soboleva等人，2023年）数据集进行持续预训练。这个数据集是LLaMA（Touvron等人，2023a）预训练数据混合的开源复现版本，由82%的网络数据（67%来自CommonCrawl，15%来自C4），4.5%的代码（Github），4.5%的维基百科，4.5%的书籍，2.5%的Arxiv，以及2.0%的StackExchange组成。

• SlimPajama数据集中，书籍和Github代码是最长的数据来源，其次是Arxiv。像C4和StackExchange这样的网页往往较短。直接上采样长数据会改变领域数据混合比例。 例如，上采样超过100K的序列会增加书籍领域的比例。同样，领域混合比例的变化也会导致长度分布的偏移。

总结:

文章探讨了一种轻量级（只需5B数据）post-training的方法，通过这种方式将模型从较短上下文（如4K）顺利扩展至128K上下文的能力。同时还发现简单地在某些领域（例如书籍）增加更长的训练数据的常见做法性能并不会更好，而一个各领域数据均衡的混合数据非常重要。研究还展示了在10亿至50亿token的这种数据上对完整模型进行持续预训练是增加语言模型上下文长度至128K的一种有效且可行的策略。

3 Conclusion

• 通过适当的数据混合和轻量级持续训练，可以有效地扩大语言模型处理的上下文长度至128K。 该方法在性能上超越了其他强大的开源长上下文模型，并缩小了与如GPT-4 128K这样的前沿模型的差距。

• 持续预训练的数据量和质量要求：
  • 对于数量，实验表明使用5亿至50亿token足以使模型能够在128K上下文内任何位置检索信息；
  • 对于质量，实验结果同样表明领域平衡和长度上采样是非常重要的。

4 Limitation

本文只涉及了持续预训练阶段如何做数据工程，而关于在100K上下文长度的任务上对语言模型进行指令微调的研究还很有限。希望本文的工作能为未来在100K级别的长上下文指令微调研究提供基础。

二、实验

1 实验结论

Table 3展示了作者提出的方法在长上下文任务（Needle-in-a-Haystack 测试） 和短上下文任务（MMLU 测试） 上的性能。这个表格比较了不同的语言模型在这两种任务上的表现，包括非LLaMA模型（如 GPT-4 Turbo 128K 和 GPT-3.5 Turbo 16K） 以及基于LLaMA-2 7B和13B 的模型。从这个表格中，我们可以得出以下结论：

1. 长上下文性能显著提升：LLaMA-2 7B模型在80K上下文长度上达到了88.0%的准确率，这与GPT-4 Turbo 128K的87.1%相当，显示出作者方法的有效性。

2. 短上下文性能保持稳定：LLaMA-2 7B模型在MMLU测试上的准确率为43.3%，这表明模型在处理短上下文任务时仍然保持了良好的性能。

3. 模型泛化能力：作者的方法不仅在训练时使用的上下文长度（80K）上表现良好，而且能够泛化到更长的上下文（128K），这表明模型在处理长距离依赖时具有较好的泛化能力。

4. 与现有模型的比较：作者的方法在长上下文任务上超越了现有的开源长上下文模型，如YaRN Mistral 7B 128K和LongLoRA 7B 100K，并且与前沿模型GPT-4 Turbo 128K的差距显著缩小。

5. 模型大小的影响：作者还展示了基于LLaMA-2 13B的模型在长上下文任务上的性能，这表明模型大小对于提升性能也是有帮助的。

总结: 实验证明提出的方法在长上下文建模方面是有效的，并且能够在不牺牲短上下文性能的情况下，显著提升模型在长上下文任务上的表现。这为未来在长上下文语言模型的研究和应用提供了有价值的参考。

2 大海捞针整体表现随训练数据量的变化情况

Figure 3 在文档中展示了随着训练数据量的增加，模型在Needle-in-a-Haystack测试中的检索性能如何变化，以及与此相关的验证损失（validation loss）的变化。这个图表通过不同数量的训练token（从100M到10B）来展示模型性能的提升情况。

从Figure 3中，我们可以得出以下结论：

1. 数据量对性能的影响：随着训练数据量的增加，模型在Needle-in-a-Haystack测试中的准确率逐渐提高。这表明更多的数据有助于模型更好地学习和泛化到更长的上下文。

2. 性能饱和点：在大约5B（十亿）训练token之后，模型的准确率似乎达到了一个饱和点，进一步增加数据量（如增加到10B）并没有显著提升模型在长上下文任务上的性能。这可能意味着模型已经从数据中学习到了足够的信息来处理128K上下文的任务。

3. 验证损失与性能的关系：随着训练的进行，模型的验证损失逐渐降低，这与模型在Needle-in-a-Haystack测试中性能的提升相一致。这表明验证损失是一个有用的指标，可以用来衡量模型在长上下文任务上的性能。

4. 数据量与性能提升的关系：图表显示，500M（五亿）训练token足以使模型解锁大部分检索准确性。这支持了作者的假设，即模型已经通过大规模预训练获得了精确检索长距离上下文的内在能力，而持续预训练只是通过轻量级的数据扩展这一能力。

总结： 模型性能随着数据量增加而提升的趋势，但是模型在长上下文任务中的性能可能存在一个饱和点，超过这个点，额外的数据量可能不会带来显著的性能提升。

3 消融实验：按数据源的长度进行上采样显著提升了检索能力

Figure 4 在文档中比较了在不同数据策略下，模型在验证损失（validation loss）和Needle-in-a-Haystack测试性能方面的表现。这个图表展示了原始数据混合（Original Datamix）与按源长度上采样（Per-source Length Upsampling）数据混合策略之间的差异。

要点：

1. 损失函数并不完全反映上下文检索能力：原始数据混合（没有进行长度上采样）尽管在验证损失上表现接近，但在Needle-in-a-Haystack测试中的表现不佳。这表明仅使用验证损失作为评估标准可能无法充分揭示模型在长上下文检索任务上的能力。
2. 按源长度上采样直观重要：按源长度上采样的数据混合策略显著提高了模型在Needle-in-a-Haystack测试中的精确检索能力。其可以在不改变领域混合比例的情况下，增加长序列数据的比例。这种方法在保持短上下文性能的同时，有效地提升了长上下文性能。

**结论：**在长上下文语言模型的研究中，数据策略的选择对于模型性能有着显著影响。例如本文发现长度上采样在提升模型在长上下文任务上的表现至关重要，并且这种提升并不以牺牲短上下文任务的性能为代价。

4. 消融实验：均衡的上采样措施至关重要

4.1 不同数据混合策略的数据长度和领域分布情况

作者考虑了以下几种方法：

• 在4K处切割：这种方法将文档截断为4K长度的块。这种方法继承了原始数据混合，并且被大多数预训练工作使用。由于大约30%的文档自然长度超过4K，这种方法破坏了这些自然存在的长距离依赖关系。
• 在128K处切割：这种方法保留了大多数自然存在的长距离依赖关系，而不改变领域混合。LongLoRA遵循这种方法。然而，作者后来展示了仅使用自然存在的长依赖关系是不够的。
• 按源上采样：这种方法保留了领域混合，然后在每个领域内上采样长文档。**这种方法在不改变领域混合的情况下上采样长文档。**作者推荐这种方法，因为实验表明这提供了最平衡的性能提升。
• 全局上采样：这种方法在不考虑它们来源的情况下上采样长文档，从而**稍微改变了领域混合。**Together使用了这种方法。
• 上采样Arxiv/书籍/Github数据：这种方法有意上采样Arxiv/书籍/Github数据，假设这些领域更有可能包含自然发生的长文档。例如，YaRN Mistral和MPT-storyteller上采样了书籍。这种方法改变了领域和长度分布。

Figure 2 通过可视化的方式展示了这些不同的数据策略如何影响数据的长度和领域分布，这对于理解如何通过数据工程来扩展语言模型的上下文长度至128K至关重要。

4.2 均衡性采样消融实验结果

Table 5 这个表格展示了在不同领域（如C4、CC、Stack Exchange、Wiki、Books、Github和Arxiv）中，通过改变原始数据混合来实现长上下文数据混合的结果，以及这些变化对验证集损失的影响。持续预训练实验都采用7B LLaMA-2模型时，使用5B（五十亿）token的数据。

从Table 5中，我们可以得出以下结论：

1. 领域特定上采样的影响不能迁移到其他领域：对特定领域（如书籍、代码、Arxiv）进行上采样，虽然在这些领域内提高了短上下文和长上下文的性能，但这些改进并不总是能够转移到其他领域。例如，对代码领域的上采样可能会损害书籍领域的表现。
2. 平衡的数据混合策略对提升上下文检索能力最有效：按源长度上采样（Per-source length upsampling）策略在不同领域之间提供了最平衡的性能提升，几乎没有显著增加任何领域的损失。这表明在保持领域混合比例不变的情况下，适当上采样长序列是提高模型在长上下文任务上性能的有效方法。
3. 全局上采样也存在一定局限性：全局上采样（不考虑数据来源）或有意上采样特定领域（如书籍、代码、Arxiv）会改变领域混合和长度分布，这可能导致在某些领域内性能提升，但在其他领域内性能下降。

总结：数据工程策略对模型性能非常重要，正确的数据策略可以帮助模型在长上下文任务上取得更好的性能，同时保持在短上下文任务上的表现。

三、总结

论文主要研究了如何通过post-training和数据工程的方法，有效地将语言模型的上下文长度扩展到128K。研究主要关注于数据的数量和质量，发现通过在领域平衡和长度上采样的基础上进行连续预训练是一个有效的策略。实验结果显示他们的方法优于其他开源的长上下文模型，并可以与最新的前沿模型（如GPT-4 128K）相媲美。该方法不仅效果显著，而且成本可控，为语言模型的长上下文处理能力的发展提供了新的思路。

另外，该工作也展示了数据工程在大型语言模型研究中的重要性，可能与建模和算法创新同等重要。这里面许多数据细节对于长上下文性能至关重要，但这些细节可能容易被忽视。

后续优化上，由于本文只涉及了持续预训练阶段如何做数据工程，而关于在100K上下文长度的任务上对语言模型进行指令微调的研究还很有限。希望本文的工作能为未来在100K级别的长上下文指令微调研究提供基础。

四、参考

[1] Xiong, W., Liu, J., Molybog, I., Zhang, H., Bhargava, P., Hou, R., Martin, L., Rungta, R., Sankararaman, K. A., Oguz, B., et al. Effective long-context scaling of founda- tion models. arXiv preprint arXiv:2309.16039, 2023.

👇关注公众号NLP PaperWeekly，对话框输入“128K”，即可领取上述论文👇

• 进技术交流群请添加我微信（FlyShines)
• 请备注昵称+公司/学校+研究方向，否则不予通过