BiMix: Bivariate Data Mixing Law for Language Model Pretraining

(Q1) The design choice of different entropy-based measures and how to derive the mixing proportions from them are hard to follow. Could the authors please elaborate more on this?

In Appendix A, we have provided the specific formulas for the four entropy-based mixtures. Here is a more detailed explanation of the entire process:

1. During model training, we need to sample a batch that includes data from different domains, which requires determining the proportion of each domain.
2. Entropy is commonly understood as a measure of the inherent information contained within data. In probability theory, higher information content indicates greater uncertainty. A domain with higher information content/uncertainty is more challenging to learn and thus requires a larger proportion in the mixture.
3. For each entropy definition employed, we count the frequency of tokens for each domain and compute the overall information represented by that domain using mathematical formulas.
4. Finally, we normalize the entropy values of each domain using the softmax function (as shown in Equation 13) to obtain proportions that sum to one, thereby generating the corresponding mixture.

(Q2) As shown in Figure 6, the downstream task performances of the model trained by different mixtures are fairly low，e.g. mostly under 10%.

This result is expected. We utilized the evaluation code provided by DoReMi and achieved results comparable to those reported in their original paper. For instance, in Table 5 of DoReMi Appendix B, even an 8B model achieves an accuracy of only 4.35% on NaturalQuestions and 6.74% on WebQuestions. It is important to note that these results were obtained using the Exact Match evaluation metric and a one-shot prompt setting. You may have observed higher performance in other papers, but discrepancies in metrics and few-shot settings across different studies make such comparisons inadvisable.

[1] Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance

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

[3] DoGE: Domain Reweighting with Generalization Estimation

We hope that our responses have effectively addressed your concerns and you will consider re-evaluating our paper favorably. If you have any further suggestions or questions, please feel free to contact us. Thank you once again for your time and effort.

We would like to express our sincere gratitude for your comments and queries. Below, we address your feedback regarding Weaknesses (W) and Questions (Q) point by point, and we have also included additional experiments in response to your suggestions.

(W1) The validness of the underlying domain-independent assumption.

We would like to clarify that the domain-independent assumption is not arbitrary; it is grounded in observations of real-world scaling behaviors, which we elaborate upon in Section 3.2. In this subsection, we fixed one of the independent variables (either training steps or domain proportions) and examined the scaling relationship between validation loss and the other variable. We find that the resulting scaling patterns could be effectively modeled using a domain-independent approach. This finding is supported by the visualizations in Figures 1 and 2, as well as experimental results in Sections 5.1 and 5.2.

Methods such as those proposed by Ye et al. [1] employ domain-dependent modeling approaches, which we have shown to involve redundant parameters that adversely affect their practicality. Through theoretical analysis and experimental validation, we confirm the advantages of our modeling in terms of both efficiency and performance:

• Computational Complexity: In Appendix C, we analyzed the relationship between the number of parameters required for fitting the two mixing laws and the number of domains. Our method clearly exhibits a complexity that is an order of magnitude lower than that of Ye et al. This significantly affects the scalability of Ye et al.'s approach when dealing with multiple domains. For example, if we consider mixing 10 data domains ($m=10$), Ye et al.'s mixing law requires fitting 120 parameters, while our BiMix only requires 20 parameters. Consequently, Ye et al.'s method would necessitate training hundreds of models to gather sufficient observation data points for fitting, which is often impractical.

• Mixture Optimization: In Appendix D, we compared the performance of models trained on the mixtures optimized by the two mixing laws. Our BiMix-driven models not only achieve faster convergence (requiring only 50% of the iterations to reach comparable levels) but also demonstrate higher accuracy in downstream task evaluations.

Our findings reveal an intriguing phenomenon: the interference between domains may not be as critical as previously assumed, provided that strict deduplication among domains is maintained. The output of SlimPajama involved rigorous deduplication across the RedPajama domains, thus aligning with our expectations. Additionally, as noted in Table 2, the prediction error of the law on the Pile dataset is slightly larger compared to SlimPajama, attributed to less effective deduplication among the Pile domains.

(W2) Missing baselines.

We have carefully examined the methodologies in the two papers you mentioned. However, the code related to "Perplexed by Perplexity" [2] is not publicly available, so we are only able to include experiments based on the DoGE [3] method. We have updated the comparison results in Figures 4(b) and 7, which continue to demonstrate the advantages of our approach.

(W3) The claim of theoretical insights into data mixing dynamics seems to be an overclaim because there is no theoretical justification.

We apologize for any misunderstanding caused by the language used. Our intent was to convey that the proposed BiMix equation mathematically describes the relationship between validation loss, training steps, and domain proportions. To avoid any potential confusion, we have removed the term "theoretical" to provide a more precise expression.

Dear Reviewer EKGu,

We hope this message finds you well. We have thoroughly addressed each of the weaknesses and questions you raised, and as per your suggestions, we have included additional experiments with more baselines. As the deadline approaches, we wanted to check in to ensure that our rebuttal thoroughly addresses your concerns.

If you find our responses helpful, we sincerely request that you consider increasing your score accordingly and take our efforts and improvements into account in your final decision.

Thank you once again for your invaluable time and effort. Your positive evaluation and constructive feedback are immensely important to us.

Best regards,
Authors

(Q4) The introduction (Line 59) claims improved convergence speed. Could you provide concrete evidence to support this?

Our conclusion regarding faster convergence is based on the analysis presented in Figures 5 and 8(a), where models trained on CE or BiMix mixtures exhibited a more rapid decrease in validation loss compared to those trained using other mixtures. For example, as shown in Figure 8(a), our BiMix-optimized data mixture achieved the same log-perplexity as the Baseline using only 50% of the training steps (80,000 vs. 160,000) required by Ye et al. [1]

(Q5) What is the log perplexity of DoReMi?

The log perplexity values are equivalent to the validation loss; the former is simply a different nomenclature for the evaluation perspective.We chose to present the results this way to maintain consistency with the metrics used in the DoReMi paper and to avoid any potential misunderstandings related to the evaluation metrics. We will clarify this point in the manuscript.

(Q6) Beyond optimizing data proportions, what other practical use cases could the proposed law have?

We have illustrated three practical applications of the proposed law in Sections 5.1 to 5.3. Section 5.1 showcases a fundamental application of scaling laws by allowing us to make reasonable predictions about the model's future performance based on current observations before committing additional computational resources. Sections 5.2 and 5.3 explore further applications directly related to "mixing"—one focuses on predicting the model's performance under novel data mixtures, while the other aims at optimizing the mixture proportions directly.

Minor Suggestions:

1. We explicitly defined $s$ in Equation (2) within the manuscript.
2. We have reviewed the literature you recommended and added comparisons to RegMix [2] and DoGE [3] in Figure 4 and Section 5.3.

[1] Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance

[2] Regmix: Data mixture as regression for language model pre-training

[3] DoGE: Domain Reweighting with Generalization Estimation

We hope that our responses have adequately addressed your concerns and encourage you to consider raising the score. If you have any additional questions or need further clarification, please do not hesitate to reach out. We greatly value your time and effort and are grateful for your support.

We sincerely appreciate the valuable time you have dedicated to reviewing our work and thank you for your insightful feedback. We address your questions one by one and hope that you will consider re-evaluating our work.

(Q1) The logical flow between Sections 3.1 and 3.2.

In Section 3.1, we prominently present our proposed mixing law to provide clarity and prevent readers from getting lost in complex derivations. Section 3.2, on the other hand, offers a detailed discussion of the interpretability of this modeling rather than empirical experiments. This subsection systematically unfolds our observations, thoughts, and the inductive process leading to BiMix, effectively demonstrating that our proposed equation is a credible description of the scaling behavior rather than being an arbitrary or baseless claim.

(Q2) In Sections 5.1 to 5.3, how many and which data points were collected to fit the scaling law and predict which target data point?

As you correctly summarized in Q3, we collect data points from four entropy-based data mixtures to fit the mixing law, as indicated in Lines 407 and 471 of the manuscript. The specific proportion ratios of these mixtures are presented in Tables 4 and 5 of Appendix E.

The calculation of the number of data points is detailed in the “Fitting Details” paragraph (Lines 292 to 299), where the formula is provided as $m \times n \times k$. Here, $m$ represents the number of domains (22 for Pile and 7 for SlimPajama), $n$ is the number of evaluations (20), and $k$ is the number of mixtures. Due to different experimental setups across the three sections, the specific data points used vary slightly:

• In Section 5.1, the focus is on the law's extrapolation capabilities concerning validation loss. For each mixture in Table 1, we hold out the loss at the maximum training step of 200,000 as the validation set, while prior data points are used to fit the law. Consequently, the number of data points used to fit the law is $22 \times (20 - 1)$ for Pile and $7 \times (20 - 1)$ for SlimPajama; the remaining 22 for Pile and 7 for SlimPajama serve as the target predictions.

• In Section 5.2, we assess the law’s generalization to new mixtures. We use all data points from the four entropy-based data mixtures to fit the law, totaling $22 \times 20 \times 4$ for Pile and $7 \times 20 \times 4$ for SlimPajama. The target predictions come from the data points collected during training on the other two mixtures (Baseline and DoReMi), specifically $22 \times 20 \times 2$ for Pile and $7 \times 20 \times 2$ for SlimPajama. For visual clarity, we downsampled the target data points in Figure 3 during presentation.

• In Section 5.3, we also used all data points from the four entropy-based data mixtures to fit the law (there was no need to set aside target data points). The fitted law was subsequently utilized to derive optimized domain proportions.

(Q3) Clarification on entropy-based data mixtures:

(Q3.1) Why not employ a more diverse range of proportions rather than relying solely on entropy-based methods?

We computed the Pearson correlation coefficients among the four entropy-based mixtures. Notably, aside from the correlation between CE and JE, which is only slightly above 0.9, the correlations among SE, CE, JE, and VNE are relatively dispersed. In contrast, DoReMi and Baseline exhibit a correlation of 0.8882, while the correlations among our four entropy-based mixtures are almost all below this value.

(Q3.2) The purpose of Section 5.4.

The use of entropy-based mixtures serves two key purposes. First, they are employed to gather the necessary data points for fitting our law. Second, we aim to extract valuable insights during this process. Given the significant computational resources required for training language models and studying scaling laws, we preferred not to waste resources on random mixtures. In Section 5.4, we contribute our findings to the community: entropy-driven data mixtures can serve as an efficient alternative that, while potentially slightly inferior to the fitted mixing law, offers a solid starting point at a significantly lower computational cost, especially in resource-constrained scenarios.

We sincerely appreciate your acknowledgment of our work and the valuable feedback you provided. Below, we address the identified Weaknesses (W) and Questions (Q) point by point and includes additional experiments in accordance with your suggestions.

(W1) The downstream evaluation includes only PPL.

In addition to PPL, we have incorporated evaluations on downstream tasks such as WebQuestions, LAMBADA, and TriviaQA, maintaining consistency with the DoReMi paper. These results are presented in Figures 4, 6, 7, and 8(b).

(W2) While the authors aim to reduce computational costs, the paper lacks a detailed comparison of its efficiency. For example, how much GPU hours do we need in order to get the good mixture.

In Appendix C, we have provided a comprehensive comparison of the fitting complexity between our method and that of Ye et al. [1] The analyses indicate that our approach is an order of magnitude more efficient. It is important to clarify that once the mixing law is fitted, it no longer relies on GPU resources; instead, mixture optimization can be achieved with just a few seconds of CPU computation.

(W3) The paper primarily contrasts BIMIX with DoReMi and baseline mixtures, without comparisons to other recent advancements, such as REGMIX [2] or Online Data Mixing [3].

We have carefully examined the methodologies in these two papers and have included a comparison with RegMix in Figures 4(a) and 6. While we observe that it demonstrates comparable performance, it is crucial to note the differences and emphasize that the focus of our study is on the mathematical explainability of mixing behavior. We have proposed a bivariate mixing law and illustrated its broader applicability, in contrast to RegMix, which primarily concentrates on optimizing mixtures. Additionally, the Online Data Mixing approach employs a fundamentally different online learning framework, which we were unable to adapt to our framework within the limited time available for this rebuttal.

Questions about section 5.4

The primary objective of the entropy-based mixtures is to gather the necessary data points for fitting our law. However, given the significant computational resources required for training language models and studying scaling laws, we aim to maximize resource utilization to extract valuable insights during this process. In Section 5.4, we contribute our findings to the community to assist practitioners by demonstrating that entropy-driven data mixtures can serve as an efficient alternative. While they may be slightly inferior to the fitted mixing law, they offer a solid starting point at a significantly lower computational cost. We believe this has positive implications for the "data mixing" community, particularly in resource-constrained scenarios.

[1] Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance

[2] RegMix: Data Mixture as Regression for Language Model Pre-training

[3] Efficient Online Data Mixing For Language Model Pre-Training

We sincerely hope that our responses adequately address your comments and encourage you to consider a favorable acceptance of our paper. Please do not hesitate to contact us if you have any further questions. Thank you once again for your helpful assessment!

Dear Reviewer VKW2,

We greatly appreciate your time in reviewing our work and truly value your thoughtful consideration. We have thoroughly addressed the misunderstandings you initially identified. Specifically, we clarified that our evaluation is not limited to PPL but also includes downstream tasks (W1), and we have detailed in the appendix how both theoretical and experimental evidence support the efficiency of our method (W2). In response to your suggestions, we have also conducted more experimental comparisons (W3). Furthermore, we have provided detailed explanations to resolve your concerns regarding Section 5.4.

We have made every effort to address each of your comments and sincerely hope that you will consider these improvements when making your final decision. If you have any further questions or concerns, please do not hesitate to let us know, and we will do our utmost to respond promptly.

Best regards,
Authors