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

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

2024-07-02

Abstract: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 <a class="link-external link-https" href="https://github.com/sail-sg/regmix" rel="external noopener nofollow">this https URL</a>.

Computation and Language,Artificial Intelligence

What problem does this paper attempt to address?

The paper mainly explores how to effectively select data mixtures in large-scale language model pre-training, which is a key factor affecting model performance. A method called REGMIX is proposed to automatically identify high-performing data mixtures by transforming the data mixture problem into a regression task. Specifically, REGMIX trains a set of small models with different data mixtures and fits a regression model using the performance of these models and their corresponding mixtures to predict the impact of other possible data mixtures on performance. Through this method, the optimal data mixture can be found to train large-scale models without requiring a large amount of computational resources.

Experimental results show that REGMIX outperforms manual selection and performs equally or better than the DoReMi method with less computational budget. In addition, the study also found that:

1. Data mixtures have a significant impact on model performance, with differences in performance of individual tasks reaching up to 14.6%.
2. Compared to Wikipedia, general web corpora (such as CommonCrawl) have a stronger positive correlation with performance improvement in downstream tasks.
3. The interaction between data domains is complex and requires automated methods (such as REGMIX) to handle.
4. The effect of data mixture goes beyond the law of scale, and REGMIX captures this complexity by considering all domains together.

In conclusion, the paper attempts to address the problem of how to effectively and automatically select data mixture strategies to improve model performance in large-scale language model pre-training.