Ingegneria dei Dati per il Ridimensionamento dei Modelli Linguistici a Contesti di 128K

Data Engineering for Scaling Language Models to 128K Context

February 15, 2024

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

cs.AI

Abstract

Studiamo la ricetta per il preaddestramento continuo per scalare la lunghezza del contesto dei modelli linguistici fino a 128K, con un focus sull'ingegneria dei dati. Ipotesizziamo che la modellazione di contesti lunghi, in particolare la capacità di utilizzare informazioni in posizioni arbitrarie dell'input, sia una capacità già in gran parte acquisita attraverso il preaddestramento su larga scala, e che questa capacità possa essere facilmente estesa a contesti sostanzialmente più lunghi di quelli osservati durante l'addestramento (ad esempio, da 4K a 128K) attraverso un preaddestramento continuo leggero su una miscela di dati appropriata. Investigiamo la quantità e la qualità dei dati per il preaddestramento continuo: (1) per la quantità, dimostriamo che da 500 milioni a 5 miliardi di token sono sufficienti per consentire al modello di recuperare informazioni ovunque all'interno del contesto di 128K; (2) per la qualità, i nostri risultati enfatizzano ugualmente l'equilibrio del dominio e il campionamento di lunghezze maggiori. Nello specifico, troviamo che il campionamento ingenuo di dati più lunghi su determinati domini come i libri, una pratica comune dei lavori esistenti, fornisce prestazioni subottimali, e che una miscela bilanciata di domini è importante. Dimostriamo che il preaddestramento continuo dell'intero modello su 1-5 miliardi di token di tali dati è una strategia efficace e conveniente per scalare la lunghezza del contesto dei modelli linguistici a 128K. La nostra ricetta supera i forti modelli open-source a contesto lungo e riduce il divario rispetto ai modelli all'avanguardia come 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.