LongRoPE: Extensión de la Ventana de Contexto de LLM Más Allá de 2 Millones de Tokens

LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens

February 21, 2024

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

cs.AI

Resumen

Una ventana de contexto grande es una característica deseable en los modelos de lenguaje de gran escala (LLMs). Sin embargo, debido a los altos costos de ajuste fino, la escasez de textos largos y los valores catastróficos introducidos por nuevas posiciones de tokens, las ventanas de contexto extendidas actuales están limitadas a alrededor de 128k tokens. Este artículo presenta LongRoPE, que, por primera vez, extiende la ventana de contexto de LLMs preentrenados a un impresionante 2048k tokens, con un máximo de solo 1k pasos de ajuste fino dentro de longitudes de entrenamiento de 256k, manteniendo el rendimiento en la ventana de contexto corta original. Esto se logra mediante tres innovaciones clave: (i) identificamos y aprovechamos dos formas de no uniformidades en la interpolación posicional a través de una búsqueda eficiente, proporcionando una mejor inicialización para el ajuste fino y permitiendo una extensión de 8x en escenarios sin ajuste fino; (ii) introducimos una estrategia de extensión progresiva que primero ajusta un LLM de 256k de longitud y luego realiza una segunda interpolación posicional en el LLM extendido ajustado para alcanzar una ventana de contexto de 2048k; (iii) reajustamos LongRoPE en una longitud de 8k para recuperar el rendimiento de la ventana de contexto corta. Experimentos extensos en LLaMA2 y Mistral en diversas tareas demuestran la efectividad de nuestro método. Los modelos extendidos mediante LongRoPE conservan la arquitectura original con modificaciones menores en la incrustación posicional y pueden reutilizar la mayoría de las optimizaciones preexistentes.

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.