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

Dear Reviewer dZPn,

Thank you for your effort in reviewing our paper and providing valuable suggestions. We summarize your questions and answer them as follows.

Q1: Lack of a detailed analysis on computation overheads

Thank you for your suggestion on analyzing the computation overheads. We include detailed experiments and discussion in Appendix F.3 of our updated manuscript.

The extra computation required in our method is the number of fitting mixture multiplied by the costs to fit scaling laws for each mixture. We aim to optimize the data mixture of a $N_{target}$-sized model trained for $S_{target}$ steps, whose training computation can be estimated as $6N_{target}S_{target} B$ [1], where $B$ is the batch size. To predict the final results, we train small models of scale $\\{N\_i\\}\_\{i=1\}^K$ on N mixtures for $S_0$ steps. Which involves training computation $6BNS_0\sum_{i=1}^K{N_i}$. Therefore, the additional compuataion overhead would be $\frac{6BNS_0\sum_{i=1}^K{N_i}}{6N_{target}S_{target} B} = N\frac{\sum_{i=1}^K N_i}{N_{target}}\frac{S_0}{S_{target}}$, where $N$ is the multiple introduced by data mixing laws, typically 20 in our experiments, while $\frac{\sum_{i=1}^K N_i}{N_{target}}\frac{S_0}{S_{target}}$ is the computation saved with scaling laws. With this fraction, we see that the compuation budget of our method mostly depends on the extrapolation ability of scaling laws.

Regarding the concrete numbers in practice, state-of-the-art scaling predictions enable 100x to 1000x extrapolation [2,3]. This makes the full pipeline only involve 1/5 to 1/50 extra computation, which is much smaller than the extra compute for the default mixture to match our optimized mixture (i.e., 48%). Admittedly, we do not have an excellent scaling law prediction technique at the time we finish this study, which costs as much as 1/10 of the final computation for each data mixture. Nonetheless, our latest experiment also confirms that L(S) is able to achieve at least 6.25x scaling and L(N) 72.5x scaling, with results we update in Appendix F.3. We plan to update these analyses with more thorough study in later versions.

Note that scaling prediction remains a developing research area[4,5,6], the accuracy and efficiency of our method can be further improved with more investigations into the practice of scaling-law prediction, which would also be a focus of our future work.

[1] Scaling Laws for Neural Language Models

[2] GPT4 Technical Report

[3] DeepSeek LLM Scaling Open-Source Language Models with Longtermism

[4] Unraveling the Mystery of Scaling Laws: Part I

[5] Scaling Laws with Learning Rate Annealing

[6] A Hitchhiker’s Guide to Scaling Law Estimation

Dear Reviewer uTuD,

Thank you for your constructive comments. We summarize the questions you concern about and answer them as follows.

Q1: The validness of the underlying domain-independent assumption. No feature interaction term like r_i*r_j models the interactions between domain proportions

We would like to clarify that we have not made domain-independent assumptions.

Theoretically, Eqn. 1 in fact includes the feature interaction term like $r_ir_j$ as you mention given the linear combination serves as an exponent. Specifically, by Taylor expansion, we have $L_i(r_{1\dots M}) = c_i + k_i\exp{\left(\sum_{j=1}^M{t_{ij}r_j}\right)}=c_i+k_i\left(1 + \sum_{j=1}^M{t_{ij}r_j}+\frac{1}{2}\sum_{j=1}^{M}\sum_{k=1}^Mt_{ij}t_{ik}r_jr_k+o^2\right),$

where there exists cross terms $r_jr_k (j\not=k)$ modeling the interaction between different domains as you mention.

Empirically, we conduct a synthetic experiment, where the two experimented domains are (1) 50% Wikipedia + 50% CommomCrawl and (2) 50% Wikipedia + 50% ArXiv. The two domains are obviously correlated with the 50% wiki data. In this case, as shown in Fig. 16 in Appendix F.2, Eqn. 1 can still model the relationship between mixture proportion and validation losses. Here are the specific results.

We hope the above discussion helps explain the efficacy of our formula better. We are grateful for your comments and include these discussions in Appendix F.2 to improve the clarity of our manuscript.

Q2: Missing baselines. There does not seem to be a sufficient comparison between the proposed method and other data selection methods.

We have compared our method to DoGE[1] and DoReMi[2], which you mention, in Fig. 9. For a more comprehensive comparison, as also suggested by other reviewers, we update the results of DoGE and DoReMi in Fig. 9 with more strictly optimized settings.

For DoGE, we further include their OOD version algorithm that aligns more with our setting. For DoReMi, we include another data mixture that is adapted from the mixture it optimizes specifically for the Pile.

We include the details in Appendix F.4 in our updated manuscript. Here are the results of training 1B models for 100B tokens, which shows the data mixture our data mixing laws provide is still optimal.

$~$

Tab. 1 Comparison of different data mixture optimization methods. Default: the mixture originally used by LLaMa. DoGE (Universal): DoGE for universal generalization, which optimizes on a validation set i.i.d. to training data, we obtain the mixture from [1]. DoGE (OOD): DoGE optimized for a given validation data. DoReMi (RedPajama): DoReMi optimized for RedPajama. DoReMi (Pile): DoReMi mixture adapted from the results of optimizing on the Pile.

$~$

For another paper you mention [3], it targets on sample selection problem instead of data mixing, which are two orthogonal problems. It can be applied together with our as well as other data mixing methods.

[1] DOGE : Domain Reweighting with Generalization Estimation

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

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

Thanks for the answers, especially the Taylor expansion part. This mitigates my main concern regarding the modeling of domain interactions. The additional experiments for more baselines and downstream tasks also make some sense to me. Therefore, I have increased my score to 6.

Dear Reviewer 1wD6,

Thank you for your insightful comments. We are pleased that you find our analysis and thought process enjoyable to read and is convincing. We summarize your concerns and try addressing them below.

Q1: The law is empirical and must be validated quite extensively in terms of their predictive power. The experiments in the paper are reasonable, but are conducted in exactly one setup.

We appreciate your recognition of the thought process we present and suggest it reasonably convinding. We do understand your concerns about limited empirical setups. To further confirm the generality of data mixing laws, we include additional experiments on more data domains.

In specific, we further experiment on Wikipedia, ArXiv, and StackExchange, which are three domains different from our experiments in Sec. 3.2. We train 70M models for 10k steps on different mixtures and fit Eqn. 1 with 7 mixtures and validate on the other 3 mixtures. We visualize our results in Appendix F.1 and present the validation accuracy as follows. The result confirms data mixing laws still hold in these domains. We hope this can further support the generality of our data mixing laws.

$~$

Tab. 1 Loss prediction errors of data mixing laws on the mixture of Wikipedia, ArXiv, and StackExchange.

Q2: Why would we care about a particular form of this law?

Thanks for insightful discussion. We believe many function forms (e.g., polynomials and MLP) can predict the model performance regarding data mixing well as we attribute the prediction accuracy of data mixing laws to that it is actually predicting through interpolation.

Despite many other potential alternatives, a specific form of data mixing laws provide the following benefits:

(1) Simplicity. Including fewer parameters helps using fewer samples to fit the law. It is easier to analyze. For example, our exponential law indicates that the benefit of increasing the proportion of a domain decreases exponentially, which implies the benefits of using a balanced mixture.

(2) Interpretability. The function form of data mixing laws can provide some interpretability on the parameters. For instance, we show the interaction between different domains (Fig. 5) with $t_ij$ in our data mixing law.

For predictive powers, we have also experimented with fitting with MLP directly but does not find an obvious difference in the accuracy compared our empirical formulas. Under the setting of Fig. 7, the validation error with data mixing laws (Eqn. 8) and an MLP are 0.01713 and 0.01727, respectively.

We hope our response helps clarify our ideas and consolidate our work better. We look for your feedback.

Q3: More comparison to the OOD setting of DoGE and DoReMi.

Thanks for your advice to make our comparison more comprehensive. We update further results of DoGE and DoReMi here along with the experimental setup and results in Sec. 4.2 with details in Appendix F.4.

For DoGE, we follow the original paper and train a small proxy model (160M) for 10k steps with their OOD version algorithm and obtain the data mixture as follows. We include the optimization process in Fig. 19.

For DoReMi, our original experiments were conducted with optimization on RedPajama training data. As the reviewers indicate, this may not align optimally with validation on the Pile given that DoReMi serves as a universal optimization method without awareness of validation data.

To provide a more comprehensive evaluation, we leveraged the domain overlap between RedPajama and the Pile datasets to adapt the optimized mixture for the Pile, provided by the DoReMi paper, to our setting. We map the mixture proportions as follows. (1) We borrow the proportion of GitHub, StackExchange, Wikipedia, and ArXiv. (2) We let CommonCrawl and C4 share the proportion of Pile-CC and (3) let the proportion of Books as the sum of BookCorpus2 and Books3. We normalized the proportion to ensure they sum up to 1 and train with RedPajama with the normalized mixture proportions. The specific data mixture is as follows.

The results are as follows, which shows that data mixture given by ours remains optimal.

$~$

Tab. 1 Comparison of different data mixture optimization methods. Default: the mixture originally used by LLaMa. DoGE (Universal): DoGE for universal generalization, which optimizes on a validation set i.i.d. to training data, we obtain the mixture from [1]. DoGE (OOD): DoGE optimized for a given validation data. DoReMi (RedPajama): DoReMi optimized for RedPajama. DoReMi (Pile): DoReMi mixture adapted from the results of optimizing on the Pile.

[1] DOGE : Domain Reweighting with Generalization Estimation

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

Dear Reviewer c3fa,

We sincerely appreciate your recognition of our work, with highlights in novelty, quality, clarity, and significance. We would also like to acknowledge the thorough and constructive feedback provided which help strengthen our work, which we summarize and respond as follows.

Q1: Clarity issues on Algorithm 1.

Algorithm 1 demonstrates how we predict language modeling performance across target model sizes, training steps, and data mixtures, as illustrated in Fig. 1. Formally, our goal is to predict the performance of a model with scale $N_{target}$ trained for $S_{target}$ steps on data mixture $r_{target}$.

The process begins with $N$ different data mixtures for fitting, where for each mixture, we train $K$ models of different scales for $S_0$ steps.

1. For each trained model, we first fit $L(S)$ using its losses from the first $S_0$ steps to predict its performance at $S_{target}$.
2. Then, for each mixture, we use the predicted performance at $S_{target}$ across different model scales to fit $L(N)$ and predict the performance of a model with $N_{target}$.
3. Finally, we use these predictions from $N_{target}$ models trained for $S_{target}$ steps on different mixtures to fit data mixing laws and estimate model performance for any unseen data mixture at the given scale.

We appreciate your suggestion for making the laws used each step more explicit and accordingly updated the notations in Alg. 1 with $L(S)$ and $L(N)$ to improve clarity.

Q2: What does resampling mean in L456?

Regarding L456 in the original manuscript (L458 in the updated version), we detail our process for selecting data mixtures to run Algorithm 1. Since different fitting mixtures can yield varying prediction accuracies on data mixing laws, we use a cross-validation approach with the smallest models (which are computationally efficient) to identify effective data mixtures which helps fit a data mixing law that predict well.

Our selection process involves: (1) initially training on 40 diverse mixtures, (2) sampling 20 mixtures and validating their predictions on the remaining 20, and (3) repeating step (2) and selecting the set of 20 mixtures that minimizes prediction errors on the validation set after fitting.

We then use the selected mixtures to run Algorithm 1 to predict language modeling performance of large scale training on different data mixtures.