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>.

Reinforced mixture learning

Model-based clustering and classification using mixtures of multivariate skewed power exponential distributions

Big Learning Expectation Maximization

Bayesian mixtures of common factor analyzers: Model, variational inference, and applications

On the Behavior of the Expectation-Maximization Algorithm for Mixture Models

Nearly Minimax Optimal Reinforcement Learning for Linear Mixture Markov Decision Processes

Dynamically regularized maximum likelihood learning of Gaussian mixtures

Learning Infinite-Horizon Average-Reward Linear Mixture MDPs of Bounded Span

An Efficient Em Approach To Parameter Learning Of The Mixture Of Gaussian Processes

Conditional mixture modeling and model-based clustering

Learning Mixtures of Ranking Models

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

Neural Mixture Models with Expectation-Maximization for End-to-end Deep Clustering

An Mcmc Based Em Algorithm For Mixtures Of Gaussian Processes

Improved Algorithm for Adversarial Linear Mixture MDPs with Bandit Feedback and Unknown Transition

EM for Mixture of Linear Regression with Clustered Data

A BYY Scale-Incremental EM Algorithm for Gaussian Mixture Learning

A Greedy Merge Learning Algorithm for Gaussian Mixture Model

MPMA: Mixture Probabilistic Matrix Approximation for Collaborative Filtering

Learning Mixed Multinomial Logits with Provable Guarantees

Robust Mixture Learning when Outliers Overwhelm Small Groups