Ingénierie des données pour la mise à l'échelle des modèles de langage à un contexte de 128K

Data Engineering for Scaling Language Models to 128K Context

February 15, 2024

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

cs.AI

Résumé

Nous étudions la méthode de pré-entraînement continu pour étendre la longueur de contexte des modèles de langage à 128K, en nous concentrant sur l'ingénierie des données. Nous émettons l'hypothèse que la modélisation de contextes longs, en particulier la capacité à utiliser des informations à des emplacements arbitraires de l'entrée, est une compétence principalement acquise lors d'un pré-entraînement à grande échelle, et que cette compétence peut être facilement étendue à des contextes nettement plus longs que ceux rencontrés pendant l'entraînement (par exemple, de 4K à 128K) grâce à un pré-entraînement continu léger sur un mélange de données approprié. Nous examinons la quantité et la qualité des données pour le pré-entraînement continu : (1) en termes de quantité, nous montrons que 500 millions à 5 milliards de tokens suffisent pour permettre au modèle de récupérer des informations n'importe où dans le contexte de 128K ; (2) en termes de qualité, nos résultats mettent également l'accent sur l'équilibre des domaines et le suréchantillonnage de la longueur. Concrètement, nous constatons que le suréchantillonnage naïf de données plus longues dans certains domaines comme les livres, une pratique courante des travaux existants, donne des performances sous-optimales, et qu'un mélange équilibré de domaines est important. Nous démontrons que le pré-entraînement continu du modèle complet sur 1 à 5 milliards de tokens de telles données est une stratégie efficace et abordable pour étendre la longueur de contexte des modèles de langage à 128K. Notre méthode surpasse les modèles open-source à contexte long performants et réduit l'écart avec les modèles de pointe comme 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.