LongRoPE: Het uitbreiden van het contextvenster van LLM's tot meer dan 2 miljoen tokens

LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens

February 21, 2024

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

cs.AI

Samenvatting

Een groot contextvenster is een gewenste eigenschap in grote taalmodellen (LLM's). Vanwege de hoge kosten voor fine-tuning, de schaarste aan lange teksten en catastrofale waarden die worden geïntroduceerd door nieuwe tokenposities, zijn huidige uitgebreide contextvensters beperkt tot ongeveer 128k tokens. Dit artikel introduceert LongRoPE, dat voor het eerst het contextvenster van vooraf getrainde LLM's uitbreidt naar een indrukwekkende 2048k tokens, met maximaal slechts 1k fine-tuningstappen binnen een trainingslengte van 256k, terwijl de prestaties binnen het oorspronkelijke korte contextvenster behouden blijven. Dit wordt bereikt door drie belangrijke innovaties: (i) we identificeren en benutten twee vormen van niet-uniformiteit in positionele interpolatie via een efficiënte zoektocht, wat een betere initialisatie voor fine-tuning biedt en een 8x uitbreiding mogelijk maakt in scenario's zonder fine-tuning; (ii) we introduceren een progressieve uitbreidingsstrategie die eerst een LLM van 256k lengte fine-tunt en vervolgens een tweede positionele interpolatie uitvoert op het gefinetunde uitgebreide LLM om een contextvenster van 2048k te bereiken; (iii) we passen LongRoPE aan op een lengte van 8k om de prestaties van het korte contextvenster te herstellen. Uitgebreide experimenten op LLaMA2 en Mistral over diverse taken demonstreren de effectiviteit van onze methode. Modellen die via LongRoPE zijn uitgebreid, behouden de oorspronkelijke architectuur met minimale aanpassingen aan de positionele embedding en kunnen de meeste bestaande optimalisaties hergebruiken.

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.