Adaptive Data Optimization: Dynamic Sample Selection with Scaling Laws

Thank you for the thoughtful response and we are glad that you find our contribution novel and paper well-written. We hope our revision and response will address your concerns.

W1: We apologize for any confusion and have updated the text to provide better motivation for these choices. We agree with the reviewer that the usefulness of scaling laws, while heuristic, has been validated by plenty of prior work. Exponential moving averages to smooth noisy online estimates (e.g., Eq 5) are also commonly used, for example in optimizers like Adam. Hence, the novel choices here are our definition of the preference distribution and credit assignment, which boil down to the following intuition given in L301: prioritize sampling from domains where we are making fast progress (learning quickly), but only if the progress can be attributed to the data from that domain. We have added more discussion to the sections to elaborate on our design choices, and would be happy to revise more if you have any further questions.

W2: We note that the smoothing we do ($\gamma_1$ and $\gamma_2$) is largely in line with optimization literature and it adds almost no cost. While it is infeasible to exhaustively ablate all hyperparameters given the cost of training, we have conducted a study sweeping over the smoothing parameter $s$ in the appendix (Table 3). We find that multiple values of $s$ work and that even not using any smoothing at all (s=1) still outperforms the baselines, though having some smoothing is even better. Having a small $s$ (e.g., 0.1) corresponds to having almost no credit assignment which does perform as well as having larger $s$, confirming the benefit of credit assignment. Please see our overall response for more details.

Q1: The blue dots are indeed the actual validation loss. We have added this to the caption. I think this question is largely an open problem and we unfortunately do not have the resources to do this study. However, anecdotally, the data mixtures used for models of different sizes in the industry labs are significantly different, which would suggest that mixtures that work well for small scales might not work well for larger models. Note that ADO finds different curricula for models of different sizes, which would corroborate this. A concurrent work [Ye et al, 2024] also reaches the same conclusion with empirical verification.

Q2: $\gamma_1$ and $\gamma_2$ are fairly standard temporal smoothing parameters used in the literature. We did not have the resources to do an extensive search but the values we picked are standard [Defazio et al, NeurIPS2024] and we did not actually tune them at all. For the smooth parameter $s$, we added a sweep to the revised paper (see Table 3 in the appendix), where we find that multiple values of $s$ work well.

Q3: This assessment was made based on comparing the predictions of the scaling laws to the actual observed loss curve (Figure 6). We see that the scaling law predictions are fairly accurate up until near the end of training.

We hope our response, revision, and new ablation address your concerns. We'd also be happy to continue the discussion.

Reference

[1] Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance. Ye et al, 2024.

[2] The Road Less Scheduled. Defazio et al, NeurIPS 2024.

We thank the reviewer for their thoughtful review and insightful questions.

W1: To clarify, the right half of Figure 4 shows perplexity on training domains, all within The Pile (the training dataset). The observed phenomenon--that ADO does not typically achieve the lowest perplexity on training domains--is interesting but does not indicate a limitation in generalization capability. In fact, ADO achieves the lowest perplexity of all methods on the two datasets that are not used in training: SlimPajama and FineWeb, and both are large, diverse, and high-quality text datasets that are considered to be better than The Pile. Hence, Figure 4 in some sense shows that ADO actually has superior generalization compared to other methods. Of course, there is a much more philosophical and open question of what generalization really means for LLMs.

W2: Both $\gamma_1$ and $\gamma_2$ are akin to the momentum coefficients in the optimization literature [Defazio et al, NeurIPS2024] and 0.9 is a fairly common value. We did not actually tune these values but we believe they should be fairly robust. The smoothing parameter s has a larger effect on the behavior of the algorithm so we have conducted an ablation on the effect of s at the 125M scales in the appendix (Figure 9). Please see our overall response for more details..

W3: $\gamma_1$ controls the strength of decay for the exponential moving average on the credit assignment, which gives history $h$. This is analogous to the $\beta_1,\beta_2$ parameters in the Adam optimizer. $\gamma_2$ controls the linear combination between the temporal averaged policy and the current preference distribution--a larger value of $\gamma_2$ makes the policy more adaptive and less stationary. This is similar to Defazio et al. We apologize for any confusion here and have clarified the text accordingly.

W4: This is an interesting point of comparison. Consider the comparison with methods that use proxy models. Empirically, training a 124M model (as a proxy for 1.3B) takes over 10 hours excluding the post processing time. On the other hand, ADO adds less than 20 minutes to training time, which it uses to fit scaling laws. Note also that the cost of fitting scaling laws is independent of model size.

Theoretically, the computation cost of ADO depends on the rate of convergence of L-BFGS, which can be problem-dependent but is generally efficient. The complexity of fitting per-domain scaling laws scales with the number of domains and the granularity of the initialization search grid. Suppose $G$ is the granularity of discretization for a parameter (e.g., 7), and $K$ is the number of domains, the complexity of fitting would be $O(G^3 K)$, which grows linearly with the number of domains. However, we can easily parallelize multiple L-BFGS runs on TPUs (or GPUs), which is why ADO does not add much time in practice (as we mentioned above, 0.04% training time in a 3 day run which is less than 20 minutes, a number that can likely be optimized further). We have added this analysis to Appendix A.1.

Q1: It is a good point and ADO is in fact designed to prevent extreme over- or under- representation. In particular, the power transformation in Eq 2 with smoothing parameter $s<1$ smooths out the credit assignment and leads to a more balanced weighting across domains. We also clip the probabilities to avoid under-representation all together (see Appendix C). The prior distribution $\mu$ also offers another way to assign more emphasis on domains that are considered important.

Q2: While training on another dataset of the same size as The Pile would require significant additional time and compute resources, we note that we did not tune any hyperparameters differently between the 124M and the 1.3B which shows robustness across model scale. Of course, like all hyperparameters, the optimal values could be different for different datasets. Also, we added a sweep over values of the smoothing parameter $s$ in Table 3 of the Appendix and found that multiple values of $s$ work well. Even not using any smoothing at all ($s=1$) outperforms the baselines, though having some smoothing is even better.

Q3: This is a good question and an important direction for future investigation. One approach could be to use scaling laws that model interactions between domains, for example by introducing parameters to the scaling laws that account for pairwise interaction. The challenge of such an approach is the increase in parameters could make fitting the scaling laws more challenging or impossible. In our preliminary exploration, we found the signal-to-noise ratio for this estimation to be too high to get a reliable estimate.

Thanks for authors' detailed responses. I think that they have adressed my concerns. I will remain my positive score. Good Luck!

We thank the reviewer for their thoughtful review and questions. We hope our response below addresses your concerns.

W1: We apologize for the lack of motivation in the paper. We have added more motivation in the text, but could you please elaborate on which parts you think lack motivation so we could better address them?

W2: Could you elaborate on which results do not support the proposed method and claims we made? The evaluations in Table 1 represent a strong improvement for ADO in comparison with prior literature. For example, ADO improves over the balanced baseline by 3.3% (1.3B) and 1.3% (124M) on average over benchmarked tasks. In Table 2 of [Fan et al, ICML 2024], both DOGE and DoReMi [Xie et al, NeurIPS 2023] improve over the balanced baseline by 1.7% and 0.9%, respectively. We note the benchmarked tasks between these two papers are not exactly the same, but this gives a sense of the scale of improvement. More importantly, ADO is completely online whereas almost all prior methods require proxy models and additional computation. We believe this is a major appeal of ADO.

Q1: We are not sure if we understand the question. Regardless of the value of h, this transformation always smooths it. Concretely, higher probability domain will have lower probability and lower probability domain will receive higher probability. In the extreme for $s=0$, the $\lambda_k(t) = h_k(t)^0 = 1$ for all $k$, so the it would be uniform across all domains. The motivation is similar to having a higher temperature for a softmax distribution, but in our preliminary exploration we found power transformation worked better. The reason why we choose to smooth the distribution is because our estimation is unlikely to be perfect so we want to avoid assigning either very high probability or very low probability to the domains.

Q2: This is a great question. There are certainly many choices one can make about this function. Following your previous question, we have conducted an ablation study on the effect of s in Figure 9 of the appendix. We can see that larger s are usually better. This is reasonable since smaller $s$ leads to less credit assignment. On the other hand, $s=1$ which would make $\lambda_k(t) = h_k(t)$ leads to slightly worse performance. This aligns with our intuition that some degree of smoothing is beneficial since using history to do credit assignment is not perfect.

Q3: We fully agree that it would be interesting to test ADO on larger models. Unfortunately, pretraining a 8B model is out of reach given the resources available to us, and arguably, most academics. For reference, our 1.3B model takes 3 days to train. An 8B model would not only take more time per step but also requires more data to train. We believe that it could take much more than 8 times as long to train a single 8B model with the resource we have (i.e., > 24 days). On the other hand, we believe that a 1.3B model is well above the standard practice for data selection methods at an academic conference. For example, DoGE [Fan et al, ICML 2024] ran experiments on models up to 684M.

Q4: Thank you for this comment. Both $\gamma_1$ and $\gamma_2$ are akin to the momentum coefficients in the optimization literature [Defazio et al, NeurIPS 2024] and 0.9 is a fairly common value for momentum. We did not actually tune these values at all, but we believe they should be fairly robust. As you have pointed out, $s$ has a larger effect on the behavior of the algorithm so we have conducted a new ablation on the effect of $s$ at the 125M scales in Table 3 of appendix. Please see our overall response for more details.

Q5: Thank you for the suggestion. We have included these plots for $\lambda_k(t)$ (Figure 10) and $\frac{\partial}{\partial n} \hat{L}_k(n)$ (Figure 11) in the appendix.

We hope our response, revision and new ablation address your concerns. We'd also be happy to continue the discussion.

Reference

[1] DoGE: Domain Reweighting with Generalization Estimation. Fan et al, ICML 2024.

[2] DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining. Xie et al, NeurIPS 2023.

[3] The Road Less Scheduled. Defazio et al, NeurIPS 2024.

We thank the reviewer for their careful analysis and thoughtful feedback. We hope our response below addresses your concerns.

W1: We agree that it would be interesting to test ADO on larger models. Unfortunately, pretraining a 8B model is out of reach given the resources available to us, and arguably, most academics. For reference, our 1.3B model takes 3 days to train. An 8B model would not only take more time per step but also requires more data to train [Hoffman et al, 2022]. We believe that it could take more than 8 times as long to train a 8B model with the resource we have. On the other hand, we believe that a 1.3B model is well above the standard practice for data selection methods at an academic conference. For example, DoGE [Fan et al, ICML 2024] ran experiments on models up to 684M.

W2: While ADO’s improvement over the "Natural" baseline is relatively modest (0.7% at 124M and 0.5% at 1.3B), as we discuss in the paper, the “natural” baseline turns out to be a strong baseline that was not commonly tested in prior work.

When we compare against the “balanced” (i.e., each domain receives the same weight) baseline, as does most prior work in this area, the magnitude of improvement for ADO is significant in comparison with existing literature. Per Table 2 in [Fan et al, ICML 2024], both DOGE and DoReMi [Xie et al, NeurIPS 2023] improve over the balanced baseline by 1.7% and 0.9%, respectively. Meanwhile, ADO improves over the balanced baseline by 3.3% (1.3B) and 1.3% (124M). While the benchmarked tasks differ somewhat and the settings are different (5-shot vs 0-shot), this shows that the scale of improvement is in line with or larger than prior work. In fact, ADO’s improvement relative to the “balanced” baseline increases with scale.

In addition, by the nature of scaling law, most, if not all, performance metrics would become more difficult to improve with larger scales because the room to improve becomes smaller, so it is unrealistic to expect that absolute performance gain would stay the same at different scales. Concretely, the error for “natural” at 124M is $1-0.463 = 0.537$ and ADO improved over “natural” by $0.07$ so the relative improvement is $0.007/0.537=0.013$. The error for “natural” at 1.3B is $1-0.585 = 0.415$ and ADO improved over “natural” by $0.05$ so the relative improvement is $0.005/0.415=0.012$. Based on this calculation, we can see that the relative improvement of ADO at different scales is consistent.

W3: The effect of learning rate schedule on the learning curve’s functional form is a very new research direction. In fact, the first paper on the topic (that we are aware of) [Tissue et al, 2024] only came out very recently after we had finished our experiments. There are many things that we don’t fully understand about these methods, and we do not want to introduce too many moving components to our method because these methods alone would be worth one or more stand alone papers. As such, we believe these directions are better left for future works.

Q1: Good question--we have added this information in Figure 9 of the appendix. An interesting observation is that ADO’s training loss actually does not decrease faster than baselines because ADO prioritizes the domains with higher learning potential, which does not necessarily mean low loss. In fact, we see that ADO tends to select higher quality data which naturally have higher entropy which would lead to higher training loss.

Q2: We believe that it should be possible to apply ADO or ADO-like methods to other domains as long as the loss curve looks like a well-behaved power law. Without other interpretations, ADO can be seen as a method that prioritizes data points whose loss is going down faster than the other ones. However, the information-theoretic interpretation might not apply depending on the exact training objective. It would be very interesting to explore this in future works.

Reference

[1] DoGE: Domain Reweighting with Generalization Estimation. Fan et al, ICML 2024.

[2] DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining. Xie et al, NeurIPS 2023.

[3] Scaling Law with Learning Rate Annealing. Tissue et al, 2024.

[4] Training Compute-Optimal Large Language Models. Hoffman et al, 2022.

Thank you for addressing the concerns raised in my review. I would like to elaborate on these points.

On Experiment Scale (W1)

• I understand that training an 8B model may be infeasible given your resource constraints. However, I must respectfully disagree with the justification that a 1.3B model is sufficient for relevance. While it is true that academic research often operates under resource limitations, the field of LLMs is evolving rapidly, and what was considered “state-of-the-art” or “adequate” a year ago may no longer hold true. For example, DoGE's experiments, very likely completed in 2023, represent a landscape that has shifted quite significantly by now, nearing 2025. Models like Llama have become widely used benchmarks and hold substantial weight in the research community.

• To make a lasting contribution, I believe your method should demonstrate its practicality and relevance to the contemporary scale and settings where it could realistically be applied. This would elevate the significance of your work beyond being an interesting but narrowly scoped academic exercise.

Magnitude of Improvement and Task Performance (W2)

• Your clarification on the relative improvement over the “natural” and “balanced” baselines is appreciated. However, I remain unconvinced by your argument that a modest improvement of 0.5-0.7% over “natural” is consistent. The weakness in ADO’s performance at larger scales is not solely due to diminished relative improvements. I noted that while ADO with the 124M model outperforms baselines on 6 out of 7 downstream tasks, this drops to 4 out of 7 tasks with the 1.3B model. This raises concerns about ADO’s consistency and robustness as model scale increases. This trend is particularly important since it would be problematic if ADO’s benefits do not consistently increase with scale relative to the “balanced” baseline. However, the task-level performance metrics raise concerns and further analysis or discussion of this discrepancy would be valuable.

While I still see value in sharing this work with the community, I stand by my assessment that the current contribution remains marginally above the acceptance threshold. The primary concerns—namely, the scale of experiments, the practical relevance of ADO, and the inconsistency in task-level performance at larger scales—prevent the paper from being a strong acceptance. To elevate this work beyond “barely acceptable,” addressing these concerns in meaningful ways (e.g., experiments with more contemporary model architectures, stronger demonstrations of generalization, and a discussion on task-level performance trends) would be necessary.

I have additional concerns after further reflection, particularly regarding comparisons to prior work and the lack of theoretical contributions in this paper. I outline these below.

On Comparisons to DoReMi and DoGE

• It is surprising that the rebuttal does not discuss DoReMi (NeurIPS 2024) when addressing conventional model sizes for academic conferences. DoReMi explicitly targets the training of 8B models, as stated clearly in their abstract, and evaluates their method at the 8B scale, showcasing their relevance in the current landscape.

• Furthermore, both DoReMi and DoGE provide significant theoretical contributions. For instance, DoReMi includes a rigorous theoretical contribution (Appendix D), while DoGE complements its experimental results with theoretical justification (Appendix B). These contributions enhance the broader impact and rigor of their work, even when experiments are limited to smaller scales (DoGE). Experiments in DoReMI are more comprehensive, testing 8B models on more datasets than this work does.

• In contrast, ADO neither performs experiments at larger scales nor offers meaningful theoretical contributions. While the empirical results are interesting, the absence of theoretical justification or deeper analysis significantly weakens the paper, particularly when compared to these prior works. This is especially problematic for a method that claims to introduce an original approach but does not delve into why ADO (or some of its elements) works at a fundamental level.

On Downstream Task Performance and the “Natural” Baseline

• While the authors argue that both “natural” and ADO are original contributions of this work, the focus of this paper is on ADO. As such, ADO must consistently outperform the “natural” baseline to demonstrate its value. For the 1.3B model, ADO only outperforms “natural” on 4 out of 7 tasks, compared to 6 out of 7 tasks at 124M. In particular, the “natural” baseline outperforms ADO on ARC-E, SciQ, and LogiQA2.

Empirical Nature of the Paper

• The paper is heavily empirical and lacks theoretical justification for why ADO works. While simplicity and practicality are valuable attributes, they are not sufficient to compensate for the absence of theoretical insights or deeper analysis. In this regard, the paper falls short of the standards set by prior works such as DoReMi and DoGE, both of which balance empirical results with robust theoretical contributions.

After reevaluating the paper, I have decided to lower my score to 5 (“below acceptance threshold”). While ADO demonstrates merit in terms of compute efficiency and modest empirical improvements, the paper falls short in its theoretical and experimental rigor relative to existing literature. Specifically:

• ADO’s inability to consistently outperform the “natural” baseline at larger scales undermines its claim of scalability.
• The lack of theoretical contributions or deeper analysis makes the paper overly reliant on its empirical results, which are not sufficiently comprehensive.
• Comparisons to prior work reveal significant gaps in scope, relevance, and contribution.

While ADO has potential, the paper’s current form does not meet the standard of significant contribution required for acceptance. I hope the authors will consider addressing these concerns in future revisions, as there is a foundation here for meaningful work if the necessary rigor and scope (as shown in DoReMI) are added.

Thank you for the additional comments and continued engagement. We will try to address these new comments below.

DoReMi did indeed contain experiments using far more resources than most academics would have access to since it is from one of the largest industry labs. We already acknowledge in the limitation of the original draft that we do not know how ADO will perform on larger models, and we have revised the limitation section to explicitly mention the 8B model of DoReMi and would be happy to revise this section more if you have further suggestions.

The lack of theoretical contributions or deeper analysis makes the paper overly reliant on its empirical results, which are not sufficiently comprehensive.

[...] prior works such as DoReMi and DoGE, both of which balance empirical results with robust theoretical contributions.

Our paper is indeed empirical in nature and we do not claim any theoretical contribution. However, we believe that the existing literature in this area is also heavily empirical in nature. The two prior works you mention here (DoReMi and DoGE) provide interesting derivations and analyses, but these do not necessarily amount to “theoretical justification” for why their methods work.

For example, in DoGE (Appendix B), the authors derive the DoGE update rule as an optimization problem--but it is not a proof for why DoGE works. In fact, from DoGE (Appendix C.5):

“the performance of the proxy model falls behind the base model with resampled dataset with DOGE domain sampling weights. It is even worse than the baseline with uniform domain weights.“

If the derivation actually explained DoGE’s efficacy, then the proxy model should do better than uniform weights because it is actually trained with the derived update rule. The fact that the proxy model does worse indicates that the motivation does not match what the algorithm is actually doing. As such, it stands to reason that DoGE’s success is also empirical. In terms of empirical results’s comprehensiveness, we believe that the scope of our results is comparable to DoGE and we also conduct experiments on larger models.

In DoReMi (Appendix D), the authors constructed a toy data generating process for which there exists a set of weights that is better than uniform, but this is not a proof that DoReMi would actually find these weights. The fact that DoReMi works on this toy model was also an empirical observation.

ADO’s inability to consistently outperform the “natural” baseline at larger scales undermines its claim of scalability.

While the authors argue that both “natural” and ADO are original contributions of this work, the focus of this paper is on ADO. As such, ADO must consistently outperform the “natural” baseline to demonstrate its value. For the 1.3B model, ADO only outperforms “natural” on 4 out of 7 tasks, compared to 6 out of 7 tasks at 124M. In particular, the “natural” baseline outperforms ADO on ARC-E, SciQ, and LogiQA2.

In our opinion, the contribution of a work should be measured against the existing literature, and we have demonstrated that ADO outperforms the existing techniques and popular baselines in the literature while being much more efficient. Furthermore, for LAMBADA, natural outperforms ADO in 124M but underperforms in 1.3B, indicating that there can be variability in individual tasks. In terms of the more robust average performance, the relative improvement of ADO over natural is consistent across scales (1.3% and 1.2% respectively). We believe that this supports ADO’s value.

Further, consistent improvement in all downstream tasks over prior works is not the standard of the literature for pretraining to the best of our knowledge. For example, DoGE does not consistently outperform the DoReMi in downstream performance (Table 2 of DoGE). In comparison, ADO outperforms DoReMi in 6/7 tasks and performs the same in 1/7 at both scales.

We hope this clarifies these points and would be happy to continue the discussion.