[论文解读] Data Engineering for Scaling Language Models to 128K Context

Yao Fu, Rameswar Panda+5 more|arXiv (Cornell University)|Feb 15, 2024

Network Packet Processing and Optimization被引用 4

一句话总结

本文提出了一种数据工程方法，通过在每源数据上进行长度上采样，对LLaMA-2语言模型进行持续预训练，将上下文长度扩展至128K。实验表明，通过精心设计的数据混合策略——在保持各领域比例的同时，提升每个领域内长序列的权重——可显著提升模型在Needle-in-a-Haystack检索任务上的表现，使性能接近GPT-4 128K，同时保持短上下文下的性能。

ABSTRACT

We study the continual pretraining recipe for scaling language models' context lengths to 128K, with a focus on data engineering. We hypothesize that long context modeling, in particular extit{the ability to utilize information at arbitrary input locations}, is a capability that is mostly already acquired through large-scale pretraining, and that this capability can be readily extended to contexts substantially longer than seen during training~(e.g., 4K to 128K) through lightweight continual pretraining on appropriate data mixture. We investigate the extit{quantity} and extit{quality} of the data for continual pretraining: (1) for quantity, we show that 500 million to 5 billion tokens are enough to enable the model to retrieve information anywhere within the 128K context; (2) for quality, our results equally emphasize extit{domain balance} and extit{length upsampling}. Concretely, we find that naively upsampling longer data on certain domains like books, a common practice of existing work, gives suboptimal performance, and that a balanced domain mixture is important. We demonstrate that continual pretraining of the full model on 1B-5B tokens of such data is an effective and affordable strategy for scaling the context length of language models to 128K. Our recipe outperforms strong open-source long-context models and closes the gap to frontier models like GPT-4 128K.

研究动机与目标

• 探究如何通过数据工程在不改变模型架构的前提下，有效将语言模型扩展至128K上下文长度。
• 解决在模型原始预训练长度之外任意位置实现精确信息检索的挑战，尤其是在长上下文中的检索。
• 确定在经过大规模预训练后，是否可通过在精选数据上进行最小量的持续预训练，解锁模型已部分具备的长上下文能力。
• 评估在长上下文适应过程中，数据量、领域平衡与序列长度分布之间的权衡。
• 提供一种实用、经济且高效的配方，用于将开源模型扩展至128K上下文，以支持未来的指令微调。

提出的方法

• 在1–50亿token的精选数据上，对LLaMA-2 7B和13B模型进行持续预训练，使用全注意力机制，上下文窗口为64K–80K。
• 应用RoPE位置嵌入基值调整，以在不修改架构的前提下支持更长的上下文长度。
• 设计了一种基于每源的长度上采样策略，在保持原始领域混合比例的同时，增加每个领域内长序列的比例。
• 使用FlashAttention降低显存占用，使在学术级硬件上进行长上下文训练成为可能。
• 使用Needle-in-a-Haystack基准测试评估模型在4K至128K上下文长度范围内的精确检索性能。
• 与YaRN-Mistral 128K和LongLoRA 100K等强基线模型进行比较，以隔离数据混合设计的影响。

Figure 1: Needle-in-a-Haystack test (Kamradt, 2023 ) performance comparison. The x-axis denotes the length of the document (the “haystack”); the y-axis indicates the position that the “needle” (a short sentence) is located within the document, from 1K to 128K tokens. For example, 50% indicates that

实验结果

研究问题

• RQ1 在仅用4K上下文预训练的模型上，仅通过在精心筛选的数据上进行少量持续预训练，是否能解锁128K上下文能力？
• RQ2 在高长度领域（如书籍）中对长序列进行简单上采样，是否会显著降低在代码或网页内容等其他领域中的性能？
• RQ3 在每个领域内对长序列进行上采样的同时保持原始领域混合比例，是否比基于领域的上采样策略更有效？
• RQ4 数据混合设计如何影响长上下文与短上下文性能之间的权衡？
• RQ5 在长上下文建模方面，数据驱动的方法是否能在效率和性能上超越大规模预训练配方？

主要发现

• 仅在1–50亿token的每源长度上采样数据上进行持续预训练，即可使LLaMA-2 7B和13B模型在128K上下文长度的Needle-in-a-Haystack测试中取得优异表现。
• 基于每源的长度上采样策略——在保持领域比例的同时提升各领域内长序列的权重——在各领域间实现了最均衡的性能提升，优于简单的基于领域的上采样方法。
• 简单的基于领域的上采样（如过度代表书籍）会导致代码和网页内容等其他领域性能显著下降，表明数据混合设计存在明显缺陷。
• 采用本研究配方训练的模型在长上下文检索任务上的表现可与GPT-4 128K相媲美，显著优于YaRN-Mistral 128K和LongLoRA 100K等强开源基线模型。
• 仅依赖验证损失不足以检测性能差异；Needle-in-a-Haystack测试揭示了测试损失无法捕捉的关键长上下文检索能力差异。
• 本研究证明，数据工程在实现顶尖长上下文性能方面，其重要性不亚于模型架构或算法创新。

Figure 2: Length and domain distributions of the various data mixture strategies. We use SlimPajama (Soboleva et al., 2023 ) , an open-source reproduction of LLaMA (Touvron et al., 2023a ) pretraining data mixture, as the source dataset. The original data mixture has about 30% of documents that are