PaperlayerView exact revision on arXiv

Source-linked AI summary

How to Train Long-Context Language Models (Effectively)

Tianyu Gao, Alexander Wettig, Howard Yen, Danqi Chen

arXiv:2410.02660v4cs.CLcs.LG

TL;DR

The paper asks how to train short-context language models to use long-context information effectively without relying on inadequate evaluation signals. It establishes broad downstream evaluation after SFT, then studies continued-training and instruction-tuning choices. ProLong achieves strong 128K performance with 5% of Llama-3.1's long-context data budget and supports up to 512K tokens.

  • Problem

    Adapting language models to extremely long inputs is challenging, and existing perplexity or needle-in-a-haystack evaluations provide limited guidance for development.

  • Method

    The paper evaluates continued-training recipes with diverse long-context downstream tasks after SFT, then ablates data mixtures, sequence lengths, instruction datasets, and related design choices.

  • Results

    ProLong achieves the best 128K performance among 10B-parameter models using 5% of Llama-3.1's long-context data budget and supports up to 512K tokens.

  • Takeaways & Limitations

    Long-context training benefits from combining long and high-quality short-context data, training beyond evaluation length, and using short-context instruction data for SFT.

  • Takeaways & Limitations

    The study is limited to the 10B-scale regime and Llama-3 models, which may limit generalizability.

Abstract

from arXiv · show

We study continued training and supervised fine-tuning (SFT) of a language model (LM) to make effective use of long-context information. We first establish a reliable evaluation protocol to guide model development -- instead of perplexity or simple needle-in-a-haystack (NIAH) tests, we use a broad set of long-context downstream tasks, and we evaluate models after SFT as this better reveals long-context abilities. Supported by our robust evaluations, we run thorough experiments to decide the data mix for continued pre-training, the instruction tuning dataset, and many other design choices such as position extrapolation. We find that (1) code repositories and books are excellent sources of long data, but it is crucial to combine them with high-quality short-context data; (2) training with a sequence length beyond the evaluation length boosts long-context performance; (3) for SFT, using only short instruction datasets yields strong performance on long-context tasks. Our final model, ProLong-8B, which is initialized from Llama-3 and trained on 40B tokens, demonstrates state-of-the-art long-context performance among similarly sized models at a length of 128K. ProLong outperforms Llama-3.1-8B-Instruct on the majority of long-context tasks despite using only 5% as many tokens during long-context training. Additionally, ProLong can effectively process up to 512K tokens, one of the longest context windows of publicly available LMs.

1 Introduction

The paper studies how to adapt short-context language models for effective long-context use, arguing that reliable downstream evaluation is needed to guide training choices. Comprehensive experiments identify data, sequence-length, and SFT practices, culminating in ProLong's strong 128K performance and 512K capacity.

  • Long-context adaptation is challenging because it requires substantial infrastructure and data, while many design choices remain poorly understood.
  • Minimal-training context extensions can fail even on needle-in-a-haystack, motivating continued training on billions of long-document tokens.
  • The evaluation protocol uses diverse downstream applications and evaluates after SFT because some gains emerge only after instruction tuning.
  • Experiments find that long-only training hurts long-context performance, longer training sequences help, and short-context instruction data can suffice for long-context performance.
  • ProLong achieves the best 128K performance among 10B-parameter models using 5% of Llama-3.1's long-context data budget and supports up to 512K tokens.

2 Guiding Model Development With Meaningful Evaluations

The paper develops long-context models using broad, realistic evaluations rather than relying on perplexity or synthetic recall alone. It finds that evaluation after SFT reveals more useful signals, while long-context training must also preserve short-context capabilities.

  • Evaluating on diverse and realistic tasks: HELMET evaluates recall, retrieval, reranking, in-context learning, book question answering, and long-document summarization across diverse long-context abilities.
  • Evaluating on diverse and realistic tasks: NIAH and perplexity can miss important differences: strong models saturate on recall while other HELMET categories continue distinguishing long-context ability.
  • Evaluating on diverse and realistic tasks: Using 100% long data improves perplexity but significantly hurts downstream long-context performance, showing that perplexity is not an adequate development signal.
  • Evaluating after supervised fine-tuning: After SFT, continued training produces clearer RAG and reranking trends, and enables realistic QA and summarization evaluations that otherwise fail completely.
  • Evaluating after supervised fine-tuning: Short-context instruction tuning remains surprisingly competitive, while adding synthetic long-context instruction data can hurt performance even at 1%.
  • Checking that short-context performance is preserved: Position extrapolation and an existing long-data mixture deteriorate Llama-3-8B's short-context performance, motivating data sources that preserve it.

3 Long-Context Data Curation

The study finds that books and code repositories are strong long-context data sources, but performance depends on mixing them with high-quality short-context data. A 60% long-data and 40% short-data mixture, supported by a curated ShortMix, achieves the best overall results.

  • Long data sources: Code repositories provide abundant long-context data by concatenating all files from each repository into a single document.This construction increases distances between dependent files and reduces recency bias.
  • Long data sources: An equal mix of books and code repositories achieves the best overall performance across long-context tasks.Code alone performs best on stress-test recall, while books are broadly beneficial for in-context learning, summarization, and re-ranking.
  • Short/long data ratio: 60% long data and 40% short data produces the best average performance and is adopted for ProLong.Increasing long data consistently decreases short-context performance, while long-context trends vary by task and SFT stage.
  • Short-context mix: A curated ShortMix outperforms other short-data sources on both short- and long-context tasks.Its domains are particularly important for retaining Llama-3-8B’s mathematical-reasoning performance.
  • Short-context mix: The data mix significantly surpasses Fu et al. (2024) on both long-context and short-context tasks.The comparison is presented as evidence for the effectiveness of ProLong’s data curation.

4 Scaling the Size and Length of the Training Data

Scaling continued training improves long-context performance, and training at 512K can improve results evaluated at 64K. However, longer sequences are more expensive and may not be computationally optimal.

  • Training scale: ProLong scales training to 40B tokens, including 20B tokens at a maximum sequence length of 64K and an additional 20B at 512K.The recipe incorporates findings from the ablation experiments and uses longer sequences and more tokens.
  • Training scale: Increasing training steps produces positive trends on recall, RAG, re-ranking, and summarization.Short-context performance initially drops, recovers with training, drops again when switching from 64K to 512K, and then recovers.
  • Training length: Training at 512K substantially improves long-context performance evaluated at 64K.The comparison uses identical hyperparameters and data mixtures after initializing from 20B tokens trained at 64K.
  • Training length: A simplified dependency-counting model provides intuition for why longer documents can expose more dependencies of a given distance.The authors explicitly note that the model’s assumptions do not hold in practice.
  • Limitations: Longer-sequence training is more expensive and may not be computationally optimal.The paper notes this cost even while demonstrating empirical benefits from longer training sequences.

5 Supervised Fine-Tuning for Long-Context LMs

The paper studies supervised fine-tuning for long-context instruction following using standard and synthetic instruction datasets. It finds that short-context UltraChat data alone is sufficient and that adding synthetic long-context data can hurt performance.

  • Scope: The study focuses on supervised fine-tuning and leaves reinforcement learning and preference optimization for future work.All SFT experiments use the ProLong base model trained for 40B tokens with a maximum sequence length of 512K.
  • Synthetic data: Synthetic QA, RAG, and summarization data are generated from long documents or books to create long instruction examples.The procedures include question-answer generation, retrieved-chunk presentation, and recursive summarization.
  • SFT data: Short-context UltraChat data yields strong long-context results and is used for final ProLong SFT.UltraChat outperforms Tulu-v2 and ShareGPT in the initial comparison.
  • Synthetic data: Adding synthetic instruction data does not improve very long-context tasks, and even 1% synthetic data hurts performance in this setting.The finding persists when Llama-3-70B is used as the data generator.
  • Synthetic data: The paper hypothesizes that synthetic data may help previous systems because they had insufficient long-context training or more extensive instruction datasets.These explanations are presented as hypotheses for differences from prior work, not as established findings.

6 The ProLong Model: Recipe and Results

ProLong combines continued training at 64K and 512K with targeted design choices and SFT, achieving strong long-context performance at 128K and extending effectively to 512K.

  • Recipe: ProLong is trained for 20B tokens at 64K, then 20B tokens at 512K, before SFT on UltraChat.The base model uses sequence parallelism for the 512K stage.
  • Design choices: Changing the RoPE frequency base improves long-context performance during position extrapolation.The recipe uses bases of 8 × 10^6 at 64K and 1.28 × 10^8 at 512K.
  • Design choices: Disabling cross-document attention benefits both short- and long-context performance while increasing training throughput.The masking prevents attention from crossing document boundaries.
  • Evaluation: ProLong outperforms Llama-3.1-8B-Instruct on all evaluated categories except summarization using only 5% of its long-context data budget.The comparison uses 40B versus 800B long-context training tokens.
  • Evaluation: ProLong continues improving as its context window increases from 32K to 512K on QA and summarization stress tests.The evaluation truncates documents at each tested length.
  • Evaluation: On NoCha, ProLong achieves the best accuracy among 10B-scale models on instances longer than 180K tokens and exceeds random guessing below 75K.It is the only model above the 25% random-guessing level in the shorter category.

7 Related Work

Related work extends context through position methods, attention changes, continued training, alternative architectures, and specialized benchmarks; this paper follows continued training with diverse evaluation.

  • Adapting existing LMs: Prior methods extend context through position extrapolation or modified attention patterns, often with minimal additional training.The paper distinguishes these approaches from its continued-training setting.
  • Adapting existing LMs: Llama 3.1 and Jamba insert long-context continued training between standard pre-training and supervised fine-tuning.This is the training-stage arrangement adopted in the paper.
  • Adapting existing LMs: The paper reports that standard short-context instruction data performs better for long-context results in its setting than synthetically generated long SFT data.This contrasts with prior work using synthetic long data during SFT.
  • Efficient long-context architectures: Efficient alternatives include linear-attention, recurrent, and other attention architectures, but many require training from scratch and have long-context recall limitations.These approaches differ from adapting existing pretrained LMs.
  • Long-context evaluation: Long-context benchmarks cover broad evaluations and focused properties such as positional bias, in-context learning, and book-length summarization.The paper follows HELMET for diverse application coverage and reliable evaluation.

8 Conclusion

The paper studies how to turn short-context pretrained LMs into long-context models through continued training and SFT, and releases ProLong and its development resources.

  • Conclusion: The paper studies how to continually pre-train and SFT a short-context pretrained LM for long-context use.It conducts ablations across components of the training recipe.
  • Conclusion: ProLong is presented as a new state-of-the-art long-context LM based on the paper’s combined findings.The authors release code, data, and models publicly.
  • Limitations: The study is limited by resource constraints that prevent exhaustive ablations of optimization hyperparameters and additional data mixtures.The authors also restrict the study to the 10B-scale Llama-3 regime, limiting potential generalizability.
  • Evaluation: HELMET evaluates long-context models at 128K and beyond with results reported across different context lengths.The protocol provides fine-grained control over evaluation lengths.
  • Evaluation: HELMET replaces saturated NIAH recall with JSON KV retrieval, which requires matching a value to a key in a large JSON file.This creates a more challenging synthetic recall task.
  • Evaluation: HELMET uses class-balanced demonstrations, abstract labels, and model-based scoring for long-context ICL, QA, and summarization.QA outputs receive GPT-4o scores on a 0–3 scale.

A.2 Data processing

The data-processing and ablation setup combines diverse sources, document-aware packing and attention, position extrapolation, initialization comparisons, and short-context SFT dataset tests.

  • Data processing: The study explores web, code, educational, mathematical, textbook, and book-related sources for long-context training.The listed sources include the Stack, SlimPajama, FineWeb, FineWeb-Edu, Tulu-v2, OpenWebMath, textbooks, and Dolma.
  • Data processing: Short training and SFT data are packed into 64K chunks, while long-context data is filtered for documents meeting the target length.Truncation is handled differently for short training and SFT data.
  • Data processing: At 512K, the final mixture combines 64K and 512K documents, preserving shorter long documents to reduce sequence-parallel communication overhead.The code mixture uses a 50%/50% ratio, while the book mixture uses 83%/17%.
  • Attention and batching: Variable-length attention computes attention separately for concatenated documents, reducing the cost of masked full-sequence attention.This exploits the quadratic complexity of attention with sequence length.
  • Ablation setting: The ablation baseline trains Llama-3-8B for 5B tokens at a maximum sequence length of 64K.The default ablations remove textbooks and use ShortMix proportions.
  • Position extrapolation: RoPE frequency-base ablations compare settings at 64K and 512K to determine suitable position extrapolation values.The experiments find that the original Llama-3 base degrades performance and larger scaling can improve it.
  • Attention masking: Document masks prevent attention across document boundaries and improve both long- and short-context results in continued training.The implementation compares masked attention with standard full attention over concatenated documents.
  • Initialization: Using Llama-3-8B-Instruct for the final model slightly improves long-context performance and substantially strengthens short-context performance.The ablations otherwise initialize from the base Llama-3-8B model.

B.5 Synthetic data with a stronger data generator

The stronger synthetic-data generator does not change the conclusion that short-context SFT data performs best in this setting. ProLong’s carefully designed data mix also outperforms Fu et al.’s strategy on both short- and long-context tasks.

  • B.5 Synthetic data with a stronger data generator: Synthetic data generated by either Llama-3-8B-Instruct or Llama-3-70B-Instruct does not improve long-context performance.The stronger generator was tested to rule out low synthetic-data quality as the explanation.
  • B.5 Synthetic data with a stronger data generator: Short-context SFT data achieves the best result even when synthetic data is generated by Llama-3-70B-Instruct.The stronger generator performs slightly better than the weaker one, but using only short-context SFT data remains best.
  • B.5 Synthetic data with a stronger data generator: Under a fair comparison using the same initialization, 5B tokens, and hyperparameters, ProLong’s data mix significantly outperforms Fu et al. on short- and long-context tasks.The main difference is ProLong’s carefully selected mixture of short and long documents rather than proportional up-sampling of long documents by domain.
  • B.5 Synthetic data with a stronger data generator: ProLong largely regains Llama-3.1-8B-Instruct’s short-context performance after SFT, except on MMLU and GSM8K.The comparison uses ProLong’s SFT setup and includes an UltraChat-based Llama-3-8B-Instruct baseline.

B.8 Evaluation on more benchmarks

ProLong is evaluated on RULER and ∞Bench, but these benchmarks have narrow domain coverage and noisy metrics that can produce unintuitive model rankings. Despite this limitation, ProLong performs best among 10B-scale models on ∞Bench.

  • B.8 Evaluation on more benchmarks: RULER and ∞Bench cannot reliably reflect long-context performance because their domain coverage is narrow and evaluation metrics are noisy.These limitations produce unintuitive rankings, including larger models performing worse than Llama-3.1-8B.
  • B.8 Evaluation on more benchmarks: ProLong performs best on ∞Bench among 10B-scale models.This result is reported despite the benchmarks’ reliability limitations.
  • B.8 Evaluation on more benchmarks: Gemini-1.5-Pro and Llama-3.1 70B perform worse than Llama-3.1 8B on these benchmarks, illustrating their unintuitive model rankings.The passage attributes this pattern to narrow domain coverage and noisy evaluation metrics.

C Full Evaluation Results

The paper provides complete evaluation results and additional long-context evidence across 128K and 512K settings. These materials include benchmark results, long-document QA examples, and links to dataset- and length-specific scores.

  • C Full Evaluation Results: Detailed evaluation results for each dataset and evaluation length are provided through an external results repository.The repository contains the specific numbers underlying the paper’s evaluations.
  • C Full Evaluation Results: At 128K, ProLong achieves the best ∞Bench performance among all 10B-scale models, although RULER and ∞Bench have noisy metrics and narrow domain coverage.The table reports results on both benchmarks and cautions against treating them as fully reliable measures of long-context ability.
  • C Full Evaluation Results: At 512K, ProLong answers questions effectively from a 300K-token document while Llama-3.1 and GPT-4o fail on the example.The examples also include comparisons with Llama-3.1-8B-Instruct at 128K and GPT-4o through the ChatGPT interface.
Loading 2410.02660v4…