Data Engineering for Scaling Language Models to 128K Context
Proposes a data engineering strategy using 80K sequence continual pretraining to extend language models' context to 128K, achieving near GPT-4 performance.
Key Findings
Methodology
This work employs a continual pretraining approach based on the LLaMA-2 architecture with full attention. Key components include:• Using 80K-long sequence data for incremental training;• Maintaining domain balance via multi-source data mixing;• Applying long sequence upsampling to enhance long-range dependency learning. Experiments conducted on 8×80G A100 GPUs over about a week, with 5B tokens, demonstrate significant improvements in long context retrieval, approaching GPT-4 128K performance on Needle-in-a-Haystack. The approach leverages lightweight data adjustments rather than large-scale retraining, making it resource-efficient.
Key Results
- The model achieves 88% accuracy on Needle-in-a-Haystack at 128K context, outperforming open-source counterparts like YaRN-Mistral 7B 128K (57.4%) and LongLoRA 100K, and nearing GPT-4 128K (87.1%).
- Training on 500M to 5B tokens suffices to unlock precise long-range retrieval, confirming that the base model already possesses this capability, which can be extended with minimal data.
- Balancing data sources and applying length upsampling within each domain significantly improves performance, avoiding degradation seen with naive long sequence upsampling.
Significance
This research demonstrates a practical, low-cost pathway for extending open-source language models' context length to 128K, closing the performance gap with proprietary models like GPT-4. It addresses the long-standing challenge of enabling models to utilize long-range dependencies effectively, which is crucial for multi-document understanding, long dialogue modeling, and autonomous agents. The methodology offers a scalable solution that balances computational efficiency with performance, opening avenues for broader academic and industrial adoption.
Technical Contribution
The key technical contribution is a data-centric approach combining domain-balanced multi-source data mixing with long sequence upsampling, enabling models trained on 80K sequences to generalize to 128K contexts. This approach reduces training costs and hardware demands, leveraging existing architectures with minimal modifications (e.g., positional encoding adjustments). It shifts focus from architecture-centric solutions to data engineering, providing a new paradigm for long context modeling.
Novelty
This work is the first to systematically validate that limited long-sequence data, when properly balanced and upsampled, can activate and transfer long-range retrieval capabilities in pre-trained models. Unlike prior efforts relying on large-scale from-scratch training, it emphasizes data engineering as a cost-effective means to extend context length, marking a significant departure from the mainstream paradigm.
Limitations
- Despite promising results at 128K, performance at even longer sequences (e.g., 200K+) remains limited due to hardware constraints and architecture scalability. Future work must explore more efficient parallelism and model modifications.
- Data balancing strategies, while effective, may not fully address domain-specific performance drops, especially in highly specialized or out-of-distribution scenarios.
- The approach still requires substantial computational resources, and further optimization is needed for broader accessibility.
Future Work
Future research will focus on developing more scalable parallel attention mechanisms, such as sequence parallelism, to support sequences beyond 200K. Additionally, integrating instruction tuning and multi-task learning could further enhance generalization and robustness across diverse long-text tasks. Exploring unsupervised data augmentation and adaptive sampling strategies may also improve long-range dependency modeling.
AI Executive Summary
Long-range dependency modeling has become a pivotal frontier in advancing language understanding, especially for complex tasks like multi-document comprehension and long dialogue management. Historically, models like LLaMA and GPT-3 were limited to 4K-8K contexts, constraining their ability to process extended texts. The emergence of GPT-4 128K has set a new benchmark, but its closed-source nature limits widespread research. This study introduces a data engineering approach that leverages lightweight continual pretraining on carefully balanced, long-sequence data to extend open-source models’ context length to 128K.
The core innovation lies in combining multi-source data mixing with long sequence upsampling, enabling models to learn long-distance dependencies without massive retraining. Experiments show that training on 5B tokens with 80K sequences achieves 88% accuracy on the Needle-in-a-Haystack benchmark, closely approaching GPT-4’s 87.1%. This approach significantly reduces computational costs and hardware demands, making long-context modeling accessible for academic research.
The results demonstrate that a small amount of targeted long-sequence data can unlock substantial capabilities, challenging the notion that large-scale from-scratch training is necessary. The methodology emphasizes data quality and diversity, maintaining domain balance to prevent performance degradation across different content types. Looking ahead, further innovations in parallel attention mechanisms and multi-task training are expected to push the boundaries of sequence length even further, enabling more sophisticated long-text understanding in practical applications.
Deep Analysis
Background
长上下文建模经历了从4K到8K的逐步发展，代表性工作包括Touvron等的LLaMA系列、Xiong等的LongLORA、Peng等的YaRN。尽管如此，长距离信息检索和多文档理解仍受限于模型训练数据和架构优化。GPT-4 128K的出现开启了新一轮研究热潮，但其封闭性限制了学术界的验证。现有开源模型如LongLoRA和YaRN在长序列支持上表现有限，主要受制于训练数据的长度和多样性。
Core Problem
核心问题在于如何在有限资源下，利用有限长序列数据实现长距离依赖能力的迁移。传统方法依赖大规模从零训练，成本高昂且难以普及。模型在128K上下文下的检索精度不足，限制了多文档问答、长历史对话等应用的性能。硬件限制、数据偏差和模型架构的瓶颈，成为实现长序列建模的主要障碍。
Innovation
创新点包括：• 采用80K长序列数据进行逐步预训练，突破4K限制；• 结合多源数据平衡策略，避免单一领域偏差；• 通过长序列上采样，增强模型长距离依赖捕获能力；• 在有限硬件条件下实现128K上下文扩展，降低成本。此方案区别于大规模从零训练的路径，强调数据工程在模型长距离能力中的核心作用。
Methodology
- �� 选择LLaMA-2架构，调整RoPE位置编码以适应长序列；• 利用SlimPajama数据集，保持领域多样性；• 采用80K序列长度逐步预训练，训练时间约7天；• 通过多源数据混合，保持领域平衡；• 实施长序列上采样，增强长距离依赖；• 使用GPU offloading和FlashAttention技术优化训练效率；• 训练过程中监控针芯入草堆测试，调整数据比例和训练策略。
Experiments
在8块80G A100 GPU上，使用5B tokens数据进行预训练，模型在针芯入草堆测试中达88%的准确率。对比Baselines包括YaRN-Mistral和LongLoRA，结果显示本方法在长序列检索和多文档理解任务中明显优越。还通过长书问答任务验证模型在实际场景中的应用潜力。多次消融实验确认长序列上采样和数据平衡对性能提升的关键作用。
Results
模型在128K长度下的针芯入草堆测试准确率达88%，超越大部分开源模型，逼近GPT-4 128K的87.1%。500M到5B tokens的数据量已足以激活模型的长距离检索能力，验证了少量长序列数据的有效性。平衡多源数据比例，避免偏向单一领域，显著提升模型在多任务、多场景中的表现。
Applications
该技术适用于多文档问答、长历史对话、自动化知识库检索等场景，能显著提升模型理解复杂长文本的能力。实现条件包括：丰富多样的数据集、强大的硬件支持和优化的训练流程。未来可结合指令调优，拓展到更长序列和多任务学习，推动智能助理、知识管理等行业应用。
Limitations & Outlook
硬件资源仍是限制因素，长序列训练对GPU内存和计算能力要求较高。数据平衡策略在多领域迁移中效果有限，模型在极端长序列（超过128K）上的表现尚待提升。未来需探索更高效的并行技术和模型架构优化，以实现更大规模的长距离建模。
Plain Language Accessible to non-experts
想象你在一个非常长的故事书里阅读，普通的模型就像只记住前几页，难以理解后面发生的事情。而这项研究就像给你一支超级笔，可以把整本书的内容都记下来，不仅能记住每个细节，还能快速找到任何你想知道的部分。通过合理整理和反复复习，你可以在阅读超长的故事时，依然记得重要的线索。这就像给模型配备了超级强大的记忆力，让它能理解和检索超长文本中的信息。
ELI14 Explained like you're 14
想象你在玩一个超级长的冒险游戏，里面有很多关卡和秘密任务。普通的游戏只能记住前几关的内容，但这个新方法就像给游戏装上了一个超级大脑，能记住上百关的细节。这样，你可以随时找到隐藏的宝藏或解开谜题，不用担心忘记之前的线索。它通过让游戏中的角色多读一些长长的故事，把记忆变得更强大。虽然需要一些特别的技巧和设备，但最终让游戏变得更有趣、更聪明。
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 \textit{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 \textit{quantity} and \textit{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 \textit{domain balance} and \textit{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.