LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens

Yiran Ding; Li Lyna Zhang; Chengruidong Zhang; Yuanyuan Xu; Ning Shang; Jiahang Xu; Fan Yang; Mao Yang

Source-linked AI summary

LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens

Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, Mao Yang

arXiv:2402.13753v1Feb 21, 2024cs.CL

TL;DR

Pretrained LLMs are difficult to extend far beyond 128k because new positions introduce catastrophic values, long training texts are scarce, and long-sequence fine-tuning is expensive. LongRoPE searches non-uniform RoPE interpolations and progressively extends models through 256k fine-tuning plus a second interpolation, reaching 2048k. Across LLaMA2 and Mistral experiments, it maintains long-context performance and comparable short-context benchmark accuracy.

Problem
Context extension beyond roughly 128k is limited by untrained-position failures, scarce corresponding long texts, and expensive fine-tuning.
Method
LongRoPE searches RoPE rescale factors across dimensions and token positions, fine-tunes at 256k, then applies a second interpolation to reach 2048k.
Results
LongRoPE maintains low perplexity from 4k to 2048k, achieves over 90% passkey retrieval accuracy, and delivers comparable accuracy on 4096-token benchmarks.
Takeaways & Limitations
The method extends LLaMA2 and Mistral to 2048k while retaining the original architecture with minor positional-embedding modifications.
Takeaways & Limitations
A 512× extension crowds positions within the original 4k window and can degrade short-context performance.

Abstract

from arXiv · show

Large context window is a desirable feature in large language models (LLMs). However, due to high fine-tuning costs, scarcity of long texts, and catastrophic values introduced by new token positions, current extended context windows are limited to around 128k tokens. This paper introduces LongRoPE that, for the first time, extends the context window of pre-trained LLMs to an impressive 2048k tokens, with up to only 1k fine-tuning steps at within 256k training lengths, while maintaining performance at the original short context window. This is achieved by three key innovations: (i) we identify and exploit two forms of non-uniformities in positional interpolation through an efficient search, providing a better initialization for fine-tuning and enabling an 8x extension in non-fine-tuning scenarios; (ii) we introduce a progressive extension strategy that first fine-tunes a 256k length LLM and then conducts a second positional interpolation on the fine-tuned extended LLM to achieve a 2048k context window; (iii) we readjust LongRoPE on 8k length to recover the short context window performance. Extensive experiments on LLaMA2 and Mistral across various tasks demonstrate the effectiveness of our method. Models extended via LongRoPE retain the original architecture with minor modifications to the positional embedding, and can reuse most pre-existing optimizations.

1. Introduction

LongRoPE addresses the obstacles limiting pretrained LLM context extension beyond roughly 128k tokens by exploiting positional non-uniformities and progressively extending context to 2048k. Experiments show low perplexity across 4k–2048k, strong passkey retrieval, and comparable short-context benchmark accuracy.

Motivation: Beyond roughly 128k tokens, context extension is constrained by catastrophic values from untrained positions, scarce corresponding long texts, and expensive long-sequence fine-tuning.Extending from 4k to over 1000k introduces more than 90% new positions.
Method: LongRoPE searches dimension- and token-position-dependent RoPE rescale factors to preserve important positional information and improve fine-tuning initialization.An evolutionary search addresses the exponentially expanding search space.
Method: A progressive strategy fine-tunes at 256k, then performs a second positional interpolation to reach a 2048k context window without direct fine-tuning on extremely long texts.The second stage exploits the method’s 8× non-fine-tuning extension capability.
Method: LongRoPE readjusts RoPE at shorter lengths to mitigate degradation within the original context window.Inference sequences below 8k use updated searched rescale factors.
Results: LongRoPE maintains low perplexity from 4k to 2048k, exceeds 90% passkey retrieval accuracy, and achieves comparable accuracy on benchmarks designed within 4096 tokens.The method applies across different RoPE-based LLMs, including LLaMA2 and Mistral.

2. Non-uniformity in Positional Interpolation

The study finds that positional interpolation is non-uniform across RoPE dimensions and token positions, and that searching these patterns improves context extension. These findings support better perplexity without fine-tuning and better initialization for fine-tuning.

2.1. Preliminary: RoPE positional interpolation methods rescale rotation frequencies differently: PI applies uniform linear interpolation, while NTK and YaRN vary interpolation across frequency dimensions.PI can crowd position information at high extension ratios.
2.1. Preliminary: Current interpolation rules are sub-optimal because RoPE dimensions exhibit substantial non-uniformities that human-designed methods do not fully exploit.Evolution search is used to discover improved non-uniform interpolations.
2.1. Preliminary: Searched rescale factors improve LLaMA2 perplexity at 8k and 16k without fine-tuning by preserving original RoPE information, especially in key dimensions.The searched factors differ from PI, NTK, and YaRN’s fixed or group-wise scaling.
2.2. Study on Non-uniform Positional Interpolation: Initial token positions benefit from less interpolation because they receive large attention scores and are important to attention layers.Retaining starting tokens without interpolation improves both PI and Dynamic-NTK.
2.2. Study on Non-uniform Positional Interpolation: Non-uniform positional interpolation extends context effectively in both fine-tuning and non-fine-tuning settings.At 64k, the searched method significantly outperforms PI and YaRN before and after fine-tuning LLaMA2-7B.
Summary: The study concludes that varying RoPE dimensions and token positions are the two non-uniformities most useful for improving context extension.Using both provides better initialization for fine-tuning and enables an 8× extension without fine-tuning.

3. LongRoPE

LongRoPE searches non-uniform RoPE rescale factors across dimensions and token positions, then progressively extends pretrained LLMs to 2048k context with limited fine-tuning. It also achieves an 8× extension without fine-tuning, while addressing short-context degradation after extreme scaling.

Search formulation: LongRoPE searches dimension- and position-dependent RoPE rescale factors to minimize perplexity at a target context length.The formulation preserves original RoPE for initial token positions and applies rescaling after a learned threshold.
Search formulation: The search space assigns each RoPE dimension a rescale factor and each solution a threshold for retaining initial positions without interpolation.Factors λ_i range from direct extrapolation to values larger than standard position interpolation, while the threshold is searched over discrete candidates.
Search algorithm: Evolutionary search combines seeded initialization, mutation, crossover, and a monotonic constraint to reduce perplexity evaluations in the exponentially large solution space.The initial population includes PI, NTK, and YaRN factors, and valid solutions require λ_i ≤ λ_i+1 across RoPE dimensions.
Non-fine-tuning extension: 8× extension without fine-tuning lets LongRoPE extend LLaMA2 from 4k to 32k while existing methods show perplexity spikes after 2× extension.The method attributes this result to preserving key dimensions and token positions through non-uniform rescaling.
Progressive extension: 2048k context is reached progressively by extending to 256k, fine-tuning for 1,000 steps within 256k training length, and performing a secondary search without further fine-tuning.The staged fine-tuning uses 400 steps with 128k factors followed by 600 steps with 256k factors.
Shorter context window recovery: Extreme 512× interpolation crowds positions within the original 4k window and causes short-context performance degradation.This motivates subsequent shorter-window recovery through adjusted RoPE rescale factors.

4. Experiments

Experiments on LLaMA2-7B and Mistral-7B evaluate LongRoPE across perplexity, passkey retrieval, and short-context benchmarks. The method reaches 2048k contexts while preserving strong long- and short-context performance, with results supported by ablations of its positional interpolation strategy.

Long sequence language modeling within 256k: LongRoPE-2048k models outperform state-of-the-art baselines within 256k context length despite a 16× longer context window.This pattern is reported across Proof-pile and PG19 perplexity evaluations.
Long sequence language modeling beyond 2000k: LongRoPE extends LLaMA2-7B and Mistral-7B to 2048k, with perplexity comparable or superior to baselines from 8k to 128k.LLaMA2 perplexity generally decreases with longer context, while Mistral exceeds 7 beyond 256k.
Long sequence language modeling beyond 2000k: Mistral performs better with 128k rather than 256k fine-tuning because both settings use 16k training lengths following YaRN.The authors attribute this to the effect of the training length on further context extension.
Passkey retrieval: Existing models’ passkey accuracy drops to 0 beyond 128k, whereas LongRoPE-LLaMA2-2048k maintains at least 90% accuracy from 4k to 2048k.LongRoPE-Mistral-2048k retains 100% accuracy through 1800k and reaches 60% at 2048k.
Standard benchmarks within original context window: LongRoPE models achieve comparable accuracy on original 4096-context benchmarks and outperform the original Mistral on TruthfulQA by +0.5%.The LLaMA2 model fine-tuned at 256k shows somewhat more degradation but remains reasonable on most tasks.
Ablation results: Non-uniform positional interpolation sustains perplexity during secondary extension, while PI and YaRN perplexity quickly increases with extension ratio.The ablation evaluates extending a fine-tuned LLaMA2-256k model to 512k, 1024k, and 2048k.
Ablation results: Readjusting RoPE factors with less interpolation significantly improves perplexity at 4k and 8k and average benchmark accuracy.This recovery targets short-context performance after extending LongRoPE to 2048k.
Ablation results: RoPE-dimension non-uniformity reduces perplexity versus linear interpolation, while token-position non-uniformity helps at 16k and 32k but not 2048k.The study also finds that preserving only initial tokens without interpolation is not useful.

5. Related Works

Related work extends context through positional interpolation, attention manipulation, retrieval memory, or long-context fine-tuning. LongRoPE differs by using minor positional-embedding modifications and is described as flexible and complementary to efficient fine-tuning methods.

Retrieval-based approaches: Retrieval-based approaches add external memory and retrieval modules, whereas LongRoPE modifies positional embeddings and supports tasks beyond retrieval.The paper specifically mentions long-document summarization and few-shot applications.
Attention-based context window extensions: Attention-based methods manipulate attention mechanisms and masks to address attention explosion from new positions; these methods are complementary to positional interpolation.The cited approaches retain the original LLM context window length while changing attention behavior.
Fine-tuning based approaches: Other fine-tuning approaches use large RoPE bases or target-length fine-tuning, while LongRoPE supports flexible target lengths and extensions beyond 2M.LongRoPE is presented as orthogonal to efficient fine-tuning methods such as LongLoRA and PoSE.

6. Conclusion

The paper presents LongRoPE as a method that extends LLM context to 2048k while maintaining shorter-context capabilities through non-uniform positional interpolation and progressive extension.

Conclusion: LongRoPE reaches a 2048k context window while maintaining capabilities within the original shorter context window.The method uses minor positional-embedding modifications rather than changing the original architecture.
Conclusion: Efficient evolutionary search exploits non-uniformities across RoPE dimensions and token positions, improving fine-tuning initialization and enabling an 8× extension without fine-tuning.The search identifies rescale factors for RoPE rotation angles based on token positions.
Conclusion: A progressive strategy uses 256k-length fine-tuned models to reach 2048k without extra fine-tuning.The conclusion frames this strategy as the basis for the final context-window extension.

A.1. Settings

The settings describe the passkey-retrieval prompt and computational environment used for long-context evaluation. The task repeats irrelevant text around a randomly sampled passkey whose location and document length vary.

Environments: Experiments run on 16 A100 GPUs using Flash Attention-2 and an internal CUBE platform to reduce long-sequence training and inference costs.The paper notes that memory and computation time increase exponentially with sequence length, making contexts beyond 512k challenging to serve.
Passkey prompt: The passkey prompt asks the model to memorize important information hidden inside irrelevant text and answer a later query.The template frames the task as finding and retaining a concealed fact.
Passkey prompt: The prompt repeats sentences around a hidden five-digit passkey, with separate repetition counts controlling document length before and after the key.The example uses 17865 as the passkey.
Passkey prompt: The document length varies with x and y, while the passkey number is randomly sampled and changes at each test.The query asks the model to retrieve the concealed number.

A.2. Additional details on fine-tuning

Fine-tuning experiments compare progressive and direct strategies for extending LLaMA2 and Mistral to 128k and 256k contexts. Starting 256k fine-tuning from a 128k checkpoint is reported as most effective, with substantial resource requirements.

128k fine-tuning incurs a large initial loss from the 32× extension, but the loss rapidly decreases afterward.
Using 128k text lengths for a 256k context causes a sharp initial loss increase, while direct fine-tuning to 256k decreases loss relatively slowly.
Fine-tuning a 256k model from a 128k checkpoint is more effective than the tested direct and truncated-text alternatives according to Proof-Pile perplexity.Table 12 reports test perplexity for checkpoints from the three settings.
LLaMA2-128k requires 8 A100 GPUs for one week and 400 steps, while LLaMA2-256k uses 16 GPUs for two weeks and 600 steps.Mistral 128k and 256k use 4 A100 GPUs for two days with a 16k training length.

A.3. Additional details on the search

The evolution search efficiently finds non-uniform RoPE rescale factors, substantially lowering validation perplexity for 256k extension. At 2048k, the 8× extension from a 256k model performs better than the 16× extension from a 128k model, while evaluation cost limits search iterations.

After one iteration, the 256k search finds solutions significantly better than PI and YaRN, and further iterations reduce validation perplexity from 273.27 to 118.47.
At 2048k, the 16× extension has higher perplexity than the 8× extension from the fine-tuned 128k and 256k models, respectively.
A single perplexity evaluation at 2048k takes about 50 minutes, constraining the number of search iterations.The passage states that better results are possible if more search time is allowed.
Search for context lengths up to 256k takes within 3 days on one A100 GPU, while 1024k and 2048k searches use 4 and 8 GPUs within a 5-day limit.A 512k search uses 2 A100 GPUs.