Optimizing Pre-Training Data Mixtures with Mixtures of Data Expert Models
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2025
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866912241646829568 |
|---|---|
| author | Belenki, Lior Agarwal, Alekh Shi, Tianze Toutanova, Kristina |
| author_facet | Belenki, Lior Agarwal, Alekh Shi, Tianze Toutanova, Kristina |
| contents | We propose a method to optimize language model pre-training data mixtures through efficient approximation of the cross-entropy loss corresponding to each candidate mixture via a Mixture of Data Experts (MDE). We use this approximation as a source of additional features in a regression model, trained from observations of model loss for a small number of mixtures.
Experiments with Transformer decoder-only language models in the range of 70M to 1B parameters on the SlimPajama dataset show that our method achieves significantly better performance than approaches that train regression models using only the mixture rates as input features. Combining this improved optimization method with an objective that takes into account cross-entropy on end task data leads to superior performance on few-shot downstream evaluations.
We also provide theoretical insights on why aggregation of data expert predictions can provide good approximations to model losses for data mixtures. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2502_15950 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Optimizing Pre-Training Data Mixtures with Mixtures of Data Expert Models Belenki, Lior Agarwal, Alekh Shi, Tianze Toutanova, Kristina Machine Learning Computation and Language We propose a method to optimize language model pre-training data mixtures through efficient approximation of the cross-entropy loss corresponding to each candidate mixture via a Mixture of Data Experts (MDE). We use this approximation as a source of additional features in a regression model, trained from observations of model loss for a small number of mixtures. Experiments with Transformer decoder-only language models in the range of 70M to 1B parameters on the SlimPajama dataset show that our method achieves significantly better performance than approaches that train regression models using only the mixture rates as input features. Combining this improved optimization method with an objective that takes into account cross-entropy on end task data leads to superior performance on few-shot downstream evaluations. We also provide theoretical insights on why aggregation of data expert predictions can provide good approximations to model losses for data mixtures. |
| title | Optimizing Pre-Training Data Mixtures with Mixtures of Data Expert Models |
| topic | Machine Learning Computation and Language |
| url | https://arxiv.org/abs/2502.15950 |