RegMix: Datenmischung als Regression für das Vortraining von Sprachmodellen

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

July 1, 2024

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

cs.AI

Zusammenfassung

Die Datenmischung für das Vortrainieren großer Sprachmodelle beeinflusst die Leistung erheblich, aber wie man eine effektive Mischung bestimmt, bleibt unklar. Wir schlagen RegMix vor, um automatisch eine leistungsstarke Datenmischung zu identifizieren, indem wir sie als Regressionsaufgabe formulieren. RegMix beinhaltet das Training einer Reihe von kleinen Modellen mit unterschiedlichen Datenmischungen und das Anpassen eines Regressionsmodells, um ihre Leistung basierend auf ihren jeweiligen Mischungen vorherzusagen. Mit dem angepassten Regressionsmodell simulieren wir die am besten bewertete Mischung und verwenden sie, um ein groß angelegtes Modell mit Größenordnungen mehr Rechenleistung zu trainieren. Um RegMix empirisch zu validieren, trainieren wir 512 Modelle mit 1M Parametern für 1B Tokens unterschiedlicher Mischungen, um das Regressionsmodell anzupassen und die optimale Mischung zu finden. Unter Verwendung dieser Mischung trainieren wir ein 1B Parameter-Modell für 25B Tokens (d.h. 1000-mal größer und 25-mal länger), das sich als das beste unter 64 Kandidatenmodellen mit 1B Parametern und anderen Mischungen erweist. Darüber hinaus zeigt unsere Methode eine überlegene Leistung im Vergleich zur menschlichen Auswahl und erzielt Ergebnisse, die mit oder über DoReMi übereinstimmen, während nur 10% des Rechenbudgets genutzt werden. Unsere Experimente zeigen auch, dass (1) Datenmischungen die Leistung signifikant beeinflussen mit Leistungsunterschieden von bis zu 14,6 %; (2) Webkorpora anstelle von als hochwertig wahrgenommenen Daten wie Wikipedia die stärkste positive Korrelation mit der nachgelagerten Leistung aufweisen; (3) Domänen auf komplexe Weise interagieren und oft dem gesunden Menschenverstand widersprechen, weshalb automatische Ansätze wie RegMix erforderlich sind; (4) Die Effekte der Datenmischung überschreiten Skalierungsgesetze, und unser Ansatz erfasst die Komplexität, indem er alle Domänen zusammen betrachtet. Unser Code ist unter https://github.com/sail-sg/regmix verfügbar.

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.