Dynamic Gradient Alignment for Online Data Mixing

We thank the reviewer for the feedback and insightful suggestions for improving our work! We provide corresponding clarifications as follows:

W1: Proof of convergence

As the reviewer points out, the theoretical understanding of data mixing (i.e. domain reweighting) algorithms is a valuable topic that could significantly enhance the theoretical contributions of our paper. We have updated the paper to include a convergence proof of DGA in a simple convex scenario (Theorem 2). To the best of our knowledge, this is the first theoretical proof of convergence for a data reweighting algorithm. For instance, prior works such as [1,2] do not prove the convergence of their algorithms.

W2: Potential explanations for the limited improvement in reasoning accuracy

We agree with the reviewer that despite a substantial improvement above the uniform baseline, DGA performs comparably importance sampling baselines in reasoning accuracy scores, as shown in Table 1.

As for the considerable gap between language modeling ability and reasoning accuracy, we provide two potential explanations in Appendix B.3.1: (1) since the target objective in the outer-loop ($L_{spe}$ in Equ. 3) is the next token prediction loss for language modeling, the language modeling performance is expected to be improved. However, it does not necessarily translate to improved reasoning accuracy (2) MMLU is considered to be a challenging reasoning task, where the accuracies can hardly be improved when the model capacity is below some specific threshold (Figure 11. (G) in [3]). Understanding the loss-accuracy gap and deriving an optimization objective that directly benefits the reasoning capability would be a compelling future direction.

W3: Comparisons to more baselines

We have updated the paper and included the comparison between DGA and DoGE in Appendix F (Fig. 15). The results demonstrate that DGA outperforms DoGE in both specialized loss and general loss.

Q1: Insights on other tasks

In this work, we mainly focus on the language modeling tasks. However, we think this gradient-based method could be generally applied to various tasks, e.g. image classification and vision-language modeling. Since the preprocessing (embedding and clustering) of the pretraining corpus takes intensive effort, we leave more broad benchmarking as one future extension of our current work.

W4&Q2: What factors influence DGA's stability when EMA is not used?

According to the domain weight evolution trajectories shown in Figure 3(a), DGA without EMA consistently up-weights a specific domain for a certain number of steps before shifting to another. This cyclical behavior amplifies overfitting over time and optimizes the model parameters to an ill-positioned location, which could be fatal for the online optimization. Also, the first identified relevant domain to be significantly down-weighted when the model shifts to the second domain. Notably, the spikes in domain weights align closely with those observed in the validation loss curves, highlighting the correlation between these dynamics and the model's performance fluctuations.

Q3: Future directions on sample-level data selection

While the sample-level selection is not computationally friendly on pretraining because of the gigantic scale of the pretraining corpus, the gradient-based method could be applied to sample-level data selection in continual pretraining or task-specific finetuning stages. Additionally, the clustering-based method could be applied to balance the diversity and relevance during data selection. We consider it as a promising direction for future research.

We hope our answers can resolve most of the concerns. If you have any further questions, don’t hesitate to contact us. We will improve the manuscript according to your valuable suggestions!

[1] Doge: Domain reweighting with generalization estimation [2] Doremi: Optimizing data mixtures speeds up language model pretraining [3] Emergent Abilities of Large Language Models

Dear reviewer 9LhJ,

Thank you again for your efforts reviewing our paper and providing valuable feedback! Since tomorrow is the last day of the author-reviewer discussion phase, we hope you can take some time to look into our rebuttals and additional results. Let us know if you have other concerns and questions!

Sincerely,

the authors

Thanks for your reply.

Indeed, in some settings, DGA and importance sampling perform quite similarly. However, we also highlight several settings where DGA outperforms importance sampling, for instance, in Figure 1 and Figure 2. So, in our view, there are several scenarios where DGA brings something new compared to importance sampling.

Finally, we would appreciate if the reviewer could clarify what they mean by “only an alternative to importance sampling”. Indeed, > both DGA and importance sampling pursue the same goal.

As both solve the same problem and DGA performs very comparable to Importance Sampling there's unfortunately little that makes it more attractive except for slightly different application properties. So, it must be viewed as equal alternative to Importance Sampling.

We thank the reviewer for the valuable feedback and insightful questions! We provide the corresponding clarifications as follows:

Q1: Clarification on DGA with EMA (Figure 1,3)

According to the domain weight evolution trajectories shown in Figure 3(a), DGA without EMA consistently up-weights a specific domain for a certain number of steps before shifting to another. This cyclical behavior not only amplifies overfitting over time but also optimize the model parameters to an ill-positioned location. Also, the first identified relevant domain to be significantly down-weighted when the model shifts to the second domain. Notably, the spikes in domain weights align closely with the spikes observed in the validation loss curves, highlighting the correlation between these dynamics and the model's performance fluctuations.

W1: Potential explanations for the limited improvement in reasoning accuracy

We agree with the reviewer that despite a substantial improvement above uniform baseline, DGA only matches or does slightly worse than importance sampling baselines in reasoning accuracy scores, as shown in Table 1.

As for the gap between language modeling ability and reasoning accuracy, we provide two potential explanations in Appendix B.3.1: (1) since the target objective in the outer-loop ($L_{spe}$ in Equ. 3) is the next token prediction loss for language modeling, the language modeling performance is expected to be improved. However, it does not necessarily translate to improved reasoning accuracy (2) MMLU is considered to be a challenging reasoning task, where the accuracies can hardly be improved when the model capacity is below some specific threshold (Figure 11. (G) in [1]). Understanding the loss-accuracy gap and deriving an optimization objective that directly benefits the reasoning capability would be a compelling future direction.

Finally, we would appreciate if the reviewer could clarify what they mean by “only an alternative to importance sampling”. Indeed, both DGA and importance sampling pursue the same goal.

W2: Confusion in the understanding of "online"

We call this method “online” because, unlike proxy-based method (e.g. DoGE), it estimates the domain weights on the fly during a single optimization loop. We have clarified this in the manuscript in L57-L59.

W3: Assumption of training a LLM with gradient

We agree with the reviewer that DGA requires the ability to compute gradients of the model. As the reviewer says, we believe that this is not a strong requirement because DGA is used during pre-training or continued pre-training, which are phases typically requiring the ability to compute gradients of the model.

We hope our answers can resolve most of the concerns. If you have any further questions, don’t hesitate to contact us. We will improve the manuscript according to your valuable suggestions!

[1] Emergent Abilities of Large Language Models

We thank the reviewer for the thoughtful feedback! We summarize them into 7 questions and provide our clarifications as follows.

W1&Q1: Proof of convergence

As the reviewer points out, the theoretical understanding of data mixing (i.e. domain reweighting) algorithms is a valuable topic that could significantly enhance the theoretical contributions of our paper. We are happy to report that we could prove that DGA converges in a simple setting. We have updated the paper to include a convergence proof of DGA in a simple convex scenario (Theorem 2). To the best of our knowledge, this is the first theoretical proof of convergence for a data reweighting algorithm. For instance, prior works such as [1,2] do not prove the convergence of their algorithms.

Extending this proof to more realistic scenarios on large language model training would be an intriguing and complex direction for future research. As discussed in the updated version, we believe it would be challenging due to the inherent non-convexity of the problem.

W2&W3&Q2&Q6&Q7: Comparison between DGA and DoGE

We have updated the paper and included the comparison between DGA and DoGE in Appendix F (Fig. 15). The results demonstrate that DGA outperforms DoGE in both specialized loss and general loss. This is because Doge requires training a proxy model before retraining. We believe that this illustrates nicely the merits of our online method.

W4&Q3: Detailed complexity comparison between distribution reweighting and domain reweighting

We agree with the reviewer that the original domain reweighting algorithm cannot handle the case when the number of domains is much larger than the update frequency ($k>>T_r$), as we mentioned in the text L342-362. The distribution reweighting algorithm (Sec. 3.2, Equ. 9) is proposed to resolve this challenge by approximating domain weights $\alpha_{domain} \in R^{k}$ as a linear combination of $N$ distributions, where ($N<<k$). In our experiments, we use $k\in \{64,4096,262k\}, N=23, T_r=100$ which induces $(N+1)/T_r = 24$ more computation overheads. In comparison, the standard domain reweighting will lead to $(k+1)/T_r = 262k$ overhead factor, which is impossible for LLM training.

We add one more sentence for clarification in L358-359.

W5&Q4: Ablations on hyperparameters

We provide the ablations on the EMA factor $\beta$ in Appendix G.2, under the same setting in Sec. (3.1). We applied $\beta={0.1,0.9}$, where we find $\beta=0.1$ is the optimal value across all token limits. In comparison, a larger $\beta$ will slow down the learning of domain weights, which hurts training efficiency.

Besides, we are running ablation experiments on the step size ($lr_{\alpha}$) and update frequency ($T_r$) of domain weights $\alpha$ on language modeling task. We will have results by the end of the weekend. For each combination of hyperparameters, we will report the GPU time and the specialized loss on the target domain in in Appendix G.2.

Q5: Comparison between DGA and Importance sampling

We agree with the reviewer that despite a marginal improvement above uniform baseline, DGA shows comparable or even lower reasoning accuracy compared to importance sampling baselines, as shown in Table 1. However, we provided the evidence in Figure 1 and Figure 2 that DGA demonstrates superior stability in language modeling than the importance sampling method. We also included a discussion in Appendix B.1, Figure 12, which shows that DGA accelerates the task-adaptive training by a large margin compared to importance sampling.

As for the limited improvement in reasoning accuracy, we included two potential explanations in Appendix B3.1: (1) since the target objective in the outer-loop ($L_{spe}$ in Equ. 3) is the next token prediction loss for language modeling, the language modeling performance is expected to be improved. However, it does not necessarily translate to improved reasoning accuracy (2) MMLU is considered to be a challenging reasoning task, where the accuracies can hardly be improved when the model capacity is below some specific threshold (Figure 11. (G) in [3]). How to derive an optimization objective that directly benefits reasoning accuracy would be a compelling direction for future research.

Q8: DGA performance across model scales

We present language modeling perplexity scores across three model scales (125M, 350M, 750M) in Appendix G.2, Figure 17. While the experiments on the largest models (750M) are still running, our current results on Arxiv demonstrate that DGA consistently outperforms uniform sampling by a significant margin across all three model scales. While it exhibits some degradation in general knowledge compared to uniform sampling, the gap narrows on the largest model (750M), indicating the potential scalability benefits of DGA in larger model regimes.

Q9: Potential future direction of data mixing

As the reviewer pointed out, the reasoning capability would not necessarily increase if we apply the next token prediction loss as the objective for data mixing optimization. How to derive an optimization objective that directly benefits the reasoning accuracy would be a compelling future direction.

We hope our answers can resolve most of your concerns. If you have any further questions, don’t hesitate to contact us. We will improve the manuscript according to your valuable suggestions!

[1] Doge: Domain reweighting with generalization estimation [2] Doremi: Optimizing data mixtures speeds up language model pretraining [3] Emergent Abilities of Large Language Models

Dear reviewer TRYF,

Thank you again for your efforts reviewing our paper and providing valuable feedback! Since tomorrow is the last day of the author-reviewer discussion phase, we hope you can take some time to look into our rebuttals and additional results. Let us know if you have other concerns and questions!

Sincerely,

the authors

We thank the reviewer for the feedback and insightful questions! We provide corresponding clarifications as follows:

W1: Limited improvement in reasoning accuracy

We agree with the reviewer that despite a substantial improvement above uniform baseline, DGA performs comparably as importance sampling baselines in reasoning accuracy scores, as shown in Table 1.

As for the considerable gap between language modeling ability and reasoning accuracy, we included an discussion in Appendix B.3.1 with two potential explanations: (1) since the target objective in the outer-loop ($L_{spe}$ in Equ. 3) is the next token prediction loss for language modeling, the language modeling performance is expected to be improved. However, it does not necessarily translate to improved reasoning accuracy (2) MMLU is considered to be a challenging reasoning task, where the accuracies can hardly be improved when the model capacity is below some specific threshold (Figure 11. (G) in [1]). How to derive an optimization objective that directly benefits the reasoning accuracy would be a compelling future direction.

In addition, we report the perplexity scores as the common measurement on language modeling capability. As shown in Appendix B.1, Figure 12, DGA achieves up to $6.5\times$ acceleration on task-adaptive training on MMLU benchmark.

Q1&W2: Performance comparison on pretraining + finetuning

We thank the reviewer for the great suggestion! Though the current work mainly focus on the task-adaptive pretraining, finetuning could have a great impact in task adaptation performance in real cases. We are currently running the finetuning experiments and we expect to present the results by the end of the weekend.

We hope our answers can resolve most of the concerns. If you have any further questions, don’t hesitate to contact us. We will improve the manuscript according to your valuable suggestions!

[1] Emergent Abilities of Large Language Models

Dear reviewer yT5T,

Thank you again for your efforts reviewing our paper and providing valuable feedback! Since tomorrow is the last day of the author-reviewer discussion phase, we hope you can take some time to look into our rebuttals and additional results. Let us know if you have other concerns and questions!

Sincerely,

the authors

We thank the reviewer for their insightful questions and suggestions which help us improve our work! We provide our clarifications as follows:

W1: Model architecture and scales

We want to clarify that we do not pre-train models on the entire redpajama-v2 (30T tokens), which would induce an intensive computation cost and also be wasteful for these model sizes according to neural scaling law. We train our largest model (N=750M) on ~25B tokens, which is above the standard 20 tokens per parameter derived from scaling laws [1]. As discussed in the paper, L204-L210 (computational cost and memory overhead), there is in principle no obstacle to scale DGA to larger models given sufficient computation resources. However, pre-training 3B or larger models sadly exceeds our computational budget.

As for the architecture, we adopt a standard decoder-only transformer, which is commonly used in most LLM papers as well as prior works on data mixing [2,3,4]. Since the proposed DGA method is derived from general optimization theory, it should be architecture-agnostic and applicable to OpenLLaMA as well. Seeing how it adapts to other generic pretraining sets with various data quality is an interesting future research direction.

W2: More restricted token constraints

We consider 30M tokens per domain as a strict restriction since this gives 30M*64 = 1.92B tokens in total to train the uniform baseline, which utilizes all tokens across 64 domains. This is a sufficient amount of tokens to train a 125M model without overfitting. A stricter token constraint, such as 1M tokens per domain, would lead to inevitable overfitting, even from the uniform baseline, which utilizes all available tokens with the best diversity.

Since we aim to investigate whether DGA could strategically sample from various domains to achieve better performance on the target domain without overfitting, we believe that using 1M tokens per domain is too far from a realistic pre-training scenario.

W3: Catastrophic forgetting

We thank the reviewer for this judicious remark. Indeed, catastrophic forgetting is a serious issue for language models, especially in the task-adaptive training.

In response, we summarized our observations in a detailed discussion in Appendix H. We investigate the catastrophic forgetting in various contexts: (1) how does DGA preserve general knowledge comparing to importance sampling and DoGE baselines; (2) how does distribution reweighting algorithm performs comparing to domain reweighting algorithm; (3) how would DGA suffer from catastrophic forgetting when applied on various scale of models. We highlight that DGA effectively mitigates catastrophic forgetting compared to other baselines (Importance sampling and DoGE). Specifically, DGA achieves a well-balanced trade-off: it produces a specialized model with low specialization loss on the target domain while maintaining general knowledge, as evidenced by its low loss on the Pile domains (representing broad, diverse domain knowledge) and the direct generic loss from RedPajama-v2.

We hope our answers resolve most of your concerns. If you have any further questions, don’t hesitate to contact us. Your valuable suggestions have helped us improve the manuscript!

[1] Scaling Laws for Neural Language Models

[2] Doge: Domain reweighting with generalization estimation

[3] Doremi: Optimizing data mixtures speeds up language model pretraining

[4] Specialized language models with cheap inference from limited domain data

Dear reviewer HaQd,

Thank you again for your efforts reviewing our paper and providing valuable feedback! Since tomorrow is the last day of the author-reviewer discussion phase, we hope you can take some time to look into our rebuttals and additional results. Let us know if you have other concerns and questions!

Sincerely,

the authors

We have finished the following experiments to complete our previous response:

Q4: Lack of ablation studies on hyperparameters

We have included the complete ablation results on two hyperparameters, the step size ($lr_{\alpha}$) and reweight frequency ($T_r$), in Appendix G.3 and Figure 18. The results demonstrate that using a step size that is too small will result in slow updates to the domain weights. On the other hand, applying a large step size ($\eta$) can accelerate the learning of domain weights but may also lead to training instability due to overly up-weighted domains. Since EMA can effectively stabilize learning, we recommend that practitioners choose $\eta$ values between 0.1 and 0.5. Regarding the reweight frequency, using a smaller $T_r$ generally improves the final specialized loss but increases computation costs. To balance cost and performance, we recommend setting $T_r$ between 30 and 100. However, these values should be tailored to specific use cases and levels of domain granularity.

Q5: Lack of a comparison of model performance after finetuning

We have included the results after fine-tuning in Appendix G.1and Figure 16, where we show that DGA consistently outperforms the uniform sampling and importance sampling baselines across various scales of fine-tuning tokens from the downstream task. In other words, the gain incurred by using DGA during pre-training is kept during fine-tuning.