RegMix: Miscelazione di Dati come Regressione per il Pre-Addestramento di Modelli Linguistici

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

July 1, 2024

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

cs.AI

Abstract

La miscela di dati per il pre-addestramento di modelli linguistici di grandi dimensioni influisce significativamente sulle prestazioni, ma come determinare una miscela efficace rimane poco chiaro. Proponiamo RegMix per identificare automaticamente una miscela di dati ad alte prestazioni formulandola come un task di regressione. RegMix prevede l'addestramento di un insieme di piccoli modelli con diverse miscele di dati e l'adattamento di un modello di regressione per prevedere le loro prestazioni date le rispettive miscele. Con il modello di regressione adattato, simuliamo la miscela con il punteggio più alto e la utilizziamo per addestrare un modello su larga scala con un ordine di grandezza maggiore di risorse computazionali. Per validare empiricamente RegMix, addestriamo 512 modelli con 1 milione di parametri per 1 miliardo di token di diverse miscele per adattare il modello di regressione e trovare la miscela ottimale. Utilizzando questa miscela, addestriamo un modello da 1 miliardo di parametri per 25 miliardi di token (cioè 1000 volte più grande e 25 volte più lungo) che risulta essere il migliore tra 64 modelli candidati da 1 miliardo di parametri con altre miscele. Inoltre, il nostro metodo dimostra prestazioni superiori rispetto alla selezione umana e ottiene risultati che eguagliano o superano DoReMi, utilizzando solo il 10% del budget computazionale. I nostri esperimenti mostrano anche che: (1) Le miscele di dati influiscono significativamente sulle prestazioni, con variazioni delle prestazioni su singoli task fino al 14,6%; (2) I corpora web, piuttosto che dati percepiti come di alta qualità come Wikipedia, hanno la correlazione positiva più forte con le prestazioni downstream; (3) I domini interagiscono in modi complessi spesso contrari al senso comune, rendendo necessari approcci automatici come RegMix; (4) Gli effetti della miscela di dati trascendono le leggi di scala, e il nostro approccio cattura la complessità considerando tutti i domini insieme. Il nostro codice è disponibile all'indirizzo 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.