General Response

We sincerely thank the reviewers for their thoughtful feedback and careful evaluation of our work. We are grateful that our paper was well received with positive scores (5, 5, 4, 3). Particularly, we are encouraged that the reviewers appreciate the novelty of our approach (6aWf, H7gH, bchD, FwDJ), the importance of the problem tackled in this work (6aWf, H7gH, bchD, FwDJ), and the performance of MFMS-GP (H7gH, bchD, FwDJ).

Reviewer bchD raised valid concerns regarding our use of the simulator and existing Bayesian Optimization (BO) methods. In response, we would like to take this opportunity to clarify and highlight what we view as the main contributions of our work:

• (a) providing the simulator as a testbed for methodological development (echoing Reviewer 6aWf’s comment)
• (b) proposing a sequential decision making framework that naturally describes data mixture optimization.

(a) A clear bottleneck of developing methods for data mixture optimization is the high cost of training new LLMs each time a method proposes a different mixture. To make progress, the scientific community needs a high-fidelity but cheap setting in order to develop effective methodologies. Therefore, we decided to build a novel high-fidelity simulator suitable for testing and benchmarking principled algorithms for data mixing. We believe this simulator will be a useful resource for the community, supporting future advancements in this area. We would also like to note that the use of a simulator is a common approach in the study of hyper-parameter optimization where evaluation is expensive [1, 2].

(b) Our motivation for proposing the sequential decision framework stems from the current status-quo in the literature: the existing approaches to data mixing optimization rely on (i) ad-hoc data curation approaches and (ii) standard scaling laws that perform deterministic extrapolation. While using BO for the sequential decision framework may seem natural, its actual effectiveness has not been established in prior work, leaving practitioners either unaware of its potential or hesitant to adopt it. We are among the first to rigorously demonstrate the effectiveness of BO for data-mixing optimization, and this verification is far from trivial—it required training 472 models up to 1B parameters. We hope these promising results will spur further exploration of BO for data mixing and other pretraining optimizations, potentially leading to significant savings in both cost and compute.

We hope to address your questions and concerns satisfactorily this round, and are truly thankful for your feedback, which has helped improve the paper in various ways.

[1] Zela et al., Surrogate NAS Benchmarks: Going Beyond the Limited Search Spaces of Tabular NAS Benchmarks, International Conference on Learning Representations, 2022

[2] Eggensperger et al., Efficient benchmarking of hyperparameter optimizers via surrogates, Proceedings of the Twentynineth National Conference on Artificial Intelligence, 2015

Point-by-Point Responses to Reviewer FwDJ

We are thankful for Reviewer FwDJ’s positive perspective on the paper, and are grateful that the reviewer appreciates the novelty of the idea and the performance of our method.

It was unclear to me why the authors used MLPs as predictor models for estimating the performance of particular combinations of model size, dataset proportions etc. MLPs are not commonly used in BO

We apologize for the confusion caused by our exposition. The MLP in fact serves only as a surrogate benchmark on which we validate our BO method (MFMS-GP), and is not used in the BO. Since training a new language model for every hyperpameter sets that BO methods recommend would be prohibitively expensive, we chose to train a MLP model. We then benchmarked MFMS-GP against baselines on the output generated by the MLP.

While the overall R^2 seems good, the results for individual cases vary quite a bit (Table 3)... A (small) additional experiment on real data …

We would like to highlight that the results in Table 3 artificially limits the training dataset available to the MLPs (to demonstrate the predictive value of smaller models’ results), and do not reflect the quality of the MLP that was used as the surrogate benchmark. For the main results of the paper, we used the MLP trained on 422 real language model training runs, and validated the MLP on 50 real validation runs. The MLP achieves a R^2 greater than 0.95 across all metrics validated on real data.

General Response

We sincerely thank the reviewers for their thoughtful feedback and careful evaluation of our work. We are grateful that our paper was well received with positive scores (5, 5, 4, 3). Particularly, we are encouraged that the reviewers appreciate the novelty of our approach (6aWf, H7gH, bchD, FwDJ), the importance of the problem tackled in this work (6aWf, H7gH, bchD, FwDJ), and the performance of MFMS-GP (H7gH, bchD, FwDJ).

Reviewer bchD raised valid concerns regarding our use of the simulator and existing Bayesian Optimization (BO) methods. In response, we would like to take this opportunity to clarify and highlight what we view as the main contributions of our work:

• (a) providing the simulator as a testbed for methodological development (echoing Reviewer 6aWf’s comment)
• (b) proposing a sequential decision making framework that naturally describes data mixture optimization.

(a) A clear bottleneck of developing methods for data mixture optimization is the high cost of training new LLMs each time a method proposes a different mixture. To make progress, the scientific community needs a high-fidelity but cheap setting in order to develop effective methodologies. Therefore, we decided to build a novel high-fidelity simulator suitable for testing and benchmarking principled algorithms for data mixing. We believe this simulator will be a useful resource for the community, supporting future advancements in this area. We would also like to note that the use of a simulator is a common approach in the study of hyper-parameter optimization where evaluation is expensive [1, 2].

(b) Our motivation for proposing the sequential decision framework stems from the current status-quo in the literature: the existing approaches to data mixing optimization rely on (i) ad-hoc data curation approaches and (ii) standard scaling laws that perform deterministic extrapolation. While using BO for the sequential decision framework may seem natural, its actual effectiveness has not been established in prior work, leaving practitioners either unaware of its potential or hesitant to adopt it. We are among the first to rigorously demonstrate the effectiveness of BO for data-mixing optimization, and this verification is far from trivial—it required training 472 models up to 1B parameters. We hope these promising results will spur further exploration of BO for data mixing and other pretraining optimizations, potentially leading to significant savings in both cost and compute.

We hope to address your questions and concerns satisfactorily this round, and are truly thankful for your feedback, which has helped improve the paper in various ways.

[1] Zela et al., Surrogate NAS Benchmarks: Going Beyond the Limited Search Spaces of Tabular NAS Benchmarks, International Conference on Learning Representations, 2022

[2] Eggensperger et al., Efficient benchmarking of hyperparameter optimizers via surrogates, Proceedings of the Twentynineth National Conference on Artificial Intelligence, 2015

Point-by-Point Responses to Reviewer H7gH

We thank the reviewer for your thoughtful comments and for appreciating the novelty of our approach and the significance of the problem.

How many simulation runs are required to obtain a relatively accurate estimation on the optimal data mixture?

We confirmed by grid search that MFMS-GP indeed converged near the optimal data mixture as shown in Figures 3 and 4. Performances on most metrics converged within 500 numbers of evaluations, and all converged near the optimum by 800 evaluations.

... could you provide the comparison results to Regmix [1] and data mixing law [2]...

We would like to highlight that the sequential decision framework under which our method operates affords more flexibility than the explore-then-commit paradigm implicitly adopted by the two referenced approaches. Concretely, it’s evident that in the two mentioned approaches, one needs to arbitrarily determine the scale and the number of the models to be trained to fit a function, and then train one model at the target scale in the end, without clear guidance on which mixture to try next should one have access to more compute. In contrast, our framework naturally considers the selection of samples sequentially from start to end. Therefore, it is more suitable to view the two approaches as possible acquisition functions that one can employ under the framework proposed in this work. That said, we provide direct comparison with the two methods below.

Comparison with Regmix:

Bolded values indicate the best-performing method, with alternatives falling outside one standard deviation from 15 runs. Note that the Regmix results on the three accuracy metrics are considerably worse than Hyperband and MFMS-GP. We suspect the method only works when fitting LightGBM models on losses.

Comparison with data mixing law: We admit that we have difficulties extracting the exact setup (i.e. how many models trained, and at what scales) for the 1B model experiment in the paper. We also tried to look for its exact implementation in code, but to no avail. Our best guess from their description in Sec. 4.2 is: 40 models trained at 70M, and 20 models trained at 160M, 305M, and 410M scales.

Assuming the interpretation is correct, the corresponding number of evaluations comes out to be ~820, which is far beyond the point MFMS-GP (full-scale) converged for most metrics.

...scalability and effectiveness of the MFMS-BO method on larger number of domains has not been explored

Indeed scalability with respect to the number of domains is a key concern. Bayesian optimization approaches are known to work well on the order of 20 dimensions, and we believe the same limitation applies to our case. We note that a high number of domains would likely introduce problems for all optimization methods.

...discussion on the computation complexity of the proposed method

In fact, the computational cost of Bayesian optimization (BO) here is minuscule compared to the cost of training language models (LM). Our BO experiments typically take under 10 minutes on a laptop, whereas training a 1B model with 4x H100s took 42 hours. We believe the cost of learning a regression model and fitting scaling law for Regmix and DML to be similarly small compared to LM training. All in all, we expect the cost of most methods for identifying promising data mixtures to be orders of magnitude smaller than the cost of LLM trainings.

That said, the computational complexity of our method is the same as standard GP-based Bayesian optimization: O(N^3) where N is the number of data points (trained models).

out-of-domain extrapolation results… could help to improve the understanding on its OOD performance

We want to thank the reviewer for highlighting this aspect in which we can improve the experimental setup. We plan to run more experiments to showcase the framework under a different setting, and we will incorporate the suggestion there.

General Response

We sincerely thank the reviewers for their thoughtful feedback and careful evaluation of our work. We are grateful that our paper was well received with positive scores (5, 5, 4, 3). Particularly, we are encouraged that the reviewers appreciate the novelty of our approach (6aWf, H7gH, bchD, FwDJ), the importance of the problem tackled in this work (6aWf, H7gH, bchD, FwDJ), and the performance of MFMS-GP (H7gH, bchD, FwDJ).

Reviewer bchD raised valid concerns regarding our use of the simulator and existing Bayesian Optimization (BO) methods. In response, we would like to take this opportunity to clarify and highlight what we view as the main contributions of our work:

• (a) providing the simulator as a testbed for methodological development (echoing Reviewer 6aWf’s comment)
• (b) proposing a sequential decision making framework that naturally describes data mixture optimization.

(a) A clear bottleneck of developing methods for data mixture optimization is the high cost of training new LLMs each time a method proposes a different mixture. To make progress, the scientific community needs a high-fidelity but cheap setting in order to develop effective methodologies. Therefore, we decided to build a novel high-fidelity simulator suitable for testing and benchmarking principled algorithms for data mixing. We believe this simulator will be a useful resource for the community, supporting future advancements in this area. We would also like to note that the use of a simulator is a common approach in the study of hyper-parameter optimization where evaluation is expensive [1, 2].

(b) Our motivation for proposing the sequential decision framework stems from the current status-quo in the literature: the existing approaches to data mixing optimization rely on (i) ad-hoc data curation approaches and (ii) standard scaling laws that perform deterministic extrapolation. While using BO for the sequential decision framework may seem natural, its actual effectiveness has not been established in prior work, leaving practitioners either unaware of its potential or hesitant to adopt it. We are among the first to rigorously demonstrate the effectiveness of BO for data-mixing optimization, and this verification is far from trivial—it required training 472 models up to 1B parameters. We hope these promising results will spur further exploration of BO for data mixing and other pretraining optimizations, potentially leading to significant savings in both cost and compute.

We hope to address your questions and concerns satisfactorily this round, and are truly thankful for your feedback, which has helped improve the paper in various ways.

[1] Zela et al., Surrogate NAS Benchmarks: Going Beyond the Limited Search Spaces of Tabular NAS Benchmarks, International Conference on Learning Representations, 2022

[2] Eggensperger et al., Efficient benchmarking of hyperparameter optimizers via surrogates, Proceedings of the Twentynineth National Conference on Artificial Intelligence, 2015

Point-by-Point Responses to Reviewer bchD

We thank the reviewer for finding our approach novel and practical. The limitations raised by the reviewer are indeed compromises we made under the constraint of academic budget. We hope that our responses clarify the rationale behind these choices and demonstrate how they do not undermine the key scientific contributions of the work.

...raising questions about generalizability to larger, production-scale models and diverse datasets

To build the simulator that supports the study, we had to train a total of 472 language models, including 68 models at the 1B scale. Unfortunately, training a comparable number of significantly larger models was beyond our available computational budget. Similarly, expanding the number of domains would have required substantially more training runs at each model scale to cover the larger space of data mixtures. We chose to focus on the RedPajama dataset with five training categories as a result. That said, we expect our approach to scale effectively to ~20 domains, as commonly observed for a Bayesian optimization method.

might be a bit too much to … perform BO on the full LLM training ...simulator for validation has serious limitation…

We agree with the reviewer that performing BO without a simulator would be prohibitively expensive, and we also recognize that using a simulator introduces limitations. However, we would like to highlight that use of simulators, sometimes referred to as surrogate benchmarking [1, 2], is a well-established strategy commonly used to study hyperparameter-tuning of expensive-to-train models.

[1] Zela et al., Surrogate NAS Benchmarks: Going Beyond the Limited Search Spaces of Tabular NAS Benchmarks, International Conference on Learning Representations, 2022

[2] Eggensperger et al., Efficient benchmarking of hyperparameter optimizers via surrogates, Proceedings of the Twentynineth National Conference on Artificial Intelligence, 2015

It is kind of natural to treat the data mixing ratio as hyperparameters… current approach largely combines existing BO…

We in fact view the naturalness of the BO approach as a strength. While natural, the approach has been overlooked in the current literature on data mixture optimization. By formulating and empirically validating this framework without developing a tailored BO method, we demonstrate its potential and aim to encourage its broader adoption in future research. We believe this is a meaningful contribution in its own right.

No. The authors claim they have included Limitations in the last section "Conclusions and Future Work", but to me it's not really a discussion of the current work's limitation.

We thank the reviewer for the feedback that has inspired valuable discussions. We'll make sure to include the discussions and thoroughly address the aforementioned limitations with the extra page afforded to us in the camera ready version.

We thank the reviewer for their consideration, and we’re happy that the response addressed your concerns. We plan to release the trained models, as well as all the accuracies, losses and other metadata obtained during training.

General Response

We sincerely thank the reviewers for their thoughtful feedback and careful evaluation of our work. We are grateful that our paper was well received with positive scores (5, 5, 4, 3). Particularly, we are encouraged that the reviewers appreciate the novelty of our approach (6aWf, H7gH, bchD, FwDJ), the importance of the problem tackled in this work (6aWf, H7gH, bchD, FwDJ), and the performance of MFMS-GP (H7gH, bchD, FwDJ).

Reviewer bchD raised valid concerns regarding our use of the simulator and existing Bayesian Optimization (BO) methods. In response, we would like to take this opportunity to clarify and highlight what we view as the main contributions of our work:

• (a) providing the simulator as a testbed for methodological development (echoing Reviewer 6aWf’s comment)
• (b) proposing a sequential decision making framework that naturally describes data mixture optimization.

(a) A clear bottleneck of developing methods for data mixture optimization is the high cost of training new LLMs each time a method proposes a different mixture. To make progress, the scientific community needs a high-fidelity but cheap setting in order to develop effective methodologies. Therefore, we decided to build a novel high-fidelity simulator suitable for testing and benchmarking principled algorithms for data mixing. We believe this simulator will be a useful resource for the community, supporting future advancements in this area. We would also like to note that the use of a simulator is a common approach in the study of hyper-parameter optimization where evaluation is expensive [1, 2].

(b) Our motivation for proposing the sequential decision framework stems from the current status-quo in the literature: the existing approaches to data mixing optimization rely on (i) ad-hoc data curation approaches and (ii) standard scaling laws that perform deterministic extrapolation. While using BO for the sequential decision framework may seem natural, its actual effectiveness has not been established in prior work, leaving practitioners either unaware of its potential or hesitant to adopt it. We are among the first to rigorously demonstrate the effectiveness of BO for data-mixing optimization, and this verification is far from trivial—it required training 472 models up to 1B parameters. We hope these promising results will spur further exploration of BO for data mixing and other pretraining optimizations, potentially leading to significant savings in both cost and compute.

We hope to address your questions and concerns satisfactorily this round, and are truly thankful for your feedback, which has helped improve the paper in various ways.

[1] Zela et al., Surrogate NAS Benchmarks: Going Beyond the Limited Search Spaces of Tabular NAS Benchmarks, International Conference on Learning Representations, 2022

[2] Eggensperger et al., Efficient benchmarking of hyperparameter optimizers via surrogates, Proceedings of the Twentynineth National Conference on Artificial Intelligence, 2015

Point-by-Point Responses to Reviewer 6aWf

We want to thank the reviewer for your encouraging feedback and strong support of our paper. We are especially glad that you found our demonstration of the use of smaller models and earlier training steps interesting.

how does this differ from prior work on multi-information source optimization… if the differences are minimal, consider avoiding the term framework

In making this point, the reviewer points to a paper discussing Multi-Information Source optimization, dealing with a situation in which multiple sources inform the objective function of interest but where each of those sources may provide biased estimates of the true function.

We agree with the reviewer that the idea of having multiple sources informing the true value of interest is not new for Bayesian optimization. At this level of abstraction, the idea indeed encompasses our approach. However, we believe the problem of data mixture optimization introduces a unique and interesting structure that deserves emphasis.

Specifically, there are two key dimensions in which fidelities are introduced: model scale and training steps. The two dimensions are different in a few ways (e.g. significantly more training steps than commonly considered model scales). But most notably, querying an information source at a later training step also yields outputs from all earlier steps as well. However, this feature is not present in the model scale dimension. This asymmetry poses an interesting structure that can potentially be better exploited by developing tailored acquisition functions. We believe this can be an exciting direction for future works.

That said, we are happy to adopt more suitable terminology to reflect better the relationship to prior works while also highlighting the interesting distinctions. We would welcome any suggestions the reviewer may have on alternative phrasing.

... offer the simulator stand out as a contribution

We thank the reviewer for this suggestion. We also think the simulator is one of the key contributions, and will be sure to highlight it in the codebase for the camera-ready version.

We wholeheartedly agree with all the points the reviewer raised. (1) Indeed the novelty lies in the application, and we will adjust the phrasing when listing contributions accordingly. (2) We agree that the phrase “traditional BO” is misleading. Our initial hope was that we could prevent misunderstanding by mentioning multi-fidelity BO in the next paragraph (starting on line 83), but we recognize that it is not optimal. We will revise our exposition in those sections.

For the camera-ready version, we will use the extra page available to more carefully discuss our work in relation to the multi-fidelity BO literature. Thank you for pointing out where our exposition could improve. The discussions brought up by the reviewer are extremely helpful for the revision.