# ADMIRE-BayesOpt: Accelerated Data MIxture RE-weighting for Language Models with Bayesian Optimization

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 74 references that are in the map.
- First posted on arXiv 15 Aug 2025.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [Overcoming catastrophic forgetting in neural networks](https://mlsanity.com/paper/1612.00796/) (2016)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Experience Replay for Continual Learning](https://mlsanity.com/paper/1811.11682/) (2018)
- [Tulu 3: Pushing Frontiers in Open Language Model Post-Training](https://mlsanity.com/paper/2411.15124/) (2024)
- [Quantifying Memorization Across Neural Language Models](https://mlsanity.com/paper/2202.07646/) (2022)
- [A Tutorial on Bayesian Optimization](https://mlsanity.com/paper/1807.02811/) (2018)
- [Phi-4 Technical Report](https://mlsanity.com/paper/2412.08905/) (2024)
- [Hyperband: A Novel Bandit-Based Approach to Hyperparameter Optimization](https://mlsanity.com/paper/1603.06560/) (2016)
- [Deep Learning on a Data Diet: Finding Important Examples Early in Training](https://mlsanity.com/paper/2107.07075/) (2021)
- [LESS: Selecting Influential Data for Targeted Instruction Tuning](https://mlsanity.com/paper/2402.04333/) (2024)
- [Beyond neural scaling laws: beating power law scaling via data pruning](https://mlsanity.com/paper/2206.14486/) (2022)
- [Progress & Compress: A scalable framework for continual learning](https://mlsanity.com/paper/1805.06370/) (2018)
- [What Neural Networks Memorize and Why: Discovering the Long Tail via Influence Estimation](https://mlsanity.com/paper/2008.03703/) (2020)
- [Deep Kernel Learning](https://mlsanity.com/paper/1511.02222/) (2015)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [DataComp-LM: In search of the next generation of training sets for language models](https://mlsanity.com/paper/2406.11794/) (2024)
- [SemDeDup: Data-efficient learning at web-scale through semantic deduplication](https://mlsanity.com/paper/2303.09540/) (2023)
- [Prioritized Training on Points that are Learnable, Worth Learning, and Not Yet Learnt](https://mlsanity.com/paper/2206.07137/) (2022)
- [Max-value Entropy Search for Efficient Bayesian Optimization](https://mlsanity.com/paper/1703.01968/) (2017)
- [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)

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

- [MergeMix: Optimizing Mid-Training Data Mixtures via Learnable Model Merging](https://mlsanity.com/paper/2601.17858/) (2026)
- [AutoLLMResearch: Training Research Agents for Automating LLM Experiment Configuration - Learning from Cheap, Optimizing Expensive](https://mlsanity.com/paper/2605.11518/) (2026)
- [Thomson: Continual Learning of Frontier Models for SovereignAI](https://mlsanity.com/paper/2608.27147/) (2026)
- [Pitfalls and Remedies for Multi-Task Bayesian Optimization](https://mlsanity.com/paper/2607.09073/) (2026)
- [Scaling Near-Optimal SFT-RL Annotation Budget Allocation from Small to Large LLMs](https://mlsanity.com/paper/2609.01573/) (2026)

## Abstract (arXiv)

Determining the optimal data mixture for large language model training remains a challenging problem with an outsized impact on performance. In practice, language model developers continue to rely on heuristic exploration since no learning-based approach has emerged as a reliable solution. In this work, we propose to view the selection of training data mixtures as a black-box hyperparameter optimization problem, for which Bayesian Optimization is a well-established class of appropriate algorithms. Firstly, we cast data mixture learning as a sequential decision-making problem, in which we aim to find a suitable trade-off between the computational cost of training exploratory (proxy-) models and final mixture performance. Secondly, we systematically explore the properties of transferring mixtures learned at a small scale to larger-scale experiments, providing insights and highlighting opportunities for research at a modest scale. By proposing Multi-fidelity Bayesian Optimization as a suitable method in this common scenario, we introduce a natural framework to balance experiment cost with model fit, avoiding the risks of overfitting to smaller scales while minimizing the number of experiments at high cost. We present results for pre-training and instruction finetuning across models ranging from 1 million to 7 billion parameters, varying from simple architectures to state-of-the-art models and benchmarks spanning dozens of datasets. We demonstrate consistently strong results relative to a wide range of baselines, resulting inspeed-ups of over 500% in determining the best data mixture on our largest experiments. In addition, we broaden access to research by sharing ADMIRE IFT Runs, a dataset of 460 full training & evaluation runs worth over 13,000 GPU hours, greatly reducing the cost of conducting research in this area.