Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance
Authors: Jiasheng Ye
Peiju Liu
Tianxiang Sun
Jun Zhan
Yunhua Zhou
Xipeng Qiu
Venue: ICLR 2025 Poster (2025)
Keywords: data mixtures
pretraining
post-training
scaling laws
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose a method to optimize data mixtures by quantitatively predicting language modeling performance of different data mixture with low cost.

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 laws to predict the performance of large models trained on massive data under various mixtures with only small-scale training.  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.

Published as a conference paper at ICLR 2025
DATA MIXING LAWS: OPTIMIZING DATA MIXTURES BY PREDICTING LANGUAGE MODELING PERFORMANCE
Jiasheng Ye" Peiju Liu Tianxiang Sun', Jun Zhan', Yunhua Zhou? Xipeng Qiu
Fudan University, SShanghai Al Labortory (jsye23, Pij iu23) @m. fudan. edd. .n zhouyunhua@pji www.pppppla..g..........om. org.cn xpqi uê Fudan. cn
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 pre- dictability of model performance regarding the mixture proportions function forms, which we refer to as the data mixing laws. Fitting such functions 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 laws to predict the performance of large models trained massive data under various mix- tures with only small-scale training. Experimental results verify that our method effectively optimizes the training mixture of a IB 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 continual training accurately predicts the critical mixture proportion that avoids catastrophic forgetting and outlooks the potential for dynamic data schedules
INTRODUCTION
Pretraining data for large language models (LLMs) are typically a mixture of multiple domains, varying from English to minority languages (Doddapaneni et al. [021 Li al. 2023), from casual dialogs to formal academic writings (Taylor al. 2022). and from texts to modalities like images and speeches (Zhan et al. 2024), among others. These data interplay with each other, showing complex relationships including facilitation, being unrelated, conflict (Guo al. (024). This necessitates adjusting the mixture proportions of training data balance the model capabilities while harnessing synergies across domains, thus enhancing the competence of the outcome models, highlighted by
extensive practices (Gururangan et al.) 2020 Zhou et al.) 2023: Xie et al. 2024a Fan et al. (024).
Nonetheless, it remains elusive to figure out an ideal training data mixture. Most existing practices tune the mixture through heuristics upsample a proportion of high-quality or underrepresented data without disclosing the concrete criteria in detail (Gao et al.) 2020; Touvron et al. [203aa] Bai et al.
2023.
Bi et al. 2024). While some studies propose automatic algorithms to qualitatively optimize data
mixture (Xie et 2024a Fan et al., 2024), it is hard to predate the effect of these strategies before the actual training run. In contrast, encouraged by advances in scaling laws that show model losses on given set of evaluation data are quantitatively predictable for wide range of variables (Kaplan et al. 2020; Hoffmann et al. 2022), 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. The intuition is that predicting the performance of unseen data mixture only involves interpolating among seen mixtures because the proportions are bounded between 0 and 1. For this reason, numerous function forms can lead to descent prediction accuracy, among which we try to find a simple one. In particular, we find that, given a mixture of M
'Equal contribution.
1 'Corresponding authors. Codess and data are available at: http:////gggghubb.bocee//ett ng1 aws

---

Published as a conference paper at ICLR 2025
Small Steps, Models, Seen Mixture
Samplex
Training Stop Laws
Mixting
Large Steps, Small Models, Seen Mixture
Lossas
Loises
Large Steps, Large Models, Seen Mixture
Misissan Lous
Medel Sine Laws
Large Steps, Large Models, Unseen Mixture
Traiiin Steps
Mosial Skats
H midhors
Propartion (Domain
Properton (Domain
Trainig Steps awes Modill moo: (8 Mikinng aave (ours)
Figure I: 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 shree-step pipeline to predict model performance on the target model size, training step, and mixtures.
domains, an exponential function over the linear combination of the proportions, i.e.,
(i(r...M) c: + K2 exp
Σ 4O
(I)
can predict the validation loss L; on any of the training domains accurately under a fixed model size and amount of training data, where T the proportions of the M domains and c,, k2 by 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 validation set end-to-end without explicitly decomposing 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 unreasonably expensive. Fortunately, fruitful rescarch 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., 2020 Henighan et al.. 2020 Hoffmann et al. 2022 Alabdulmohssnn et al., 2022, OpenA|| [ Muennighoff et al. [2024; [Bi et al. (2024). 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. !
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 1B model trained on 100B tokens and achieve performance comparable to a model trained on default mixture for 48% more steps. The prediction domain-speciffi validation sets also offers plausible references to the balance of model capabilities. Further applying our data mixing law to continual pretraining accurately find the proportion that avoids catastrophic forgetting (French) 1999 Kirkpatrick et al.) 2017 Luo et al., 2023), revealing the prospect of applying data mixing laws to guide a multi-stage pertaining, and thus 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 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 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, show- ing effectiveness in optimizing model performance, balancing model capabilities, and the prospects of guiding the design of the data schedule.
BACKGROUND
We briefly review the pretraining process of large language models and summarize key findings from neural scaling laws, then formalize the problem we study. Further related works are in App.
2

---

Published as a conference paper at ICLR 2025
Pretraining large language models. We consider the task of pretraining an autoregressive language model pe via next-token predictions (Radford et al. 2018). The training dataset Prriin = {D;}:-- composes M domains with mixture proportions r E A M In each training step, the task first samples 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.,
Le = Einr. Dr
Σ log Po Po (10y(=o..)-1))
(2)
To evaluate the learned model, we compute the loss on validation data Dval-
Scaling laws. For wide range of factors z, scaling laws (Kaplan et al. [2020] Henighan et al. 2020 Hoffmann al. 2022) show that their effect on the losses L follows power laws
+ kx.,
where c, k, and a are parameters to and x can be model sizes, numbers of training data, training steps, and the amount of computation. Previous experience (Alabdulmohsin et al. 2022 OpenAl 2023; Bi et al. 2024 Su et al., 2024) highlights the impressive predictability of scaling laws. Specifically, Eqn. fitted collection of small models, training data, computation to precisely predict the test loss of larger cases over orders of magnitudes. This enables practitioners to estimate the performance of 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. [20TT) sparse architectures (Frantar et al. 2023). and repeated data (Muennighoff et al. (024), consolidating the predictability language model performance.
extrapolate
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 tackle mixture optimization by establishing a quantitative framework that predicts the loss given any mixture proportion. Formally, for training dataset comprising M domains, we parameterize the function
L= fr))
under the fixed model sizes and number of training steps, where T Ar 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.,
= arg (r.).
(5)
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 refer to 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 - degrees of freedom in the mixture proportions and, correspondingly. K 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-mone 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 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
3

---

Published as a conference paper at ICLR 2025
from one of the training domains to circumvent the non-monotonicity (Sec. 3.
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.
Setups We train 70M and 160M language mod- els on the mixture of Github and Pile-CC subset from the Pile dataset (Gao al.) 2020) 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.
Reducible Loss on Pil-CC
2w:10"
Reducitle
10
0.0
02
04
US
08 00
04
288
08
14
Findings Results in Fig. reveal the quanti- tative 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 shared constan 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 pretrained language model regarding the domain proportion precisely fit into an exponential law
Proportion of Github
Proportion of Pilo-CC
7001 20calge
70U1 30cakpa
1004 20h 30pa
(0094 DDe staa
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 sepa- rately. We train 70M and 160M models on five different mixtures of Github and Pile-CC and ob- tain 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 the domain proportions.
Figure 2:
(r.) c, ky exp (tur:),
(6)
where L1 and are validation loss and training mixture proportion of domain i, respectively, while Ca k4, and tu are learnable parameters
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.
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. 0.875, 1 and retain valid ones in which three proportions sum up and do not use up all tokens in any of the domains, which results in 32 mixtures in total.
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. if the number of domains M = 2.
Symmetry. Any exchanging of variables should not change the functional form as we should not incorporate any domain-specifee bias.
"The constant term, known as irreducible loss, arises from finite training data and the entropy of the evaluation data theoretically (Bishop) 2006 Henighan et al. (020).
'While power laws are more common in existing studies on scaling laws (Kallan et al. 2020 Hoffmann
et al. 2022). we do not consider it for its ill po-po properties that the function value blows up when the variable, mixture proportion in our case, approaches 0. This contradicts the observations that the losses remain low (e.g.. no more than 10) with the generalization between domains.
'Despite a simple case, our findings on two domains have practical applications to continue pretraining Gu rurangan et al, (020), where aim to enhance pretrained model on given domain by training it mixture of the original pretraining data and upcoming domain data. Please see Sec. 5 for details.
4

---

Published as a conference paper at ICLR 2025
Table 1: Mean absolute emrors (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 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
Together, the two principles lead to candidate functions that replicate the exponential term in Eqn. 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:
M
MI: [((r) =G +
Bigg exp (145T))), M2: Tx(r) =G + Aig
exp (Li5rs),
#=1
M
M3: L,(r) =c; + h;s exp
II EigTy
M4: L,(r) =cs + his exp
tift's
We summarize their fitting errors on three target domains in Tab.
Findings The results in Tab. suggests both
Pile-CC
Github
M1 and M4 gives reliable predictions while M4 has fewer coefficients. Therefore we adopt M4
fitting validation
fitting validation
3.0
observation
observation
M Σ
8.2
2.5
Tr(r...as) = Gs kis exp
8.1
1.8
as the function forms of our data mixing law, where Li is the validation loss on i-th validation domain, ry is the proportion of the j-th training domain, and c₁ k;, tiy are learnable parameters. The fitting results are in Fig. and Fig. demon-
8.0
1.0
3.0
3.2
prediction
prediction
Books3
Total
fitting validation
3.1
fitting validation
strates the prediction accuracy. The results indi-
abservation
observation
cate that Eqn. fits the given samples well and estimates the unseen ones accurately.
2.9
2.8
Meanings of the coefficients. To provide more intuition, we discuss the meanings of the co- efficients in Eqn. 7 In general, k; > 0, thus the exponential term is always positive and the prediction loss is strictly greater than the con- stant c. Hereby, C; represents losses that are reducible by adjusting the data mixture. lij. depending on both training domain and valida- tion domain j, shows the interaction between them. A negative try indicates that training data of domain helps reduce validation loss on domain and vice versa.
2.
3.0
52
5.5
25
2.6
prediction
prediction
Figure 3: Prediction results on the domain losses and overall losses in the three-domain experiment. The domain losses are fitted with Eqn. / and we obtain the total losses through explicit domain ag- gregation of Eqn. g
Patterns of the coefficients. We visualize normalized &ij of training and validating on the 5 domains mixture of the Pile in Fig. 4 The relationship between domains can be categorized into 3 types.
'The Pile contains 22 fine-grained domains which collected into five coarse-grained domains, i.c., academic, internet, prose, dialogues, and misc, where misc include Github and the DeepMind Mathematics
5

---

Published as a conference paper at ICLR 2025
(A)
Lonone on Caimuu
(B)
Loceeo Booked
[c)
Loooes on PieoO
(D)
Total Loooes
Perpartnn Boose
Propartion of Piyc0
Proportin of Bossd
Rigcorioon Pf D0 50
propation Ganuta
Rroporton of Booka]
Propertion Cansen
Progeston Glied
Figure 5: 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 mixed with the three domains. x: validation samples. the predicted minimum loss on the overall validation set.
Being unrelated: The figure shows a highly sparse pattern where most of the domains have little relationship to each other and the validation performance nated by training data of the same domain, which supports the intuitive progressive mixture tuning strategy that adds data for underperforming capability during training (Team) 2023). Meanwhile, we also observe facilitation (e.g., training dialogue for the internet) and conflict (e.g.. train- ing symbolic data for prose) between domains, which
Insemee Symbois o slop. Froc Azzdems
a domain is domi-
Validation Domain
indicates the room for leveraging domain interaction enhance model performance.
Pross
Training Domain
3.3 PREDICTING LANGUAGE MODELING PERFORMANCE ANY VALIDATION MIXTURE
Figure 4: The interaction between dif- ferent training and validation domains on the Pile. Each boxes are fitted nor- malized lij from Eqn. 7 We normalize the value by lij with the maximum abso- lute value for each validation set (i.e.. tgg ///ti,arg max; [t,y)) to compare the val- ues intuitively.
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 composi- tion 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 claborate as follows.
larger value (greener) indicates the training domain helps learn the validation domain more.
Explicit domain aggregation. Considering a validation set made up of K domains with the propor- tions as $1..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 training mixture. Therefore, the functional relationship of the overall validation loss to the training mixture proportions expands into
K G ki exp
L(ri..M) Σ s.L.(ri...M)
Σ 8;
Σ LigTh
i=1
Using Eqn. S we can the loss on each validation domain L; and sum them up to obtain the prediction of overall loss.
Implicit domain aggregation. A remaining limitation is that still need to acquire the components of validation data $1..K in advance. This can be inconvenient if the validation set is collected separately from the training queries that cover unknown compositions of various domains. To remove the constraint validation components, we assume that we can decompose the validation data into K implicit domains whose losses are predictable with Eqn. and we treat their proportions in the validation data $1... as
For instance, the validation data may come from real-world user
Dataset which are symbolic content. We do not experiment fine-grained domains for their limited number of tokens available.

---

Published as a conference paper at ICLR 2025
learnable parameters, leading to the final form of our data mixing laws With this perspective, we fit a data mixing law with the overall losses end to end.
Introducing implicit domains may draw concerns about the number of fitring samples exploding with the number parameters to fit and questions on deciding the number of implicit domains without knowing the oracle. We study and discuss their impact, respectively.
Do we need quadratic number of samples to fit Eqn. aas the number of domain grows? No.
Mean Abookte Erross
domaine
The parameters in Eq 8 scale as OMM x K), where M and K represent training and implicit validation domains. Nevertheless, as shown in 6 the quadratic growth in the number of pa- rameters does not translate quadratic growth in sample requirements. We attribute this to the high sparsity in the parameters as fig 4 re- veals, which allows us to fit the equation with substantially fewer samples when using appro- priate regularization. While using more samples decreases prediction errors, the number of sam- ples that reach similar precision level does not grow dramatically. This may pave the way for applying implicit domain aggregations for cases with more training domains. Although concluding the exact number of samples required can be challenging due to the differences among training data, can tune the fitting mixtures the smallest experimented models, which is cheap and works well in practice (see Sec. 4.2 and App. D.3).
0.08
domains
domaine
0.04
0.02
10
12 16 20
Figure 6: The mean absolute validation errors of Eqn. fiiteed with different numbers of samples for training mixtures containing different numbers of training domains. For each number, we resample and select the batch of fitting mixtures that reach lowest errors.
How to choose the number of implicit do- mains? Set larger than the oracle one. Fig. shows our experiments where we train language models on the 5 coarse-grained domains of Pile and evaluate validation set mixed with these domains. We compare the errors obtained by implicit domain aggregation with different num- bers of implicit domains to those obtained by explicit domain aggregation. We find that apply- ing implicit domain aggregation and setting the number of implicit domains 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 large number without careful tuning"
Impkot (K-1)
Implat (K-5)
Impket
(8-5)
(K-10)
(K-50)
Figure 7: Prediction errors of the five-domain data mixing laws fitted with explicit and implicit do- main aggregation. Explicit domain aggregation: we fit Eqn. 7 for five domains respectively and sum them up according to their weight the over- all validation sets. Implicit domain aggregation: we fit the losses on overall validation with Eqn. g assuming different numbers of implicit domains K and treating the proportion different implicit domains learnable parameters.
NESTED SCALING LAWS PREDICT LOSSES TRAINED ON VARIOUS MIXTURES USING ONLY SMALL-SCALE EXPERIMENTS
4.1 A PIPELINE FOR Loss PREDICTIONS
While data mixing laws enable to predict the performance of models trained on unseen mixtures, fitting the laws directly on target scales is unaffordably expensive. Firstly, fitting the laws involves training multiple models across diverse mixtures with model sizes and token counts identical to the
We note that the final forms of our data mixing law resemble multilayer perception (see the computation graph Fig. [4].. We include further discussion and implementation details in Appendix
We set K 30 if not stated in later experiments.
7

---

Published as a conference paper at ICLR 2025
Algorithm 1 A pipeline to predict losses of different mixture proportions on large models trained on massive data through small-scale training
Input: Validation data Dvat, training data of M domains {Dm 1M m==1, target training steps Stargets target
model size Naarget. target mixture predict Plarget, training steps to fit the step laws So, model sizes to fit the size laws {N₂ ),, and data mixtures {rt] to fit.
Output: The validation loss of Naarget model trained for Storget steps on mixture Plargets i.c. ([Naarget:Starren Ptarget).
1: for Each mixture r; do
2: for Each model size N, do
3: Train model of size N, on mixture r, for So steps to obtain L S < Su, r.)
4: Fit training step scaling law L(S) with L(N,, < So,r:)
Predict L(N Starget, r:)) <- L(S = Starget)
6: end for
7: model size scaling law L(N) with T.(NI...K. K: Starget,
8: Predict (NNtagget, Starget, ri) LIN Ntarget)
9: end for
10: Fit the data mixing law L(r) with L(Nearget, Starget, Fi...N)
11:
predict [(Nargget, Stargets Psarget) +- L(Paarget
target ones. Furthermore, we must repeat this process for each target model size and training datase Such expensive costs hinder the practical value of our data mixing laws.
We thus wonder whether can obtain the losses of different mixture proportions without training at large scales. Fortunately, this idea gains endorsement from existing experiences that verify the im- pressive extrapolation of scaling laws of training steps and model sizes. particular, OpenAl (2023) predicts the loss of the target model with merely 1, 000x-10, 000x less compute. Consequently, 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.
We illustrate the proposed pipeline in Fig. [ with details depicted in Alg. T. Scaling laws in our pipeline are largely based the function forms of Chinchilla Scaling Laws (Hoffmann et al. 2022). i.e. L(N, D) = E + No DJ. where N is the model size and D is the number of training data. Under fixed batch sizes,
A B
treat the number of training data the number of training steps S as well. Notably, we do not directly fit the complete Chinchilla Scaling Law with two variables as we find it practically unstable to fit such many parameters simultancously in our preliminary study, similar to the findings in Besiroglu et al. (2024). Instead, we decompose the law into two power laws A respectively. We first fit power laws of training steps to predict model performance with more training data then fit power laws of model sizes to predict the performance when scaling up models. We empirically find this routine
for training steps L(S) E1 s and model sizes L(N) = E2 N
stable
4.2 EXPERIMENT
We verify the effect of our pipeline with an experiment to minimize the validation loss of 1B model trained on 100B tokens.
Setups. We train our models on the mixture of RedPajama and validate the validation set of the Pile 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 series of 70M, 160M, 305M, and 4 410M models for 30B tokens. For all the models, we set the batch size as 1M tokens thus translating into
*An idea is to transfer the optimized training mixture on small models trained with few tokens the training of large models and large volumes of training data. Nevertheless, as recent studies (Goyal et al.) 22044] Kang et al.) 22044) Covert et al.) highlight, the rankings of the data mixture vary as the model size and number of trained tokens change (Appendix C). Therefore, the optimal mixture at experimented scales can be suboptimal at
the target scale.
'See Appendix D.2 for our preliminary verification. We notice some recent efforts try to investigate democratizing the implementation of predictions with scaling laws to facilitate applications (Su et al.) 2024]] Porian et al.] (2024) While we illustrate our pipeline with the nested use of scaling laws, other implementations of scaling law predictions are also feasible and complementary to our method.
8

---

Published as a conference paper at ICLR 2025
Defaut
Optiizzdd
16.5
Peppecyy
16.0
Coom
Domairs
Default Mastuee
0.770
Optimized Misture
C4
0.150
0.2000
Cithab
0.450
6.180
Tooks
0 0088
0.73x
Starsfhrchunge
Wilepedia
(U/G1)
0.0000
0/1156
61280
1.48x
14.5
1.4
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 the original number training steps and eventually achieves performance comparable to 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.
Figure 8:
100k steps for the IB models and 30k steps for small models. We apply a cosine learning rate decay with warmup of 2k steps which decays 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 emrors all 40 samples as our final set of mixtures to run pipeline. For more details, please refer to Appendix D33
Results. Our pipeline optimizes language modeling performance effectively.
Fig.
g
shows
the default mixture of RedPajama (Touvron al. 2023a in and the optimized mixture obtained from Alg. with their performance on the validation data. The model trained on the optimized mixture can achieve performance comparable to the one trained on the default mixture with only 73% steps. It even- tually attains a performance that requires 48% more steps if trained using the default mixture. This in- dicates the effectiveness of our pipeline in mixture optimization 10
Delautt
DoGE
DoGE (OOD)
DoRaMi (RedPajama) (Plb)
DoñaMi
Ours
(Unhvoraal)
The language modeling perfor- mance of different data mixtures. All models are IB models trained for 100B tokens with the same hyperparameters and validated on the validation set of the Pile. Specific propor- tions are in Fig. 2T
Figure 9:
We also compare our optimized data mixture to pre- vious optimization algorithms, which provide qualita- tive optimization. Specifically, we compare our method DoGE (Fan et al. 2024) and DoReMi (Xie et al. 2024a). For DoGF, we compare both their universal generalization setting which assumes i.i.d. between training and validation data, and the OOD setting which optimizes for a given validation set, similar to ours. For DoReMi, which only works for universal optimization that ignores the validation data, we experiment on both a mixture optimized exactly on RedPajama and a mixture adapted from the one optimized on the Pile using the domain overlap between RedPajama and the Pile. More specific details on obtaining these data mixtures are in Appendix[FA. [F4] As shown in Fig. our method finds the mixture that reaches the lowest losses for the same model sizes trained with the same data budgets. This further verifies the effectiveness of our method.
10 loss predictions are in Fig. 22 which shows the predictions are plausibly accurate and are consistent with the rankings of actual runs.

---

Published as a conference paper at ICLR 2025
(C)
26v10
Lo0k inTePPie
2.02
Proportion 0.375
Pie Propotion 0.3
Reducible Loss on Python
2.08
Criics Propertin
P.L. wil Contrnaa Pritran
Recucibee Loss the
2.0
abx1e"
2030
4000
6500
5000
10000
(D)
1.5
Pytoon Proporiin 6462
Pyter Poppotio 6.
Crtice Propertin
vel Condinal Pickiah
Pile Proportion
Python Proportion
30
2000
4.000
10000
Figure 10: Loss predictions and the training curve continual pretraining Pythia-70M 70M on a mixture of the Pile and python code. (A) Loss prediction 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).
DATA MIXING LAWS HELP AVOID CATASTROPHIC FORGETTING IN CONTINUAL PRETRAINING
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 pre- trained models. injects up-to-date knowledge into the model, avoiding performance degradation due distribution shifts (Gururangan et al., 2020; Xiong al.) (023). In addition, researchers also apply continual pretraining to reuse existing model parameters to build models of different architecture (Komatsuzaki al. 2022).
We experiment typical scenario of continual pretraining, where we train the model on the mixture of original pretraining data and upcoming data of 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 mixtures and fit the Eqn. 6 on losses of the Pile and python codes. Results in Fig. 10 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 the original data. A representative mixture optimization target is maintain the general-purpos ability (losses on the Pile) unchanged. To this end, using the fitted data mixing laws, we predict the critical proportion leading the same loss before continual pretraining. Fig. I0 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. suggest continual pretraining is significant for connection to the design of data schedules (Albalak et al. 2023: Chen et al. (024b). Usually, continual pretraining applies pretrained model, while it is natural to further continually pretrain the continual pretrained models, ie., multi-stage pretraining (Chen al. 2024b). In each stage, the mixture proportions 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, more comprehensive data curating paradigm.
the
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. Using data mixing laws allows practitioners to quantitatively estimate the model performance on unseen mixture proportions before the actual training, allowing low-cost tuning of data mixture together with scaling laws. Given the burgeoning interest in data engineering, we hope that our study paves the way for further quantitative inquiries and theoretical analyses in this research area.
10

---

Published as a conference paper at ICLR 2025
ACKOWLEDGEMENT
This work was supported by the National Key Research and Development Program of China N1o22022ZD0100100 The computations in this research were performed using the CFFF plat- form of Fudan University. We thank Botian Jiang, Shiduo Zhang, and anonymous reviewers for their constructive feedback.
REFERENCES
Ibrahim M Alabdulmohsin, Behnam Neyshabur, and Xiaohua Zhai. Revisiting neural scaling laws in language and vision. Advances in Neural Information Processing Systems, 35:22300-22312, 2022.
Alon Albalak, Liangming Pan, Colin Raffel, and William Yang Wang. Efficient online data mixing for language model pre-training. arXiv preprint arXiv:2312.02406,, 2023.
Shun-ichi Amari, Naotake Fujita, and Shigeru Shinomoto. Four types of learning curves. Neural Computation, 4(4):605-618 618, 1992.
Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jachoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. arXiv preprint arXiv: 100.06711,, 2021.
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et al. Qwen technical report. arXiv preprint arXiv:2309. 16609, 2023.
Tamay Besiroglu, Ege Erdil, Matthew Barnett, and Josh You. Chinchilla scaling: A replication attempt. arXiv preprint arXiv:2404.10102, 2024.
Xiao Bi, Deli Chen, Guanting Chen. Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek Ilm: Scaling open-source language models with longtermism. arXiv preprini arXiv:2401.00 102954,,0044 2024.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
Christopher Bishop. Pattern recognition and machine learning. Springer google schola, 2:5-43, 2006.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D) Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Ethan Caballero, Kshitij Gupta, Irina Rish, and David Krueger. Broken neural scaling laws. arXiv preprint arXiv:2210. 14891, 2022.
Daoyuan Chen, Yilun Huang, Zhijian Ma, Hesen Chen, Xuchen Pan, Ce Ge, Dawei Gao, Yuexiang Xie, Zhaoyang Liu, Jinyang Gao, Yaliang Li, Bolin Ding, and Jingren Zhou. Data-juicer: A one-stop data processing system for large language models. In International Conference on Management of Data, 2024a.
Mayee Chen, Nicholas Roberts, Kush Bhatia, Jue Wang, Ce Zhang, Frederic Sala, and Christopher Ré. Skill-it! data-driven skills framework for understanding and training language models. Advances in Neural Information Processing Systems, 36, 2024b.
Ian Connick Covert, Wenlong Ji, Tatsunori Hashimoto, and James Zou. Scaling laws for the value of individual data points in machine learning. In Forty-first International Conference Machine
Learning.
Sumanth Doddapaneni, Gowtham Ramesh, Mitesh Khapra, Anoop Kunchukuttan, and Pratyush Kumar. primer pretrained multilingual language models. arXiv preprint 2rXiv:2107.000 2021.
Harris Drucker. Improving regressors using boosting techniques. Icml, volume 97, pp. e115. Citeseer, 1997.
11

---

Published as a conference paper at ICLR 2025
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts. In International Conference Machine Learning, pp. 5547-5569. PMLR, 2022.
Simin Fan, Matteo Pagliardini, and Martin Jaggi. Doge: Domain reweighting with generalization estimation, 2024.
Elias Frantar, Carlos Riquelme, Neil Houlsby, Dan Alistarh, and Utku Evci. Scaling laws for sparsely-connected foundation models. arXiv preprint arXiv:2309008520, 2023.
Robert M French. Catastrophic forgetting in connectionist networks. Trends in cognitive sciences, (4):128-135, 1999.
Paul Friedl. Dis/similarities in the design and development of legal and algorithmic normative systems: the case of perspective api. Law, Innovation and Technology, 15(1):25 2023.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.000
2020.
Sachin Goyal, Pratyush Maini, Zachary C Lipton, Aditi Raghunathan, and J Zico Kolter. Scaling laws for data filtering-date data curation cannot be compute agmostic. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22702-227 2, 2024.
Yuxian Gu, Li Dong, Yaru Hao, Qingxiu Dong, Minlie Huang, and Furu Wei. Towards optimal learning language models. arXiv preprint arXiv:2402. 17759, 2024.
Shangmin Guo, Yi Ren, Stefano V Albrecht, and Kenny Smith. Sample relationship from learning dynamics matters for generalisation. arXiv preprint arXiv:2401.0888 2024.
Suchin Gururangan, Ana Marasović, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, and Noah A Smith. Don't stop pretraining: Adapt language models to domains and tasks. arXiv preprint arXiv:2004. 10964, 2020.
Suchin Gururangan, Margaret Li, Mike Lewis, Weijia Shi, Tim Althoff, Noah Smith, and Luke Zettlemoyer. Scaling expert language models with unsupervised domain discovery. arXiv preprint arXiv:2303. 14177, 2023.
Tatsunori Hashimoto. Model performance scaling with multiple data sources. In International Conference on Machine Learning, pp. 4107-4116. PMLR, 2021.
David Haussler. Quantifying inductive bias: Ai learning algorithms and valiant's learning framework.
Artificial intelligence, 36(2):177-2211 1988.
Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B Brown, Prafulla Dhariwal, Scott Gray, et al. Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010. 14701, 2020.
Danny Hernandez, Jared Kaplan, Tom Henighan, and Sam McCandlish. Scaling laws for transfer. arXiv preprint arXiv:2102.0129.0 2021.
Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou. Deep learning scaling is predictable, empirically. arXiv preprint arXiv: 2017.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556 2022.
Kurt Homik,, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, 265):359-366, 1989.
Mareus Hutter. Learning curve theory. arXiv preprint 1rX04:2102.00074, 2021.
12

---

Published as a conference paper at ICLR 2025
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023.
Feiyang Kang, Yifan Sun, Bingbing Wen, Si Chen, Dawn Song, Rafid Mahmood, and Ruoxi Jia. Autoscale: Automatic prediction of compute-optimal data composition for training llms. arXiv preprint urXiv:2407.20117, 2024.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint ar arXiv:2001.08.0 2020.
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho. Agnieszka Grabska- Barwinska, et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114 (13):3521-3526, 2017.
Aran Komatsuzaki, Joan Puigcerver, James LeeTTh Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upcycling: Training mixture-of- experts from dense checkpoints. arXiv preprint arXiv:2212.05055. 2022.
Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Mare Marone, Christopher Akiki, Jia Li, Jenny Chim, et al. Starcoder: may the source be with you! arXiv preprint arXiv:2305.0616/1 2023.
Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing Jiang, and Min Lin. Regmix: Data mixture regression for language model pre-training. arXiv preprint arXiv:2407.00 0.9422, 2024.
Xiaoran Liu, Hang Yan, Chenxin An, Xipeng Qiu, and Dahua Lin. Scaling laws of rope- based extrapolation. In The Twelfth International Conference on Learning Representations, 2023.
Shayne Longpre, Gregory Yauney, Emily Reif, Katherine Lee, Adam Roberts, Barret Zoph, Denny Zhou, Jason Wei, Kevin Robinson, David Mimno, et al. A pretrainer's guide to training data: Measuring the effects of data age, domain coverage, quality, & toxicity. arXiv preprint WrXiv:2305.13165 2023.
Yun Luo, Zhen Yang. Fandong Meng, Yafu Li, Jie Zhou, and Yue Zhang. An empirical study of catastrophic forgetting in large language models during continual fine-tuning. arXiv preprint AXXv:23008088747 2023.
Roger Mead. The design of experiments: statistical principles for practical applications. Cambridge university press, 1990.
Eric Michaud, Ziming Liu, Uzay Girit, and Max Tegmark. The quantization model of neural scaling. Advances in Neural Information Processing Systems, 36, 2024.
Sören Mindermann, Jan M Braumer, Muhammed T Razzak, Mrinank Sharma, Andreas Kirsch, Winnie Xu, Benedikt Höltgen, Aidan N Gomez, Adrien Morisot, Sebastian Farquhar, et al. Prioritized training on points that are learnable, worth learning, and not yet learnt. In International Conference on Machine Learning, pp. 15630-15649, PMLR, 2022.
Niklas Muennighoff, Alexander Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin A Raffel. Scaling data-constrained language models. Advances in Neural Information Processing Systems, 36, 2024.
OpenAl. Gpt-4 technical report. arXiv preprint
2023.
Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay. The refinedweb dataset for falcon Ilm: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.011 / 2023.
13

---

Published as a conference paper at ICLR 2025
Tomer Porian, Mitchell Wortsman, Jenia Jitsev, Ludwig Schmidt, and Yair Carmon. Resolving discrepancies in compute-optimal scaling of language models. arXiv preprint arXiv:2406.19144 19146, 2024.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing with unsupervised learning. 2018.
Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112. 11446, 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21 21(1):5485-555) 2020.
Yunfan Shao, Linyang Li, Zhaoye Fei, Hang Yan, Dahua Lin, and Xipeng Qiu. Balanced data sampling for language model training with clustering. arXiv preprint wrXiv:2402.14526 2024.
Hui Su, Zhi Tian, Xiaoyu Shen, and Xunliang Cai. Unraveling the mystery of scaling laws: Part i. arXiv preprint urXiv22403.06563. 2024.
Tianxiang Sun, Xiaotian Zhang, Zhengfu He, Peng Li, Qinyuan Cheng, Xiangyang Liu, Hang Yan, Yunfan Shao, Qiong Tang. Shiduo Zhang, Xingjian Zhao, Ke Chen, Yining Zheng. Zhejian Zhou, Ruixiao Li, Jun Zhan, Yunhua Zhou, Linyang Li, Xiaogui Yang, Lingling Wu, Zhangyue Yin, Xuanjing Huang, Yu-Gang Jiang, and Xipeng Qiu. Moss: An open conversational large language model. Machine Intelligence Research, 2024. ISSN 2731-5398. doi: 10.1007//11633-0221111
URL htt.httts::/// http:///ttttbbboooo com
wttp:aac..ooopoeooco/c//c/ /ttpoo///eMMMSS//MO//M MMOSS
Ross Taylor, Marcin Kardas, Guillem Cucurull, Thomas Scialom, Anthony Hartshorn, Elvis Saravia, Andrew Poulton, Viktor Kerkez, and Robert Stojnic. Galactica: A large language model for science. arXiv preprint arXiv:2211.0908 2022.
InternLM Team. Internlm: A multilingual language model with progressively enhanced capabilities, 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13977 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023b.
Leslie G Valiant. A theory of the learnable. Communications of the ACM, 27(11):111 11344111344114,, 1984.
VN Vapnik and A Ya Chervonenkis. On the uniform convergence of relative frequencies of events to their probabilities. Theory of Probability and its Applications, 16(2):264, 1971.
Ashish Vaswani, Noam Shazeer, Niki Parmar. Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.. 2017. URL http://///pprceeeiinnneeeeeeeeie FLL-ffaeesssteeeeifoossaaaaaaaaa
Pablo Villalobos. Scaling laws literature review, 2023. URL hwt.hhtsss//pccch..coo
scaling laws literatuue--evviat Accessed: 2024-02-27.
Guillaume Wenzek, Marie-Anne Lachaux, Alexis Conneau, Vishrav Chaudhary, Francisco Guzmán, Armand Joulin, and Edouard Grave. Cenet: Extracting high quality monolingual datascts from web crawl data. arXiv preprint
2019.
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu. Percy S Liang, Quoc V Le, Tengyu Ma. and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. Advances in Neural Information Processing Systems, 36, 2024a.
14

---

Published as a conference paper at ICLR 2025
Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models via importance resampling. Advances Neural Information Processing Systems, 36, 2024b.
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajiwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309. 16039, 2023.
Jun Zhan, Junqi Dai, Jiasheng Ye, Yunhua Zhou, Dong Zhang, Zhigeng Liu, Xin Zhang, Ruibin Yuan, Ge Zhang, Linyang Li, Hang Yan, Jie Fu, Tao Gui, Tianxiang Sun, Yugang Jiang, and Xipeng Qiu. Anygpt: Unified multimodal IIm with discrete sequence modeling, 2024.
Chunting Zhou, Pengfei Liu, Puxin Xu, Srinivasan Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, Lili Yu, al. Lima: Less is more for alignment. Advances in Neural Information Processing Systems, 36, 2024.
Kun Zhou, Yutao Zhu, Zhipeng Chen, Wentong Chen, Wayne Xin Zhao, Xu Chen, Yankai Lin, Ji-Rong Wen, and Jiawei Han. Don't make your Ilm evaluation benchmark cheater. arXiv preprint arXiv:2311 1.01964 2023.
15

---

Published as a conference paper at ICLR 2025
A
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. 2020 OpenAl 2023 Jiang et al. 2023
Touvron et al., 2023b] Sun et al., 2024)). 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. 2024). Such critical impact motivates the development of data curating strategies to reduce computation costs and enhance knowledge (Longpre et al.) 2023). 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. 2019 Penedo et al. 2023). A mass of endeavors in existing practice has involved multifarious filters. scoring the documents with from superficial features on characters (Rae et al. [202]] Xie et al.] [2024b] Raffel et al. 2020) to semantics including similarity to the high-quality reference corpus (Wenzek) jt al.] 2019) and toxicity (Longpre et al. 2023] |Fiedd] 2023). 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. 2022 Touvron et al..
2023a Xie et al. 2024a) and devising data schedules (Mindermann et al.) [0222 Albalak et al. [2023]
Chen ct al. 2024b; Fan et al. 2024). 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 quantitative solution. As
concurrent work, Liu et al. (2024) also proposes to predict optimal data mixtures through regression but assuming rank invariance across training scales.
Scaling laws are functional relationships between the properties of interests (e.g., test loss or other performance metries) and the scales of controllable factors regarding the optimization process or
architecture (e.g., model sizes and numbers of training samples) (Villalobos) (023). Along with
the development of machine learning, characterizing scaling behaviors has garnered great research interest under the context of learning thcorics, bounding the generalization error given the number of
training samples in the form of power laws (Vapnik & Chervonenkis 1971 Valiant 1984 Haussler
1988 Amari et al.) 1992). Nevertheless, overly strict assumptions hinder their practical applications. In recent years, statistical estimation scaling gained fast progress for deep neural networks and
spawns the introduction of scaling laws. Hestness et al. (2017) pioneers the trend and demonstrates power-law generalization error scaling across breadth of factors but the power-law exponents differ from previous theoretical analysis. Kaplan et al. (2020): Hoffmann et al. (2022); Hemighan et al. (2020) conduct more comprehensive investigationn on Transformer architecture (Vaswani et al. [0117). 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.) 2020 Hoffmann
eet al., 2022; OpenAl, 2023). Lately, progressive investigations propose amendments to existing scaling laws (Caballero et al.. 2022 Alabdulmohsin et al. 2022). seeking theoretical explanations on
the empirical formulas Bahri et al. (021) Hutter (2021): Michaud al. (2024), and exploring the
functional relationships in broader scenarios (Hernandez et al. 2021 Frantar et al. 2023 Liu et al.,
2023). The most relevant study ours is Hashimoto (2021) which explores performance prediction under multiple data sources but is limited to small-scaled supervised learning tasks.
B LIMITATIONS AND DISCUSSIONS
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 dellede this paper, similar to related studies (Xie et al.. 2024a] [Che et al.) 2024b; Allalak et al.] 2023 Fan et al., 2024). 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.) 2023 Shao et al. 2024), could further benefit the performance of outcome models. For the validation domains, 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
16

---

Published as a conference paper at ICLR 2025
acknowledge that implicit domain aggregation may be less interpretable compared 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, requires careful design of experiments (Mead] [990] 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. [2] we notice our predictions are smaller than the actual loss, which we attribute to the underestimation from the step laws and model size laws fit. Further practical experience demystifies the technical details of scaling laws (Su et al. 2024 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 joint law of training steps, model sizes, and mixture proportions. Although we can predict the losses with our pipeline, 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. 2017 Kaplan et al. 2020)] Hashimoto 2021] Hoffmann et al., 2022, Frantar et al., 2023). 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 pro- portions for pretraining. More sophisticated data curating can include dynamic proportions (Albalak) it al. 2023) and even a curriculum that changes data domains (Chen et al. (2024b). 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 dynamic data mixing law.
On theoretical understandings. Our data mixing laws, similar to most scaling laws, are empirical findings. We believe theoretical understanding of the training dynamics that form the laws provides a more solid justification. A potential perspective is understanding the target of tuming mixture proportion through gradient estimation (Guo et al., 2024 Gu et al.. (024). 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.) 2022)) and the relationship between data samples is formalized with their gradient directions (Guo al. (0024).
We believe that further investigation into these issues could facilitate more sophisticated quantitative methods for data engineering. We leave them as future work.
17

---

Published as a conference paper at ICLR 2025
C THE RANKING OF DATA MIXTURES DEPEND ON MODEL SIZES AND TRAINING STEPS.
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 com- pare the relative performance of models different sizes and trained with different numbers of steps in Fig.
(A)
70M
15
180M
Ranking
10
305M
10
15
20
Mixture
(B)
Results show that the relative performance fluctuates despite a relatively consistent trend across sizes and training steps. This indicates that a mixture is better small scales but does not always perform better at large scales, consistent with findings Goyal et al. (2024); Covert al. Kang al. (2024). The longest common sequence the partial orders among the 20 mixtures in Fig. [T] and Fig. B) only reaches lengths of 10 and 11, respectively.
10k steps
15
20k steps
Planking
10
30k steps
10
15
20
Mixture Index
Figure 11: 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.
D IMPLEMENTATION DETAILS
D.1 MODEL TRAINING
Throughout this study, we employ the Pythia Suit (Biderman et al. 2023) 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 I IM tokens and a maximum learning rate of le-4. We warm up the learning rates for 2000 steps and decay it 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 learning rate warmup. For the costs of our experiments, it takes around 3.5/8/166 hours to train a 70M/160M/305M/4100 model for 30B tokens on A100 GPUs on our infrastructure.
Table 2: Model architectures for experiments in this paper.
70M
160M
305M
410M
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
12
16
16
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 The Pile contains coarse-grained domains, which are further decomposed into 22 fine-grand tne-ggrined domains. Our experiment in Sec. 33. 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 domains in Sec. 3.3 we adopt the five coarse-graings domains, i.e., academic, internet, prose, dialogues, and misc, where misc include Github and the Deep.ind Mathematies Dataset which are symbolic content. We the coarse-graines domains because it is hard to find five fine-graing domains with sufficient tokens. For the RedPajama, download the version produced and shared by Chen et al. (002a)).
18

---

Published as a conference paper at ICLR 2025
D.2 PREDICTING LANGUAGE MODELING PERFORMANCE WITH SCALING LAWS
In our prediction pipeline in Sec. we adopt nested use scaling laws of training steps and model sizes, which are both power laws, to predict language modeling performance at scale. To fit the laws, we follow Hoffmann et al. 2022) to search over of initialized parameters and fit the samples by minimizing the Huber errors between predictions and observations with LBFGS.
We present our results on verifying the feasibility of applying scaling laws to predict language modeling performance. Our prediction pipeline (described in Sec. f employs two scaling laws: one related to training steps and another to model sizes, to extrapolate performance with increased training data and larger models. We evaluate the precision of predictions for each of these scaling laws, respectively.
Scaling laws of training steps. Fig. 12 shows the training curve 70M models on three different data mixtures. We fit power laws within 30K steps (marked with circles) and extrapolate to predict model performance to as many as 100k steps (marked with stars). On all validation sets, the fitted curves give descent prediction precision, with low mean absolute error of 0.02.
Aratomo
Pissn
Dakge
Symbolo
Ieame:
Total
M
Valand
Moxana
Vakaso
Mekava
1
Vanast
Mosust
Figure 12:
Verification on predicting language modeling performance with scaling laws of training
steps. We train 70M models on three different mixtures up to 100k steps and validate them on 5 validation domains as well as the overall validation mixture. All curves are fitted within 30k steps (+)
and and extrapolated to predict model performance to 100k steps (+)
Scaling laws of model sizes. Fig. 13 shows the results where we fit power laws on 70M, 160M, and 305M models (marked with circles) and extrapolate the curve to predict 410M model performance (marked with stars) at different training steps and under different data mixtures. The results are positive, showing that we can precisely predict the 410M model performance at different training steps, with a mean absolute error of 0.003.
Moien1
Mean?
Mesent Miture
Milin Mitue
Me
Micass
22
Sax
Model dae
Modill dize
Moten
Menee Mirue
Meni
Figure 13: Verification on predicting language modeling performance with scaling laws of model
sizes. We train models of 70M, 160M, and 406M on three different mixtures and validate them on 5 validation domains as well as the overall validation mixture. All curves are fitted with models of 70M, 160M, and 305M (+) and extrapolated to predict the performance of 410M models (+). We verify the predictability at different numbers of training steps.
Overall, we consider fitting power laws to predict model performance for more training steps and larger models are feasible. Therefore we adopt them to implement the scaling law predictions in our pipeline (Sec. 4).
19

---

Published as a conference paper at ICLR 2025
Algorithm 2 Sampling mixture proportions for fitting mixture laws.
Input: Maximum proportions of M domains Pmaz = [fmaz,ly... Mm) where Tmazzi Deargel with D; and Diarget being numbers of available tokens in i-th domain and target number of training tokens, at), minimum proportion respectively, sorted in descending orders (i.c., Iman.1 Fmmor;22 grid size 8, number of mixture to run experiment N.
Output: set mixtures experiment {rn } Candidate mixtures C + GETALLCANDIDATESS ())
1:
Split mixtures with 0 proportion C into and others into C1
3:
Samples (ra m=m
[X/4] from Co and {ra 1 [N/44] from C1
4:
procedure GETALLCANDDATTEXOOmman index i. proportions first domains ri...i-))
6: Candidate mixtures C = 0
if M then
8:
i 1 1
j= rj Taaz, then
9:
nh... 1 Σj
10:
C CU
11:
end
12:
else
+' 5+ "maz.i
14:
for = 0 To [log 5
P- do
Ii t- max 00 /
16:
C C=CCUGETALCANDDDATTE
+ (r....)
17:
end for end if
18:
19:
return
20:
end procedure
D.3
FITTING DATA MIXING 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 extreme case where all sample mixtures have proportions around a small region, it is hardly possible to fit 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 couble-diminisiie 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 domain proportion in any of the training domains take up 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) (997) 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.
the proportion of each domain
E CONNECTIONS BETWEEN IMPLICIT DOMAIN AGGREGATION AND MLP
We first repeat our final mixture law (Eqn.
g here for convenience:
K Σ
M Σ
L(ri..M)
Σ 85 (rL,(ri...M)
88 Gi ki exp
LigTy
where F1...M are mixture proportions on M training domains, L, are validation loss on K implicit domains with s; as their weight in the overall validation set, and ct, byy are other parameters fit.
The mixture law boils down to computation graph in Fig. 14 which contains two layers. The first layers predict the domain losses, while the second sums up the domain losses to ob- tain the overall validation loss. In this way, the mixture law becomes a multilayer perception (MLP) with exponential activation function. In practice, fit the mixture laws with implicit
20

---

Published as a conference paper at ICLR 2025
domain aggregation by fitting multilayer per- ception with exponential activation and applying softmax to the output layer weights. Addition- ally, considering the high variance of ML.,, we further employ AdaBoost Regressor (Drucker) 1997) for fitting the mixture laws to stabilize the predictions and improve their accuracy.
Mintue Proporton
Impket Losass
Es
1860
Inspired by this perspective, we attribute the suc- cessful fitting of data mixing laws two aspects. First, the MLP with a sufficiently large hidden dimension is universal approximator (Hornik ct al. 1989p thus being able to the relation- ships between losses and mixture proportions. Second, the mixture proportions are bounded between 0 and For this reason, predicting an unseen mixture is an interpolation problem, which is usually easier than extrapolation.
The computation graph of mixture law with implicit domain aggregation. We take an case of 3 training domains and 4 implicit validation domains example. The parameters correspond the notations in Eqn. g
Figure 14:
21

---

Published as a conference paper at ICLR 2025
F SUPPLEMENTED RESULTS
F.1 PREDICTION RESULTS ON MORE DOMAINS
To further consolidate the efficacy of data mixing laws and show that they are general for different data, we experiments on domains different from those in Sec. 3.2
We train and validate on Wikipedia, ArXiv, and StackExchange of RedPajama. which are three domains different from those in Sec. 3.2 All samples are from 70M models trained for 10k steps. The prediction accuracy is in Fig. [5] The result shows the predicted and observed losses are consistent for different mixtures. This confirms that data mixing laws also work on domains besides those in the main experiments.
Wikipedia
ArXiv
StackExchange
fitting
fitting
fitting
233
observation
validation
observation
validation
observation
21
validation
0.2
2.1
0.0
24
22
4.
46-
46
13
4.50
475
5/0
52t
650
575
800
200
22
20
22
244
prediction
prediction
prediction
Figure 15: Prediction results domain losses with Eqn. We train 70M models mixtures of Wikipedia, ArXiv, and StackExchange for 10k steps. We fit on mixtures and validate on 3 other mixtures.
F.2 DATA MIXING LAWS MAKE NO DOMAIN-INDEPENDENN INDEPENDENT ASSUMPTIONS
Although our data mixing laws combine the terms with the proportion of different domains through linear combination, we make no Gomain-indeppndent assumption that different domain affects the losses independently. This is because the linear combination serves exponent in Eqn. and Eqn. 7 Specifically, by Taylor expansion, we have
M
M
M
(.(r...ar) cs + hi exp
tasTs
= + Ki
1
tistaryyk + +o²
2
j=1
A=
where there exists interaction terms (jr + k) of different mixture proportions.
Empirically, we evaluate the effectiveness of our data mixing laws in modeling domain in-
Wikipeda.Cormmin Commommrravw
4x.10%
2:175 10*
Reducible Loss on
Reducible Loss on Wikipedi+
teractions by examining their ability to predict
2.15x
language modeling performance when mixing two correlated domains. Specifically, we con-
21 1xPP
struct two synthetic data domains with deliber- ate overlap. The first domain consists of 50% Wikipedia and 50% CommonCraw) data, while
3.7
2,075 10
10
2.05 109
the other domain comprises 50% Wikipedia and
0.4
0.6
0.6 Proportion Wikipedia+AAXiv
0.4
Proportion Wilkiedia-Coommmmrre
50% ArXiv content. this case, increasing the proportion of one domain necessarily increases the shared Wikipedia component. Therefore, the contribution of training domain on target losses coupled with the proportion of the other domain given their joint contribution through Wikipedia. As demonstrated in Fig 16 our pro- posed mixing law (Eqn successfully models the language modeling performance across var- ious mixing ratios of these correlated domains.
Figure 16: Data mixing laws model the lan- guage modeling performance of mixing correlated domains with different proportions. We train 70M models on the mixtures of "Wikipedia+ Common- Crawl" and "Wikipedia+ArXiv" for 15k steps. We validate on the two domains separately and fit the relationship between mixture proportions and vali- dation losses with Eqn. 6
F.3 EXTRA VALIDATION ON SCALING LAWS PREDICTION
22

---

Published as a conference paper at ICLR 2025
We discuss the computation that our prediction method with nested scaling laws requires. In particular, the cost primarily depends on how much scaling laws can accurately extrapolate.
Fifiin Samp
Loss
60000
100000
150000
200000
260000
Sleps
Specifically, we need to train N different data mixtures on model sized N1. .2 .N...., .Nx for So steps to predict the model performance of dif- ferent data mixtures trained with a model with Nearget parameters for Starges steps. The to- tal extra computational overhead relative to di- rect training is N So Σ N; Starget Ne Niarget" where the frac-
0.00
Ener
-0.01
50000
100000
150000
200000
260000
Slep
Figure 17: The scaling law of training steps ac- curately extrapolates to 6.25x more steps. We fit L(S) E + B/S with 40k training steps of IB model and validate the prediction on language modeling performance up to 250k steps.
tion 8o 2 N₁
Starget represents computation saved through scaling law predictions. State-of-the- art scaling law prediction demonstrates that this
fraction can be 1/100 to 1/1000 (OpenAI 2023: Bi et al. 2024 Together with the typical value of
N, which is 20 in experiments, the overall method should require extra 1/5 to 1/50 training computation expectedly.
Given that achieving accurate scaling law pre- dictions remains a developing area, we would like to provide our preliminary investigation to support 100x to 1000x scaling. Fig. [7 shows the scaling prediction of training steps with the scaling law of training steps L(S). where we fit with the first 40k steps and predict the model performance up to 250k steps. This shows that fitting with 40k steps accurately predicts the language modeling performance on 250k steps. which is 6.25x scaling. Additionally, Fig. [8 shows the scaling prediction of model sizes with LNN,, where we fit with models smaller than 100M and find it accurately predicts model per- formance up to 7.25B, which is 72.5x scaling. Combining L(S) and L(N), we may achieve 450x scaling.
(4. BBa+ (4./
Filting Samples
Validation, N-1.51B
3.5-
Validation, N-7,25B
10a
10'
10a 10P 1010
Model size
Figure 18: The scaling law of model sizes accu- rately extrapolates to 70x larger models. We fit language modeling performance at convergence following (Kaplan et al.) 2020) with L(N) = B/N E. The language modeling performance
[18
of 1,5B and 7.25B models are predicted with L(S).
F.4 COMPARISON OTHER DATA MIXING METHODS
We compare method representative data mixing methods, DoGE (Fan et al. 2024) and DoReMi (Xie et al. 2024a). As our experiment in Sec. the Pile.
we train on RedPajama and validation on
DoGE (Fan et al., 2024) contains universal
generalization setting, which assumes validat- ing on the same data as training, and an OOD setting which targets any validation data. We experiment with both of them. For universal
Rocks
Proportio 1
generalization, we refer the data mixture pro- vided by Fan et al. (2024). For the OOD setting, we follow the original paper train a small proxy model (160M) for 10k steps and apply their online updating rule adjust the data mix- ture, shown in Fig. [9| We also follow Fan et al.) (2024) to calculate the average proportions along the training steps of the proxy model the final optimized mixture.
2000
4000
8000
16.K
Step
Figure 19: The evolution of mixture proportions when training the proxy model with the updating rule in the OOD setting of DoGE.
For DoReMi (Xie et al. 2024a), which is only designed for general optimization without awareness of the validation data, we experiment on both its mixture proportion optimized with RedPajama and the Pile. For the mixture optimized with RedPajama, adopt the result of DoReMi 10k from Fan et al.
23

---

Published as a conference paper at ICLR 2025
(2024). For the mixture optimized the Pile, refer to the optimized Pile mixture in the original pa- per (Xie et al. 2024a) and adapt the mixture to the one for RedPajama according to the domain overlap. Specifically, for ArXiv, Wikipedia, Github, and StackExchange, we di- rectly borrow the mixture proportion. CommonCrawl and C4 equally share the proportion of Pile-CC. The propor- tion of Books is obtained as the sum of Books3 and CookCorpus2 in the Pile. We renormalize the proportions of these domains to ensure they sum up to 1.
2.33
2.79
2.75
Default
DoGE
DoGE
DoReMI DoReMi
Ours
(Universal)
(OOD) (RedPajama) (Pile)
Figure 20: Comparisons of the language modeling perfor- mance of different data mixtures. All models are models trained for 100B tokens with the same hyperparameters and validated on the validation set the Pile. Default: original data mixture from Touvron al. (023a). DoGE (Univer- sal): DoGE for universal generalization, obtained from Fan et al. 2024). DoGE (OOD): OOD generalization setting of DoGE optimized with validation set of the Pile. DoReMi (RedPajama): DoReMi mixture optimized by training proxy model on RedPajama. DoReMi (Pile): DoReMi mixture op- timized by training proxy model on the Pile and adapted for our training on RedPajama through the domain overlaps between two dataset. Specific proportions are Fig. 21
Fig. 21 summarizes the final mixture proportion we use for different setups. We train all models for 100B tokens the model size IB. The outcome per- formance is in Fig. 20 which shows the mixture provided by our data mix- ing law indeed archives the lowest val- idation loss.
Default
DoGE (Universal)
DoGE (OOD)
cD (28.90%))
(0755%)
(67
Books (4.50%)
Books (1.85%)
Books 14,33 (4,50%)
Arkiv (8.80%)
(12.11%)
Wispedia (4.50% (21.40%)
Anthv
Wilbedia (0.38%)
StackExchangg Citbub (4-00%
Wikpedia
GlackExchange (11.72%
G4 (15.00%)
Grhub (7 (7%0
BtaokExchange (16.50%)
(29.30%)
Gilbub (7 (7.00%)
DoReMi (RedPajama)
DoReMi (Pile)
Ours
(38.10%)
(41.02%)
Books (30.75%)
(8.20%)
(25
ArXi (4.28%)
Wikipedia 11 (1.56%)
(11.4155
(9.38%%)
StackExchange (12.50)
Wikpedia (23.05%)
cannch Eyrggog
Cithub (14.08%)
udEechhange (8.47%)
C4 (41.02%)
CC (12.00%))
C4 (26.00%))
Figure 21: Specific mixture proportions on Redpajama from different data mixture optimization methods.
24

---

Published as a conference paper at ICLR 2025
G LOss PREDICTION RESULTS WITH NESTED SCALING LAWS
Fig. shows the prediction results nested use of scaling laws in Sec. 4.2 The result demonstrates plausible reference the relative scale of losses both the overall validation data and different validation domains. The optimized mixtures perform better in most domains. While the overall loss helps optimize the overall performance, losses on different domains show model capabilities in various aspects, Our result indicates that, by tuning data mixtures, it is possible to improve specific model capabilities without sacrificing others, consistent with the findings of Xie (2024a)
Overall
Internet
Academic
Observetion
3.3
2.0
3.2
2.6
2.9
2.6
2.8
3.0
3.0
3.2
2.4
2.8
3.0
Prediction
Prediction Dialogue
Prediction
3.3
2.2
3.2
34-
Obsesvation
3.1
3.0
1.8
2.9
3/2
344
344
1.6
2.2
Prediction
fitting data
extrapolation (default mixture)
extrapolation (optimized mixture)
Figure 22: 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 models and 100B tokens.
25


Reviews:
[{"review_id": "JxdhPxOG1N", "reviewer": "Reviewer_c3fa", "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.", "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.\n2. Can you comment on the mixing law fitting well in other settings (such as different subsets of RedPajama/Pile, or a new dataset overall)?\n3. 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?\n\nOverall 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.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "Originality: \n- 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.\n\nQuality: \n- 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.\n- 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.\n\nClarity: \n- Paper was well-structured in starting with results on the simplest mixing law and then gradually moving towards more practical settings.\n\nSignificance:\n- 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.", "weaknesses": "Clarity: \n- 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\".\n- 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. \n- 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.  \n\nQuality: \n- 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.\n- 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). \n- Table 15 in the appendix suggests to me that the mixing laws are not that accurate for the nested approach.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OympsYxPyN", "reviewer": "Reviewer_1wD6", "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”.", "questions": "I would appreciate the authors’ responses as well as plans to address two weaknesses I have described above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "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.", "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.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TYjb9MlYuv", "reviewer": "Reviewer_uTnD", "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.\n\nKey contributions include:\n\n1. 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.\n\n2. 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.\n\n3. Applications in Continual Training: The framework also provides insights for maintaining performance in continual pretraining by dynamically adjusting mixture proportions, thus preventing catastrophic forgetting.", "questions": "See the Weaknesses.\n\nI am open to raising my score if the authors address the above Weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and easy to understand.\n\n2. The research topic on pre-training data mixture is timely and important for the LLM community.\n\n3. 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.", "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.\n\n2. 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].\n\n3. Missing downstream performances. Only perplexity is evaluated, which is not sufficient to demonstrate the effectiveness of the proposed method.\n\n[1] BiMix: Bivariate Data Mixing Law for Language Model Pretraining\n\n[2] Perplexed by Perplexity: Perplexity-Based Data Pruning With Small Reference Models\n\n[3] DoGE: Domain Reweighting with Generalization Estimation\n\n[4] DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tzkbEYQbd3", "reviewer": "Reviewer_dZPn", "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.", "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?\n2. **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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written, where the methodology and experimental results are clearly demonstrated.\n2. The effectiveness from the method is well-justified from extensive experiments on LLM pretraining and continual pretraining.", "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?\n\n2. **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.\n\n3. trivial points: \n    - the footnotes (9,10) are not properly displayed on the paper.\n\n[1] DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining\n[2] DoGE: Domain Reweighting with Generalization Estimation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)