RegMix : Mélange de données comme régression pour le pré-entraînement de modèles de langage

RegMix: Data Mixture as Regression for Language Model Pre-training

July 1, 2024

Auteurs: Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing Jiang, Min Lin

cs.AI

Résumé

Le mélange de données pour le pré-entraînement des grands modèles de langage a un impact significatif sur les performances, mais la manière de déterminer un mélange efficace reste incertaine. Nous proposons RegMix pour identifier automatiquement un mélange de données performant en le formulant comme une tâche de régression. RegMix implique l'entraînement d'un ensemble de petits modèles avec des mélanges de données diversifiés et l'ajustement d'un modèle de régression pour prédire leurs performances en fonction de leurs mélanges respectifs. Avec le modèle de régression ajusté, nous simulons le mélange le mieux classé et l'utilisons pour entraîner un modèle à grande échelle avec une puissance de calcul plusieurs ordres de grandeur supérieure. Pour valider empiriquement RegMix, nous entraînons 512 modèles de 1 million de paramètres sur 1 milliard de tokens de différents mélanges pour ajuster le modèle de régression et trouver le mélange optimal. En utilisant ce mélange, nous entraînons un modèle de 1 milliard de paramètres sur 25 milliards de tokens (c'est-à-dire 1000 fois plus grand et 25 fois plus long), qui s'avère être le meilleur parmi 64 modèles candidats de 1 milliard de paramètres avec d'autres mélanges. De plus, notre méthode démontre une performance supérieure par rapport à la sélection humaine et obtient des résultats qui égalent ou surpassent DoReMi, tout en utilisant seulement 10 % du budget de calcul. Nos expériences montrent également que (1) Les mélanges de données ont un impact significatif sur les performances, avec des variations de performance sur une seule tâche allant jusqu'à 14,6 % ; (2) Les corpus web, plutôt que des données perçues comme de haute qualité comme Wikipédia, ont la plus forte corrélation positive avec les performances en aval ; (3) Les domaines interagissent de manière complexe, contredisant souvent le bon sens, ce qui rend nécessaire des approches automatiques comme RegMix ; (4) Les effets du mélange de données transcendent les lois d'échelle, et notre approche capture cette complexité en considérant tous les domaines ensemble. Notre code est disponible à l'adresse https://github.com/sail-sg/regmix.

English

The data mixture for large language model pre-training significantly impacts performance, yet how to determine an effective mixture remains unclear. We propose RegMix to automatically identify a high-performing data mixture by formulating it as a regression task. RegMix involves training a set of small models with diverse data mixtures and fitting a regression model to predict their performance given their respective mixtures. With the fitted regression model, we simulate the top-ranked mixture and use it to train a large-scale model with orders of magnitude more compute. To empirically validate RegMix, we train 512 models with 1M parameters for 1B tokens of different mixtures to fit the regression model and find the optimal mixture. Using this mixture we train a 1B parameter model for 25B tokens (i.e. 1000x larger and 25x longer) which we find performs best among 64 candidate 1B parameter models with other mixtures. Further, our method demonstrates superior performance compared to human selection and achieves results that match or surpass DoReMi, while utilizing only 10% of the compute budget. Our experiments also show that (1) Data mixtures significantly impact performance with single-task performance variations of up to 14.6%; (2) Web corpora rather than data perceived as high-quality like Wikipedia have the strongest positive correlation with downstream performance; (3) Domains interact in complex ways often contradicting common sense, thus automatic approaches like RegMix are needed; (4) Data mixture effects transcend scaling laws, and our approach captures the complexity by considering all domains together. Our code is available at https://github.com/sail-sg/regmix.