Dynamic Gradient Alignment for Online Data Mixing

Authors: Simin Fan, David Grangier, Pierre Ablin
Primary area: foundation or frontier models, including LLMs
Keywords: LLM pretraining, data reweighing, data mixing

Abstract:
The composition of training data mixtures is critical for effectively training large language models (LLMs), as it directly impacts their performance on downstream tasks. Our goal is to identify an optimal data mixture to specialize an LLM for a specific task with access to only a few examples. Traditional approaches to this problem include ad-hoc reweighting methods, importance sampling, and gradient alignment techniques.
This paper focuses on gradient alignment and introduces Dynamic Gradient Alignment (DGA), a scalable online gradient alignment algorithm. DGA dynamically estimates the pre-training data mixture on which the models' gradients align as well as possible with those of the model on the specific task.
DGA is the first gradient alignment approach that incurs minimal overhead compared to standard pre-training and outputs a competitive model, eliminating the need for retraining the model. Experimentally, we demonstrate significant improvements over importance sampling in two key scenarios: (i) when the pre-training set is small and importance sampling overfits due to limited data; and (ii) when there is insufficient specialized data, trapping importance sampling on narrow pockets of data.
Our findings underscore the effectiveness of gradient alignment methods in optimizing training data mixtures, particularly in data-constrained environments, and offer a practical solution for enhancing LLM performance on specific tasks with limited data availability.

Decision: reject
Reviewer scores: 5, 5, 5, 6, 5

Reviews and discussion:
<thread>
creation date:
1730576913740

thread_type:
official_review

<thread_object>
creator: reviewer_9lhj

<content>
summary:
The paper introduces Dynamic Gradient Alignment (DGA), an online data reweighting algorithm aimed at optimizing data mixtures for training large language models (LLMs) on specific tasks. Unlike static or previously established gradient-alignment methods, DGA dynamically updates domain weights during training by aligning model gradients on the target task, leading to improved performance and reduced overfitting, especially in data-constrained environments. The method is scalable, incurs minimal computational overhead, and leverages exponential moving averages (EMA) for stabilizing updates.

strengths:
- **Innovative Approach**: The paper presents DGA as an online, scalable alternative to existing domain reweighting methods, offering significant advantages in flexibility and adaptability.
- **Empirical Evidence**: Detailed experiments show DGA's effectiveness compared to importance sampling, particularly under token-constrained conditions and fine-grained data domains.
- **Theoretical Rationale**: The method is well-anchored in optimization principles, connecting gradient alignment and bilevel optimization frameworks.
- **Practical Relevance**: DGA demonstrates strong utility for real-world data scenarios with limited resources, showcasing improved balance between exploration and exploitation in domain weighting.

weaknesses:
- **Theoretical Assumptions**: While the approach connects well to optimization theories, the convergence proofs are limited due to non-convexity in practice.
- **Specificity vs. Generality**: The results on downstream reasoning (MMLU) indicate that while DGA improves language modeling, the gains do not always translate to enhanced reasoning capabilities, suggesting potential limits in generalizability.
- **Sparse Baseline Comparisons**: The paper could benefit from more varied baseline methods beyond standard importance sampling to contextualize DGA's performance.
- **Stability Under Extreme Data Constraints**: Although EMA helps with stability, the paper notes challenges when DGA operates without it, implying potential reliability issues under certain configurations.

questions:
- Can the authors provide insights on the performance of DGA on tasks beyond language modeling to substantiate its generalization claims?
- What are the primary factors influencing DGA's stability when EMA is not used?
- Are there any plans to explore hybrid approaches combining DGA with more computationally intensive, sample-level methods?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721193556

thread_type:
official_review

<thread_object>
creator: reviewer_ejdq

<content>
summary:
The submission proposes a new method to improve training of LLMs on datasets that have an increased diversity / size ratio, e.g. mix of many small datasets. The method is a combination of increased (local) adaptability via dynamic gradient alignment (DGA) and smoothing via a moving average (EMA). Results report effects on training loss and accuracy in comparison to prior art of "Importance Sampling" with the same target.

strengths:
The submission is well-written and a good mixture of motivation, analysis of previous work, detailing theoretic foundations, experimental verification and conclusions.

weaknesses:
Experimental results only show loss, but no accuracy improvements over Importance Sampling (Table 1). Tables 2 & 3 also show no consistency in the pattern of first and second best of the compared approaches and alternating "wins" by tiny margins. This suggests that the proposed technique is just possibly only an alternative to Importance Sampling.

Initially, I had some issues correctly understanding the term "online" in the name of the proposed technique. Similar to me it may mislead other readers to think that this postpones the mixing to inference time (that would be an awesome new result), but it clearly doesn't.

A slight weakness lies in the fact that the technique is based on the assumption of LLMs being trained via gradients. This limits applicability to the current state-of-the-art. In contrast, most sampling techniques are more generally applicable.

questions:
Lines 067-069 ("Since the domain weights... leads to snow-balled errors."): this claim is vague and should be verified by experiment(s). Or did I miss that (maybe Figure 1(a)? Can you please add clarity to this?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730729363189

thread_type:
official_review

<thread_object>
creator: reviewer_tryf

<content>
summary:
This paper introduces Dynamic Gradient Alignment (DGA) as a stable and scalable data mixing method for LLM pretraining to address two key challenges of online domain reweighting. DGA is an online algorithm that adjusts the training data distribution according to the current model status, relying on gradient alignments to estimate progress on the target task. The paper demonstrates that DGA with EMA can significantly mitigate overfitting and yield superior performance on the end-task by balancing exploitation and exploration under limited tokens within generic domains. Additionally, a novel distribution reweighting strategy is proposed, enabling DGA to scale up to extremely fine-grained data domains without incurring intensive computations. Experiments on MMLU show that DGA effectively leverages fine-grained domain knowledge to balance specialty and diversity during training, demonstrating the scalability and efficiency of gradient-alignment-based data reweighting methods in data-constrained settings.

strengths:
1. The paper introduces Dynamic Gradient Alignment (DGA), a novel online gradient alignment algorithm that dynamically adjusts training data mixtures.
2. The experiments are well-designed, covering key scenarios of small pretraining sets and insufficient specialized data, using the MMLU benchmark.
3. The results clearly demonstrate the advantages of DGA over importance sampling and uniform baselines, with detailed analysis of different model scales and domain granularities

weaknesses:
1. The paper provides a brief description of the theoretical foundation of DGA, lacking in-depth mathematical derivation and theoretical analysis. There is no rigorous proof of convergence and stability for DGA, which raises concerns about the robustness and reliability of the proposed method.
2. The paper does not compare DGA with other advanced domain reweighting methods. This omission makes it difficult to assess the relative performance and advantages of DGA compared to existing methods.
3. Benchmarking against other methods is essential to establish the novelty and superiority of the proposed approach. The absence of such comparisons weakens the paper's contribution to the field.
4. The paper lacks a detailed complexity analysis of the experiments, particularly in Section 3.2. A thorough analysis of the computational complexity and resource requirements would provide better insights into the practical implications of implementing DGA.
5. The paper does not provide a comprehensive analysis of the impact of different hyperparameters on the performance of DGA.
Understanding how sensitive DGA is to various hyperparameters is important for its practical implementation and optimization.

questions:
1. Can you provide more mathematical derivation and theoretical analysis of DGA, including proofs of convergence and stability? This would help in understanding the robustness and reliability of the proposed method.
2. How does DGA compare with other advanced domain reweighting methods (e.g., DoGE)? A detailed comparison would help in understanding the relative advantages and limitations of DGA.
3. Can you provide a detailed complexity analysis of the experiments in Section 3.2?
When k is much larger than T_r, 1) is DGA much more complicated than importance sampling method and 2) no better than DoGE in terms of both performance and complexity? Understanding the computational complexity and resource requirements is crucial for evaluating the practical implications of implementing DGA.
4. Can you provide a comprehensive analysis of the impact of different hyperparameters on the performance of DGA? Such ablation study is important for practical implementation and optimization.
5. In Table 1, DGA does not show significant improvement above importance sampling. However, DGA is much more complicated under this experiment setting?
6. Do you have more benchmarking results against other methods? it is essential to establish the novelty and superiority of the proposed approach.
7. In experiments section, there should be a comparison between DGA and DoGE in terms of performance and computational complexity? As DGA is heavily inspired by DoGE, a detailed comparison would provide insights into the relative strengths and weaknesses of both methods.
8. How does DGA perform across different model scales? Can you provide more detailed results and analysis for various model sizes to understand its scalability and effectiveness?
9. What are the potential directions for future work and improvements on DGA? Are there any planned extensions or enhancements to address the current limitations and expand its applicability?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730773364472

thread_type:
official_review

<thread_object>
creator: reviewer_yt5t

<content>
summary:
The authors propose a novel pre-training data mixing method, called DGA, which enables LLMs to adapt to downstream tasks with just a few examples. By alternately updating model parameters and domain weights, DGA achieves exploration and exploitation on the training data, thereby solving the overfitting problems of importance sampling. DGA provides an approach for the data mixing problem in pre-training for specific downstream tasks. Additionally, the authors validated their method on the Pile and MMLU datasets, demonstrating DGA’s effectiveness in language modeling for downstream tasks.

strengths:
1,DGA effectively addresses the overfitting issue of importance sampling in scenarios with limited data.

2, By incorporating the EMA term, DGA effectively addresses the issue of training instability.

3, By combining with importance sampling, DGA resolves the issue of "scaling linearly with the number of domains k". Even when there are a large number of domains, it can also demonstrate its advantages.

weaknesses:
1, DGA's reasoning accuracy on datasets like MMLU is not good, but the authors only emphasize that DGA has strong capabilities in language modeling. The performance on downstream tasks should be related to reasoning accuracy, rather than the capability of language modeling.

2, The authors did not compare the performance differences between DGA and other methods under the paradigm of pre-training + fine-tuning or pre-training + ICL, but LLMs typically adapt to downstream tasks using these paradigms, which significantly differs from actual usage scenarios.

questions:
1, How would DGA perform under the paradigms of pre-training + fine-tuning or pre-training + ICL?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730867780470

thread_type:
official_review

<thread_object>
creator: reviewer_haqd

<content>
summary:
This paper proposes Dynamic Gradient Alignment (DGA), an online domain reweighting approach for the purpose of identifying an optimal data mixture for training an LLM for a downstream task with access to only a few examples (denoted specialized set).  DGA estimates stepwise optimal domain weights during model training.  The traditional importance sampling approaches suffer from some limitations, such as overfitting into small-scale pretraining data in each domain, difficulty to handle large domain granularity, and easily trapping into quite suboptimal local optimum with limited specialized set. The authors empirically show that DGA alleviates these limitations, offers an effective approach for optimizing training data mixture for downstream tasks with limited specialized set.

strengths:
(1)	The proposed DGA shows advantages over traditional importance sampling approaches in two challenging scenarios, that is, the number of training samples in each domain is limited instead of the unrealistic assumption of unlimited number of tokens, and the domain granularity is quite large making traditional domain reweighting methods cause intractable computational overheads.

(2)	The paper proposed two major innovations in DGA, including incorporating EMA term into domain reweighting to mitigate overfitting and improve stability, and also introducing a distribution reweighting scheme, with the number of distributions smaller than the number of domains, for scaling up to very fine-grained domains (very large number of domains). Though DGA is built upon DoGE, the differentiators from DGA over DoGE are clearly explained in the paper, as the lower variance, smaller overhead, and the novel EMA in DGA, compared to DoGE.

(3) Overall, the paper is clearly written. The Appendix provides useful additional results for the two challenging scenarios, hyperparameter information, and algorithm details.

weaknesses:
(1)	For empirical validations, the paper used Redpajama-v2 as the pre-training data, which has 30T tokens after filtering and deduplication. This is non-trivial amount of data. However, the paper pre-trained models with 125M, 350M, and 750M number of parameters. In comparison, some SOTA LLMs are trained on smaller or comparable amount of data with much larger model sizes, for example, Llama 3 405B trained on 15.6T tokens, QWen2.5-8B trained on 18T tokens, OpenLLaMA 3B, 7B, and 13B trained on 1T tokens, etc. The model sizes explored in this paper may be too small and too limited model capacity to sufficiently learn from 30T tokens data. It would be useful to pre-train models with larger sizes, e.g., to 3B and 7B, which will correspondingly impact the experimental results and analysis.  A related question is, why not use OpenLLaMA setup including its open-source data for the experiments?

(2)	Figure 1 analyzes the effect of DGA and DGA with EMA. However, 30M tokens/domain are still quite large. How about the comparison when the number of tokens/domain is further reduced, e.g., less than 1M, less than 100K, 10K?

(3) One critical question when adapting pre-trained LLM for downstream tasks is how much the resulting model can achieve a good balance between general capabilities and performance on downstream tasks. For example, for instruction following capabilities on a downstream task, how much the model can perform for general instruction following tasks with DGA compared to the baselines. This issue of catastrophic forgetting has not been discussed in the paper.

questions:
(1)	Please address the points listed under Weaknesses.

(2)	There are some typos and grammar errors in the paper. For example, “obtain obtain”.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732308682169

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response: Summary of the proposed changes to address the reviewers' concerns.

comment:
We sincerely thank all reviewers for their efforts in reading our paper and providing insightful feedback, which helps us improve our work. We appreciate that the novelty and superiority of DGA as an online data mixing method have been recognized by all the reviewers. 

We summarize the main concerns raised by the reviewers with our proposed changes as follows:
1. **Q**: Lack of a theoretical convergence analysis of the proposed method. [TRYF, 9LhJ]
    
    **A: We have updated the paper to include a convergence proof of DGA in a simple quadratic 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.
2. **Q**: Lack of comparison to other domain reweighting baselines (e.g. DoGE). [TRYF, 9LhJ]
    
    **A: We have updated the paper and include 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.
3. **Q**: Lack of discussions on catastrophic forgetting of generic knowledge. [HaQD]
    
    **A: We provide an in-depth discussion on catastrophic forgetting in Appendix H**, where we show that DGA can preserve more general knowledge during task-adaptive pretraining, while other baselines (Importance sampling, DoGE) are more vulnerable to catastrophic forgetting. 
4. **Q**: Lack of ablation studies on hyperparameters. [TRYF] 
    
    **A: We provide the ablations on the EMA factor $\beta$ in Appendix G.2** in domain reweighting with $k=64$, under the same setting in Sec. (3.1). We applied $\beta={0.1,0.9}$, where we find DGA with $\beta=0.1$ can effectively prevent overfitting without sacrificing the learning efficiency on the specilized knowledge. In comparison, a larger $\beta$ will slow down the learning of domain weights, which hurts training efficiency. Meanwhile, **we are running ablation experiments on the step size ($lr_{\alpha}$) and update frequency ($T_r$) of domain weights $\alpha$ on one 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 Appendix G.2. However, we found that the hypeparameters discussed in the main text worked well for a broad range of tasks/model sizes.

For other questions and points of clarification, we provide detailed responses correspondingly under each individual review.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes Dynamic Gradient Alignment (DGA), an online domain re-weighting algorithm designed to optimize training data mixtures for large language models (LLMs) in downstream tasks with limited examples. The method addresses limitations of traditional importance sampling approaches, such as overfitting to small-scale data, difficulty handling fine-grained domains, and suboptimal local minima. The paper claims following contributions: 1) DGA dynamically adjusts domain weights during training by aligning gradients with the target task, enabling effective exploration and exploitation of training data. Exponential moving averages (EMA) is used to stabilize weight updates and improve training scalability. The method is capable of handling extremely fine-grained data domains without significant computational overhead, and adapting to diverse datasets, including mixes of small-scale datasets with high diversity. 2) ablations and empirical results on benchmarks such as MMLU and the Pile showed the efficacy of proposed method; overfitting is mitigated with balances specialty and diversity during training.


Strength of this paper
- Tackles challenges in recently popular topics in LLMs area such as overfitting and instability when limited training samples are available. The problem has strong utility for real-world scenarios with limited resources.
- Empirically validate the efficacy of proposed method. Using benchmarks like MMLU to validate DGA’s effectiveness compared to importance sampling and uniform baselines, across varying model scales, domain granularities, and data limitations.
- The proposal of using online gradient alignment of dynamically adjusting training data mixtures , and exponential moving average to stablize is reasonable


Weakness of this paper

Several reviewers raised few concerns/limitations of this paper. By addressing these limitations, the paper could strengthen its experiment and expand impact.

- Experiment scope is limited. The paper only compare with standard importance sampling, but not advanced domain reweighting methods or other variations. The model sizes (125M, 350M, 750M parameters) are too small to fully leverage the 30T tokens in the Redpajama-v2 dataset. Larger models, such as 3B or 7B, should be explored to evaluate the method's effectiveness, and verify whether the effectiveness can be translated to models with sizes commonly used in the community. Results also lack exploration of smaller token/domain scenarios (e.g., <1M, <100K, or 10K tokens/domain), which would better illustrate DGA's adaptability under extreme data constraints. The paper also need to evaluate how DGA impacts instruction-following capabilities or balances general and task-specific performance without forgetting general capabilities. Comparisons with common paradigms such as pre-training + fine-tuning or pre-training + in-context learning (ICL) should also be conducted.
- Some concerns about generalizability. E.g., the paper shows DGA improves language modeling but does not consistently enhance downstream reasoning capabilities (e.g., MMLU); the paper does not provide sensitivity analysis of hyperparameters, which is essential for optimization in real-world scenarios. There is limited mathematical derivation or discussion about training convergence to validate stability and robustness of proposed method; the assumptions behind DGA, such as reliance on gradient-based LLM training, limit its applicability compared to more general sampling techniques.