Ingeniería de Datos para Escalar Modelos de Lenguaje a Contextos de 128K

Data Engineering for Scaling Language Models to 128K Context

February 15, 2024

Autores: Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, Hao Peng

cs.AI

Resumen

Estudiamos la receta de preentrenamiento continuo para escalar la longitud de contexto de los modelos de lenguaje a 128K, con un enfoque en la ingeniería de datos. Planteamos la hipótesis de que el modelado de contextos largos, en particular la capacidad de utilizar información en ubicaciones arbitrarias de la entrada, es una habilidad que en su mayoría ya se adquiere a través del preentrenamiento a gran escala, y que esta capacidad puede extenderse fácilmente a contextos sustancialmente más largos que los vistos durante el entrenamiento (por ejemplo, de 4K a 128K) mediante un preentrenamiento continuo ligero en una mezcla de datos apropiada. Investigamos la cantidad y calidad de los datos para el preentrenamiento continuo: (1) en cuanto a la cantidad, demostramos que entre 500 millones y 5 mil millones de tokens son suficientes para permitir que el modelo recupere información en cualquier lugar dentro del contexto de 128K; (2) en cuanto a la calidad, nuestros resultados enfatizan por igual el equilibrio de dominios y el aumento de longitud. Concretamente, encontramos que el aumento ingenuo de datos más largos en ciertos dominios como libros, una práctica común en trabajos existentes, ofrece un rendimiento subóptimo, y que una mezcla equilibrada de dominios es importante. Demostramos que el preentrenamiento continuo del modelo completo en 1B-5B tokens de dichos datos es una estrategia efectiva y asequible para escalar la longitud de contexto de los modelos de lenguaje a 128K. Nuestra receta supera a modelos de contexto largo de código abierto robustos y reduce la brecha con modelos de vanguardia como GPT-4 128K.

English

We study the continual pretraining recipe for scaling language models' context lengths to 128K, with a focus on data engineering. We hypothesize that long context modeling, in particular the ability to utilize information at arbitrary input locations, is a capability that is mostly already acquired through large-scale pretraining, and that this capability can be readily extended to contexts substantially longer than seen during training~(e.g., 4K to 128K) through lightweight continual pretraining on appropriate data mixture. We investigate the quantity and quality of the data for continual pretraining: (1) for quantity, we show that 500 million to 5 billion tokens are enough to enable the model to retrieve information anywhere within the 128K context; (2) for quality, our results equally emphasize domain balance and length upsampling. Concretely, we find that naively upsampling longer data on certain domains like books, a common practice of existing work, gives suboptimal performance, and that a balanced domain mixture is important. We demonstrate that continual pretraining of the full model on 1B-5B tokens of such data is an effective and affordable strategy for scaling the context length of language models to 128K. Our recipe outperforms strong open-source long-context models and closes the gap to frontier models like GPT-4 128K.