Extensión de la Ventana de Contexto de los LLMs con 100 Muestras

Extending LLMs' Context Window with 100 Samples

January 13, 2024

Autores: Yikai Zhang, Junlong Li, Pengfei Liu

cs.AI

Resumen

Se sabe que los Modelos de Lenguaje de Gran Escala (LLMs, por sus siglas en inglés) tienen una capacidad limitada de extrapolación más allá de su ventana de contexto preentrenada, lo que restringe su aplicación en tareas posteriores con entradas extensas. Estudios recientes han buscado extender la ventana de contexto de los LLMs modificando el método de codificación de posición conocido como rotary position embedding (RoPE), adoptado por modelos populares como LLaMA, PaLM y GPT-NeoX. Sin embargo, trabajos previos como Position Interpolation (PI) y YaRN son intensivos en recursos y carecen de experimentos comparativos para evaluar su aplicabilidad. En este trabajo, identificamos la necesidad inherente de que la entropía de atención de los LLMs (es decir, la entropía de la información de los puntajes de atención) mantenga estabilidad e introducimos una extensión novedosa de RoPE que combina el ajuste de la frecuencia base de RoPE y la escalación de los logits de atención para ayudar a los LLMs a adaptarse eficientemente a una ventana de contexto más amplia. Validamos la superioridad de nuestro método tanto en el rendimiento de ajuste fino como en la robustez en diferentes tamaños de ventana de contexto en diversas tareas que demandan contexto. Notablemente, nuestro método extiende la ventana de contexto de LLaMA-2-7B-Chat a 16,384 con solo 100 muestras y 6 pasos de entrenamiento, demostrando una eficiencia extraordinaria. Finalmente, también exploramos cómo las composiciones de datos y los planes de entrenamiento afectan la extensión de la ventana de contexto para tareas específicas posteriores, sugiriendo que el ajuste fino de LLMs con conversaciones extensas es un buen punto de partida. Publicamos nuestro código y datos de SFT en https://github.com/GAIR-NLP/Entropy-ABF.

English

Large Language Models (LLMs) are known to have limited extrapolation ability beyond their pre-trained context window, constraining their application in downstream tasks with lengthy inputs. Recent studies have sought to extend LLMs' context window by modifying rotary position embedding (RoPE), a popular position encoding method adopted by well-known LLMs such as LLaMA, PaLM, and GPT-NeoX. However, prior works like Position Interpolation (PI) and YaRN are resource-intensive and lack comparative experiments to assess their applicability. In this work, we identify the inherent need for LLMs' attention entropy (i.e. the information entropy of attention scores) to maintain stability and introduce a novel extension to RoPE which combines adjusting RoPE's base frequency and scaling the attention logits to help LLMs efficiently adapt to a larger context window. We validate the superiority of our method in both fine-tuning performance and robustness across different context window sizes on various context-demanding tasks. Notably, our method extends the context window of LLaMA-2-7B-Chat to 16,384 with only 100 samples and 6 training steps, showcasing extraordinary efficiency. Finally, we also explore how data compositions and training curricula affect context window extension for specific downstream tasks, suggesting fine-tuning LLMs with lengthy conversations as a good starting point. We release our code and SFT data at https://github.com/GAIR-NLP/Entropy-ABF.