LongRoPE: 200만 토큰 이상의 LLM 컨텍스트 윈도우 확장

LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens

February 21, 2024

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

cs.AI

초록

대규모 언어 모델(LLMs)에서 넓은 컨텍스트 창은 바람직한 기능이다. 그러나 높은 파인튜닝 비용, 긴 텍스트의 부족, 그리고 새로운 토큰 위치에 의해 도입된 치명적인 값들로 인해, 현재 확장된 컨텍스트 창은 약 128k 토큰으로 제한되어 있다. 본 논문은 LongRoPE를 소개하며, 이는 사전 훈련된 LLMs의 컨텍스트 창을 처음으로 2048k 토큰까지 확장하고, 256k 훈련 길이 내에서 최대 1k 파인튜닝 단계만으로 원래의 짧은 컨텍스트 창에서의 성능을 유지한다. 이는 세 가지 주요 혁신을 통해 달성된다: (i) 효율적인 탐색을 통해 위치 보간에서 두 가지 형태의 비균일성을 식별하고 활용하여, 파인튜닝을 위한 더 나은 초기화를 제공하고 비파인튜닝 시나리오에서 8배의 확장을 가능하게 한다; (ii) 256k 길이의 LLM을 먼저 파인튜닝한 후, 파인튜닝된 확장 LLM에 두 번째 위치 보간을 수행하여 2048k 컨텍스트 창을 달성하는 점진적 확장 전략을 도입한다; (iii) 8k 길이에서 LongRoPE를 재조정하여 짧은 컨텍스트 창 성능을 회복한다. LLaMA2와 Mistral을 대상으로 한 다양한 작업에서의 광범위한 실험은 본 방법의 효과를 입증한다. LongRoPE를 통해 확장된 모델은 위치 임베딩에 대한 사소한 수정만으로 원래의 아키텍처를 유지하며, 대부분의 기존 최적화를 재사용할 수 있다.

English

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.