# Data Mixture Optimization: A Multi-fidelity Multi-scale Bayesian Framework

arXiv: https://arxiv.org/abs/2503.21023 · HTML: https://mlsanity.com/paper/2503.21023/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 49 references that are in the map.
- First posted on arXiv 26 Mar 2025.

## Builds on (most cited first, 30 of 31)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [A Tutorial on Bayesian Optimization](https://mlsanity.com/paper/1807.02811/) (2018)
- [Hyperband: A Novel Bandit-Based Approach to Hyperparameter Optimization](https://mlsanity.com/paper/1603.06560/) (2016)
- [BOHB: Robust and Efficient Hyperparameter Optimization at Scale](https://mlsanity.com/paper/1807.01774/) (2018)
- [TinyStories: How Small Can Language Models Be and Still Speak Coherent English?](https://mlsanity.com/paper/2305.07759/) (2023)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [Estimating the Carbon Footprint of BLOOM, a 176B Parameter Language Model](https://mlsanity.com/paper/2211.02001/) (2022)
- [SMAC3: A Versatile Bayesian Optimization Package for Hyperparameter Optimization](https://mlsanity.com/paper/2109.09831/) (2021)
- [Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance](https://mlsanity.com/paper/2403.16952/) (2024)
- [RegMix: Data Mixture as Regression for Language Model Pre-training](https://mlsanity.com/paper/2407.01492/) (2024)
- [Scaling Laws for Generative Mixed-Modal Language Models](https://mlsanity.com/paper/2301.03728/) (2023)
- [Observational Scaling Laws and the Predictability of Language Model Performance](https://mlsanity.com/paper/2405.10938/) (2024)
- [Multi-fidelity Bayesian Optimisation with Continuous Approximations](https://mlsanity.com/paper/1703.06240/) (2017)
- [Practical Multi-fidelity Bayesian Optimization for Hyperparameter Tuning](https://mlsanity.com/paper/1903.04703/) (2019)
- [Freeze-Thaw Bayesian Optimization](https://mlsanity.com/paper/1406.3896/) (2014)
- [Multi-Information Source Optimization](https://mlsanity.com/paper/1603.00389/) (2016)
- [Efficient Online Data Mixing For Language Model Pre-Training](https://mlsanity.com/paper/2312.02406/) (2023)
- [Scaling Laws for Data Filtering -- Data Curation cannot be Compute Agnostic](https://mlsanity.com/paper/2404.07177/) (2024)
- [YAHPO Gym -- An Efficient Multi-Objective Multi-Fidelity Benchmark for Hyperparameter Optimization](https://mlsanity.com/paper/2109.03670/) (2021)
- [Does your data spark joy? Performance gains from domain upsampling at the end of training](https://mlsanity.com/paper/2406.03476/) (2024)
- [Adaptive Data Optimization: Dynamic Sample Selection with Scaling Laws](https://mlsanity.com/paper/2410.11820/) (2024)
- [Improving Pretraining Data Using Perplexity Correlations](https://mlsanity.com/paper/2409.05816/) (2024)
- [Cost-aware Bayesian Optimization](https://mlsanity.com/paper/2003.10870/) (2020)
- [The Depth-to-Width Interplay in Self-Attention](https://mlsanity.com/paper/2006.12467/) (2020)
- [(Mis)Fitting: A Survey of Scaling Laws](https://mlsanity.com/paper/2502.18969/) (2025)

## Built on by (most cited first, 9 of 9)

- [ADMIRE-BayesOpt: Accelerated Data MIxture RE-weighting for Language Models with Bayesian Optimization](https://mlsanity.com/paper/2508.11551/) (2025)
- [AC-ODM: Actor--Critic Online Data Mixing for Sample-Efficient LLM Pretraining](https://mlsanity.com/paper/2505.23878/) (2025)
- [Support-Proximity Augmented Diffusion Estimation for Offline Black-Box Optimization](https://mlsanity.com/paper/2605.11246/) (2026)
- [BoLT: A Benchmark to Democratize Black-box Optimization Research for Expensive LLM Tasks](https://mlsanity.com/paper/2605.17000/) (2026)
- [Embedding by Elicitation: Dynamic Representations for Bayesian Optimization of System Prompts](https://mlsanity.com/paper/2605.19093/) (2026)
- [The Chicken and Egg Dilemma: Co-optimizing Data and Model Configurations for LLMs](https://mlsanity.com/paper/2602.08351/) (2026)
- [Evidence-Gated LLM Priors for Multi-Objective Bayesian Optimization](https://mlsanity.com/paper/2606.01730/) (2026)
- [Scaling Near-Optimal SFT-RL Annotation Budget Allocation from Small to Large LLMs](https://mlsanity.com/paper/2609.01573/) (2026)
- [Selecting Repetition Counts Across Model Scales in Data-Constrained Pretraining](https://mlsanity.com/paper/2610.05126/) (2026)

## Abstract (arXiv)

Careful curation of data sources can significantly improve the performance of LLM pre-training, but predominant approaches rely heavily on intuition or costly trial-and-error, making them difficult to generalize across different data domains and downstream tasks. Although scaling laws can provide a principled and general approach for data curation, standard deterministic extrapolation from small-scale experiments to larger scales requires strong assumptions on the reliability of such extrapolation, whose brittleness has been highlighted in prior works. In this paper, we introduce a textitprobabilistic extrapolation framework for data mixture optimization that avoids rigid assumptions and explicitly models the uncertainty in performance across decision variables. We formulate data curation as a sequential decision-making problemunicodex2013multi-fidelity, multi-scale Bayesian optimizationunicodex2013where data mixtures, model scale, training steps are adaptively selected to balance training cost and potential information gain. Our framework naturally gives rise to algorithm prototypes that leverage noisy information from inexpensive experiments to systematically inform costly training decisions. To accelerate methodological progress, we build a simulator based on 472 language model pre-training runs with varying data compositions from the SlimPajama dataset. We observe that even simple kernels and acquisition functions can enable principled decisions across training models from 20M to 1B parameters and achieve 2.6x and 3.3x speedups compared to multi-fidelity BO and random search baselines. Taken together, our framework underscores potential efficiency gains achievable by developing principled and transferable data mixture optimization methods.