Research map / Large language models / Language, safety and interpretability
Data Engineering for Scaling Language Models to 128K Context
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Studies a data recipe for continually pretraining LLaMA-2 7B/13B to 128K context, finding that 1-5B tokens of per-source length-upsampled, domain-balanced SlimPajama suffice to pass Needle-in-a-Haystack retrieval.
Problem
Extending language models' context windows to 128K while retaining precise retrieval at arbitrary positions and short-context ability, at academic-scale compute; prior long-context continual pretraining used hundreds of billions of tokens and ad hoc data mixtures.
Approach
Long-context retrieval is mostly acquired during standard pretraining, so light continual pretraining on 1-5B tokens suffices; the key data choice is per-source length upsampling, which increases long sequences within each domain while keeping the original domain mixture fixed.
Key results
- Ours LLaMA-2 7B (80K) scores 88.0 on Needle vs 70.0 for LongLoRA 7B and 57.4 for YaRN Mistral 7B; GPT-4-Turbo 128K scores 87.1.
- Ours LLaMA-2 13B (64K) reaches 90.0 Needle and 52.4 MMLU, outperforming LongLoRA 13B (54.1 Needle, 50.1 MMLU).
- On InfiniBench BookQA at 128K, Ours LLaMA-2 13B scores 31.1, above YaRN Mistral 7B (26.3) and close to GPT-4-Turbo (37.4).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Continue pretraining LLaMA-2 7B on 5B tokens of per-source length-upsampled data | LLaMA-2 7B · Needle-in-a-Haystack | Needle performance: 70.0 -> 88.0 | Helped |
| Continue pretraining LLaMA-2 13B on 5B tokens of per-source length-upsampled data | LLaMA-2 13B · Needle-in-a-Haystack | Needle performance: 54.1 -> 90.0 | Helped |
| Continue pretraining LLaMA-2 7B on 5B tokens of per-source length-upsampled data | LLaMA-2 7B · MMLU | accuracy: 37.9 -> 43.3 | Mixed |
| Continue pretraining LLaMA-2 13B on 5B tokens of per-source length-upsampled data | LLaMA-2 13B · InfiniBench BookQA, 128K | BookQA: 24.6 -> 31.1 | Helped |
| Increase continual-pretraining data from 500M to 5B tokens (ablation) | LLaMA-2 7B · Needle-in-a-Haystack, 128K | retrieval performance | Helped |
| Increase continual-pretraining data from 5B to 10B tokens (ablation) | LLaMA-2 7B · Needle-in-a-Haystack, 128K | length generalization | Hurt |
| Add per-source length upsampling while retaining the original domain proportions (ablation) | LLaMA-2 7B · Needle-in-a-Haystack, 128K | precise retrieval | Helped |
| Replace per-source length upsampling with the original mixture without length upsampling (ablation) | LLaMA-2 7B · Needle-in-a-Haystack, 128K | precise retrieval | Hurt |
| Upsample book data instead of keeping a balanced domain mixture (ablation) | LLaMA-2 7B · Book and Github domains, 0-4K context | validation loss: + .029 Github; -- .175 Book | Mixed |
| Upsample code data instead of keeping a balanced domain mixture (ablation) | LLaMA-2 7B · Book and Github domains, 0-4K context | validation loss: + .030 Book; -- .023 Github | Mixed |
| Upsample Arxiv data instead of keeping a balanced domain mixture (ablation) | LLaMA-2 7B · Arxiv and Book domains, 0-4K context | validation loss: -- .060 Arxiv; + .040 Book | Mixed |
Surprising findings
- Validation loss can be nearly identical between data recipes while retrieval performance differs substantially, so loss is an insufficient proxy for long-context ability.
- Naively upsampling book data (a common practice) hurts other domains such as code and web text, while per-source length upsampling gives the most balanced results.
Limitations
- Only continual pretraining is studied; no instruction tuning for 100K context.
- Training context limited to 80K/64K by the HuggingFace-DeepSpeed framework; 200K not explored.
- Relies on the Needle-in-a-Haystack test and a single BookQA benchmark for long-context evaluation.
Experimental setup
Datasets: SlimPajama InfiniBench (BookQA) MMLU
Benchmarks: Needle-in-a-Haystack MMLU InfiniBench BookQA
Metrics: retrieval accuracy MMLU accuracy BookQA score validation loss
Baselines: YaRN Mistral 7B 128K LongLoRA Together LLaMA-2 32K LongChat v1.5 32K GPT-4 Turbo 128K
Builds on: LLaMA 2 SlimPajama FlashAttention DeepSpeed ZeRO YaRN
Closest papers by meaning
- Scaling Instruction-tuned LLMs to Million-token Contexts via Hierarchical Synthetic Data Generation (ICLR 2025 · similarity 0.90)
- Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance (ICLR 2025 · similarity 0.90)
- Bridging Compute- and Data-Optimal Pretraining (NeurIPS 2026 · similarity 0.90)
- Scaling Retrieval-Based Language Models with a Trillion-Token Datastore (NeurIPS 2024 · similarity 0.89)
- Data-Constrained Language Model Pretraining: Improved Regularization and Scaling Laws (NeurIPS 2026 · similarity 0.89)
- L$^2$M: Mutual Information Scaling Law for Long-Context Language Modeling (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.