Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance

## Abstract

Abstract Pretraining data of large language models composes multiple domains (e.g., web texts, academic papers, codes), whose mixture proportions crucially impact the competence of outcome models.
While existing endeavors rely on heuristics or qualitative strategies to tune the proportions,
we discover the quantitative predictability of model performance regarding the mixture proportions in function forms, which we refer to as the data mixing laws .
Fitting such functions on sample mixtures unveils model performance on unseen mixtures before actual runs, thus guiding the selection of an ideal data mixture. Furthermore, we propose nested use of the scaling laws of training steps, model sizes, and our data mixing law to enable predicting the performance of large models trained on massive data under various mixtures with only small-scale training.
Moreover, experimental results verify that our method effectively optimizes the training mixture of a 1B model trained for 100B tokens in RedPajama, reaching a performance comparable to the one trained for 48% more steps on the default mixture.
Extending the application of data mixing laws to continual training accurately predicts the critical mixture proportion that avoids catastrophic forgetting and outlooks the potential for dynamic data schedules. 1 1 1 Codes and data are available at: https://github.com/yegcjs/mixinglaws .

## 1 Introduction

Pretraining data for large language models (LLMs) are typically a mixture of multiple domains, varying from English to minority languages (Doddapaneni et al., ; Li et al., ), from casual dialogs to formal academic writings (Taylor et al., ), and from texts to modalities like images and speeches (Zhan et al., ), among others.
These data interplay with each other, showing complex interchangeable, unrelated, or contradictory relationships (Guo et al., ).
This necessitates adjusting the mixture proportions of training data to balance the model capabilities while harnessing synergies across domains, thus enhancing the competence of the outcome models, as highlighted by extensive practical experience (Gao et al., ; Gururangan et al., ; Zhou et al., ; Xie et al., ).

Nonetheless, it remains elusive to figure out an ideal training data mixture.
Most existing practices tune the mixture through heuristics to upsample a proportion of high-quality or underrepresented data without disclosing the concrete criteria in detail (Gao et al., ; Touvron et al., ; Bai et al., ; Bi et al., ) and it is hard to predate whether these data strategies are effective before finishing the training run. Encouraged by advances in scaling laws that show model losses on a given set of evaluation data are quantitatively predictable for a wide range of variables (Kaplan et al., ; Hoffmann et al., ), we wonder whether this also holds for mixture proportions,
so that we can estimate the outcome model performance given any mixture before actually training on them, including the desired one that reaches minimum loss.

In this paper, we answer this proposition affirmatively.
We find that, given a mixture of $M$ domains, an exponential function over the linear combination of the proportions, i.e.,

$$ $L_{i}(r_{1\dots M})=c_{i}+k_{i}\exp{\left(\sum_{j=1}^{M}{t_{ij}r_{j}}\right)},$ (1) $$

can predict the validation loss $L_{i}$ on any of the training domains $i$ accurately under a fixed model size and amount of training data, where $r_{1\dots M}$ are the proportions of the $M$ domains and $c_{i},k_{i},t_{ij}$ are parameters to fit.
Fitting such functions on all the evaluated domains and calculating the weighted sum according to their proportions in the validation data leads to the prediction of final validation loss.
Further, treating the validation proportions as learnable parameters allows fitting the estimated losses on a validation set end-to-end without explicitly decomposing it into known domains.

Despite the predictability, fitting the function between mixture proportions and validation loss, or the data mixing laws for simplicity, requires samples of numerous runs with different mixtures.
Running these experiments with the same model size and the amount of training data as the target model is unreasonably expensive.
Fortunately, fruitful research on scaling laws demonstrates impressive results that
fitting power laws with small models and small data effectively predicts the losses on larger models and data over orders of magnitudes (Kaplan et al., ; Henighan et al., ; Hoffmann et al., ; Alabdulmohsin et al., ; OpenAI, ; Muennighoff et al., ; Bi et al., ).
On this basis, we propose a pipeline to nested utilize the scaling laws of training steps, model sizes, and our data mixing law, so that we can study the impact of mixture proportions for the target model sizes and data amount with only experiments at the affordable scales, illustrated in Fig. 1.

Figure: Figure 1: Illustration on our pipeline to optimize data mixture. Left: Our pipeline takes three steps. Starting from small-scale training results, the three steps use the scaling laws of training steps, model sizes, and data mixing laws to predict model performance on large steps, large models, and unseen mixtures, respectively. Right: Visualization of the three-step pipeline to predict model performance on the target model size, training step, and mixtures.

Experimental results verify the reliability of our data mixing law and prediction pipeline.
By predicting the overall validation loss, we optimize the training mixture of RedPajama for a 1B model trained on 100B tokens and achieve performance comparable to a model trained on default mixture for 48% more steps.
The prediction on domain-specific validation sets also offers plausible references to the balance of model capabilities.
Further applying our data mixing law to continual pretraining can accurately find the proportion that avoids catastrophic forgetting (French, ; Kirkpatrick et al., ; Luo et al., ), revealing the prospect of applying data mixing laws to guide a multi-stage pertaining, and thus a dynamic data schedule.

Overall, our contributions and findings are as follows:

- •
We discover the quantitative predictability of model performance regarding data mixture, and summarize this into a functional relationship, namely the data mixing laws.
- •
We propose a pipeline to predict model performance of large-scale training on different mixture proportions but only experiments on small models with few training data through nested use of scaling laws of training steps, model sizes, and data mixing laws.
- •
We experiment to verify the reliability of our data mixing laws and prediction pipeline, showing its effectiveness in optimizing model performance, balancing model capabilities, and the prospects of guiding the design of the data schedule.

## 2 Background

We briefly review the pretraining process of large language models and summarize key findings from neural scaling laws, then we formalize the problem we study.
For more related work, please refer to Sec. 6.

#### Pretraining large language models.

We consider the task of pretraining an autoregressive language model $p_{\theta}$ via next-token predictions (Radford et al., ).
The training dataset $\mathcal{D}_{\mbox{train}}=\{\mathcal{D}_{i}\}_{i=1}^{M}$ composes $M$ domains with mixture proportions $\bm{r}\in\Delta^{M-1}$.
In each training step, the task first samples a batch of domain indices according to the mixture proportions and then sample sequences of $L$ tokens from the sampled domains.
Using the sampled data, it learns to optimize the negative log-likelihood of sampled data, i.e.,

$$ $\mathcal{L}_{\theta}=\mathbb{E}_{i\sim\bm{r},x_{0\dots L}\sim D_{i}}\left[- \sum_{j=1}^{L}\log P_{\theta}(x_{j}|x_{0\dots j-1})\right].$ (2) $$

To evaluate the learned model, we compute the loss on validation data $\mathcal{D}_{\mbox{val}}$.

#### Scaling laws.

For a wide range of factors $x$, scaling laws (Kaplan et al., ; Henighan et al., ; Hoffmann et al., ) show that their effect on the loss $L$ of a pretrained model follows power laws

$$ $L=c+kx^{\alpha},$ (3) $$

where $c$, $k$, and $\alpha$ are parameters to fit and $x$ can be model sizes, numbers of training data, training steps(^2^22 The training step law (Kaplan et al., ) refers to the function between the intermediate training step and its corresponding loss in one single training run. For clarity, we highlight its difference from the data scaling law that predicts the early exit loss (Kaplan et al., ) or the loss obtained with optimal hyperparameters for each data budget (Hoffmann et al., ).), and the amount of computation.
Previous experience (Alabdulmohsin et al., ; OpenAI, ; Bi et al., ; Su et al., ) highlights the impressive predictability of scaling laws.
Specifically, Eqn. 3 fitted on a collection of small models, training data, or computation can extrapolate to precisely predict the test loss of larger cases over orders of magnitudes.
This enables practitioners to estimate the performance of a pretrained large language model without actually finishing the expensive runs.
Recent development further shows various functional relationships between the performance of language models and a broader range of factors, including transfer learning (Hernandez et al., ), sparse architectures (Frantar et al., ), and repeated data (Muennighoff et al., ), consolidating the predictability of language model performance.

#### Problem formalization.

We study optimizing the mixture proportions of pretraining data for large language models.
Motivated by the impressive predictability of existing scaling laws, we try to tackle mixture optimization by establishing a quantitative framework that predicts the loss given any mixture proportion.
Formally, for a training dataset comprising $M$ domains, we parameterize the function

$$ $L=f_{\theta}(\bm{r}),$ (4) $$

under the fixed model sizes and number of training steps, where $\bm{r}=r_{1\dots M}$ is the proportion of the $M$ domains.
Harnessing this function, we seek a mixture that achieves the desired performance.
Without loss of generality, we search for the mixture that reaches minimum validation loss, i.e.,

$$ $\bm{r}^{*}={\arg\min}_{\bm{r}}f_{\theta}(\bm{r}).$ (5) $$

## 3 The proportions of data mixtures influence model losses in a quantitatively predictable way

In this section, we present our findings on the predictability of model losses regarding data mixtures, which boils down to functional relationships we refer to as the data mixing laws.

To discover the data mixing laws, we encounter two challenges posed by their characteristics.

- (i)
Multi-variables.
For a data mixing law for $K$ domains, we should consider $K-1$ degrees of freedom in the mixture proportions and, correspondingly, $K-1$ variables in the target function.
The increase of variables considerably enlarges the scope of potential functions thereby complicating the identification of the function form.
- (ii)
Nonmonotonicity. A monotonic relationship between losses and the proportion of any domain indicates that a lopsided mixture can achieve minimum loss without endeavors to balance domain proportions, which contradicts the practice.
Therefore, differing from existing scaling laws that loss monotonically decreases with the scale of concerning factors, the data mixing law we study should accommodate non-monotonic functions.
This nonmonotonic nature adds another layer of complexity to our analysis.

To navigate these challenges, we initially simplify the problem by studying a scenario where the relationship between loss and mixture proportions fits into a univariate monotonic function then retracts the simplifications progressively.
In specific, we begin our study on the case where we only train on two domains thus avoiding multi-variables, and only consider the validation data coming from one of the training domains to circumvent the non-monotonicity (Sec. 3.1).
Subsequently, we broaden our framework to encompass training on multiple domains (Sec. 3.2) and explore the predictability of losses on general validation data that also comprises various domains (Sec. 3.3).

### 3.1 Pilot Study on Domain Losses under Two-domain Mixtures

We begin our exploration with the simplest case where we only learn on mixtures of two domains and evaluate our model on the two domains respectively.

Figure: Figure 2: Quantitative predictability of domain losses on two domains, which are Github and Pile-CC. We train on the mixtures of these two domains and validate the outcome models on them separately. We train 70M and 160M models on five different mixtures of Github and Pile-CC and obtain the reducible losses by subtracting the original losses with a constant shared across models of the same sizes and trained for the same number of steps. The reducible losses in log scale show linear correlations to the domain proportions.
Refer to caption: extracted/2403.16952v1/figs/2mix.png

#### Setups

We train 70M and 160M language models on the mixture of Github and Pile-CC subset from the Pile dataset (Gao et al., ) with five different mixture proportions, which are {0.25, 0.375, 0.5, 0.625, 0.75} for Github.
We train all models with a batch size of 1M tokens for 30k steps, which is 30B tokens in total, and evaluate checkpoints at different steps on the validation set of GitHub and Pile-CC.

#### Findings

Results in Fig. 2 reveal the quantitative predictability of domain losses given the domain proportions.
We encouragingly find that, for checkpoints with the same size and trained with the same number of steps, after subtracting a shared constant(^3^33The loss function of language models is lower bounded by the entropy of the evaluation data theoretically (Bishop, ; Henighan et al., ). Thus, it is common to have constant terms in scaling laws, known as irreducible loss.), their domain losses in the log scale demonstrate a linear relationship to the domain proportion.
This holds for both domains in our experiments.
The result indicates that with other factors fixed, the domain losses of a pretrained language model regarding the domain proportion precisely fit into an exponential law

$$ $L_{i}(r_{i})=c_{i}+k_{i}\exp{\left(t_{ii}r_{i}\right)},$ (6) $$

where $L_{i}$ and $r_{i}$ are validation loss and training mixture proportion of domain $i$, respectively, while $c_{i}$, $k_{i}$, and $t_{ii}$ are learnable parameters (^4^44 Despite a simple case, our findings on two domains have practical applications to continue pretraining (Gururangan et al., ), where we aim to enhance a pretrained model on a given domain by training it on a mixture of the original pretraining data and upcoming domain data. Please see Sec. 5 for details.).

### 3.2 Extension to Domain Losses Trained on Multi-domain Mixtures

To accommodate real-world pretraining data that mostly contains more than two domains, we extend our investigation into multiple domains.
For simplicity and the ease of visual aids, we start with the case of three domains.

**Table 1: Mean absolute errors (MAE) of different candidate functions for predicting the target domain losses. We also include random guesses that randomly predict between the maximum and minimum loss of the training samples for reference. In specific, we report the MAE of the expectation of this random guess which predicts the median of the maximum and minimum loss. The training data contain $M=3$ domains and we fit each function with the same 24 mixtures and validate on 8 other mixtures. The split is random. The lowest error for each target domain are in bold while the second lowest are underlined.**
|  |  | GitHub | Books3 | Pile-CC |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Method | # Coeff. | Train | Validation | Train | Validation | Train | Validation |
| Random | - | 0.8895 | 0.8758 | 0.1291 | 0.1331 | 0.0768 | 0.1045 |
| M1 | 2M+1 | 0.0292 | 0.0312 | 0.0082 | 0.0121 | 0.0045 | 0.0050 |
| M2 | M+2 | 0.1558 | 0.3327 | 0.0114 | 0.0119 | 0.0072 | 0.0083 |
| M3 | M+2 | 0.3389 | 0.2177 | 0.0914 | 0.0465 | 0.0746 | 0.0947 |
| M4 | M+2 | 0.0298 | 0.0365 | 0.0062 | 0.0074 | 0.0036 | 0.0078 |

#### Setups

We train on the mixtures of GitHub, Pile-CC, and Books3 subset from the Pile for a total of 30B tokens and evaluate the model on the three domains, respectively.
For specific mixtures, we grid search from $\{0,0.125,0.25,\dots,0.875,1\}^{3}$ and retain valid ones in which three proportions sum up to 1 and do not use up all tokens in any of the domains(^5^55The GitHub and Books3 subset in the Pile do not contain as many as 30B tokens after deduplication.), which results in 32 mixtures in total.

We utilize the losses on these experimented mixtures to identify the function forms between losses and mixture proportions through conjecture and then verification.
In specific, we base our conjecture of possible forms on the following two principles.

- •
Compatibility. The form can reduce to Eqn. 6 if the number of domains $M=2$.
- •
Symmetry. Any exchanging of variables should not change the functional form.

The second principle stems from the intuition to avoid introducing any domain-specific bias.
Together, the two principles lead to candidate functions that replicate the exponential term in Eqn. 6 for each training domain and combine them through operations that adhere to commutative law.

According to the two principles, we experiment with the following candidate functions:

$$ M1: $\displaystyle L_{i}(\bm{r})=$ $\displaystyle c_{i}+\sum_{j=1}^{M}\left[k_{ij}\exp\left(t_{ij}r_{j}\right) \right],$ M2: $\displaystyle L_{i}(\bm{r})=$ $\displaystyle c_{i}+k_{i}\sum_{j=1}^{M}\exp\left(t_{ij}r_{j}\right),$ M3: $\displaystyle L_{i}(\bm{r})=$ $\displaystyle c_{i}+k_{i}\exp\left(\prod_{j=1}^{M}t_{ij}r_{j}\right),$ M4: $\displaystyle L_{i}(\bm{r})=$ $\displaystyle c_{i}+k_{i}\exp\left(\sum_{j=1}^{M}t_{ij}r_{j}\right).$ $$

We summarize their fitting errors on three target domains in Tab. 1.

Figure: Figure 3: Quantitative predictability of domain losses on three domain mixtures, Github, Books3, and Pile-CC. We train on the mixture of these three domains and validate the outcome models on them as well. The surfaces show the predicted losses on (A) Github; (B) Books3; (C) Pile-CC; and (D) the overall validation set mixed with the three domains. $\times$: validation samples. ${\star}$: the predicted minimum loss on the overall validation set.

Figure: Figure 4: Prediction results on the domain losses and overall losses in the three-domain experiment. The domain losses are fitted with Eqn. 7 and we obtain the total losses through explicit domain aggregation of Eqn. 8.

#### Findings

The results in Tab. 1 suggests both M1 and M4 gives reliable predictions while M4 has fewer coefficients.
Therefore we adopt M4

$$ $L_{i}(r_{1\dots M})=c_{i}+k_{i}\exp{\left(\sum_{j=1}^{M}{t_{ij}r_{j}}\right)}$ (7) $$

as the function forms of our data mixing law, where $L_{i}$ is the validation loss on $i$-th validation domain, $r_{j}$ is the proportion of the $j$-th training domain, and $c_{i},k_{i},t_{ij}$ are learnable parameters.
The fitting results are in Fig. 3 and Fig. 4 demonstrates the prediction accuracy.
The results indicate that Eqn. 7 fits the given samples well and estimates the unseen ones accurately.

To provide more intuition, we discuss the meanings of the coefficients in Eqn. 7. In general, $k_{i}>0$, thus the exponential term is always positive and the prediction loss is strictly greater than the constant $c$.
Hereby, $c_{i}$ represents losses that are not reducible by adjusting the data mixture.
$t_{ij}$, depending on both training domain $i$ and validation domain $j$, shows the interaction between them.
A negative $t_{ij}$ indicates that training data of domain $j$ helps reduce validation loss on domain $i$ and vice versa.
Besides, rewriting the coefficient $k_{i}$ into $\exp(\log k_{i})$ makes it a bias term on the exponent which adjusts the aggregate effect of all training domains.

### 3.3 Predicting Losses of Any Validation Mixture

We further loosen constraints in Sec. 3.1 and Sec. 3.2 that the validation data are from one of the training domains.
We first consider the validation set to be a known composition of the training domains and then free this requirement for more general cases of arbitrary validation sets.
These correspond to the two strategies we fit the data mixing laws, which we elaborate on as follows.

#### Explicit domain aggregation.

Considering a validation set made up of $K$ domains with the proportions as $s_{1\dots K}$, the validation loss can be written into the weighted sum of domain losses.
Thanks to the discovery of Eqn. 7, we can apply the equation to predict domain losses herein given a training mixture.
Therefore, the functional relationship of the overall validation loss to the training mixture proportions expands into

$$ $L(r_{1\dots M})=\sum_{i=1}^{K}s_{i}L_{i}(r_{1\dots M})=\sum_{i=1}^{K}s_{i} \left[c_{i}+k_{i}\exp{\left(\sum_{j=1}^{M}t_{ij}r_{j}\right)}\right].$ (8) $$

Using Eqn. 8, we can fit the loss on each validation domain $L_{i}$ and sum them up to obtain the prediction of overall loss.

#### Implicit domain aggregation.

A remaining limitation is that we still need to acquire the components of validation data $s_{1\dots K}$ in advance.
This can be inconvenient if the validation set is collected separately from the training ones.
For instance, the validation data may come from real-world user queries that cover unknown compositions of various domains.
To remove the constraint on validation components, we assume that we can decompose the validation data into $K$ implicit domains whose losses are predictable with Eqn. 7, and we treat their proportions in the validation data $s_{1\dots K}$ as learnable parameters.
This leads to the final form of our data mixing laws(^6^66We note that the final forms of our data mixing law resemble a multilayer perception (see the computation graph Fig. 10). We include further discussion and implementation details in Appendix C.).
With this perspective, we fit a data mixing law with the overall losses end to end.

Figure: Figure 5: Prediction errors of the five-domain data mixing laws fitted with explicit and implicit domain aggregation. Explicit domain aggregation: we fit Eqn. 7 for five domains respectively and sum them up according to their weight in the overall validation sets. Implicit domain aggregation: we fit the losses on overall validation with Eqn. 8, assuming different numbers of implicit domains $K$ and treating the proportion of different implicit domains as learnable parameters.
Refer to caption: extracted/2403.16952v1/figs/5mix_end2end.png

Fig. 5 shows our experiments where we train language models on the 5 coarse-grained domains of Pile(^7^77The Pile contains 22 fine-grained domains which are collected into five coarse-grained domains, i.e., academic, internet, prose, dialogues, and misc, where misc include Github and the DeepMind Mathematics Dataset which are symbolic content. We do not experiment on fine-grained domains for their limited number of tokens available.) and evaluate a validation set mixed with these 5 domains.
We compare the errors obtained by implicit domain aggregation with different numbers of implicit domains to those obtained by explicit domain aggregation.
We find that applying implicit domain aggregation and setting the number of implicit domains no smaller than the actual one (5 in the experimented case) results in lower errors than explicit domain aggregation.
Moreover, the error remains low as we set the number of implicit domains much larger.
This verifies the prediction accuracy of our implicit domain aggregation strategy for data mixing law and the number of implicit domains $K$ can be a large number without careful tuning(^8^88We set $K=30$ if not stated in the rest of this paper.).

## 4 Nested scaling laws predict losses trained on various mixtures using only small-scale experiments

Figure: Algorithm 1 A pipeline to predict losses of different mixture proportions on large models trained on massive data through small-scale training

### 4.1 A Pipeline for Loss Predictions

While data mixing laws enable us to predict the performance of models trained on unseen mixtures, the requirement to fit the laws involves training multiple models across diverse mixtures with model sizes and token counts identical to the target ones.
Furthermore, we must repeat this process for each target model size and training dataset(^9^99An idea is to transfer the optimized training mixture on small models trained with few tokens to the training of large models and large volumes of training data. Nevertheless, in Appendix A, we show that the partial order of data mixture concerning model performance varies as the model size and number of trained tokens change. Therefore, the optimal mixture at experimented scales can be suboptimal at the target scale.).
This results in expensive costs thus hindering the practical value of our data mixing laws.

We thus wonder whether we can obtain the losses of different mixture proportions without training at large scales.
Fortunately, this idea gains endorsement from existing experiences that verify the impressive extrapolation of scaling laws of training steps and model sizes.
In particular, OpenAI () predicts the loss of the target model with merely $1,000\times$–$10,000\times$ less compute.
As a consequence, we can train small models with few training steps on different mixtures, and fitting scaling laws on them to estimate the losses of the target model size and the target number of training steps.
We can then use the predicted losses to fit a data mixing law and search for the optimal mixture.
The proposed pipeline is illustrated in Fig. 1 with details depicted in Alg. 1.

We note that the nested use of scaling laws is still not a widely adopted strategy in existing literature to the best of our knowledge.
Most studies try connecting newly found scaling laws to existing ones by incorporating variables into the latter (Hashimoto, ; Frantar et al., ; Aghajanyan et al., ).
Although unifying of different scaling laws is attractive, the added complexity to the interaction among multiple variables makes identifying the function form extremely challenging and hard to verify.
Pioneering efforts that derive scaling laws with multiple concerning factors (model sizes and the numbers of training data) (Kaplan et al., ; Hoffmann et al., ) make assumptions upon the theory of error decomposition (Harville, ; Domingos, ) but their outcomes are still not widely verified in large-scale practice.
From this perspective, we consider that our nested use of scaling laws offers a new technical routine of using scaling laws, which circumvents the huge challenge of complicating existing laws, making the discovery and application of scaling laws more accessible.

### 4.2 Experiment

Figure: Figure 6: The validation perplexity on the Pile validation set for 1B models trained on the default mixture and the optimized mixture of RedPajama for 100B tokens. Our optimized mixture achieves the performance of the default mixture only using 0.73 of the original number of training steps and eventually achieves a performance comparable to a default mixture trained with 1.48 times more tokens (estimated by the scaling law of training steps, shown as the dashed line). The specific mixture proportions are in the right table.
Refer to caption: extracted/2403.16952v1/figs/optimize_loss_2.png

We verify the effect of our pipeline with an experiment to minimize the validation loss of a 1B model trained on 100B tokens.

#### Setups.

We train our models on the mixture of RedPajama and validate the validation set of the Pile to mimic the scenario where validation data are collected separately from the training data.
To fit the scaling laws of training steps and model sizes, we train a series of 70M, 160M, 305M, and 410M models for 30B tokens.
For all the models, we set the batch size as 1M tokens thus translating into 100k steps for the 1B models and 30k steps for small models.
We apply a cosine learning rate decay with a warmup of 2k steps which decays to 0.1 of the maximum learning rate at the 100k-th steps.

To reach low prediction errors with a limited number of experiment runs, we select the mixtures for experimentation by leveraging the fact that mixture proportion terms are represented as exponential functions within our data mixing law.
Specifically, we enumerate candidate mixtures by double-diminishing the proportion for each training domain, starting from the maximum available one that does not use up all the domain tokens.
In this way, the losses of each (implicit) validation domain are distributed evenly.
We then sample 40 mixtures from all the candidates and train the smallest 70M models.
We resample groups of 20 mixtures from them to fit the data mixing law and select the group that reaches minimum prediction errors on all 40 samples as our final set of mixtures to run our pipeline.
For more details, please refer to Appendix B.2.

#### Results.

Figure: Figure 7: Results of our loss prediction pipelines for the overall validation loss and domain losses. Fitting data are from 70M to 410M models trained for 30B tokens, while the extrapolated points are from the default and optimized mixture for 1B models and 100B tokens.
Refer to caption: extracted/2403.16952v1/figs/predictions.png

Fig. 6 shows the default mixture of RedPajama used in Touvron et al. () and the optimized mixture obtained from Alg. 1 with their performance on the validation data.
The loss predictions are in Fig. 7.
We find that:

- •
The pipeline optimizes performance effectively. The model trained on the optimized mixture can achieve a performance comparable to the one trained on the default mixture with only 73% steps. It eventually attains a performance that requires 48% more steps if trained using the default mixture. This indicates the effectiveness of our pipeline in mixture optimization.
- •
Predictions provide references for the balance of model capabilities. Prediction results in Fig. 7 demonstrate plausible reference on the relative scale of losses on both the overall validation data and different validation domains. While the overall loss helps optimize the overall performance, losses on different domains show model capabilities in various aspects. We suppose that this facilitates further mixture curating for the nuanced balance of model capabilities.

## 5 Application to Continual Pretraining Outlooks Data Schedule Driven by Data Mixing Laws

We are also interested in applying our data mixing laws to continual pretraining, which shares the same paradigm as pertaining but begins the model with pretrained parameters instead of random initialization.
Generally, continual pretraining is a common technique to enhance existing pretrained models.
It injects up-to-date knowledge into the model, avoiding performance degradation due to distribution shifts (Gururangan et al., ; Xiong et al., ).
In addition, researchers also apply continual pretraining to reuse existing model parameters to build models of a different architecture (Komatsuzaki et al., ).

Figure: Figure 8: Loss predictions and the training curve of continual pretraining Pythia-70M on a mixture of the Pile and python code. (A) Loss prediction on the Pile; (B) Loss prediction on python; (C) training curves with losses on the Pile; (D) training curves with losses on python. We predict final losses with Eqn. 6. The law accurately finds the critical mixture proportion that maintains model performance on the original domain (i.e., the Pile).
Refer to caption: extracted/2403.16952v1/figs/continue_all.png

We experiment on a typical scenario of continual pretraining, where we train the model on the mixture of original pretraining data and upcoming data of a target domain to enhance.
For instance, we continually pretrain Pythia-70M models with a mixture of the Pile and Python codes, where the former is the original pretraining data of the base model.
To verify whether our data mixing laws apply to continual pretraining, we train the models for 10B tokens on 4 mixtures and fit the Eqn. 6 on losses of the Pile and python codes.
Results in Fig. 8 confirm that Eqn. 6 fits into the losses of continual pretraining.

During continual pretraining, a too-large proportion of the target data can hurt the performance of the original data.
A representative mixture optimization target is to maintain the general-purpose ability (looses on the Pile) unchanged.
To this end, using the fitted data mixing laws, we predict the critical proportion leading to the same loss as before continual pretraining.
Fig. 8 demonstrates the success of our prediction where the proportion we find results in similar performance compared to the model before continual pretraining while gaining improvement in the target domain.

Remarks. We suggest continual pretraining is significant for its connection to the design of data schedules (Albalak et al., ; Chen et al., ).
Usually, continual pretraining applies to a pretrained model, while it is natural to further continually pretrain the continual pretrained models, i.e., multi-stage pretraining (Chen et al., ).
In each stage, the mixture proportions or even the domain components of training data can be different.
This becomes a dynamic data schedule as the number of training stages approaches the infinite limit.
Therefore, the successful application of our data mixing laws on continual training signifies a promising prospect for using it to design dynamic data schedules, a more comprehensive data curating paradigm.

## 6 Related Work

#### Curating pretraining data for LLMs.

Training massive transformer architecture on trillions of tokens, a.k.a. pretraining, is the primary step to building modern large language models that exhibit impressive human-like generalist abilities (Brown et al., ; OpenAI, ; Jiang et al., ; Touvron et al., )).
It takes up most of the computation resources for model training and researchers believe it endows almost all the knowledge in LLMs (Zhou et al., ).
Such critical impact motivates the development of data curating strategies to reduce computation costs and enhance knowledge (Longpre et al., ).
The efforts can be categorized into two steps.
The first step focuses on obtaining a high-quality training dataset.
A typical procedure includes selecting data sources to constitute different domains, deduplication, and the most intricate filtering (Wenzek et al., ; Penedo et al., ).
A mass of endeavors in existing practice has involved multifarious filters, scoring the documents with from superficial features on characters (Rae et al., ; Xie et al., ; Raffel et al., ) to semantics including similarity to the high-quality reference corpus (Wenzek et al., ) and toxicity (Longpre et al., ; Friedl, ).
With a dataset on hold, the second step aims to make the best use of it.
This includes tuning the data mixture (Du et al., ; Touvron et al., ; Xie et al., ) and devising data schedules (Mindermann et al., ; Albalak et al., ; Chen et al., ; Fan et al., ).
Our work is among those tune data mixtures and our extension to continue pretraining signifies our prospect of guiding the schedule design.
Different from existing attempts that rely on intuition or qualitative targets, our study seeks a quantitative solution.

#### Scaling laws

are functional relationships between the properties of interests (e.g., test loss or other performance metrics) and the scales of controllable factors regarding the optimization process or architecture (e.g., model sizes and numbers of training samples) (Villalobos, ).
Along with the development of machine learning, characterizing scaling behaviors has garnered great research interest under the context of learning theories, bounding the generalization error given the number of training samples in the form of power laws (Vapnik & Chervonenkis, ; Valiant, ; Haussler, ; Amari et al., ).
Nevertheless, overly strict assumptions hinder their practical applications.
In recent years, statistical estimation on scaling gained fast progress for deep neural networks and spawns the introduction of scaling laws.
Hestness et al. () pioneers the trend and demonstrates power-law generalization error scaling across a breadth of factors but the power-law exponents differ from previous theoretical analysis.
Kaplan et al. (); Hoffmann et al. (); Henighan et al. () conduct more comprehensive investigations on Transformer architecture (Vaswani et al., ), further highlighting the power-law relationship on test loss regarding model sizes, the amount of training data and computation across orders of magnitudes.
These findings foretell the performance gain with scaling quantitatively and guide the trade-off between larger models and more training data, directing to the later development of large language models (Brown et al., ; Hoffmann et al., ; OpenAI, ).
Lately, progressive investigations propose amendments to existing scaling laws (Caballero et al., ; Alabdulmohsin et al., ), seeking theoretical explanations on the empirical formulas Bahri et al. (); Hutter (); Michaud et al. (), and exploring the functional relationships in broader scenarios (Hernandez et al., ; Frantar et al., ; Liu et al., ).
The most relevant study to ours is Hashimoto () which explores performance prediction under multiple data sources but is limited to small-scaled supervised learning tasks.

## 7 Discussions

In this work, we discover the quantitative predictability of model losses regarding the mixture proportions of training data, which boils down to the data mixing laws.
Our findings cover training data from two domains to multiple domains, and validation data from a single domain to unknown compositions.
Using data mixing laws allows practitioners to quantitatively estimate the model performance on unseen mixture proportions before the actual training, facilitating the tuning of ideal mixture proportions.
We further propose the nested use of scaling laws of training steps, model sizes, and data mixture to make predictions with only experiments at small scales, which enables the reuse of existing experiments and reduces the computation costs.
Experimental results suggest that our method effectively optimizes the data mixture, which leads to better performance in pretraining and maintaining the original abilities of pretrained models in continual pretraining.

Admittedly, how data mixtures affect model training is far from fully understood.
Our study makes preliminary attempts at a quantitative framework while leaving several limitations.

#### On the clarification of domains.

The concept of domains is not well-defined.
In this paper, similar to related studies (Xie et al., ; Chen et al., ; Albalak et al., ; Fan et al., ), we directly adopt the predefined domains in the open-source training data.
Nevertheless, we suppose that more operationally defined training domains, e.g., clustering (Gururangan et al., ; Shao et al., ), could further benefit the prediction accuracy of data mixing laws and the performance of outcome models.
For the validation domains, our implicit domain aggregation method obviates the necessity of explicitly aligning validation data with training domains.
This requirement is often encountered, given that validation data typically comprises trustworthy datasets rather than mere compilations from training domains.
However, we acknowledge that implicit domain aggregation may be less interpretable compared to the explicit approach and may raise concerns regarding its accuracy, as elaborated subsequently.

#### On the error analyses.

Leveraging scaling laws requires experiments to provide samples to fit the functions.
Consequently, it requires careful design of experiments (Mead, ) to decide the number of fitting samples to experiment with and how to distribute these samples to reduce prediction errors to the greatest extent.
In this study, we decide the number according to our affordable budget and leverage the simple rule that evenly distributes the losses of the data samples but considering more theoretically justified rules should be necessary.
Additionally, our nested use of scaling laws can introduce errors in each step.
Therefore, further analyses to mitigate the error accumulation are also demanding.
In Fig. 7, we notice our predictions are smaller than the actual loss, which we attribute to the underestimation from the step laws and model size laws we fit.
Further practical experience demystifies the technical details of scaling laws (Su et al., ) can help eliminate the errors.

#### On joint laws of multiple factors.

We propose the nested use of scaling laws for circumventing the difficulties in finding a joint law of training steps, model sizes, and mixture proportions.
Although we can predict the losses with our pipeline, a joint law unveils clear synergies of different factors.
For instance, previous studies indicate the power-law exponent in the scaling laws of model sizes and training data are insensitive to training and validation data (Hestness et al., ; Kaplan et al., ; Hashimoto, ; Hoffmann et al., ; Frantar et al., ).
Figuring out their joint laws with data mixture can further confirm this surmise.
Moreover, a joint law also implements coefficient-sharing of separate laws, reducing the number of required fitting samples.

#### On dynamic data curating.

Our study presents a pipeline to decide on a group of fixed mixture proportions for pretraining.
More sophisticated data curating can include dynamic proportions (Albalak et al., ) and even a curriculum that changes data domains (Chen et al., ).
The application of our data mixing laws in continual pretraining (Sec. 5) implies the prospect of extending our findings to these settings.
On top of this, we believe that it is promising to incorporate further analysis to pursue a dynamic data mixing law.

#### On theoretical understandings.

Our data mixing laws, similar to most scaling laws, are empirical findings.
We believe a theoretical understanding of the training dynamics that form the laws provides a more solid justification.
A potential perspective is understanding the target of tuning mixture proportion through gradient estimation (Guo et al., ; Gu et al., ).
Specifically, the mixture proportions weight data from different domains, whose effect boils down to the weight for the linear combination of gradients from different domains during training.
This perspective turns the target of tuning mixture proportions into finding an ideal gradient direction (Gu et al., ) and the relationship between data samples is formalized with their gradient directions (Guo et al., ).

To sum up, we initiate an investigation into quantitative prediction methods for pretraining data curation.
With its increasing focus on data engineering, we hope our exploration facilitates further quantitative studies and theoretical analyses in this research area.

## Acknowledgments

The computations in this research were performed using the CFFF platform of Fudan University.
We thank Botian Jiang and Shiduo Zhang for their insightful discussions.

## Appendix A The relative qualities of data mixtures relate to model sizes and training steps.

Figure: Figure 9: The rankings of the relative performance of 20 sample mixtures trained on RedPajama and validate on the Pile. (A) The rankings of models of different sizes all trained for 30k steps. (B) The rankings for 70M models trained for different steps.
Refer to caption: extracted/2403.16952v1/figs/rankings.png

One may wonder whether we can find the optimal data mixtures on small models and few numbers of steps, and then transfer the found mixture proportions to large-scale training.
To answer this question, we compare the relative performance of models in different sizes and trained with different numbers of steps in Fig. 9.

Results show that the relative performance fluctuates despite a relatively consistent trend across sizes and training steps.
This indicates that a mixture is better at small scales probably but does not always perform better at large scales.
The longest common sequence of the partial orders among the 20 mixtures in Fig. 9(A) and Fig. 9(B) only reaches lengths of 10 and 11, respectively.

## Appendix B Implementation Details

### B.1 Model Training

Throughout this study, we employ the Pythia Suit (Biderman et al., ) as our model architectures, the specific configurations are in Tab. 2.
The maximum sequence length is 4096 for pretraining from scratch and 2048 for continual pretraining, where the latter aligns with the setting of the original pretrained models.
In all our experiments, we train the model with a batch size of 1M tokens and a maximum learning rate of 1e-4.
We warm up the learning rates for 2000 steps and decay it to 0.1 of the maximum at the last training step with a cosine decay schedule.
For continual pretraining, we initialize the models with the 20k-step checkpoint of the Pythia 70M model and do not apply a learning rate warmup.

**Table 2: Model architectures for experiments in this paper.**
|  | 70M | 160M | 305M | 410M | 1B |
| --- | --- | --- | --- | --- | --- |
| Vocabulary Size | 50304 | 50304 | 50304 | 50304 | 50304 |
| Non-embedding Params | 18,915,328 | 85,056,000 | 201,541,632 | 302,311,424 | 805,736,448 |
| Layers | 6 | 12 | 16 | 24 | 16 |
| Model Dimension | 512 | 768 | 1024 | 1024 | 2048 |
| Heads | 8 | 12 | 16 | 16 | 8 |

For datasets, we mainly experiment with the Pile and RedPajama.
For the Pile, we find duplicates in the raw data, so deduplication is performed before training with it.
The Pile contains 5 coarse-grained domains, which are further decomposed into 22 fine-grained domains.
Our experiment in Sec. 3.1 is on Github and Pile-CC domains while the experiment in Sec. 3.2 is on Github, Pile-CC, and the Books.
All these are fine-grained domains.
For our experiments with 5 domains in Sec. 3.3 we adopt the five coarse-grained domains, i.e., academic, internet, prose, dialogues, and misc, where misc include Github and the DeepMind Mathematics Dataset which are symbolic content.
We use the coarse-grained domains because it is hard to find five fine-grained domains with sufficient tokens.
For the RedPajama, we download the version produced and shared by Chen et al. ().

### B.2 Fitting Mixture Laws

Fitting the mixture law requires us to first experiment on a few mixtures and obtain their losses.
The sample mixture chosen for fitting could largely affect the prediction accuracy.
Consider an extreme case where all sample mixtures have proportions around a small region, it is hardly possible to fit a law that reliably predicts the whole proportion space.

In this paper, we intuitively try evenly allocating the mixture proportions regarding their losses.
Specifically, we enumerate candidate mixtures by double-diminishing the proportion of each domain so that the losses are distributed evenly among these mixtures.
Then, according to the available computation budget, we sample a certain number of mixtures from the candidates to run experiments.
During sampling, we find candidate mixtures with a 0 domain proportion in any of the training domains take up a majority of the candidates.
To avoid these candidates making up all our samples, we specifically down-sample them.
The concrete algorithms are in Alg. 2.
Additionally, we employ AdaBoost Regressor (Drucker, ) for fitting the mixture laws to stabilize the predictions and improve their accuracy.
We encourage future studies to dive into a more careful design of candidate mixture selection with theoretical support.

Figure: Algorithm 2 Sampling mixture proportions for fitting mixture laws.

## Appendix C Connections between Implicit Domain Aggregation and MLP

We first repeat our final mixture law (Eqn. 8) here for convenience:

$$ $L(r_{1\dots M})=\sum_{i=1}^{K}s_{i}L_{i}(r_{1\dots M})=\sum_{i=1}^{K}s_{i} \left[c_{i}+k_{i}\exp{\left(\sum_{j=1}^{M}t_{ij}r_{j}\right)}\right],$ $$

where $r_{1\dots M}$ are mixture proportions on $M$ training domains, $L_{i}$ are validation loss on $K$ implicit domains with $s_{i}$ as their weight in the overall validation set, and $c_{i},t_{ij}$ are other parameters to fit.

Figure: Figure 10: The computation graph of mixture law with implicit domain aggregation. We take an case of 3 training domains and 4 implicit validation domains as example. The parameters correspond to the notations in Eqn. 8.
Refer to caption: extracted/2403.16952v1/figs/computation_graph.png

The mixture law boils down to a computation graph in Fig. 10, which contains two layers.
The first layers predict the domain losses, while the second sums up the domain losses to obtain the overall validation loss.
In this way, the mixture law becomes a multilayer perception (MLP) with an exponential activation function.
In practice, we fit the mixture laws with implicit domain aggregation by fitting a multilayer perception with exponential activation and applying softmax to the output layer weights.
Additionally, considering the high variance of MLP, we further employ AdaBoost Regressor (Drucker, ) for fitting the mixture laws to stabilize the predictions and improve their accuracy.

Inspired by this perspective, we attribute the successful fitting of mixture law to two aspects.
First, the MLP with a sufficiently large hidden dimension is a universal approximator (Hornik et al., ) thus being able to fit the relationships between losses and mixture proportions.
Second, the mixture proportions are bounded between 0 and 1.
For this reason, predicting an unseen mixture is an interpolation problem, which is usually easier than extrapolation.

## Peer Reviews

confidence: 5
contribution: 4
initial_rating: 8
presentation: 3
questions: 1. Can you explain Algorithm 1's nested structure more? It would be great if there was a longer version in the appendix that could stand on its own.
2. Can you comment on the mixing law fitting well in other settings (such as different subsets of RedPajama/Pile, or a new dataset overall)?
3. Is there a better way to evaluate DoGE and DoReMi in Figure 9 than using the mixtures reported in the DoGE paper? What do the DoGE mixtures look like in the OOD setting, which would account for the shift between RedPajama and the Pile?

Overall I think this is a very nice paper. It could benefit from more clarity in describing the algorithm and a fairer evaluation of the method.
soundness: 3
strengths: Originality: 
- The paper departs from other data mixing works by presenting the following novel methodology: let's first focus on learning the parametric form of the relationship between proportions and loss, and then use this to find the proportions that minimize the predicted loss. The idea is inspired by work on scaling laws, but mixing laws are quite different. They must capture different factors, such as cross-domain interactions, and can be multi-dimensional depending on the number of domains.

Quality: 
- Beyond just presenting the mixing laws, which are interesting as a scientific artifact, the paper does a good job of discussing practical considerations (training budget) as well as insights from fitting the laws, such as the sparsity of the interaction matrix.
- The paper explains how mixing laws are broadly applicable, presenting practical settings that are not just limited to minimizing loss on the training domains (which e.g., DoReMi focuses on). In particular, the paper considers when we have a separate validation dataset, and when we wish to continually pre-train a model to learn some new domain without forgetting the original one.

Clarity: 
- Paper was well-structured in starting with results on the simplest mixing law and then gradually moving towards more practical settings.

Significance:
- The paper provides a feasible method of doing data mixing that can be much cheaper than brute-force searching---which is often the convention when training LLMs---and is more predictable than recent algorithmic approaches. It captures the intuition of the practical trial-and-error process that practitioners would make when tuning mixtures, and formalizes it into a rigorous procedure.
summary: It is important to determine the best mixture proportions over a set of domains (i.e., code, web text, for pre-training). However, naive approaches require manual tuning over the search space, and recent algorithms that learn the proportions are not well-understood. Instead, this paper discovers a log-linear "mixing law" that relates mixture proportions to loss. Experiments are done over 2, 3, 7 domains (Pile, RedPajama) to confirm that this mixing law is accurate. This mixing law provides a way of determining the optimal data mixture in a principled manner: conduct multiple training runs over various mixtures, fit the parameters of the mixing law, and use those to compute the proportions that minimize the desired loss. One drawback of naively using this approach is that this requires many full training runs. The paper thus proposes layering on training-step and model-size scaling laws so that the optimal mixture for some target model size and target number of training steps can be inferred by training on various mixtures for smaller models and fewer number of steps. The paper also discusses various ways this approach can be used: it can be applied when we are given a validation dataset that is disjoint from training domains, and it can be applied in the continual pre-training setting too. Empirical results show that conducting shortened runs on smaller models allows them to produce a mixture over RedPajama that attains low validation loss on the Pile faster than the default mixture.
weaknesses: Clarity: 
- Minor typos: "the implementation details." on L328. Figure 6 has no x-axis. In L409, the second L(S) should be L(N). In L473 "date" should be "data".
- What does resampling mean in L456? The procedure here does not appear to match up with the algorithm since you only mention training 70M models. 
- My main clarity issue is in Algorithm 1. Presented as is, it is very difficult for someone to implement. What are the (x, y)'s you need to store in each loop and fit? What are the exact scaling laws? While L(S) and L(N) are defined in L409, this does not typecheck with the algorithm using only L(N, S, r)s.  

Quality: 
- Evaluation of DoGE and DoReMi: the mixtures taken from their paper are for the "universal generalization" setting; for DoGE these are the learned proportions that aim to minimize the average validation loss across domains, not a disjoint validation loss. For a fair comparison, you would need to use the OOD version of DoGE (their algorithm 2). Furthermore, it is unclear how much compute is used by your method vs DoGE vs DoReMi; the compute should be normalized to make a fair comparison, especially since DoGE and DoReMi require 1 extra and 2 extra training runs while this method has parameters scaling in the number of domains.
- It is unclear if the proposed mixing law holds on other domains. The 2 domain experiment is only done on Pile-CC and Github, and the 3 domain experiment is only done on Pile-CC, Github, and Books3. Is it possible that some property of these domains gives way to the linearity in the mixing law? It would also be interesting to explore if this mixing law held on tasks (i.e. in instruction-tuning). 
- Table 15 in the appendix suggests to me that the mixing laws are not that accurate for the nested approach.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: I would appreciate the authors’ responses as well as plans to address two weaknesses I have described above.
soundness: 2
strengths: I very much enjoyed reading how the authors present their thought process leading to the proposed law. Each step was reasonable, and some of the assumptions the authors put forward made sense (although it is never possible to know whether it’s post hoc or ad hoc by reading the manuscript alone.) Because of clear presentation of the steps, I was reasonably convinced of the utility of the proposed law for figuring out the data mixing ratio, despite its limited empirical validation. Furthermore, it was good to see that the authors showed that this law can be used in conjunction with so-called scaling law to improve large-scale training based on the mixing law from smaller scale training.
summary: In this paper, the authors carefully describe the steps they took to arrive at so-called data mixing law which is a simple formula that predicts the validation loss (potentially of arbitrary mix of diverse domains) of a model trained on a fixed data mixing ratio. This law has O(n_domain x n_domain + n_domain) parameters. Although there are a few arbitrary decisions here and there (e.g., choosing M4 over M1) and the validation (of the method and law) is limited to just one particular setup (training a language model on a relatively general set of domains), i enjoyed reading the authors’ thought process arriving at the proposed “law”.
weaknesses: There are two main weaknesses. First, any such law is by nature empirical and thereby must be validated quite extensively with terms of their predictive power, before they can be established as the law. The experiments in the paper are reasonable, but are conducted in exactly one setup. Although the authors consider continual pretraining as a separate setup, i don’t believe it is separate enough to show that this law generalizes beyond this particular setup. I would’ve even appreciated some synthetic experiments, as long as some diversity was introduced. This is the major weakness, since the authors claim it to be the “law” of data mixing.

The second weakness is whether this particular form is the right form. As the authors pointed out, one can think of this as a particular form of a multilayer perceptron. A natural question is then whether this particular form is the one we should prefer. For instance, why do we want to choose the form out of M1, M2, M3 and M4; can’t we simply build a much simpler MLP and fit the MLP directly? If the goal is to use this so-called law to improve the efficiency of training, why would we care about a particular form of this law? Instead, we would much rather care about the predictive power of any predictor fitted on a large number of training runs.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: See the Weaknesses.

I am open to raising my score if the authors address the above Weaknesses.
soundness: 3
strengths: 1. The paper is well-written and easy to understand.

2. The research topic on pre-training data mixture is timely and important for the LLM community.

3. The visualization of the model weights t_{ij} in Figure 5 is interesting because it shows the nuanced interaction between the validation domain and training domain proportions.
summary: This paper proposes a novel framework for optimizing data mixtures in pretraining large language models (LLMs). The authors introduce data mixing laws—predictive functions that estimate model performance based on the proportions of different data domains, allowing for the prediction of model outcomes across unseen data mixtures without extensive re-training.

Key contributions include:

1. Data Mixing Laws: The authors formalize a quantitative relationship between data mixture proportions and model losses, termed "data mixing laws." These functions enable the prediction of model performance with different data mixtures.

2. Pipeline for Scalable Predictions: By leveraging nested scaling laws for training steps, model sizes, and data mixtures, the paper proposes a pipeline that allows the prediction of large-model performance using small-scale experiments.

3. Applications in Continual Training: The framework also provides insights for maintaining performance in continual pretraining by dynamically adjusting mixture proportions, thus preventing catastrophic forgetting.
weaknesses: 1. The validness of the underlying domain-independent assumption. Similar to other recent papers that model the data mixture problem using scaling law like [1], there is also an implicit assumption of the proposed method: different domain proportions are independent. In Eq. (1), no feature interaction term like r_i*r_j models the interactions between domain proportions. The loss is modeled by an exponential function over the **linear combination** of the proportions. We would like to understand how each domain affects every other domain, this interaction might be highly complex and not local. I would like to further clarify that the feature/proportions interaction mentioned here is different from those shown in Figure 5.

2. Missing baselines. There does not seem to be a sufficient comparison between the proposed method and other data selection methods, e.g. comparing with other token-level and sample-level automatic data selection methods [2, 3, 4].

3. Missing downstream performances. Only perplexity is evaluated, which is not sufficient to demonstrate the effectiveness of the proposed method.

[1] BiMix: Bivariate Data Mixing Law for Language Model Pretraining

[2] Perplexed by Perplexity: Perplexity-Based Data Pruning With Small Reference Models

[3] DoGE: Domain Reweighting with Generalization Estimation

[4] DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. **Computation costs analysis**: how many small-scale experiments are required to get an accurate prediction of the scaling law function? If the default baseline train with these extra computation flops, how would the performance comparing with the proposed method?
2. **Fair comparison between DOREMI and DOGE**: if apply the doremi weights from original doremi paper and doge weights using the correct validation domain, how the performance of the proposed data mixing law comparing to them?
soundness: 2
strengths: 1. The paper is well-written, where the methodology and experimental results are clearly demonstrated.
2. The effectiveness from the method is well-justified from extensive experiments on LLM pretraining and continual pretraining.
summary: The paper extends the scaling law approach to optimize data mixtures for efficient language model pretraining. Specifically, the authors fit a scaling law function (i.e. data mixing laws) between the validation loss and training in-domain sampling ratios. From extensive experiments, the approach can improve language model pretraining efficiency on RedPajama dataset, as well as mitigating the catastrophic forgetting by training on mixture of generic dataset and domain specific data.
weaknesses: 1. **Lack of a detailed analysis on computation overheads**: how many small-scale experiments are required to get an accurate prediction of the scaling law function? If the default baseline train with these extra computation flops, how would the performance comparing with the proposed method?

2. **unfair comparison between some baselines**: the DOREMI [1] and DOGE [2] weights from Fan et al.'s paper are both targeting at minimizing the model's performance on SlimPajama dataset, which has a great domain shift to the target used by the authors (the Pile). I would suggest to use the DOREMI weights from their original paper [1] which also optimized on the Pile and retrain DOGE with target domain as your validation set from the Pile.

3. trivial points: 
    - the footnotes (9,10) are not properly displayed on the paper.

[1] DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining
[2] DoGE: Domain Reweighting with Generalization Estimation

Meta-review: The authors propose to estimate optimal training data mixture proportions by training a series of smaller models with varying mixture proportions and then fitting a simple model to relate those proportions to performance. They show that this leads to more efficient training (as a function of the number of training tokens) compared to baselines. Training data optimization is an important problem as large training runs are expensive; the approach here is simple, natural, and well-executed. Reviewers were unanimous in recommending acceptance and I recommend acceptance as well.

I encourage the authors to take the reviewers' feedback into account. For example, as Reviewer c3fa mentions, it would be helpful to discuss the relative computational costs of these different data mixture optimization methods, since the amount of compute spent on such methods could instead be used to train on a (possibly less optimal) data mixture for longer. In addition, please incorporate the experiments on other domains into the main text. Finally, I suggest that the authors compare with related concurrent work: https://arxiv.org/abs/2407.01492.

Decision: Accept (Poster)