Research map / Probabilistic and scientific modeling / Bayesian optimization
Data Mixture Optimization: A Multi-fidelity Multi-scale Bayesian Framework
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Frames LLM data mixture optimization as multi-fidelity multi-scale Bayesian optimization, using a GP with EI-per-unit-cost over mixture, model scale, and training steps, benchmarked on a simulator built from 472 OLMo 2 pretraining runs on SlimPajama.
Problem
Choosing data mixtures for LLM pretraining is costly and brittle: deterministic scaling-law extrapolation from small models to large ones is unreliable, and existing multi-fidelity methods do not exploit model scale as a fidelity dimension.
Approach
Treats data mixture, model scale, and training steps jointly as decision variables in a Gaussian-process Bayesian optimization with product RBF kernels, linear mean, and expected improvement per unit cost, so cheap small-model/short runs inform costly target-scale decisions.
Key results
- MFMS-GP achieves 2.6x speedup over baselines in finding the best downstream-accuracy configuration (2.6x to 3.3x range reported across plots).
- On validation cross-entropy loss, MFMS-GP achieves more than 2.7x speedup over baselines.
- MLP predictor achieves R^2 >= 0.95 across all metrics on held-out runs, and predictive R^2 is highest (0.82) when trained on many short runs (E8) versus fewer longer runs (E6, 0.69).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Optimize data mixtures jointly over mixture, model scale, and training steps rather than using Hyperband or random search | MFMS-GP · Downstream task accuracy | speedup: 2.6x to 3.3x | Helped |
| Optimize data mixtures jointly over mixture, model scale, and training steps rather than using multi-fidelity BO baselines | MFMS-GP · Validation cross-entropy loss | speedup: 2.7x | Helped |
| Include smaller-model runs when training the predictor for larger-model performance | MLP performance predictor · 1B model performance prediction | predictive accuracy | Helped |
| Use 32 shorter truncated runs instead of fewer, longer runs at the same FLOPS budget (ablation) | MLP performance predictor · Predicting final model losses | R²: 0.69 -> 0.82 | Helped |
| Use 32 shorter truncated runs instead of fewer, longer runs at the same FLOPS budget (ablation) | MLP performance predictor · Predicting final model losses | R²(log): 0.74 -> 0.85 | Helped |
| Add 700M runs to the predictor training data (ablation) | MLP performance predictor for 1B runs · 1B model performance prediction | predictive accuracy | Helped |
| Add all smaller-model runs to the predictor training data (ablation) | MLP performance predictor for 1B runs · 1B model performance prediction | predictive accuracy | Helped |
| Add 500M runs to the predictor training data (ablation) | MLP performance predictor for 700M runs · 700M model performance prediction | predictive accuracy | Helped |
| Replace the RBF data-proportion distance with Jensen–Shannon divergence (ablation) | MFMS-GP · Downstream task accuracy | accuracy | Mixed |
| Replace the RBF data-proportion distance with total variation distance (ablation) | MFMS-GP · Common Crawl, Wikipedia, and Stack Exchange | validation cross-entropy loss | Mixed |
Surprising findings
- The optimal data mixture proportion does not remain constant across model scales, contrary to naive transfer assumptions.
- For a fixed FLOPS budget, many short runs gave better predictive power than fewer full-length runs.
Limitations
- Experiments use a predictor-based simulator rather than actual pretraining runs for each BO proposal.
- Only five data categories used for training mixtures; evaluated on a single data source (SlimPajama) and small models up to 1B.
- Model scale treated as a fidelity dimension without hierarchical structure; training steps restricted to three discrete values in the GP.
Experimental setup
Datasets: SlimPajama RedPajama
Benchmarks: HellaSwag PIQA ARC Easy
Metrics: downstream accuracy validation cross-entropy loss R^2 speedup vs baselines
Baselines: Hyperband with random forest (SMAC) Random Search Multi-fidelity BO
Builds on: Hyperband SMAC BOHB Data Mixing Laws OLMo 2
Closest papers by meaning
- Adaptive Data Optimization: Dynamic Sample Selection with Scaling Laws (ICLR 2025 · similarity 0.88)
- Scaling Laws for Optimal Data Mixtures (NeurIPS 2025 · similarity 0.88)
- Large Language Models to Enhance Bayesian Optimization (ICLR 2024 · similarity 0.88)
- Compute-Constrained Data Selection (ICLR 2025 · similarity 0.88)
- DataDecide: How to Predict Best Pretraining Data with Small Experiments (ICML 2025 · similarity 0.88)
- Bayesian Optimisation with Unknown Hyperparameters: Regret Bounds Logarithmically Closer to Optimal (NeurIPS 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.