Source-linked AI summary
Effective Long-Context Scaling of Foundation Models
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, Hao Ma
TL;DR
The paper addresses effective long-context processing by building LLMs that support contexts up to 32,768 tokens. It uses continual pretraining from Llama 2 with longer sequences and a cost-effective instruction-tuning procedure, achieving improvements over Llama 2 and strong long-context benchmark performance, including surpassing gpt-3.5-turbo-16k on 7 of 10 ZeroSCROLLS tasks.
Problem
Effective long-context processing is needed for LLM applications involving dense documents, richer chatbot interactions, and iterative creation, but existing models have limited ability to model distant dependencies.
Method
The models are built by continually pretraining Llama 2 with 400 billion additional tokens in longer sequences, refining RoPE position encoding, and augmenting Llama 2 Chat data with synthetic long-document instruction examples.
Results
The models improve over Llama 2 on short- and long-context tasks, while the 70B instruction-tuned model outperforms gpt-3.5-turbo-16k on 7 of 10 ZeroSCROLLS tasks without human-annotated long-context data.
Takeaways & Limitations
Continual pretraining from a short-context model, combined with minimal positional-encoding refinement and inexpensive synthetic instruction tuning, provides a practical route to strong long-context performance.
Takeaways & Limitations
The models have not been fine-tuned for a wide range of long-context applications, and their tokenizer produces about 10% more tokens on average than GPT-3.5's tokenizer.
Abstract
from arXiv · showhide
We present a series of long-context LLMs that support effective context windows of up to 32,768 tokens. Our model series are built through continual pretraining from Llama 2 with longer training sequences and on a dataset where long texts are upsampled. We perform extensive evaluation on language modeling, synthetic context probing tasks, and a wide range of research benchmarks. On research benchmarks, our models achieve consistent improvements on most regular tasks and significant improvements on long-context tasks over Llama 2. Notably, with a cost-effective instruction tuning procedure that does not require human-annotated long instruction data, the 70B variant can already surpass gpt-3.5-turbo-16k's overall performance on a suite of long-context tasks. Alongside these results, we provide an in-depth analysis on the individual components of our method. We delve into Llama's position encodings and discuss its limitation in modeling long dependencies. We also examine the impact of various design choices in the pretraining process, including the data mix and the training curriculum of sequence lengths -- our ablation experiments suggest that having abundant long texts in the pretrain dataset is not the key to achieving strong performance, and we empirically verify that long context continual pretraining is more efficient and similarly effective compared to pretraining from scratch with long sequences.
1 Introduction
The paper targets open long-context LLMs whose downstream performance has lagged proprietary systems and has been incompletely evaluated. It presents continual pretraining and broad evaluations showing gains across long- and short-context tasks.
- Motivation: Open long-context models often underperform proprietary APIs and are evaluated mainly with language modeling loss and synthetic tasks.The paper identifies a need for an open recipe with comprehensive downstream evaluation.
- Approach: Continual pretraining from Llama 2 uses 400 billion additional tokens in long training sequences, with sequence lengths up to 32,768 tokens.The 7B/13B variants use 32,768-token sequences, while the 34B/70B variants use 16,384-token sequences.
- Evaluation: The models are evaluated on language modeling, synthetic tasks, and real-world benchmarks spanning long- and short-context settings.This evaluation scope is broader than the language-modeling and synthetic-task focus described for prior open models.
- Results: The models show significant improvements on long-context tasks and modest improvements on standard short-context tasks compared with Llama 2.The short-context gains are especially noted for coding, math, and knowledge-intensive tasks.
- Results: Context length exhibits power-law scaling, with performance continuing to improve as context increases to 32,768 tokens.The result suggests context length is an additional scaling axis for these models.
2 Method
The method combines efficient continual pretraining, a modified positional encoding, data-mix analysis, and inexpensive instruction tuning to extend context while preserving useful capabilities.
- Continual Pretraining: Continual pretraining extends Llama 2 with longer sequences while avoiding the full computational cost of long-sequence pretraining from scratch.Longer sequences incur quadratic attention overhead, motivating the continual-pretraining strategy.
- Positional Encoding: A minimal RoPE modification decreases the rotation angle to reduce distant-token decay and improve long-context modeling.The paper describes this change as necessary for enabling the model to attend longer.
- Data Mix: Pretraining-data quality often matters more than text length for improving long-context continual pretraining.Ablations found that changing the length distribution alone provides no major benefit, while a new data mix improves many short-context tasks.
- Instruction Tuning: The instruction-tuning procedure combines Llama 2 Chat data with synthetic long question-answer data generated from long pretraining documents.The synthetic data is designed to transfer diverse skills from short instruction data to long-context scenarios.
- Instruction Tuning: Language-modeling loss on long input prompts, in addition to output-token loss, consistently improves downstream task performance.This differs from standard instruction tuning, which calculates loss only on output tokens.
3 Main Results
The models retain or improve performance on standard tasks while showing stronger long-context results as usable prompt length increases. Instruction tuning further yields competitive performance across benchmarked long-context tasks and human evaluations, with evaluation caveats for proprietary-model comparisons.
- Short Tasks: Overall, the models achieve on-par or stronger results than LLAMA 2 on standard benchmarks, with notable gains in coding, math, and MMLU.The authors attribute these improvements to additional computation and knowledge from newly introduced long data.
- Long Tasks: Long-context evaluation uses real-world tasks including NarrativeQA, QuALITY, Qasper, and QMSum rather than relying mainly on perplexity or synthetic tasks.The evaluation uses 0-shot, 1-shot, or 2-shot settings depending on average sample length.
- Effective Context Utilization: Performance on each long-context task improves monotonically as the maximum prompt context length increases.The figure evaluates performance while progressively increasing the context length available to prompts.
- Instruction Tuning Results: On ZeroSCROLLS, the instruction-tuned 70B model outperforms gpt-3.5-turbo-16k on 7 of 10 long-context tasks without human-annotated long-context instruction data.The model is particularly strong on question-answering tasks, matching the main focus of its self-instruct data.
- Evaluation Caveats: Long-context evaluation remains difficult because automatic metrics may poorly reflect response quality, truncation can remove answer-relevant information, and proprietary training data can obscure leakage.These caveats affect both benchmark interpretation and comparisons involving proprietary models.
- Human Evaluation: Human evaluation compares responses across multi-turn conversation and multi-document search, using helpfulness, honesty, and harmlessness judgments from three annotators per example.The evaluation covers 2,352 examples and reports win rates with 95% confidence intervals.
4 Analysis
The analysis isolates positional encoding, data composition, instruction tuning, and training curriculum as key design choices for long-context performance. Results favor adjusted RoPE, data quality over long-text abundance, short-data instruction tuning, and continual pretraining for efficiency.
- Positional Encoding: The original Llama 2 architecture could not effectively attend beyond 4,000–6,000 tokens after long-context continual pretraining.The analysis hypothesizes that RoPE’s decay of attention scores for distant tokens creates this bottleneck.
- Positional Encoding: RoPE ABF performs best among the explored positional encoding variants and maintains performance through the full 32,768-token FIRST-SENTENCE-RETRIEVAL context.It modifies RoPE by increasing the base frequency, reducing distant-token attention decay.
- Pretraining Data Mix: Increasing the long-text ratio provides no clear, consistent advantage, while removing most long texts still preserves most performance gains over Llama 2.The ablations compare altered length distributions using Llama 2 pretraining data.
- Pretraining Data Mix: Data quality, rather than length distribution, appears to explain most improvements from the new pretraining mix, including gains on knowledge-intensive short-context tasks.The conclusion is based on both long-context and standard short-context evaluations.
- Instruction Tuning: Short instruction data blended with pretraining data provides a cost-effective route to instruction-tuning long-context models without human-annotated long instruction data.The supplied evidence introduces this strategy and its comparison table, but does not report a specific quantitative result.
- Training Curriculum: Continual pretraining from short-context models saves around 40% FLOPs while imposing almost no performance loss relative to longer-sequence training curricula.Training-loss curves also show rapid adaptation after increasing the sequence length.
5 AI Safety
The paper evaluates safety and factuality using three established benchmarks plus internal long-context red teaming. The instruction-tuned model retains safety performance comparable to Llama 2 Chat, but the evaluation remains limited by the absence of dedicated long-context safety benchmarks.
- Safety Results: The instruction-tuned model maintains similar safety performance to Llama 2 Chat and is safer and less biased than Falcon-instruct and MPT-instruct.The comparison covers the reported safety evaluation across TruthfulQA, ToxiGen, and BOLD.
- Evaluation Scope: The authors characterize the three-benchmark safety evaluation as a pilot study rather than a comprehensive assessment.They identify dedicated long-context safety benchmarks as an unmet need for future work.
- Safety Benchmarks: TruthfulQA measures the percentage of generations that are both truthful and informative, while ToxiGen and BOLD assess toxicity and demographic-group sentiment.The three metrics use different preferred directions: higher TruthfulQA, lower ToxiGen, and BOLD closer to zero.
- Long-Context Red Teaming: Internal red teaming probes long-context jailbreak vulnerability by placing long inputs before adversarial prompts targeting harmful, hateful, and criminal behaviors.The paper notes that no open-source benchmark is currently designed specifically for long-context safety.
6 Limitations
The paper identifies limitations in application coverage, tokenizer efficiency, and hallucination. These constraints affect long-form use, the amount of text represented by 32,768 tokens, code processing, and reliability with dense inputs.
- Limited Functionality: The models are not yet fine-tuned for many long-context applications, including creative writing requiring long-form outputs.Applying existing alignment recipes to varied long-context scenarios is described as expensive and nontrivial.
- Tokenizer Efficiency: The Llama tokenizer often produces about 10% more tokens than GPT-3.5’s tokenizer, reducing the amount of text represented by the nominal context window.Its 32k-symbol vocabulary contributes to this efficiency difference.
- Tokenizer Efficiency: The tokenizer handles whitespace inefficiently, making the model inefficient for long code data.This is a specific boundary on code-oriented long-context processing.
- Hallucination: The model exhibits hallucinations, which may be more pronounced for long-context inputs because they contain dense information and receive insufficient alignment.The paper presents this as an observed issue shared with other LLMs.
7 Conclusion
The paper presents long-context LLMs built through a RoPE refinement and continual pretraining from Llama 2. They report strong performance across short- and long-context tasks, favorable comparisons with open-source models and gpt-3.5-turbo-16k, and analyses of the design factors shaping results.
- Conclusion: The model series uses a position encoding refinement and continual pretraining to achieve strong long-context performance.The study analyzes positional encodings, data mix, and pretraining curriculum alongside its model results.
- Conclusion: The models outperform Llama 2 on both short- and long-context tasks and compare favorably with gpt-3.5-turbo-16k on long-context tasks after simple instruction fine-tuning without human supervision.The conclusion also reports superior performance relative to existing open-source long-context models.
A More Results
This section reports benchmark comparisons across general knowledge, reasoning, coding, mathematics, commonsense, and long-context evaluations. The supplied passages identify the evaluated task groups but do not state their numerical outcomes.
- Table 13 compares the models with open-source long-context models on research benchmarks using prompts up to 32,768 tokens.
- Table 14 decomposes the models’ MMLU results.
- Table 15 reports HumanEval, MBPP, MATH, GSM8K, NaturalQuestions, and TriviaQA-wiki results with specified shot counts.
- Table 16 decomposes commonsense reasoning results using the same shots and evaluation metrics as Llama 2.
B Theoretical Analysis of Positional Encodings
The analysis frames long-context positional encoding as a trade-off between preserving the original embedding geometry and maintaining separation between consecutive positional images. It compares Position Interpolation and Adjusted Base Frequency through geometric distances and theoretical bounds.
- RoPE maps an input vector into an embedding curve on a sphere in C^d/2 parameterized by position and base frequency.
- Position Interpolation and Adjusted Base Frequency extend a trained transformer through alternative positional embedding curves parameterized by α and β.
- The desired extended embedding minimizes distance from the original mapping while keeping the distance between its closest images sufficiently large.
- Theoretical bounds scale as (log b + log β)^-1 for ABF and α(log b)^-1 for PI under the stated high-dimensional assumptions.
- With b = 10000, experiments use PI with α = 1/4 or 1/8 and ABF with β = 50, and both methods permit rapid adaptation.
- For those settings, PI has lower granularity than ABF, with approximate bounds 0.027 and 0.076 respectively; the authors hypothesize this relates to ABF’s stronger downstream evaluation.
C Length Extrapolation Results
The paper evaluates whether the 70B model can extrapolate beyond its training sequence length. It examines both position-wise validation loss and a strict synthetic retrieval task up to 32,768 tokens.
- Extrapolation is defined as inference on input sequences longer than the model’s training sequences.
- The 70B model’s extrapolation evaluation uses position-wise validation loss and the synthetic FIRST-SENTENCE-RETRIEVAL task.
- Validation-loss evaluation averages losses across 50 batches at each position of a 32,768-token sequence, with the first 16,384 tokens as interpolation and the second half as extrapolation.
- With either RoPE ABF or XPOS ABF, the 70B model maintains loss in the extrapolation region, unlike Llama 2 with a 4,096-token context window.
D Self-Instruct Data
The paper bootstraps long-context instruction data from Llama 2 Chat despite having only short-context models available. It chunks long documents, generates question–answer pairs, and uses those pairs for instruction tuning.
- The self-instruct procedure uses Llama 2 Chat to bootstrap data for instruction fine-tuning.
- Long documents are split into chunks that fit the short-context model before question–answer generation.
- Each chunk is processed with one of several prompts selected with equal probability to generate question–answer pairs.
- Generated answers are wrapped in XML tags to enable more accurate extraction.
- The constructed data includes templates for normal answer generation and long question–answer examples extracted from Llama 2 Chat responses.