Adaptive Data Optimization: Dynamic Sample Selection with Scaling Laws

TL;DR

Adaptive Data Optimization (ADO) addresses the challenge of allocating a fixed compute budget across heterogeneous pretraining data domains by introducing an online, model-agnostic data policy that updates concurrently with training. It relies on per-domain neural scaling laws, fitting a local power-law for each domain: $\widehat{{\mathcal{L}}}_k(n) = \varepsilon_k + \beta_k n^{-\alpha_k}$, and uses the derivative $\frac{d\widehat{{\mathcal{L}}}_k(n)}{dn} = -\frac{1}{n} \alpha_k (\widehat{{\mathcal{L}}}_k(n) - \varepsilon_k)$ to forecast learning speed. The data policy combines learning-speed forecasts, credit assignment $\lambda_k(t)$, and priors to form a time-smoothed sampling distribution $\pi(t)$, with online updates and a minimum sampled probability to ensure exploration. Empirically, ADO matches or surpasses prior online data-selection methods on 124M and 1.3B parameter language models with minimal overhead (roughly $0.4\%$ additional wall-clock time) and demonstrates sensible, scale-dependent data mixtures. These results suggest a practical, scalable pathway for automatic data optimization that avoids proxy models while leveraging scaling-law insights for better pretraining efficiency.

Abstract

The composition of pretraining data is a key determinant of foundation models' performance, but there is no standard guideline for allocating a limited computational budget across different data sources. Most current approaches either rely on extensive experiments with smaller models or dynamic data adjustments that also require proxy models, both of which significantly increase the workflow complexity and computational overhead. In this paper, we introduce Adaptive Data Optimization (ADO), an algorithm that optimizes data distributions in an online fashion, concurrent with model training. Unlike existing techniques, ADO does not require external knowledge, proxy models, or modifications to the model update. Instead, ADO uses per-domain scaling laws to estimate the learning potential of each domain during training and adjusts the data mixture accordingly, making it more scalable and easier to integrate. Experiments demonstrate that ADO can achieve comparable or better performance than prior methods while maintaining computational efficiency across different computation scales, offering a practical solution for dynamically adjusting data distribution without sacrificing flexibility or increasing costs. Beyond its practical benefits, ADO also provides a new perspective on data collection strategies via scaling laws.

Figures (8)

• Figure 1: ADO is a cheap online technique for adjusting the data distribution while training large models. In contrast to prior methods, ADO tailors its data distribution to the model as it is being trained, and does not require training smaller proxy models in advance.
• Figure 2: Extrapolating larger models' loss using scaling laws up to a threshold model size, where the threshold is indicated by color. Extrapolations become more accurate with more models.
• Figure 3: The average learning curves of SGD and the meta-optimized curricula for a logistic regression (Left) and an MLP (Middle). The meta-optimized curricula consistently outperform SGD. Further, the meta loss for curricula is still decreasing after 1000 steps (Right). The first two figures are averaged over 50 runs and the shaded region corresponds to 1 standard deviation.
• Figure 4: Perplexity of 1.3B models on different domains, trained using either DoReMi, Natural, or ADO. On the left side of the red line, we have the validation perplexity on the Pile validation set, the unweighted Pile validation set, the SlimPajama validation set, and a random subset of FineWeb. On the right side of the red line, we show the validation perplexity of each domain of the Pile.
• Figure 5: The sampling distribution produced by our data policy during training on The Pile. For ease of visualization, we only highlight the top 6 largest domains. ADO produces qualitatively different strategies at different model scales and adaptively changes its weightings over time.

Theorems & Definitions (3)

• Definition 1
• Definition 2
• Definition 3