Optimizing Pretraining Data Mixtures with LLM-Estimated Utility

TL;DR

This work tackles how to optimally mix diverse pretraining data for large language models under compute and data budgets. It first shows that simple token-size heuristics can outperform manual and learned data mixes, then introduces UtiliMax, a portfolio-optimization approach that uses reduced-scale ablations to estimate dataset utility and risk, and MEDU, a cost-efficient method that uses LLMs to infer data utility descriptions from small samples. Together, UtiliMax and MEDU form a compute-efficient framework for automated data mixing that remains robust across training regimes, achieving up to a 10.6x speedup over manual baselines and up to a 200x reduction in utility-estimation cost. A key finding is the surprising effectiveness of Uniform Utility (UniMax) across budgets, suggesting that diversity and scale can be more influential than precise task-specific utility estimates in many settings. The approach provides a practical path to faster frontier-model pretraining by explicitly balancing utility, diversity, and data size through principled optimization and scalable utility estimation.

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 $\sim$200x. Together, these approaches establish a new framework for automated, compute-efficient data mixing that is robust across training regimes.

Figures (11)

• Figure 1: Scaling curves for data mixing methods for compute-optimal models trained for $6 \times 10^{19}$ to $3 \times 10^{21}$ Floating Point Operations (FLOPs). Compared to manual OLMo, heuristic UniMax, and learned DoReMi data mixes, UtiliMax leads to more compute efficient models that perform better on average across tasks.
• Figure 2: Comparison of baseline data mixing methods. The model with the top average rank at $3\times10^{21}$ FLOPs is in bold. UniMax consistently outperforms all other baselines in both settings.
• Figure 3: Comparison of utility optimization methods. The model with the top average rank at $3\times10^{21}$ FLOPs is in bold. UtiliMax outperforms alternative optimization procedures in both settings.
• Figure 4: Scaling curves comparing MEDU, Ablation Estimates, and UniMax. The model with the best mean rank at $3\times 10^{21}$ FLOPs is marked in bold. Speedup indicates cases where using MEDU improves over Ablation-Based Utility, while slowdowns indicates the opposite.
• Figure 5: A high-level overview of MEDU which describes benchmarks and classifies document utility based on this description.