LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens

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.

한 줄 요약

LongRoPE는 2048k 토큰의 컨텍스트 윈도우를 구현하면서 기존 성능을 유지하는 LLM 확장 기법이다.

핵심 기여도

LongRoPE를 통해 2048k 토큰의 컨텍스트 윈도우를 구현 (기존 128k 대비 16배 확장).
256k 토큰 길이에서 최대 1,000번의 fine-tuning 스텝으로 확장 가능.
RoPE 위치 임베딩의 두 가지 비균일성(non-uniformities)을 고려한 진화적 탐색 알고리즘 제안.
8× 확장이 fine-tuning 없이 가능하며, 8k 토큰 길이에서 성능 회복 기법 포함.

핵심 아이디어

LongRoPE는 기존 RoPE 임베딩의 비균일성을 활용해 컨텍스트 윈도우를 확장하는 새로운 접근법이다. 기존 연구는 RoPE의 회전 각도를 일정 비율로 선형 보간(linear interpolation)하거나, 특정 빈도 기반 그룹에 따라 보간을 적용하는 방식을 사용했으나, 이는 정보 손실을 초래했다. LongRoPE는 RoPE 차원과 토큰 위치의 비균일성을 동시에 고려하여, 각 차원과 위치에 맞는 최적의 스케일링 인수를 탐색한다. 이를 위해 진화적 탐색 알고리즘을 도입하여, 확장 비율에 따라 지수적으로 증가하는 탐색 공간을 효율적으로 탐색한다. 이는 fine-tuning 없이도 8×의 컨텍스트 확장을 가능하게 하며, 기존 성능을 유지하는 데 기여한다.

기술적 접근법

**RoPE 비균일성 탐색**: RoPE 차원과 토큰 위치에 따라 다른 스케일링 인수를 탐색.
**진화적 탐색 알고리즘**: 확장 비율에 따라 지수적으로 증가하는 탐색 공간을 효율적으로 탐색.
**진행적 확장 전략**: 256k 토큰 길이에서 fine-tuning 후, 2048k로의 2차 보간 수행.
**짧은 컨텍스트 성능 회복**: 8k 토큰 길이에서 RoPE 스케일링 인수 재조정.
**LLaMA2와 Mistral 모델 적용**: 기존 아키텍처 유지, 위치 임베딩에만 최소한의 수정.

주요 결과

**LLaMA2-2048k**: 4k에서 2048k 평가 길이에서 perplexity 유지, 90% 이상의 passkey retrieval 정확도 달성.
**표준 벤치마크**: 4096 토큰 내에서 기존 정확도 유지.
**비선형 보간 vs. PI/YaRN**: LongRoPE는 2048k 길이에서 perplexity가 일정 유지, PI/YaRN은 급격히 증가.
**짧은 길이 성능 회복**: 4k와 8k 길이에서 LongRoPE-2048k 모델의 perplexity가 개선됨.

의의 및 한계

LongRoPE는 기존 RoPE 기반 LLM에서 컨텍스트 윈도우를 2048k로 확장하면서 기존 성능을 유지하는 데 성공한 첫 번째 연구로, 대규모 텍스트 처리, in-context learning, LLM 에이전트 등 다양한 분야에 기여할 수 있다. 또한, 기존 최적화를 재사용할 수 있어 실용성도 높다. 그러나, 2048k 길이에서 토큰 위치 비균일성의 영향이 줄어들어 성능 개선 효과가 제한될 수 있으며, 초기 토큰만 보존하는 방식의 효과는 아직 연구 중이다.

실용적 활용

LongRoPE는 대량 텍스트 분석, 문서 요약, 다중 예시 기반 학습, LLM 에이전트 등 긴 컨텍스트가 필요한 산업 및 연구 분야에 적용 가능하다. 특히, 기존 RoPE 기반 LLM을 최소한의 수정으로 확장할 수 있어, 기존 모델 최적화를 유지하면서 새로운 애플리케이션 개발에 활용할 수 있다.