Adaptive Data Optimization: Dynamic Sample Selection with Scaling Laws

Authors: Yiding Jiang, Allan Zhou, Zhili Feng, Sadhika Malladi, J Zico Kolter
Primary area: foundation or frontier models, including LLMs
Keywords: Data selection, Pre-training, Curriculum learning, Language Models, Scaling laws

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.

Decision: accept (poster)
Reviewer scores: 5, 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730212806279

thread_type:
official_review

<thread_object>
creator: reviewer_axts

<content>
summary:
This work proposes an online adaptive data optimization (ADO) method for finding a policy (weighting) for different domains when pre-training a language model. The core idea is to use scaling law to model training loss. The derivative of the loss w.r.t. to the number of data points can be understood as loss reduction per data point locally, i.e. outline the potential for a domain. The final scoring also considers the contribution of the domain itself from the last update and makes it stable with a moving average. Their experiments show the ADO brings improved downstream performances on 7 common-sense reasoning tasks.

strengths:
- Novelty. It’s a novel idea to improve over online data optimization (ODM, Albalak et al. (2023)) on 3 aspects. 1) Using scaling law to predict training loss on each domain and the derivative of the loss w.r.t. to the number of data points has intuitive interpretation as the loss reduction per data point. 2) The weighting also considers the number of data points used for a domain. 3) A combined weighting with the above two and temporal average.
- Readability. The paper not only shows their method works but also spends a reasonable amount of space discussing perspectives, hence better readability. For example on the curriculum learning (Section 2) and on the requirements (Section 3, online and agnostic to downstream tasks) of data optimization for pretrained models.
- Besides their proposed method, the authors also discover a strong baseline called “natural” policy, that depends on the number of tokens in each domain.

weaknesses:
- Many heuristics are used such as fitting of scaling law (Section 3.1 eq1), credit assignment (Section 3.2, eq2 ), preference distribution (Section 3.3, eq3) and temporal average (eq4 and eq5). The authors try to motivate those choices from related works and intuitions, but only eq1 is adequately explained and validated. 
- Following the above, are all the heuristics (eq 1- 5) necessary? How important are they? This is a missing part in the paper. An ablation on them can validate if the invented heuristics are actually all useful.

questions:
- About Figure 2, I think the point on the variance is clear. But what’s more important is the agreement of the relative order of data strategies: if a data strategy is better than another on a smaller model, is it also better on a larger model? Did the authors study that? Also, what are the blue dots? I guess they are the actual validation loss of different model sizes, but would be nice to make it clear in the caption.
- The authors also write “they find $\gamma_1 = 0.1$, $s = 0.5$, $\gamma_2 = 0.1$” works well. How exactly is the process of finding those values?
- The authors write in line 480: “they are accurate locally for much of training and thus can act as a learning signal for the data policy”. On what basis do the claims are made?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730374264438

thread_type:
official_review

<thread_object>
creator: reviewer_7mvd

<content>
summary:
The paper introduces ADO, a novel online algorithm for optimizing data distributions during the training of large-scale foundation models. ADO addresses the challenge of efficiently managing data mixtures in pretraining by continuously adapting the data distribution based on scaling laws, without requiring proxy models or significant computational overhead. The primary contribution lies in the use of per-domain scaling laws to estimate learning potential and dynamically adjust the data mixture during model training. Experimental results demonstrate that ADO can achieve comparable or better performance than prior methods, while maintaining computational efficiency across various model scales.

strengths:
1.	ADO utilizes online scaling laws to dynamically adapt data selection during training, replacing the need for pre-trained proxy models or multi-staged processes. It requires no prior domain knowledge or external models, making it highly practical and versatile.
2.	The paper presents empirical evaluations demonstrating ADO's effectiveness across multiple benchmarks and datasets. ADO outperforms or matches existing baselines on several downstream tasks while incurring a small computational overhead (e.g., less than 0.4% additional wall-clock time).
3.	ADO offers a practical, scalable solution for data selection in large model training, with the potential to significantly reduce computational waste. The method aligns with the growing emphasis on efficient resource utilization in AI research, especially as models scale in size and cost.

weaknesses:
1.	Figure 4 reveals that fewer than half of the datasets achieve minimal perplexity under the ADO algorithm, indicating potential limitations in the algorithm’s generalization capability across diverse scenarios. The paper would benefit from a more detailed analysis of ADO's applicability.
2.	The algorithm employs heuristic choices for key parameters, such as the exponential moving average in credit assignment and smoothing coefficients, which may impact the robustness and consistency of ADO across different datasets and tasks due to the lack of a more systematic basis for parameter selection.
3.	Symbols such as $\gamma_2$ appears in key formulas but might be clarified further with more detailed context, especially for readers unfamiliar with its role in the algorithm.
4.	Although ADO highlights low computational overhead as a strength, the paper lacks a systematic computational complexity analysis and comparative experiments. The absence of quantitative and visual evidence of its efficiency limits the clarity of its computational advantages.

questions:
1.	In ADO's data selection process, some datasets might receive minimal training due to low selection probabilities, potentially leading to under-representation. Have any additional measures been implemented to balance data distribution and ensure adequate representation across datasets?
2.	How sensitive is ADO to changes in its hyperparameters, such as the smoothing coefficient (γ1) and update intervals? Would different tasks or datasets require significant parameter tuning?
3.	Given the current limitation in modeling solely intra-domain interactions, have the authors considered potential methodologies to incorporate inter-domain interactions within the ADO framework? What approaches might be feasible for capturing these cross-domain dynamics to enhance the model’s adaptability and performance in tasks where inter-domain relationships play a significant role?
4.	ADO is designed to be task-agnostic, but certain applications may benefit from targeted data selection. Could ADO be extended to incorporate task-specific objectives without compromising its efficiency?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730616467805

thread_type:
official_review

<thread_object>
creator: reviewer_monz

<content>
summary:
This paper proposes a novel method ADO to adaptively sample data from different domains based on their contribution to model training, which is estimated from scaling law in an online manner. Empirical results on different data sets demonstrate the effectiveness of the proposed method.

strengths:
The proposed method is introduced clearly and easy to understand

weaknesses:
- Several parts of ADO lack sufficient motivation
- Empirical results seem not supportive enough to fully evaluate the proposed method

questions:
- The smoothing parameter $s$ for computing the credit assignment scores seems a bit confusing. By using $s<1$, I suppose $\lambda_k$ indeed changes drastically when $h_k$ is close to 0? Should using $s<1$ encourages $\lambda_k$ to be larger than $h_k$ especially when $h_k$ is originally close to 0? Some further discussion on the effect of this hyper-parameter may be necessary. 
- Also, can we consider other types of functions to compute $\lambda_k$ from $h_k$? Some empirical comparison on different types of functions may help others better understand how the credit assignment score affects model training, and why the authors reject simply using $\lambda_k=h_k$ here. 
- Experiments are limited to medium scale models (124M and 1.3B), and the performance gain seems to decrease with increasing model size. Such a tendency can naturally make one wonder if the performance gain can be smaller or even diminish on larger models (e.g., 2.3B or 7/8B)? Some more empirical results may be necessary to better support the proposed method. 
- While the proposed method includes many hyper-parameters (e.g., the smoothing parameter $s$ mentioned above), the authors did not provide sufficient analysis on how these hyper-parameters can affect the final performance. These results should be necessary to demonstrate that ADO can be robust to different values of these hyper-parameters. 
- Despite showing the final sampling probabilities in Figure 5, the authors may also consider further showing how $\lambda_k(t)$ and $\frac{\partial}{\partial n} \hat{L}_k(n)$ changes with training iteration $t$. These results may help others better understand how these two terms affect the sampling probabilities of different domains and how they can affect model training.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730703243477

thread_type:
official_review

<thread_object>
creator: reviewer_26jp

<content>
summary:
The authors propose ADO, an online data selection method that adjusts data distribution dynamically across various domains by leveraging domain scaling laws. Without requiring a proxy model or external knowledge, ADO forecasts the model's loss on different data domains and automatically modifies the training data distribution according to each domain's learning potential. The experiments indicate that ADO surpasses baseline methods overall, with only a minimal increase in clock time.

strengths:
1. The motivations presented in Section 2 highlight significant concerns, and the proposed method properly addresses these issues.
2. ADO outperforms the baselines overall, offering online data selection without the requirement of external knowledge or proxy models.
3. The authors present several interesting observations that align with our intuition and the literature (e.g. line 428, 409, 431).

weaknesses:
1. Although the 1.3B model is not small, ADO is especially practical and relevant for LLMs with at least 8B parameters—the size typical of many popular "small" LLMs such as Llama 3.1 8B. The paper's relevance and importance might have been enhanced if the authors had conducted experiments with language models of at least 8B parameters
2. The performance improvement achieved by ADO does not seem significant enough, especially on the 1.3B model. More problematically, Table 1 shows diminishing returns as the model scale increases from 124M to 1.3B parameters. For example, 1.3B-ADO only outperforms 4 out of 7 downstream tasks, whereas 124M-ADO outperforms 6 out of 7 tasks. Additionally, the gap between the average score of ADO and the second-best baseline is less pronounced in the 1.3B model. This suggests that ADO may not scale well as the number of parameters increases.
3. Although the authors propose interesting future research directions (e.g. Section 6), incorporating some of these suggestions—such as learning rate scheduling—into the current paper would make it more complete and thorough.

questions:
1. Figures 6, 7, and 8 show the training loss for ADO only. I'd like to see how quickly ADO's training loss decreases compared to the baselines. I assume that ADO converges faster, but it would be interesting to observe how it performs relative to the baselines.

2. Can ADO be effectively applied to vision data and tasks? It would be of great importance if ADO could be used in training vision models that require multiple domains, such as large vision-language models (LVLMs) and diffusion models for image generation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732052825520

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of revision

comment:
We thank all of the reviewers for their questions and feedback, much of which we have already included in our revised manuscript. Please see the revised manuscript which answers many of the questions from the reviews--changes since the original submission are made in red. We also did some tightening so the text remains within the page limit.

As a quick summary of the changes:

* **Table 3** in the appendix: We conducted an ablation study on sweeping the hyperparameter $s \in [0,1]$ which controls the smoothing during credit assignment. We find that ADO is not particularly sensitive to $s$: of the four values we swept, which span a wide range of possible values (0.1, 0.3, 0.5, and 1.0), all four values resulted in improvements over the “Balanced” baseline and three out of four improved over the “Natural baseline.” Most notably, not using any smoothing at all ($s=1$) still outperforms the baselines, though having some smoothing is better which confirms our hypothesis. (Reviewer monZ, 7Mvd, Axts)

* **Figure 9** in the appendix: Training curves of 1.3B models under each method, including ADO and the baselines (Reviewer 26JP)

* **Figure 10** in the appendix: A plot of the credit assignment function $\lambda_k(t)$ for each domain throughout the course of training. (Reviewer monZ)

* **Figure 11** in the appendix: The estimated speed of learning $\frac{\partial}{\partial n} \hat{L}_k(n)$ on each domain given by the slope of the scaling laws. (Reviewer monZ)

We also edited Section 3 to include more intuition about the hyperparameters and how adjusting them affects ADO.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes an algorithm to optimise data distributions accros domains during model training in order to leverage domain-based scaling laws. The claims are supported and validated experimentally. The topic is timely and all reviewers agreed that the results are worth sharing with the community. One remaining question remained unanswered, namely that we do not know for sure whether the proposed method would still work with bigger models than the ones considered in the paper. While this is a valid point, it does not invalidate the results reported in this work. The authors also provided an adequate discussion of this limitation.