English

Optimizing Pretraining Data Mixtures with LLM-Estimated Utility

Computation and Language 2025-01-27 v2 Artificial Intelligence

Abstract

Large Language Models improve with increasing amounts of high-quality training data. However, leveraging larger datasets requires balancing quality, quantity, and diversity across sources. After evaluating nine baseline methods under both compute- and data-constrained scenarios, we find token-count heuristics outperform manual and learned mixes, indicating that simple approaches accounting for dataset size and diversity are surprisingly effective. Building on this insight, we propose two complementary approaches: UtiliMax, which extends token-based heuristics by incorporating utility estimates from reduced-scale ablations, achieving up to a 10.6x speedup over manual baselines; and Model Estimated Data Utility (MEDU), which leverages LLMs to estimate data utility from small samples, matching ablation-based performance while reducing computational requirements by \sim200x. Together, these approaches establish a new framework for automated, compute-efficient data mixing that is robust across training regimes.

Keywords

Cite

@article{arxiv.2501.11747,
  title  = {Optimizing Pretraining Data Mixtures with LLM-Estimated Utility},
  author = {William Held and Bhargavi Paranjape and Punit Singh Koura and Mike Lewis and Frank Zhang and Todor Mihaylov},
  journal= {arXiv preprint arXiv:2501.11747},
  year   = {2025}
}

Comments

10 pages, 8 figures

R2 v1 2026-06-28T21:11:47.966Z