RegMix: Data Mixture as Regression for Language Model Pre-training
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.
668375662ca1c52c2744c833
Qian Liu
claimed_verified
2024-07-02T07:38:41.749Z
612ee6a7b960e78c6d2319d4
https://cdn-avatars.huggingface.co/v1/production/uploads/612ee6a7b960e78c6d2319d4/2Hu9BaAyXbyh1vt0v1Qui.jpeg
Qian Liu
user
SivilTaram
668375662ca1c52c2744c834
Xiaosen Zheng
claimed_verified
2024-07-02T07:38:44.083Z
6274a2315d12b3a734adebc9
https://cdn-avatars.huggingface.co/v1/production/uploads/6274a2315d12b3a734adebc9/zLQDbszAvWh0F2BjdKock.jpeg
Xiaosen Zheng
user
xszheng2020
668375662ca1c52c2744c835
Niklas Muennighoff
admin_assigned
2024-07-03T09:26:53.541Z
5f1eb362eec0ad2a071ad6e2
https://cdn-avatars.huggingface.co/v1/production/uploads/5f1eb362eec0ad2a071ad6e2/IXMYkYKuTwn6kBdWnQeeY.png
Niklas Muennighoff
user
Muennighoff
668375662ca1c52c2744c836
Guangtao Zeng
668375662ca1c52c2744c837
Longxu Dou
admin_assigned
2024-07-03T09:27:12.212Z
6214e4ee1e35c843d42d1f88
https://cdn-avatars.huggingface.co/v1/production/uploads/6214e4ee1e35c843d42d1f88/fj-9wuIdPhvogh3BrcXTB.jpeg
Longxu Dou
user
dreamerdeo
668375662ca1c52c2744c838
Tianyu Pang
claimed_verified
2025-04-20T15:05:26.688Z
63d91b6d255ef6add20e1b38
https://cdn-avatars.huggingface.co/v1/production/uploads/1675921369867-63d91b6d255ef6add20e1b38.jpeg
Tianyu Pang
user
P2333
668375662ca1c52c2744c839
Jing Jiang
668375662ca1c52c2744c83a
Min Lin
RegMix: Data Mixture as Regression for Language Model Pre-training
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.
RegMix automatically identifies optimal data mixtures for pre-training large language models by formulating it as a regression task, significantly improving performance with reduced computational cost.
large language model pre-training data mixture regression task small models large-scale model compute budget parameter tokens Web corpora downstream performance domains scaling laws automatic approaches