Dynamic Gradient Alignment for Online Data Mixing

## Abstract

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.

## 1 Introduction

Large Language Models (LLMs) are typically pre-trained on extensive, generic corpora sourced from a variety of data domains (Brown et al., ; Touvron et al., ; Zhang et al., ), with the composition of these corpora often depending on domain availability or heuristics (Gao et al., ; Together AI Team, ).
While the diversity of natural texts allows the model to learn from various knowledge sources, not all data domains are equally beneficial according to the targeted tasks. The uncurated nature of web-crawled contents could lead to sub-optimal outcomes due to the variations in data quality (Longpre et al., ). Plus, some domains may contain misinformation and biases, as one potential source of hallucinations in language generation (Lin et al., ; Huang et al., ).

To better generalize to the downstream target tasks, it is critical to identify the most beneficial pretraining subset from large, generic pretraining corpora. While data selection per sample can be costly, domain reweighting offers an efficient group-level selection approach. Domain reweighting methods assume that samples from the same domain share similar features and search for optimal sampling weights across domains (Xie et al., ; Fan et al., ; Liu et al., ; Kang et al., ; Grangier et al., ).
The domains that most positively impact the target tasks should be assigned higher weights.

In this work, on top of a large, generic pretraining corpus, we assume we have access to a few examples representative of the downstream task on which we want the model to generalize, a so-called *specialized set*.
For this setup, Grangier et al. () recently proposed a simple and scalable *importance sampling* based method to domain reweighting, where the weight of a domain is given by the frequency of samples in the specialized set closest to the domain, where distance is measured with SentenceBert (Reimers and Gurevych, ) embeddings. This method determines the domain weights before any training and is model-agnostic.

Likewise, prior gradient-alignment methods determine a static domain weights for large-scale LM training, often relying on a small-scale proxy model (Xie et al., ; Fan et al., ) or fitting a scaling law (Liu et al., ; Kang et al., ). While these methods show improvements over training on the natural distribution of a generic corpus, they do not dynamically update domain weights during training to adapt to the current model state. In practical training scenarios, a large model may quickly overfit on certain domains with high weights. In such cases, an online weighting method can respond by shifting emphasis to other domains.

We propose Dynamic Gradient-Alignment (DGA), an online domain reweighting method that estimates step-wise optimal domain weights during model training.
Inspired by DoGE (Fan et al., ), at each reweighting step, DGA upweights the data domain whose gradient aligns more with the model’s gradient on the specific set. From the optimization perspective, training the model on the most-aligned data domain yields the greatest reduction in the targeted loss.
By incorporating an exponential-moving-average (EMA) term in online domain weights updates, DGA effectively mitigates overfitting and prioritizes the domains that currently benefit the target task the most.
Since the domain weights and model parameters are updated concurrently, inaccurate domain weights can potentially drive the model into suboptimal states, which further leads to snow-balled errors. In such cases, the EMA term serves as a correction factor, guiding the model back to a more stable state.
As an additional contribution, we scale the domain reweighting methods into extremely fine-grained domains (e.g. $262k$ domains) by introducing a novel distribution reweighting mechanism. Rather than directly reweighting $262k$ data domains, distribution reweighting reparameterizes the high-dimensional domain weights as a convex combination of weight vectors derived from a set of distributions estimated from embedding-based importance sampling (Grangier et al., ).
With the number of distributions less than the number of training data domains, it allows DGA to scale to thousands of domains and make the most of the fine-grained group-level features.

Our experiments demonstrate the effectiveness of DGA compared to standard pre-training and importance sampling baselines in two challenging cases: (1) the resource of training tokens in each domain is limited instead of infinite (§ 3.1), and (2) the domain granularity is extremely large, which introduces intractable computation overheads on the domain reweighting problem (§ 3.2).

## 2 Data Mixing with Specialized Target

### 2.1 Generic dataset and specific tasks

We consider a generic training corpus $D_{\mathrm{gen}}=\{D_{1},\ldots,D_{k}\}$, which is partitioned into $k$ distinct data domains.
We can sample from each of the $k$ domains to train a model.
Consequently, we can sample from a *mixture* of these domains and draw a batch of data following the law ${\bm{x}}\sim\mathrm{mix}({\bm{\alpha}})\triangleq\sum_{i=1}^{k}\alpha_{i}D_{i}$, where ${\bm{\alpha}}\in\mathbb{R}^{k}$ is the mixture *weights*, belonging to the *simplex* ${\bm{\alpha}}\in\Delta^{k}\triangleq\{{\bm{\alpha}}\in\mathbb{R}^{k}|\sum_{i=1
}^{k}\alpha_{i}=1\text{ and }\alpha_{i}\geq 0\text{ for all }i\}$.
Here, getting one sample from $\mathrm{mix}({\bm{\alpha}})$ means first getting a random index $i\in\{k\}$ from the categorical distribution corresponding to the vector of probabilities ${\bm{\alpha}}$, and then outputting a random sample from the domain $D_{i}$.
Sampling from this law is computationally efficient if we can efficiently sample from each domain.
Next, we consider a model, parameterized by ${\bm{\theta}}\in\mathbb{R}^{p}$, and a loss function $\ell({\bm{\theta}},{\bm{x}})$ defined for ${\bm{x}}\in D_{\mathrm{gen}}$.
To simplify notation, given a set of samples $S$ (which can be either a full dataset $D_{i}$, or a mini-batch), we denote the average over $S$ of the loss $\ell({\bm{\theta}},S)\triangleq\frac{1}{\#S}\sum_{{\bm{x}}\in S}\ell({\bm{
\theta}},{\bm{x}})$.
Since we focus on LLMs, $\ell$ is typically the next-token-prediction loss.
For a given mixture weight ${\bm{\alpha}}$, we can update ${\bm{\theta}}$ by doing optimization steps on the *generic loss*

$$ $\displaystyle L_{\mathrm{gen}}({\bm{\theta}},{\bm{\alpha}})\triangleq\mathbb{E }_{{\bm{x}}\sim\mathrm{mix}({\bm{\alpha}})}[\ell({\bm{\theta}},{\bm{x}})]=\sum _{i=1}^{k}\alpha_{i}L_{i}({\bm{\theta}})\text{ with }L_{i}({\bm{\theta}}) \triangleq\ell({\bm{\theta}},D_{i})$ (1) $$

In this paper, our goal is to use this data-mixture to train a model that performs well a *specific* task.
We assume to have access to samples from this task, split into train and test sets.
We call the train set the *specific dataset* $D_{\mathrm{spe}}$ that we use to train models.
The performance on the specific set is measured with the *specific loss*

$$ $L_{\mathrm{spe}}({\bm{\theta}})\triangleq\ell({\bm{\theta}},D_{\mathrm{spe}}).$ (2) $$

We assume that the specific set $D_{\mathrm{spe}}$ is small; hence, optimizing $L_{\mathrm{spe}}$ directly leads to overfitting: the loss on the test data would be much higher than $L_{\mathrm{spe}}$.
Instead, to get to a low specific loss, we aim to find the optimal data mixing across $k$ data domains ${\bm{\alpha}}$ at each training step to get a good model while training on the reweighted generic distribution $\mathrm{mix}({\bm{\alpha}})$.

The target specialization task can be flexible according to the application domains, ranging from reasoning, instruction following, etc., corresponding to various objective loss functions, including the next-token prediction loss and preference-based losses when applied on pair-wise datasets.
In this paper, we focus on next-token prediction on another dataset.
Next, we introduce a general bilevel formulation of the data mixing problem.

### 2.2 Bilevel formulation

Since the lack of data forbids optimizing directly $L_{\mathrm{spe}}$, we look for the mixture ${\bm{\alpha}}$ such that optimizing the generic loss $L_{\mathrm{gen}}({\bm{\theta}},{\bm{\alpha}})$ yields the smallest specific loss (Grangier et al., ).
This is formalized by the following *bilevel optimization* problem (Bracken and McGill, ; Dagréou et al., ):

$$ ${\bm{\alpha}}^{\star}\in\operatorname*{arg\,min}_{{\bm{\alpha}}\in\Delta^{k}}L _{\mathrm{spe}}({\bm{\theta}}^{\star}({\bm{\alpha}})),\text{ such that }{\bm{ \theta}}^{\star}({\bm{\alpha}})\in\operatorname*{arg\,min}_{{\bm{\theta}}}L_{ \mathrm{gen}}({\bm{\theta}},{\bm{\alpha}})$ (3) $$

This bilevel formulation is intuitive: for a given weight ${\bm{\alpha}}$, the parameters obtained by minimizing the generic loss $L_{\mathrm{gen}}$ are ${\bm{\theta}}^{*}({\bm{\alpha}})$, and we want those weights to yield a small specific loss.
Notably, if the specific loss is a mixture of generic data with an unknown weight $\tilde{{\bm{\alpha}}}$, the bilevel formulation is guaranteed to recover it. In other words:

###### Theorem 1 .

Assume that there exists $\tilde{{\bm{\alpha}}}$ such that $D_{\mathrm{spe}}=\mathrm{mix}(\tilde{{\bm{\alpha}}})$
. Then, $\tilde{{\bm{\alpha}}}$ is a solution to the bilevel problem in Equation 3.

###### Proof.

We let $\tilde{{\bm{\theta}}}$ the minimizer of $L_{\mathrm{spe}}$. Then, for all ${\bm{\alpha}}$, we have by definition that $L_{\mathrm{spe}}({\bm{\theta}}^{*}({\bm{\alpha}}))\geq L_{\mathrm{spe}}(\tilde
{{\bm{\theta}}})$. Furthermore, since $D_{\mathrm{spe}}=\mathrm{mix}(\tilde{{\bm{\alpha}}})$, we have that $L_{\mathrm{gen}}({\bm{\theta}},\tilde{{\bm{\alpha}}})=L_{\mathrm{spe}}({\bm{
\theta}})$ for all ${\bm{\theta}}$, hence minimizing this yields ${\bm{\theta}}^{*}({\bm{\alpha}})=\tilde{{\bm{\theta}}}$. Putting these results together, we have proven that for all ${\bm{\alpha}}$, it holds $L_{\mathrm{spe}}({\bm{\theta}}^{*}({\bm{\alpha}}))\geq L_{\mathrm{spe}}({\bm{
\theta}}^{*}(\tilde{{\bm{\alpha}}}))$, so that $\tilde{{\bm{\alpha}}}$ is a solution to Equation 3.
∎

We consider two types of methods to solve Equation 3. Static methods construct a single mixture weight vector ${\bm{\alpha}}$ and then minimize $L_{\mathrm{gen}}({\bm{\theta}},{\bm{\alpha}})$;
we describe in the next section how to obtain this vector ${\bm{\alpha}}$.
Online methods modify the weights dynamically during model training. They
produce a sequence of weights ${\bm{\alpha}}^{(t)}$ where $t$ is the optimization iterate. In that case, at each training step, the parameters ${\bm{\theta}}^{(t)}$ are updated by doing an optimization step — with gradient descent or Adam — on the function $L_{\mathrm{gen}}({\bm{\theta}},{\bm{\alpha}}^{(t)})$.
We now discuss methods to obtain a weight vector ${\bm{\alpha}}$ or a sequence ${\bm{\alpha}}^{(t)}$.

###### Theorem 1 .

###### Proof.

### 2.3 A Strong Baseline: Importance Sampling

A sensible strategy is to train the model on a data mixture that most resembles the composition of the targeted specialization data distribution.
This is the philosophy behind importance sampling (Kloek and Van Dijk, ).
We estimate the importance sampling weights ${\bm{\alpha}}^{\mathrm{IS}}$ using the method of Grangier et al. ().
The core idea is to embed each generic domain using SentenceBert (Reimers and Gurevych, ), and then compute the centroid of each domain ${\bm{b}}_{i}=\frac{1}{\#D_{i}}\sum_{{\bm{x}}\in D_{i}}\mathrm{Bert}({\bm{x}})$. This defines a simple and cheap to compute selection function $c({\bm{x}})\in\{1\dots k\}$, assigning ${\bm{x}}$ to its closest centroid, i.e., $c({\bm{x}})=\operatorname*{arg\,min}_{i}\|\mathrm{Bert}({\bm{x}})-{\bm{b}}_{i}\|$ for ${\bm{x}}\in D_{\mathrm{gen}}\cup D_{\mathrm{spec}}$. We use it to predict the closest generic data domain for each data instance from the specific set. The importance sampling weights are obtained as the ratio of examples falling in each bin:

$$ ${\bm{\alpha}}^{\mathrm{IS}}_{i}\triangleq\frac{\#\{{\bm{x}}\in D_{\mathrm{spe} }|c({\bm{x}})=i\}}{\#D_{\mathrm{spe}}}$ (4) $$

One of the main advantages of this method is its simplicity: the computation of the weights ${\bm{\alpha}}^{\mathrm{IS}}$ is decoupled from model optimization and can be performed before training.
It is expected to work well when the specialization set can be well approximated by the reweighted generic set, i.e., when $L_{\mathrm{spe}}({\bm{\theta}})\simeq L_{\mathrm{gen}}({\bm{\theta}},{\bm{
\alpha}}^{\mathrm{IS}}).$
When this is not the case, it might not lead to a good specific loss.
Another potential issue with this method arises when it assigns a large weight to a generic domain $D_{i}$ with little available data.
In this case, training a model on $\mathrm{mix}({\bm{\alpha}}^{\mathrm{IS}})$ will overfit on that domain $D_{i}$, and it would have been better to reduce the weight of that domain to mitigate overfitting.
A last issue arises when the number of specific examples, $\#D_{\mathrm{spe}}$, is significantly smaller than the number of domains $k$. In this situation, the importance weights become sparse, as they can have at most $\#D_{\mathrm{spe}}$ non-zero coefficients. This sparsity could be problematic, as some domains with zero weights might still be close to $D_{\mathrm{spe}}$.
We illustrate these shortcomings in our experiments and explain how gradient alignment methods — which we introduce next — overcome them.

### 2.4 DGA: Dynamic Gradient Alignment

Algorithm. We introduce the DGA: Dynamic Gradient Alignment method for data reweighting to approximately solve the bilevel problem in Equation 3.
This algorithm builds upon DoGE (Fan et al., ) and we give a precise account of their differences later.
DGA keeps track of the model’s parameters ${\bm{\theta}}^{t}$ and dynamic weights ${\bm{\alpha}}^{t}$.
Once every $T_{r}$ steps, we compute the gradient alignments ${\bm{a}}^{t}$, by doing

$$ ${\bm{a}}^{t}_{i}=\langle\nabla\ell({\bm{\theta}}^{t},{\bm{x}}_{i}),\nabla\ell( {\bm{\theta}}^{t},{\bm{z}})\rangle\text{ where }{\bm{x}}_{i}\sim D_{i}\text{ and }{\bm{z}}\sim D_{\mathrm{spe}}.$ (5) $$

and update the weights by mirror descent on the simplex (Beck and Teboulle, ) with step $\eta>0$:

$$ ${\bm{\alpha}}^{t+1}=\frac{\hat{{\bm{\alpha}}}}{\sum_{i=1}^{k}\hat{{\bm{\alpha} }}_{i}}\text{ where }\hat{{\bm{\alpha}}}={\bm{\alpha}}^{t}\odot\exp(\eta{\bm{a }}^{t})$ (6) $$

We optionally store an EMA version of the weights ${\bm{\alpha}}^{t}$ parameterized by $\beta\in[0,1]$ to stabilize the training dynamics of the model’s parameters, and define ${\bm{\alpha}}_{\mathrm{EMA}}^{t+1}=(1-\beta){\bm{\alpha}}^{t}_{\mathrm{EMA}}+
\beta{\bm{\alpha}}^{t+1}$.
Finally, at each step, we update the model’s parameters ${\bm{\theta}}^{t}$ by doing an optimization step on $L_{\mathrm{gen}}({\bm{\theta}},{\bm{\alpha}}^{t}_{\mathrm{EMA}})$. The full algorithm pseudo-code is given in Algorithm 1.

Rationale. This algorithm can be seen as a heuristic to solve the bilevel problem in Equation 3.
Indeed, each update on ${\bm{\theta}}$ optimizes the inner loss.
The update rule on ${\bm{\alpha}}$ can be seen as a mirror-descent step on $L_{\mathrm{spe}}({\bm{\theta}}^{*}({\bm{\alpha}}))$ with several approximations. The first approximation consists of approximating the solution of the inner problem with one gradient descent step with step-size $\rho$: ${\bm{\theta}}^{*}({\bm{\alpha}})\simeq{\bm{\theta}}^{t}-\rho\sum_{i=1}^{k}{\bm
{\alpha}}_{i}\nabla L_{i}({\bm{\theta}}^{t})$.
We then approximate the specific loss at ${\bm{\theta}}^{*}$ by the post-update specific loss function: $L_{\mathrm{spe}}({\bm{\theta}}^{*}({\bm{\alpha}}))\simeq f({\bm{\alpha}},\rho)
\triangleq L_{\mathrm{spe}}({\bm{\theta}}^{t}-\rho\sum_{i=1}^{k}{\bm{\alpha}}_
{i}\nabla L_{i}({\bm{\theta}}^{t}))$, that is, the drop on the specific loss after an update.
When the step size $\rho$ is small, a Taylor expansion gives

$$ $f({\bm{\alpha}},\rho)=L_{\mathrm{spe}}({\bm{\theta}}^{t})-\rho\sum_{i=1}^{k}{ \bm{\alpha}}_{i}\langle\nabla L_{i}({\bm{\theta}}^{t}),\nabla L_{\mathrm{spe}} ({\bm{\theta}}^{t})\rangle+o(\rho)$ (7) $$

Similarly, we get that the gradient of $f$ is the gradient alignment:

$$ $\frac{\partial f}{\partial{\bm{\alpha}}_{i}}({\bm{\alpha}},\rho)=-\rho\langle \nabla L_{i}({\bm{\theta}}^{t}),\nabla L_{\mathrm{spe}}({\bm{\theta}}^{t}) \rangle+o(\rho)$ (8) $$

We want to use this gradient of $f$ to implement a mirror-descent method. Unfortunately, the gradients involved in the alignment are full-batch, so we approximate them with stochastic gradients obtained from mini-batches, yielding the alignments ${\bm{a}}^{t}$ from Equation 5.
Overall, we get the approximation $\nabla_{\bm{\alpha}}L_{\mathrm{spe}}({\bm{\theta}}^{*}({\bm{\alpha}}))\simeq-
\rho{\bm{a}}^{t}$; and the update rule in Equation 6 is a mirror descent step with this approximated gradient and step $\eta/\rho$.

We have explained the link between our algorithm and the bilevel problem in Equation 3.
Proofs showing convergence of our method require assumptions violated in practice, e.g.
most theoretical work assumes that the function ${\bm{\theta}}\to L_{\mathrm{gen}}({\bm{\theta}},{\bm{\alpha}})$ is
convex (Ghadimi and Wang, ; Arbel and Mairal, ; Dagréou et al., ). Nevertheless, successful
applications of related bilevel algorithms to non-convex neural networks have been reported recently (Fan et al., ; Grangier et al., ).

Figure: Algorithm 1 Dynamic Gradient Alignment method

Computational cost and memory overhead.
The computation cost of DGA is compared to the cost of a regular pre-training run.
For a base run iteration, the main cost is $t_{g}$, the cost of computing a gradient with a mini-batch $B$.
For DGA, we need to add the cost of updating the domain weights ${\bm{\alpha}}$, which only happens every $T_{r}$ iterations. This update requires computing the $k+1$ gradients (one per domain, one for $L_{\mathrm{spe}}$).
Hence the average cost of one iteration of DGA is $(1+(k+1)T_{r}^{-1})t_{g}$.
Therefore, DGA’s compute overhead is small when $T_{r}$ is large compared to the number of domains $k$.

During training, the memory is essentially used by the optimizer state, the model gradients and its activations.
For simplicity, we assume the same precision for storing all vectors.
The optimizer state (the model parameters and the two EMA terms for Adam) and the gradients have a storage cost of $4m_{g}$, where $m_{g}$ denotes the cost of storing the model parameters.
The cost of storing the activations during backpropagation is $m_{b}$. Regular pretraining with Adam therefore costs $4m_{g}+m_{b}$. DGA computes the required gradients sequentially and does not require more memory to store activations. It simultaneously stores two gradients instead of one (one domain gradient and one specific gradient): DGA, therefore, costs $5m_{g}+m_{b}$. This means that DGA memory overhead ranges from 0 (when $m_{b}\gg m_{g}$) to $25\%$ (when $m_{g}\gg m_{b}$).
DGA scales in terms of computational cost and memory.

Comparison with DoGE.
While our method is heavily inspired by DoGE (Fan et al., ), there are several key differences.
First, DGA samples from the mixture: the weights ${\bm{\theta}}^{t}$ are updated using samples drawn from the mixture $\mathrm{mix}({\bm{\alpha}}^{t})$, with the gradient $\nabla\ell({\bm{\theta}}^{t},{\bm{x}})$ where ${\bm{x}}\sim\mathrm{mix}({\bm{\alpha}}^{t})$; this is the same gradient that one would use during pre-training with weight ${\bm{\alpha}}^{t}$.
In contrast, DoGE’s weights are updated using a reweighted gradient $\sum_{i=1}^{k}{\bm{\alpha}}^{t}_{i}\nabla\ell({\bm{\theta}}^{t},{\bm{x}}_{i})$, where each ${\bm{x}}_{i}$ are drawn from the domain $D_{i}$.
For a fixed number of samples available at each draw, DGA’s gradient estimate has a lower variance (Seiffert et al., ).
As explained above, DGA has a small overhead compared to regular pre-training, while DoGE updates the weights at each iteration.
These two key differences mean that DGA is much closer to regular pre-training than DoGE. For instance, DGA never requires retraining a model from scratch using the mixture weights estimated from a previous run, while this is the costly strategy used for DoGE.
Finally, the EMA strategy described above is novel.

## 3 Experiments

Our experiments focus on two challenging cases.
First, given limited token resources within each training domain, the model would risk overfitting with weights concentrated on a few domains. Second, given large number of training domains, applying DGA on domain reweighting could introduce intractable computation overheads linearly increasing according to the domain granularity.

Generic Datasets and Domains. For all the experiments, we use Redpajama-v2 (Together AI Team, ) as the generic training set $D_{\mathrm{gen}}$. This is one of the largest public corpus for LLM pretraining. Redpajama-v2 contains 30 trillion filtered and deduplicated tokens from web-crawled dumps. Since this corpus does not come pre-segmented into domains, we obtain obtain individual generic domains from $D_{\mathrm{gen}}$ with clustering. Specifically, we use the embedding-and-clustering pipeline from Grangier et al. ().
We first embed all the training sequences ${\bm{x}}\in D_{\mathrm{gen}}$ with SentenceBert (all-MiniLM-L6-v2), yielding a 384 dimensional embedding $\mathrm{Bert}({\bm{x}})$. We then apply $k$-means clustering on the sentence embeddings into $k=64$ clusters yielding $k$ domains $D_{1},\dots,D_{k}$.

To get fine-grained generic domains, we apply hierarchical clustering on the top of the first level of $k_{1}=64$ clusters. Specifically, each domain is further clustered once again into $64$ smaller clusters. We apply this strategy twice to get domains with granularity $k_{2}=64^{2}=4096$ and $k_{3}=64^{3}=262k$.

Model Architecture.
We train small ($125$M), medium ($350$M) and large ($750$M) models with decoder-only transformers (Vaswani et al., ). We adopt most of the training settings and architectures from (Brown et al., ). Their details are provided in Appendix C.
For optimization, we use the AdamW optimizer (Loshchilov, ).

### 3.1 Domain Reweighting with Limited Resources

Figure: (a) 30M tokens per domain
Refer to caption: x1.png

Previous works on domain reweighting implicitly assume infinite token resources from all training domains (Xie et al., ; Fan et al., ; Liu et al., ) while it is not always applicable in real-world cases.
The scenario with limited training resources is challenging for online domain reweighting.
Indeed, if the weights are concentrated on a few domains, e.g. on a single domain $D_{i}$, a large model
will quickly overfit when the number of tokens in $D_{i}$ is small.

We expect DGA to mitigate overfitting by dynamically adjusting the domain weights.
Specifically, once a model starts overfitting on $D_{i}$, the magnitude of the gradients $\nabla\ell({\bm{\theta}},D_{i})$
decreases as its training loss $\ell({\bm{\theta}},D_{i})$ is low, i.e. the domain knowledge from $D_{i}$ is well-learned.
Consequently, the corresponding gradient alignment score ${\bm{a}}_{i}=\langle\nabla\ell({\bm{\theta}},D_{i}),\nabla\ell({\bm{\theta}},D
_{\mathrm{spe}})\rangle$ decreases as well and DGA explores other domains with higher alignment scores. In other words, DGA down-weights domains once they are well-learned, thereby achieving a balance between exploration – by learning from diverse data domains – and exploitation, by intensively training on the most relevant domains.

However, with limited data per domain, we remark that
DGA without EMA demonstrates drastic changes at each domain weight update, focusing heavily on one domain at a time.
Quickly changing domain weights is problematic since we want to use the same domain
weights for $T_{r}$ steps in the future. This motivates the introduction of
the EMA update in Algorithm 1, which regularizes the model and domain weights with the previous state when it starts to overfit.

Experiment Setup.
We consider the generic domain split into $k=64$ domains. We construct three scales of generic sets,
either taking the full dataset or randomly sub-sampling $30$M, $0.1$B tokens per domain.
For the targeted specific set $D_{\mathrm{spe}}$, we use $5$ subsets from the Pile (Gao et al., ) covering common specialized data types for LM applications: Math (dm_mathematics), Code (github, stackexchange), Medical (pubmed_central), Legal (free_law) and Scientific articles (arxiv).

We implement the importance sampling baseline described in subsection 2.3. We also compare to the uniform baseline with the domain weights ${\bm{\alpha}}_{\mathrm{uniform}}$ as the natural proportion of each data domain in the generic Redpajama-v2 dataset. For importance sampling and uniform baselines, the domain weights are fixed throughout the entire training run. For both vanilla DGA and DGA with an EMA ($\beta=0.1$), we update domain weights ${\bm{\alpha}}$ every $T_{r}=100$ steps. We use $125$M models.

Results. We report the validation loss on the specialized set under various token constraints in Figure 1 for free_law and the results on other domains in Appendix A.
With $30$M tokens per domain, DGA with EMA effectively stabilizes the training, while vanilla DGA exhibits several loss spikes, suggesting a lack of robustness. Under a $0.1$B token constraint, both DGA and DGA with EMA are able to dynamically adjust domain weights to mitigate overfitting. In contrast, fixed domain weights from importance sampling consistently lead to overfitting in token-limited scenarios, demonstrating the limitations of static weighting strategies in comparison to dynamic approaches like DGA. It is worth noting that adding the EMA has no negative effect on the learning efficacy when there is no token limit, which can be used as a robust regularization in the online domain reweighting context.

Domain Weights Evolution.
In the experiments with a limited generic token budget (subsection 3.1), DGA without EMA often assigns excessive weight to one generic domain, leading to overfitting due to the restricted number of training tokens.
This iterative over-weighting pattern on generic domain weights aligns with the observed loss spikes on the specific set (3(a)). In contrast, the EMA helps to regularize the weight dynamics, effectively preventing the model from overfitting by maintaining more balanced domain weights throughout the training process.

### 3.2 Distribution Reweighting: Scaling-up Data Mixing on Extremely Fine-grained Data Domains

The computational overhead from DGA scales linearly with the number of domain $k$. This is intractable for
datasets segmented in many fine-grained domains and, consequently, prior domain reweighting methods (Xie et al., ; Fan et al., ; Liu et al., ; Kang et al., ) have not been applied in that setting.
The fine-grained setting motivates distribution reweighting as an alternative to direct domain reweighting.

Distribution reweighting leverages the strength from both embedding-based (importance sampling) and gradient-based (DGA) strategies. We consider a generic training set partitioned into $k$ domains with a large $k$ (e.g. $4096,262k$). We also have a set of $N$ auxiliary datasets $\{S_{1},\ldots,S_{N}\}$, called *basis sets*, each from a specific domain of interest. We compute the importance sampling histograms for each basis set as $P=\{{\bm{p}}_{1},\ldots,{\bm{p}}_{N}\}$, ${\bm{p}}_{i}\in\Delta^{k},P\in R^{k\times N}$.
We then use DGA to search over a reparameterized space leveraging this basis. We define the
domain weights ${\bm{\alpha}}_{\mathrm{domain}}\in\Delta^{k}$ as a convex combination of $N$ $k-$dimensional distributions derived from importance sampling,

$$ ${\bm{\alpha}}_{\mathrm{domain}}\approx P{\bm{\alpha}}_{\mathrm{dist}}=\alpha_{ \mathrm{dist},1}\cdot{\bm{p}}_{1}+\alpha_{\mathrm{dist},2}\cdot{\bm{p}}_{2}+ \ldots+\alpha_{\mathrm{dist},N}\cdot{\bm{p}}_{N}$ (9) $$

where the low-dimensional weights ${\bm{\alpha}}_{\mathrm{dist}}\in\Delta^{N}$ are learned by DGA.
This allows the use of fine-grained domain features while eliminating intensive gradient computation on each generic domain.
Importantly, this is equivalent to applying DGA with the $N$ generic domains $\tilde{D}_{1},\dots,\tilde{D}_{N}$ where $\tilde{D}_{i}=\mathrm{mix}({\bm{p}}_{i})$.
Hence, it does not require any modification to the base DGA algorithm; it suffices to be able to sample according to each $\mathrm{mix}({\bm{p}}_{i})$. We provide the pseudo-code for the distribution reweighting with DGA in Appendix D.

Experiment Setup.
We demonstrate the efficacy of distribution reweighting on the MMLU benchmark (Hendrycks et al., ). MMLU consists of 57 tasks from various knowledge fields, which serves as a testbed of multi-domain language modeling;
by measuring the downstream accuracy, we can assess whether the improvements obtained in language modeling transfer to reasoning abilities.

We construct two specific datasets with different amounts of accessible samples: (1) MMLU$_$a: we take half of the examples from each task used as $D_{\mathrm{spe}}$. We denote the other half of datapoints as MMLU$_$b, which is used for evaluation; (2) MMLU$_$dev: we randomly select 5 samples from each task, simulating the few-shot learning scenario.
MMLU$_$a has $7.1k$ samples while MMLU$_$dev only has $285$ samples, which yields sparse importance sampling histograms.
For evaluation, we assess the language modeling performance by computing perplexity on MMLU$_$b. We also measure the accuracy for multiple choice question answering on MMLU$_$b with llm-eval (Gao et al., ).

We use generic domain splits with $k=64,4096,262k$ domains.
We rely on $22$ auxiliary subdomains from The Pile (Gao et al., ) as our basis sets. For each auxiliary set,
we take $15M$ tokens and compute their importance-sampling histograms as ${\bm{p}}_{1},\dots,{\bm{p}}_{N}\in\Delta^{k}$.
To search for the optimal balance between diversity and specificity, we extend the basis sets with the importance sampling histogram from the specific set itself (i.e. MMLU$_$a or MMLU$_$dev), yielding $N=23$ distributions.
For this experiment, we use $750$M models.

DGA with distribution reweighting greatly improves language modeling. We report the loss on MMLU$_$b and average validation loss across 22 domains in the Pile in Figure 2.
Both importance sampling and DGA significantly outperform the uniform baseline.
Compared to importance sampling, we observe that DGA with distribution reweighting leads to a better Pareto-front, indicating a better balance between specialized (MMLU) and general knowledge (The Pile).
With a large domain granularity ($k=262k$ domains), training with importance sampling greatly suffers from the sparse histograms, leading to significant performance degradation. In contrast, DGA can consistently provide satisfying domain weight estimation.

Figure: (a) MMLU$_$a (half the examples)
Refer to caption: x5.png

##### Language modeling ≠ \neq ≠ reasoning accuracy.

According to Table 1, both importance sampling and DGA reweighting greatly outperform the uniform baseline, while DGA does not show significant improvement above importance sampling despite the great improvement in language modeling. It indicates that better language modeling performance may not be necessarily transferable to better reasoning abilities. We report the full results with different model scales in Appendix B.

Table: Table 1: MMLU accuracies with domain reweighting methods. Both importance sampling and DGA reweighting greatly improve the accuracy above uniform baseline, while DGA does not show significant improvement above importance sampling.

Figure: (a) DGA w. target at free_law ($30$M tokens per domain)
Refer to caption: x8.png

Weights Evolution on Distributions.
We present the evolution of domain weights for each basis distribution from DGA in Figure 3. Comparing different levels of granularity, with $k$=$262k$, the importance of the MMLU distribution is more emphasized than with $k$=$4096$, with the help of fine-grained domain features.
Additionally, with sufficient samples from the specific domain (MMLU_a, 3(b)), the MMLU distribution is consistently up-weighted across $262k$ generic domains.
In contrast, on MMLU_dev, while the distribution on MMLU is initially up-weighted, it declines gradually in the late stage of training.
Owing to the number of accessible samples from the specific set, the importance sampling distribution on MMLU_dev across $262k$ generic domains is very sparse.
During the training, the learnability of the few activated generic domains diminishes, making other distributions more beneficial to the model.

In addition to the importance sampling distribution from the specific sets (MMLU_a and MMLU_dev), DGA effectively identifies other relevant distributions from The Pile that contribute to the learning on MMLU. These influential distributions, which include phil_papers, free_law, and dm_mathematics, are all considered to contain high-quality, academic-related contents. We present detailed curves with domain labels in subsection B.2.
This ability to adaptively select beneficial distributions enhances the model’s generalization and helps mitigate overfitting by leveraging a broader yet pertinent set of data sources during pretraining.

Impact of generic domain granularity.
In Figure 4, we present the validation loss on the specific domain according to the number of clusters within the generic dataset. From $k=64$ to $4096$, both DGA and importance sampling demonstrate significant improvement in language modeling in terms of validation loss (i.e., log of perplexity). However, when the number of clusters exceeds the scale of the accessible specific set, the importance sampling method overfits the limited number of activated generic domains, failing to capture broader domain knowledge. In contrast, DGA effectively leverages extremely fine-grained domain information across $262k$ generic domains with only $7k$ samples from MMLU_a.
In the few-shot context (MMLU_dev), DGA mitigates a large performance degradation by utilizing diverse domain knowledge from other relevant distributions.

Figure: (a) $D_{\mathrm{spe}}=\texttt{MMLU_a}$
Refer to caption: x11.png

## 4 Related Work

##### Task-adaptive Data Selection for Domain-Specific LLMs.

Many works have shown that one can effectively improve the LLM’s performance on a specific downstream task with data selection according to the relevance of generic data for the targeted data domain.
Gururangan et al. () show that continued pretraining on data with high vocabulary overlap can boost its performance on the specific end-task.
On machine translation task, Aharoni and Goldberg () identify task-relevant pretraining datasets from a generic corpus using nearest neighbor of a small specialist dataset based on SentenceBert sentence representation.
Wang et al. (); Grangier et al. () train a small proxy model to give an importance weight per sample.
Xie et al. () proposed DSIR as a lexical-based importance sampling method using n-gram features.

Other than feature-based importance sampling (Grangier et al., ), influence function-based method select data points which leads to the greatest loss drop on the target from the optimization perspective (Koh and Liang, ; Kwon et al., ; Agarwal et al., ). However, these methods often introduce intensive computational overheads from the second-order gradient computations, which is not applicable on large generic pretraining corpus.

##### Data Resampling through Domain Reweighting.

Given the large scale of the generic pretraining corpus, sample-level selection strategies are hard to implement for LLM pretraining. Alternatively, domain reweighting methods (Xie et al., ; Fan et al., ; Liu et al., ; Kang et al., ) apply group-level selection by adjusting data sampling weights across different domains to reflect their importance in pretraining.
Based on the weak-to-strong generalization strategy (Burns et al., ), existing domain re-weighting methods typically estimate the optimal domain weights for a larger base model based on the preferences of a small-scale proxy model. Xie et al. () apply group distributed robust optimization to optimize the worst-case loss gap between two small-scale proxies. Fan et al. () use gradient alignment to dynamically adjust domain weights during proxy model training. Specifically, it identifies the most beneficial domains by aligning the gradients of the training data with the target task. However, it trains the proxy model on reweighted domain gradients to simulate the resampling scenario, which introduces more variance in the domain weights estimation.

## 5 Conclusion

To tackle two key challenges of online domain reweighting, we introduce Dynamic Gradient Alignment (DGA) as a stable and scalable data mixing method for LLM pretraining.
Given a target task, DGA is an online algorithm that adjusts the training data distribution according to the current model status.
This adaptation relies on an estimate of the progress on the target task from gradient alignments.
We show that under limited tokens within generic domains, DGA with EMA can notably mitigate overfitting and yields superior performance on the end-task by balancing exploitation and exploration.
We also propose a novel distribution reweighting strategy, which enables DGA to scale up to extremely fine-grained data domains without incurring intensive computations. Our experiments on MMLU show that applying distribution reweighting with DGA effectively leverages fine-grained domain knowledge to balance specialty and diversity during training.
Our work demonstrates the scalability of gradient-alignment-based data reweighting methods, as well as their efficiency in data-constrained settings.

#### Acknowledgments

We thank Angelos Katharopoulos, Skyler Seto, and Matteo Pagliardini for their help and fruitful discussions during the project.

## Appendix A Training with Limited Generic Tokens

### A.1 Validation Loss on the Targeted End-task

We present the complete results on all six target domains (arxiv, free_law, dm_mathematics, pubmed_central, github, stackexchange) as follows.
Across all six target domains, DGA with EMA ($\beta=0.1$) consistently stablize the training and yields better language modelling performance under token-limited contexts.

Figure: (a) 30M tokens per domain
Refer to caption: x13.png

Figure: (a) 30M tokens per domain
Refer to caption: x17.png

Figure: (a) 30M tokens per domain
Refer to caption: x21.png

Figure: (a) 30M tokens per domain
Refer to caption: x25.png

Figure: (a) 30M tokens per domain
Refer to caption: x29.png

Figure: (a) 30M tokens per domain
Refer to caption: x33.png

### A.2 Domain Weights Evolution

We present the domain weights evolution on 64 generic domains from DGA with and w.o. EMA regularization. With both stackexchange and free_law as the specific set, EMA effectively smoothes the spiky domain weights, which therefore stablize the training process.

Figure: (a)
Refer to caption: x37.png

## Appendix B Distribution Reweighting

### B.1 Evaluation Results on MMLU

We present the complete evaluation results on MMLU benchmark on small- ($125$M) and large- ($750$M) scale models. $k$ denotes the number of generic domains, $N$ denotes the number of reweighted importance sampling distributions from basis sets. $N$=$22$ indicates we only reweight 22 distributions from 22 The Pile subsets, while $N$=$23$ includes the importance sampling histgram from the specific set (MMLU).
Since the $125$M model shows marginal difference in accuracy because of limited capacity, we only scored $750$M model on MMLU reasoning accuracies.

**Table 2: Results on the domain reweighting experiment, with half MMLU as train set. The best results is Bolded and the second best is Underlined.**
| 125M model | MMLU loss | MMLU acc. | avg. Pile loss |
| --- | --- | --- | --- |
| Uniform | 3.56 | - | 3.27 |
| Importance S. (k=4096) | 3.32 | - | 3.16 |
| Importance S. (k=262k) | 3.22 | - | 3.19 |
| DGA domain reweighting (k=64) | 3.31 | - | 3.10 |
| DGA dist. reweighting (N=22, k=4096) | 3.34 | - | 3.13 |
| DGA dist. reweighting (N=22, k=262k) | 3.34 | - | 3.05 |
| DGA dist. reweighting (N=23, k=4096) | 3.33 | - | 3.13 |
| DGA dist. reweighting (N=23, k=262k) | 3.25 | - | 3.10 |
| 750M model | MMLU loss | MMLU acc. | avg. Pile loss |
| Uniform | 3.03 | 26.1 $\%$ | 2.82 |
| Importance S. k=4096 | 2.82 | 27.7 $\%$ | 2.75 |
| Importance S. k=262k | 2.82 | 28.4 $\%$ | 2.81 |
| DGA domain reweighting (k=64) | 2.97 | 27.1 $\%$ | 2.77 |
| DGA dist. reweighting (N=22, k=4096) | 2.86 | 27.2 $\%$ | 2.73 |
| DGA dist. reweighting (N=22, k=262k) | 2.84 | 27.0 $\%$ | 2.66 |
| DGA dist. reweighting (N=23, k=4096) | 2.83 | 26.8 $\%$ | 2.73 |
| DGA dist. reweighting (N=23, k=262k) | 2.74 | 28.0 $\%$ | 2.76 |

**Table 3: Results on the domain reweighting experiment, with 5 examples per task of MMLU as train set. We score only the 750M models.**
| 125M model | MMLU loss | MMLU acc. | avg. Pile loss |
| --- | --- | --- | --- |
| Uniform | 3.56 | - | 3.27 |
| Importance S. (k=4096) | 3.40 | - | 3.16 |
| Importance S. (k=262k) | 3.41 | - | 3.29 |
| DGA domain reweighting (k=64) | 3.46 | - | 3.19 |
| DGA dist. reweighting (N=22, k=4096) | 3.37 | - | 3.12 |
| DGA dist. reweighting (N=22, k=262k) | 3.36 | - | 3.03 |
| DGA dist. reweighting (N=23, k=4096) | 3.37 | - | 3.14 |
| DGA dist. reweighting (N=23, k=262k) | 3.35 | - | 3.08 |
| 750M model | MMLU loss | MMLU acc. | avg. Pile loss |
| Uniform | 3.03 | 26.1 $\%$ | 2.82 |
| Importance S. k=4096 | 2.96 | 27.7 $\%$ | 2.75 |
| Importance S. k=262k | 3.13 | 27.0 $\%$ | 2.99 |
| DGA k=64 | 3.01 | 27.0 $\%$ | 2.77 |
| DGA dist. reweighting from 22 domains, k=4096 | 2.89 | 26.8 $\%$ | 2.76 |
| DGA dist. reweighting from 22 domains, k=262k | 2.93 | 27.0 $\%$ | 2.68 |
| DGA dist. reweighting from 23 domains, k=4096 | 2.88 | 27.4 $\%$ | 2.75 |
| DGA dist. reweighting from 23 domains, k=262k | 2.90 | 27.0 $\%$ | 2.70 |

### B.2 Weights Evolution on Distributions

Figure: (a) DGA w. target at MMLU$_$a
Refer to caption: x39.png

## Appendix C Hyperparameters

Table 4 provides the model architectures and hyperparameters used in this paper.

Table: Table 4: Architecture hyperparameters for various model scales used in the paper. All models are vanilla Transformer decoder-only models.

## Appendix D DGA for Distribution Reweighting

Algorithm 2 explains the distribution reweighting with DGA. The implementation can be easily adapted from domain reweighting DGA with minor modifications.

Figure: Algorithm 2 Distribution Reweighting w. DGA. (Difference from domain reweighting are marked in blue)

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: (1)	Please address the points listed under Weaknesses.

(2)	There are some typos and grammar errors in the paper. For example, “obtain obtain”.
soundness: 2
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.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1, How would DGA perform under the paradigms of pre-training + fine-tuning or pre-training + ICL?
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
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?
soundness: 2
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
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
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?
soundness: 3
strengths: The submission is well-written and a good mixture of motivation, analysis of previous work, detailing theoretic foundations, experimental verification and conclusions.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
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?
soundness: 3
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.
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.
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.

Meta-review: 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.

Decision: Reject