Data Mixture Optimization: A Multi-fidelity Multi-scale Bayesian Framework
Authors: Thomson Yen
Andrew Wei Tung Siah
Haozhe Chen
C. Daniel Guetta
Tianyi Peng
Hongseok Namkoong
Venue: NeurIPS 2025 poster ()
Keywords: language models
Bayesian optimization
multi-fidelity optimization
data mixture
transfer learning
adaptive experimentation
model scaling
hyperparameter optimization
Primary area: deep_learning
TL;DR: A multi-scale, multi-fidelity Bayesian Optimization (BO) approach where {data mixtures, model scale, training steps} are adaptively selected, achieving >2.6x speedups compared to multi-fidelity BO and random search baselines.

Abstract: Careful curation of data sources can significantly improve the performance of LLM pre-training, but predominant approaches rely heavily on intuition or costly trial-and-error, making them difficult to generalize across different data domains and downstream tasks. Although scaling laws can provide a principled and general approach for data curation, standard deterministic extrapolation from small-scale experiments to larger scales requires strong assumptions on the reliability of such extrapolation, whose brittleness has been highlighted in prior works. In this paper, we introduce a probabilistic extrapolation framework for data mixture optimization that avoids rigid assumptions and explicitly models the uncertainty in performance across decision variables. We formulate data curation as a sequential decision-making problem–multi-fidelity, multi-scale Bayesian optimization–where {data mixtures, model scale, training steps} are adaptively selected to balance training cost and potential information gain. Our framework naturally gives rise to algorithm prototypes that leverage noisy information from inexpensive experiments to systematically inform costly training decisions. To accelerate methodological progress, we build a simulator based on 472 language model pre-training runs with varying data compositions from the SlimPajama dataset. We observe that even simple kernels and acquisition functions can enable principled decisions across training models from 20M to 1B parameters and achieve 2.6x and 3.3x speedups compared to multi-fidelity BO and random search baselines. Taken together, our framework underscores potential efficiency gains achievable by developing principled and transferable data mixture optimization methods. Our code is publicly available at https://github.com/namkoong-lab/data-recipes.

Data Mixture Optimization: A Multi-fidelity Multi-scale Bayesian Framework
Thomson Yen
Andrew Wei Tung Siah Decision, Risk, and Operations Division Columbia Business School
Decision, Risk, and Operations Division Columbia Business School ty2531@co 1ee311@columiaa edu
andrew siah@coluubbiii edu
Haozhe Chen
Tianyi Peng
Decision, Risk, and Operations Division Columbia Business School
Decision, Risk, and Operations Division
Columbia Business School
haozhe. chen@col cennnccolumia...mmmmm edu
tp28450columiia edu
Daniel Guetta
Hongseok Namkoong
Decision, Risk, and Operations Division Columbia Business School
Decision, Risk, and Operations Division
Columbia Business School
crg21338columbi.. edu
hongseok namkoong@oolummiiee... edu
Abstract
Careful curation of data sources significantly improve the performance of LLM pre-training, but predominant approaches rely heavily on intuition or costly trial-and-error, making them difficult to generalize across different data domains and downstream tasks. Although scaling laws can provide a principled and general approach for data curation, standard deterministic extrapolation from small-scale experiments to larger scales requires strong assumptions on the reliability of such extrapolation, whose brittleness has been highlighted in prior works. In this paper, we introduce a probabilistic extrapolation framework for data mixture optimization that avoids rigid assumptions and explicitly models the uncertainty in performance across decision variables. We formulate data curation as a sequential decision- (data
making vrbbem--mltt-fidliinn multi-scale Bayesian andimization-when mixtures, model scale, training steps) are adaptively selected to balance training cost and potential information gain. Our framework naturally gives rise algo- rithm prototypes that leverage noisy information from inexpensive experiments systematically inform costly training decisions. To accelerate methodological progress, we build simulator based on 472 language model pre-training runs with varying data compositions from the SlimPajama dataset. We observe that even simple kernels and acquisition functions can enable principled decisions across training models from 20M IB parameters and achieve 2.6x and 3.3x speedups compared to multi- multi-fidelity Bayesian optimization and random search baselines. Taken together, our framework underscores potential efficiency gains achievable by developing principled and transferable data mixture optimization methods. Our code is available at http://gtuubb.coommmmoooogggg lab/datt labddata-reciipesss. recipes.
Introduction
Data is the foundational infrastructure upon which all AI systems are built. Scaling data has been key driver of progress in machine learning, particularly in language model training (Deng et al., 2009;
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Edbily
Uniertiitty
Wiippeia
oxper ment
Fldcity reletion
Redd
Vodd
Data mbrure relddion
Optimal Training
Ora vidun
Data Distribution
Evalusto Experiment
Wkipedia
Upcate posterior
heration
Distribution with Heuristics- based Filtering
(c) Data mixture adaptive optimization problem
(c) Our multi scale multi fidelity bayesian optimization setup
Our multi-fidelity multi-scale Bayesian optimization (BO) framework. (a) Given un- known optimal training data distribution, (b) existing methods use heuristic- filtering techniques to approximate the optimal distribution. (c) Our algorithm treats data mixture optimization problem. (d) We explore data mixtures in a cost-aware fashion; when we test new data mixture, will observe. Larger models trained for more steps will result high higliidelity observations, but be more expensive. Every point we observe updates our probabilistic belief of model performance over the data mixture, model size, and training steps space, which guides subsequent parameters.
Figure I:
BO
also choose the fidelity of the observation
Hoffmann et al., 2022a: Gadre et al., 2024). While this data-centric approach has led to impressive performance gains, also meurs substantial computational and financial costs in training state-of-the art models (Hoffmann et al., 2022a; Luccioni et al., 2023). Beyond raw scale, the composition of training data has emerged as critical factor in model performance (Albalak al., 2023a; Goyal et al., 2024a). For instance, TinyStories (Eldan and Li, 2023) demonstrated that models with only 10 million parameters, when trained on a carefully chosen synthetic dataset, can generate coherent and consistent English text surpassing the capabilities of significantly larger models like GPT-2 Small (125 million parameters) (Radford et al., 2019). Similarly, Li et (2024) showed that carefully curated dataset enables training -billion-parameter model comparable in performance to that of Mistral-7Bv0.3 and Ilama 8B while using six times less compute. In practical scenarios where heterogencous data sources available, the choice of training mixture has been shown to significantly impact model performance.
Pitfalls of Prior Approaches: This growing recognition of data composition's 's importance has led institutions develop proprietary data mixtures based domain expertise and empirical observations (Radford et al., 2021; Jiang et al., 2023; OpenAI, 2024), Others have introduced heuristics, such as data selection, refine training mixtures (Thrush et al., 2024; Blakeney et al., 2024). However, these ad hoc approaches are often tailored to specific training datasets and downstream tasks, and may fail to trunsfer across domains and data types. For instance, when organizations in specialized sectors such healthcare or finance seek to train custom language models on proprictary datasets, it remains unclear whether heuristics developed for public datasets still effective. Given the substantial resources required for training eigh-performance language models, there a pressing need for principled framework to address data mixture optimization.
Wikipedia upsampling and perplexity-guide
Another line of research attempts to deterministically identify the functional relationship between data composition and model performance (Ye et al., 2024; Ge al., 2025a; Liu al., 2025). Data Mixing Laws (Ye et al., 2024), for example, proposes fitting validation losses as an exponential function of linear combinations of data proportions. However, collecting sufficient data to fit parameters such functions at the desired scale and training duration is often computationally prohibitive. As result, these methods inevitably rely on extrapolating the parameters learned from smaller models trained for fewer steps. Numerous studies have highlighted the brittleness of using such naive extrapolation

---

1000
Curvature
Proess
Oksemabore Riod Exporotia Docr
2
CE Brackrrhhnneeiii weggeee Weight
1.8
2.4
Model
400
Mede Sizee (8)
250 400
8
600
8G0
1000 1000
0.4 0.6 0.8
StackExchangeWeihht
Model Scale (8)
Figure 2: Left: The predicted validation cross-entropy loss on ArXiv data (Shen et al., 2024) as a function of data mixing coefficient and model sizes from data-driven predictor on 472 runs (see details in Sec. 2). Notice the highly non-smooth geometry. Orange dots highlight the optimal data mixture proportion for each model scale. Note that they not consistent across scales. Middle: The curvature at these points shows there are points of high irregularities, suggesting that the relationship between data mixture and model performance is unlikely take simple functional form. Right: A demonstration showing how functional forms like exponential decay fitted on a small number of points would result in high predictive error. In contrast, a probabilistic model such as a Gaussian Process can capture uncertainty over the points.
to guide hyperparameter decisions (Levine et al., 2020; Yang et al., 2021; Jiang et al., 2025). Notably, Jiang et al. (2025) demonstrated that extrapolating validation losses based on small models often leads to inaccurate results. We reinforce this finding in our empirical study, observing that the optimal data mixing proportion does not remain constant across model scales (see Figure. 2)
Existing approaches, whether based on heuristic trial-and-error or deterministic but unverifiable extrapolation, risk failing to identify the optimal data mixture as experiment scale or data domains change. Moreover, they share another fundamental limitation: when additional computational budget becomes available for training new models, these approaches offer no clear guidance on how to select the next promising data mixture to experiment with. To address these challenges, we propose viewing the problem of curating the optimal data mixture as aduptive optimization problem, where practitioners iteratively refine their mixing decisions based on empirical observations from prior experiments. This framework leverages the intuition that model performance exhibits local consistency across similar mixtures, training steps, and model scales, while avoiding rigid assumptions about the global structure of the performance landscape.
Sequential optimization of data mixtures necessitates comprehending which data compositions suffer the highest uncertainty and sharpening beliefs on performance as more observations are gathered. In particular, good adaptive policies must distinguish between aleatoric and epistemic uncertainty: epistemic uncertainty can be reduced with more data, while aleatoric uncertainty irreducible. Measurements must be planned to maximally reduce epistemic uncertainty on future runs by balancing exploration and exploitation.
Our Formulation:
We formulate this sequential optimization framework as a Bayesian optimization (BO) problem: we maintain probabilistic beliefs on the performance of various data mixtures and model scales, and use these beliefs to choose the next model scale to train, on what data mixture, and for how long. Once we fit and evaluate this new model, we use its performance to update our beliefs (Hutter et al., 2011; Falkner et al., 2018; Frazier, 2018).
In standard BO, the cost of each new observation is the same, and we aim to optimize an objective while observing the smallest number of points possible. Our setting presents additional challenges- the cost of training a new model and observing its performance is affected by (1) the number of steps for which the model is trained and (2) the scale of the model (the number of parameters therein).
The number of steps for which a model is trained affects the quality of the observation - the more steps we use to train the model, the more accurately the results will reflect the utility of training on the data mixture in question. Previous work has handled this conundrum using so-called multi- fideliy BO, in which evaluations information revealed during additional training steps will not be worth the expense (Swersky et al., 2014; Domhan al., 2015; Kandasamy al., 2017: Li et al., 2018a).
'stopped early' during the training process if it becomes clear the
Multi- Multi-Fidelity Multi-Scale (MFMS) Bayesian Optimization: Our setting is distinguished by the second factor above we also want to use data gathered on smaller model scales guide our search
3

---

for optimal parameters in larger models. We refer to this new setting as Multi-Fideling Multi-Scale (MFMS) BO. While it can be viewed a special case of the broader framework that optimizes black- box functions with access to cheaper approximate evaluations (Poloczek et al., 2016), the MFMS setting introduces distinctive and interesting structure. Specifically, the model scales dimension differs fundamentally from the usual fidelity dimension associated with training steps. For one, the number of training steps typically far exceeds the number of model scales considered. More importantly, when training for steps, we naturally obtain observations for all intermediate steps up to 2. In contrast, evaluating a model of size m provides no inherent information about the performance of smaller or larger architectures. This raises interesting questions about how appropriately treat and exploit this structure, opening new methodological directions for investigation.
Fortunately, unlike conventional hyperparameters such as learning rate or momentum, where optimal configurations exhibit complex scaling behavior across model sizes (Yang et al., 2021), recent empirical evidence suggests that optimal data mixture compositions enjoy greater transferability from smaller to larger model architectures (Ye et al., 2024; Ge et al., 2025a). This transferability property enables the strategic use of smaller-scale evaluations to identify optimal data mixture configurations that remain effective at target model scales, substantially reducing the computational cost of the optimization process.
The main contributions of the paper are as follows:
We propose a probabilistic extrapolation framework to address the problem of optimizing data mixture for training LLMs. The framework avoids rigid assumptions on the functional dependence between decision variables and model performance by explicitly modeling the performance uncertainty, thereby highlighting the need to strategically experiment with selected data mixtures to minimize the uncertainty in the optimal data mixture at the desired scale.
The framework gives rise to a Multi-Fidelity Multi-Scale BO problem, which provides a principled foundation for developing and evaluating transferable methods for optimizing data mixture. While the BO formalism that treats cheaper but noisy evaluations of the cost function of interest is well established, applying it data mixture optimization introduces distinctive structural features, stemming from the interplays between training steps and model scales, that opens up interesting avenues for future methodological advancements, including batching strategies, asynchronous optimization, custom Gaussian process kernels, and look-ahead methods.
To motivate the MFMS setting, we illustrate the use of smaller model sizes and earlier training steps for informing large-scale training decisions. We empirically show that exploit- ing cheaper evaluations on both dimensions enable more accurate predictions of model's performance at larger scale (Section 2).
To spur methodological progress, we build an empirical testbed based on a simulator trained on 472 language model pre-training runs with varying data compositions from the SlimPajama dataset. We demonstrate the promises of the MFMS BO setting by introducing a couusian-process---e Hyperband (Li et al., 2018a) 8a) and Random Search, which do not consider model scale as decision variable. We find that even this simple approach can better explore different data mixtures and model scales, and deliver the best terminal model (as measured by downstream task performance) at least 2.6x faster than baselines (Section 4).
method. We compare the method against baselines such as
2 Motivation for Multi-fidelity Multi-scale Framework
While deterministically extrapolating optimal data mixtures from small-scale experiments to large- scale models can be unreliable, intuitively, results from smaller models and earlier training steps still provide valuable information that can guide large-scale training decisions. In this section, we experimentally validate two key premises of MFMS BO: (1) smaller models help predict the performance of larger models under various data mixtures, and (2) undertrained models trained for fewer steps over different data mixtures can inform the performance of fully-trained models.
To demonstrate the predictive utility of smaller-scale experiments, we train predictors that take language model training parameters (data mixtures, model scale, training step) as input and

---

Train F1
Test
Dataset
Es
E7
Es
half of 1B runs
remaining 1B runs remaining IB runs
R²
0.69
0.77
0.82
E2
half of 1B runs
R (log)
0.74
0.82
0.85
700M runs
Es
half of IB runs
remaining IB runs
Table 2: Results averaged over runs. Notice the predictive power of our MLP is strongest when it trained on many runs for fewer steps. This provides support for the intuition that one should distribute fixed compute bud- gets across multiple shorter training sessions rather than fewer longer ones.
all smaller runs
E1
half of 700M runs
remaining 700M runs remaining 700M runs
Es
half of 700M runs
500M runs
Table 1: Model size experiments
predict either validation losses or downstream task accuracies of the resulting language model. Our experiments show that training the predictors on results from smaller-scale experiments improves the predictors' ability estimate model performance as a function of data mixture at larger scales.
is important to note that these predictors differ from the language models. The predictors take training hyperparameters input and forecast the resulting language model's performance.
2.1 Collecting Predictor's Data Through LLM Pre-training Runs
To generate the training data for our predictors, pretrained 472 language models using the OLMo 2 package (OLMo et al., 2024) with datasets derived from SlimPajama (Shen et al., 2024), a deduplicated subset of RedPajama (Weber al., 2024). We used five categories of SlimPajama data for pretraining: Wikipedia, StackExchange, Github, ArXiv, and Book. Data from the Common Cravl and C4 categories categories in each run were randomly sampled from a Dirichlet distribution. We pretrained language models varying in size from 20M to IB parameters. We recorded training losses, validation losses across all seven data categories, and evaluated performance on three downstream tasks: HellaSwag. PIQA, and Arc Easy (Zellers al.. 2019 Bisk et al., 2020; Clark al., 2018). The entire dataset was collected using 4x NVIDIA H100 80GB HBM3 for 500 compute days. Additional implementation details for pretraining are provided in Appendix A.
held to simulate out-of-distribution scenarios. The proportions of these five
2.2 Predictor Training
We trained predictors using multilayer perceptrons (MLPs) that take as input the model size, number of training steps, and dataset proportions from the five pretraining categories. These predictors output predictions for training loss, validation losses across seven categories (including the held- out CommonCrava and C4), and downstream task accuracies for HellaSwag, PIQA, and Are Easy. Predictor performance was measured using the coefficient of determination (R²). Complete training details are documented in Appendix B.
2.3 Small Models Help Predict Larger Models Outcomes
We begin by investigating the extent to which smaller model runs can inform the dynamics of larger ones, as the literature on scaling laws (Gadre et al., 2024) would suggest. Table details these experiments, and Table 3 lists the results of the experiment.
We note that, expected, information garnered from training runs smaller models seems considerably increase the accuracy of our predictions on larger models, motivating our hope that carefully crafted optimization algorithm can exploit the relationship.
Unsurprisingly, we note that Eg F. Eg > Ej and Es > Eq: the closer in scale the smaller models the larger model about which wish to make prediction, the more useful the information is. We, therefore, expect our MFMS BO algorithm to 'step through' model scales, starting with small and cheap models to identify promising data mixtures, and then progressing larger models, all the while refining the data mixtures it considers optimal.
2.4 Earlier Training Steps Help Predict Later Training Steps
5

---

30% of our model runs (evenly distributed The second central premise of our approach is that, given fixed compute budget, can be better to attempt many runs for fewer training steps than fewer runs for a larger number of training steps. To test this hy- pothesis, we carried out three additional experiments. In each of these experiments, we attempt to predict the final losses in
Dataset
E1
E2
E3
E4
Es
Wikipedia
0.75
0.96
0.94
0.73
0.88
ArXiv
0,68
0.92
0.93
0.59
0.82
Github
0.66
0.95
0.95
0.62
0.87
Book
0.83
0.97
0.97
0.79
0.92
StackExchange
0.73
0.95
0.95
0.68
0.90
CommonCrawl
0,84
0.98
0.98
0.81
0.94
C4
0,86
0.99
0.98
0.82
0.95
across model sizes). The MLP for each of (1) set of complete runs, one for each model size (2) set of 'truncated" runs, evenly dis- tributed across model sizes. In Es, we use 16 runs truncated at 19600 training steps, in E7, we use 22 runs truncated at 13000 training steps, and in És, we use 32 runs truncated at 8500 steps; thus, these exper- iments are trained on numbers generated with the same FLOPS budget. As expected, Es results in the most performant MLP model, providing evidence for our second central premise.
ArcEasy
0,92
0.94
0.94
0.88
0.90
HellaSwag
0.97
0.98
0.97
0.94
0.96
these experiments is trained
PIQA
0.94
0.96
0.96
0.90
0.93
Table 3: R values of the experiments listed in Table 1, averaged over random seeds. Notice that Eg >> E1 and E3 > Es our ability to predict the performance of larger models is considerably enhanced by insights from smaller models. Note also that Eg Ez; adding information about much smaller models does not seem help.
3 Multi-Fidelity Multi-Scale Bayesian Optimization
Knowing that both smaller models and early stopping can provide valuable insights for optimizing data mixtures, practitioners face a critical dilemma: given a fixed computational budget, how should one allocate resources to train the best model? Should one train many small models to explore different data mixtures, use early stopping at the target scale to abandon poor-performing data mixtures, or adopt hybrid approach combining both strategies? The MFMS setting articulates this dilemma, and in this section, we provide a mathematical formulation of the problem.
We consider having access to a set of datasets D = =D1, D2..... Dn}, and aim to train the best-performing language model, evaluated on a given metric, with m parameters for 2" training steps using T datapoints. The key decision variables are the fractions of the data budget T allocated to each dataset. Specifically, we sample T data points from D,. where w = fwi, W2, .... wn | E A and A denotes n n-dimensional probability simplex. Let μ(w, m, 2) denote the performance of a model with ra parameters trained for training steps using dataset proportions w.
The model's performance is represented by an unknown function y(w, m m", ;=),, which maps a given data mixture w, model size m, and training steps an evaluation metric. Our goal is to solve the optimization problem, arg max p ((w, m ,²),, where m² and 2* represent the target model size and training steps.
We have a budget B with which we can experiment using different values of w, m, and 2. Each evaluation of plw, m, ineurs cost c(m, 2). c(m, 2) is typically an increasing function of and 2, though our framework does not require this. Using m² parameters and 2" steps every time we evaluate new set of weights would quickly exhaust our budget. Instead, therefore, we might probe particular set of weights on a smaller model with m < m² parameters, or with z < x" training steps while the resulting observation u(w, m, 2) would be less informative than p(w, m², ="), it would be considerably cheaper and still provide valuable information.
However, since ux() is unknown function, it is critical to address the uncertainty in its values. BO provides a natural framework that explicitly models uncertainty in ya(-) and systematically refines estimates of uu(-) through posterior updates, one observes more evaluations at different input configurations. This technique is called multi-iddelity BO. See Figure for graphical representation of the setup.
Traditional approaches to fidelity-aware BO primarily address scenarios which the model archi- tecture w remains fixed and only the number of training steps is varied (Swersky et al., 2014; Domhan et al., 2015; Kandasamy et al., 2017; Li et al., 2018a). We add a layer of complexity by also considering model scale. One might be tempted to regard model scale as merely an additional fidelity

---

dimension, otherwise identical to training steps. However, note that there is fundamental distine- tion between the two dimensions: In the course of evaluating model trained for training steps, we must also evaluate that model for all steps z' < 2, whereas no such hierarchical relationship exists for evaluations across different model scales. This structural difference prevents direct application of the prior techniques, and suggests promising avenues for novel methodological developments in multi-fidelity optimization theory, though such extensions lie beyond the scope of our present work.
Experimental Setup
4.1 Evaluation
We propose MFMS BO as natural framework find optimal data mixture in practice. To demon- strate the potential of our framework, we adapt BO algorithm based on Gaussian processes (GPs) to optimize the data mixture. However, training a new language model for each parameter that the algorithm proposes would be prohibitively expensive. Instead, therefore, we train a predictor as described in Section 2.2, and use the predictor to benchmark our approach. We train the predictor on 472 language model training runs described in Section 2.1, where 422 of the runs are randomly selected as a training set and remaining 50 runs as a validation set. The predictor achieves R² > 0.95 across all metrics, suggesting that it is suitable as surrogate benchmark (Eggensperger ct al., 2022; Pfisterer et al., 2022).
Note: the predictor is not necessary for applying our algorithm in real-world scenarios. It serves only as an efficient benchmarking tool.
4.2 Multi-fidelity Multi-scale Gaussian Process (MFMS-GP)
Algorithm 1 Gaussian Process and Elpu
1: Input: Probability space A", model-scale space M, training-step space 2, and cost function d(.-)
2: Initialize Gaussian Process (GP) surrogate model with three RBF kernels over A", .,, and Z
and a linear mean function
3: Randomly sample points from A", .M, and Z to initialize hyperparameters of GP.
4: Initialize history H. with the randomly sampled points
5: for each optimization iteration do
6: for each (m, 2) E M x 2 do
7: Optimize the EI evaluated at these values of m and with respect to gradient descent
8: end for
9: Select the next configuration Anext (Wnext, Mnexts znext) that maximizees the Expected Im-
provement per Unit (Elpu):
EI(A)
Elpu(A) c(m, 2)
10: Evaluate ((Dnex))
Store results in
12: Update posterior of GP with H
end for
14:
Return best 20 = X [0] from configuration X arg max (AA)
We implement a GP surrogate model for our MFMS setting. The kernel of the GP a product of three separate RBF kernels for the data proportion, the model scale, and the training step dimensions. To enable learning the positive correlation between model performance and both model scales and training steps,
linear mean function.
For the acquisition function, we use Expected Improvement (EI). El aims to quantify the expected gain over the current best-observed function value, EI(x)= E [max(y" - fxx), 0)]. where the ex- pectation is taken over the posterior distribution predicted by the surrogate models, and 3 represents the current best-obeerved function value, given by y" := f(xmin) (Frazier, 2018). The El function
7

---

quantifies the expected improvement in the objective value compared to the current best, thereby encouraging the selection of points that are likely to yield better performance.
Equipped with EI, we proceed by optimizing EI over the parameter space to find the most promising point evaluate, using gradient-based methods such as L-BFGS-B (Zhu al., 1997). However, motivated by the fact that the parameter space is discrete over parameter counts (m) and training steps (z), we choose optimize El over each unique tuple (m, Then, to account for the fact that evaluation for each tuple incurs varying costs c(m, z), we evaluate the point that has the greatest EI per unit cost (EIpu) (Lee al., 2020). Additional details of our implementation are provided in Appendix C.
4.3 Baselines Overview
Multi-fidelity Bayesian Optimization: We employ Hyperband, implemented via SMAC (Lindauer et al., 2022), which leverages a random forest surrogate model and expected improvement as its acquisition function. Hyperband (Li et al., 2018b) efficiently explores multiple parameter configura- tions by early stopping poorly performing candidates. In our experiments, we fix the model scale at the target model size throughout optimization, which is not cost-aware.
Random Search: We uniformly sample data proportions and evaluate them at the largest model size and maximum training steps. See Appendix D for the detailed algorithm and implementation specifics.
Results
In this section, we demonstrate the effectiveness of our MFMS-BO framework by benchmarking the simple MFMS-GP method against baselines that do not leverage smaller model scales.
For each of the algorithms evaluated, we run 20 experiments over different seeds, and show the one standard deviation bound with shaded regions. The number of evaluations (the x-axis) was recorded by Wandb (Biewald, 2020) in the experiments we used train our predictor, and normalized so that I evaluation corresponds to the number of flops required to train a IB model for 100 training steps. As an example, in "random search" algorithm, where we randomly sample parameters and train models for 20,000 steps wtih those parameters, doing this with five different set of parameters for IB model would involve 1,000 evaluation units. In a more complex algorithm, where the number of training steps and model scales are dynamically determined, we report the exact number of evaluation steps that would have been required for the various training runs in that algorithm. The MFMS-GP method requires tuning its hyperparameters. As described in Appendix C, we randomly select few configurations to tune these parameters, and the cost of evaluating these configurations accounted for.
Since MFMS-GP relies on GP posterior and potentially noisy EI optimizations to select model scales and training steps, it may take some time before the algorithm samples points at the target scale and fidelity, even after identifying a performant data mixture. If we were to plot only the performance of the best model trained by MFMS-GP at any given evaluation unit, the result could appear deceptively poor simply because the model was smaller or trained for fewer steps. failing to reflect the true quality of the corresponding data mixture. Therefore, to allow us to compare MFMS-GP with other baselines that always train at target model scale, we add an additional curve: MFMS-GP full-scale, which shows the performance one would have gotten if one takes the best configuration MFMS- GP has observed, and simply sets the model scale and training steps to the target m and .* The additional curve can be viewed representing realistic use of our method: fixed compute budget is first allocated for exploration, allowing the algorithm to identify the best data mixture. a subsequent phase, one would take this mixture to train model to completion at the target scale. The additional computational cost of the final training is of course accounted for in these plots.
In Figures and 4, we see that both of the plots for our MFMS- GP algorithm have a 2.6x to 3.3x speedup in finding the configuration that achieves the highest accuracy. Observe that MFMS-GP rapidly explores different data mixtures and identifies promising ones earlier than Hyperband or random search. This expected, since MFMS-GP can utilize inexpensive evaluations by training on smaller models and for fewer steps. These strong results achieved with simple GP- method
8

---

He Vawww Acuraay
PIGA Accursc
ARC Easy
Rccuney (Highe is betted
0.375
0.400
0.70
1.3x spoodup
0,350
0.65
0.45
0.925
0.69
0.40
0.390
0.55
0.35
0.275
0.30
0.50
800
1000
400
0.25
200
1000
Number Evoluations
Pipprrbnn
RunderBearch
MANG-GP
MFMS-GP idl-aals
Best fin valce
Figure 3: On maximizing accuracy in the downstream tasks, our multi-scale multi-fidelity approach achieves more than 2.6x speedup and finds the hest configuration the fastest. Shaded area indicates standard deviation across random seeds.
Common Crowd Cross Ertropy
CB Cross Entropy
Waggeda Croes Eitroyy
Stuk Exchonge Gross Eieropy
Cross Emerony (Lowerr bettee
2.50
1.4
1.8
3.00
2.75
1.6
1.4
1.2
2.50
1.4
up..peddppp
Seppoccup
2,25
0.8
250
500
750
1000
1250
250
500
750
1000
1250
250
500
750
1000
1250
250
500
750
1000
1250
Number of Evoluttions
Olthub Cross Entrapy
Addy Cross Entrapy
Book Cross Eneropy
Crasa Entropy Lowwr ia batted
3.50
2.0
Hyoerburd
3.00
2.75
MFMS.GP
MFNS-GP fullaada
2.25
0.4
1.2
Number
Number
Figure 4: On minimizing the validation cross-entropy losses, our multi-scale multi-fidelity approach achieves more than 2.6x speedup and finds the best configuration the fastest. Shaded area indicates standard deviation across random seeds.
underscore the potential of the MFMS framework in efficiently gathering low-cost yet valuable information to guide large-scale language model training. We expect that thoughtful algorithmic improvements, such as better kernel design acquisition functions, will lead to even more powerful methods. include the results using several standard kernels Appendix E, and leave more extensive explorations of this new setting to future work.
Related Work
Several approaches aim to move beyond heuristic methods for data mixture by leveraging algorithmic techniques. Albalak al. (2023b) propose online data mixing strategy using a non-stochastic bandit algorithm to dynamically adjust data proportions during training, maximizing perplexity. DoReMi (Xie al., 2023) focuses on identifying and emphasizing the "hardest" datasets for base model through distributionally robust language modeling to improve training efficiency. Ge al. (2025b) model a joint scaling behavior of domain proportions and training steps; we push this further through modeling the model scale. Goyal et al. (2024b) delve into the guality-quantin tradeoff in data, exploring how data filtering and repetition affect model performance and introducing scaling laws that account for data utility decay. These works highlight the increasing interest in principled and adaptive methods for data mixture optimization, yet often focus on fixed model scales. In contrast, our multi-fidelity multi-scale approach considers the practical scenario where practitioners have a fixed budget to experiment with data mixture, and can exploit cheap information gathered from smaller models and early training steps.
Data Mixtures.

---

Scaling Laws. Scaling laws provide crucial insights into the relationship between model size, training compute, and performance in large language models (Kaplan al., 2020; Li et al., 2025; Pearce and Song, 2024; Zhao et al., 2025; Mikami et al., 2021; Aghajanyan et al., 2023). Hoffmann et al. (2022b) established foundational scaling laws demonstrating predictable performance improvements with increased compute, model parameters, and training data. (Muennighoff et al., 2023) investigates the impact of data repetition in data-constrained scenarios, showing diminishing returns beyond certain repetition threshold. Ruan et al. (2024) propose observational scaling laws based on "principal capabilities" " to explain and predict language model performance across diverse models and benchmarks, has been done in other BO settings. These scaling law studies motivate our framework by providing an empirical basis for extrapolating performance variations of different data mixtures across model scales.
Bayesian Optimization. Data mixture optimization, like hyperparameter tuning, benefits from efficient search strategies. Approaches range from full configuration selection with methods like BO to hyperband, which employs early termination of unpromising runs. Early BÔ methods (Hutter et al., 2011) used GPs to model the relationship between hyperparameters and model performance. Subsequent works explored random forests (Lindauer et al., 2022) and Parzen estimators (Bergstra et al., 2011) I) as surrogate models. Early stopping techniques like Hyperband (Li et al., 80) focus on efficiently evaluating multiple parameter configurations by progressively eliminating poorly performing candidates, and exploring many combinations with fewer resources. More recent methods like BOHB (Falkner et al., 2018) combine both the BO exploration of Parzen estimators with the multi-fidelity benefits of Hyperband. Our work, WFMSS-BO, is the first to explore multi-fidelity multi-scale approach for data mixture optimization.
7 Conclusion and Future Work
This work introduces principled framework, optimizing data mixture compositions in large language model training, a critical challenge in modern Al system development, Our framework unifies recent advances in predicting optimal data mixtures across scales with classical multi-fidelity BO techniques. Based on this unified framework, we implemented the GP using the RBF kernels and intpeeteddimmoooemeettteeeee to balance the information gain and the cost of exploring new points in the functional landscape. We find that the method achieves optimal downstream task performance 2.6 times faster than traditional multi-fidelity approaches by strategically exploring the joint space of data mixtures and model scales.
multi-scale Bayesian optimization, for
acquisition function
In addition, we empirically demonstrate two key insights that inform future efficient optimization of data mixtures. First, our analysis reveals that training runs on smaller models (below 500M parameters) provide valuable predictive signals for optimizing larger architectures (IB parameters). Second, we establish that partial training runs can effectively inform full-scale training decisions. Specifically, our results show that a combination of full and partial training runs (e.g., 5 complete and 10 heng-length runs) yields better predictive utility than an equal-compute allocation of full training runs alone (e.g., 10 complete runs).
Several promising directions emerge for future research. First, extending our framework more settings, such as language model fine-tuning, data filtering, and diverse collections of datasets with more categories, would help validate its generalizability across different data mixing scenarios. In particular, the low-cost fine-tuning regime provides an opportunity to evaluate the framework without the predictor. From a methodological perspective, incorporating domain knowledge about the positive correlation between model performance and both parameter count and training duration could enhance GP kernel design. Exploring alternative acquisition functions, such as knowledge gradient (Poloczek et al., 2016; Wu et al., 2019), could further improve efficiency in navigating the optimization landscape. The fundamental differences between model scale and training steps as fidelity dimensions also require further investigation to refine their treatment within the framework. Other realistic considerations could enhance the practical benefits of the framework, such as developing algorithms supporting batch evaluations and asynchronous updates for efficient parallel exploration. Addressing these challenges will further strengthen the framework and its applicability in large-scale language model training.
10

---

Acknowledgements
The authors thank Empire AI (Bloom et al., 2025) for generously providing extensive computational resources, without which this study would not be possible. This work was supported by the Digital Future Initiative at Columbia Business School.
References
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248-255, 2009. doi: 10. 200005559922009526665555
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rac, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models, 2022a. httptt/:///aaa htt///axivvvoo///b/// wttg/////////////o/o 15556.
Samir Yitzhak Gadre, Georgios Smymis, Vaishaal Shankar, Suchin Gururangan, Mitchell Wortsman, Rulin Shao, Jean Mercat, Alex Fang, Jeffrey Li, Sedrick Keh, Rui Xin, Marianna Nezhurina, Igor Vasiljevic, Jenia Jitsev, Luca Soldaini, Alexandros G. Dimakis, Gabriel Ilharco, Pang Wei Koh, Shuran Song, Thomas Kollar, Yair Carmon, Achal Dave, Reinhard Heckel, Niklas Muennighoff, and Ludwig Schmidt. Language models scale reliably with over-training and on downstream tasks, 2024. URL https: //appii org/ 2403. 08540.
Alexandra Sasha Luccioni, Sylvain Viguier, and Anne-Laure Ligozat. Estimating the carbon footprint of bloom, a 176b parameter language model. Journal of Machine Learning Research, 24(253): 1-15, 2023. URL btttp://jmaae org/ httpp/pppss//////2// 200///9220099000 htmml.
Alon Albalak, Liangming Pan, Colin Raffel, and William Yang Wang. Efficient online data mixing for language model pre-training 2023a. URL https: / /arxivvvvv.. org/ abs abs/2312. 2312. 02406.
Sachin Goyal, Pratyush Maini, Zachary C. Lipton, Aditi Raghunathan, and J. Zico Kolter. Scaling laws for data filtering -attring-tttaaaeo curation cannot be compute agnostic. In 2024 /EEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 22702-22711 2024a. doi: 10.1
CVPR52733.2024..221
Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english?, 2023. URL https: //arxii/.vvv//// org//og///. htt//bb/23000 07759.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Jeffrey Li, Alex Fang, Georgios Smyrnis, Maor Ivgi, Matt Jordan, Samir Yitzhak Gadre, Hritik Bansal, Etash Kumar Guha, Sedrick Keh, Kushal Arora, Saurabh Garg, Rui Xin, Niklas Muennighoff, Reinhard Heckel, Jean Mercat, Mayee Chen, Suchin Gururangan, Mitchell Wortsman, Alon Albalak, Yonatan Bitton, Marianna Nezhurina, Amro Kamal Mohamed Abbas, Cheng- Yu Hsieh, Dhruba Ghosh, Joshua Gardner, Maciej Kilian, Hanlin Zhang, Rulin Shao, Sarah M Pratt, Sunny Sanyal, Gabriel Ilharco, Giannis Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevic, Sham M. Kakade, Shuran Song. Sujay Sanghavi, Fartash Faghri, Sewoong Oh, Luke Zeulemoyer, Kyle Lo, Alaaeldin El-Nouby, Hadi Pouransari, Alexander T Toshev, Stephanie Wang, Dirk Groeneveld, Luca Soldaini, Pang Wei Koh, Jenia Jitsev, Thomas Kollar, Alex Dimakis, Yair Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. Datacomp-LM: In search the next generation of training sets for language models. In The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024. URL https: http///ppeeeeeeeeeeeeeeeeeeeeeee net/forum id=CNidWn471E.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and llya Sutskever. Learning transferable visual models from natural language supervision, 2021. URI. https:
// //arxiv. org/ /tbe//213.0000

---

Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford. Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang. Timothée Lacroix, and William Sayed. Mistral 7b, 2023. URLhtttss:/arxxi org/ abs/2310 06825.
OpenAl. Gpt-4 technical report, 2024. URL
Tristan Thrush, Christopher Potts, and Tatsunori Hashimoto. Improving pretraining data using perplexity correlations, 2024. URL https: http:: :tttppp://rxxiv..// org httg/bbb/2405 05816..
Cody Blakeney, Mansheej Paul, Brett W. Larsen, Sean Owen, and Jonathan Frankle. Does your data spark joy? performance gains from domain upsampling at the end of training, 2024. URL
attps://arxe org/ abs//406. 03476.
Jiasheng Ye, Peiju Liu, Tianxiang Sun, Yunhua Zhou, Jun Zhan, and Xipeng Qiu. Data mixing laws: Optimizing data mixtures by predicting language modeling performance, 2024. URL
https://arxiv. http://axxxx....aabb/// wwwporgg////////e///eee ....//43... 16952.
Ce Ge, Zhijian Ma, Daoyuan Chen, Yaliang Li, and Bolin Ding. Bimix: bivariate data mixing law for language model pretraining, 2025a. URL https: // /awxivv.vvv///////////a/////atv/ org/ aalaab//22 /2002.455. 14908.
Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing Jiang, and Min Lin. Regmix: Data mixture as regression for language model pre-training. In The Thirteenth International Conference on Learning Representations, 2025. URL https:// openreview
Yoav Levine, Noam Wies, Or Sharir, Hofit Bata, and Amnon Shashua. Limits depth efficiencies of self-attention. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 22640-22651 Curran As- sociates, Inc., 2020. URL https::::// //tpc//ee/n/c/eeeeeeeeeeeeeee hwwprrceeeeinnnnneeeii.p.e..ep neurips. ctt.ccccppaaa meurips.cc/paper (lleffffddd55004000ee22e-leeeeed
files/ willss/apper/2222
Paper. pdf..
Ge Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperpa- rameter transfer. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, Advances in Neural Information Processing Systems, volume 34, pages 17084-17097 Cur- neurips http:/ppprrfflseppppppaaeee Associates, Inc., 2021. URL http::::::oocceeeteete Paper. .apperr..f...... 220////11/0offf222ccc00006002220
Yiding Jiang, Allan Zhou, Zhili Feng. Sadhika Malladi, and Zico Kolter. Adaptive data optimization: Dynamic sample selection with scaling laws. In The Thirteenth International Conference on Learning Representations, 2025. URL httph:ttt////ccc htttpp://openreviest net www..et/ooommmidaaakiio0t
Zhiqiang Shen, Tianhua Tao, Liqun Ma, Willie Neiswanger, Zhengzhong Liu, Hongyi Wang, Bowen Tan, Joel Hestness, Natalia Vassilieva, Daria Soboleva, and Eric Xing. Slimpajama-de: Understand- ing data combinations for I1m training, 2024. URL ptttpp://aaxiiiorrr/bbbse.e
10818.
Frank Hutter, Holger H. Hoos, and Kevin Leyton-Brown Sequential model-based optimization for general algorithm configuration. In Proceedings of the 5th International Conference on Learning and Intelligent Optimization, LION' 05,, page 507-523, Berlin, Heidelberg, 2011. Pegringer-Veellee ISBN 9783642255656. doi: 10. 11110007978--36422 25566-5366668256 3_40.
40. URL https: /ttt:::///oio.rr//00.. org/10.
Stefan Falkner, Aaron Klein, and Frank Hutter. BOHB: robust and efficient hyperparameter optimiza- tion at scale. CoRR, abs/1807.00 80774774, 2018. URL http:///xxvvvvvo http:///wwwww..rr/aa/////0../...
Peter I. Frazier. tutorial bayesian optimization, 2018. URL URL.ttp:://axxiv.orr//b/// 02811.
Kevin Swersky, Jasper Snoek, and Ryan Prescott Adams. Freeze-thaw bayesian optimization, 2014.
URL https: //t///aaaa //arxivvvvvv////// org/ abs tabb//400 3896.
12

---

Tobias Domhan. Jost Tobias Springenberg, and Frank Hutter. Speeding up automatic hyperparameter optimization of deep neural networks by extrapolation learning curves. In Proceedings of the 24th International Conference on Artificial Intelligence, IJCAP115 page 3460-3460 AAAI Press, 2015. ISBN 9781577357384.
Kirthevasan Kandasamy, Gautam Dasarathy, Jeff Schneider, and Barnabás Póezos. Multi-fidelity Bayesian optimisation with continuous approximations. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1799-1808. PMLR, 06-11 Aug 2017. URL nttps://proceeding mlr www.rppess//vee /v00 wpwpes///T/kaaaaaaaaaa//aaaaaa/a
html.
Lisha Li, Kevin Jamieson, Giulia DeSalvo, Afshin Rostamizadeh, and Ameet Talwalkar. Hyperband: A novel bandit-based approach to hyperparameter optimization. Journal of Machine Learning Research, 18(185):1 201 URL
http://irrrrrrrrr/rp//eeeeeee wttpp///rr.pp/ppp/////ea///. 111.///66658888811 .tml..
Matthias Poloczek, Jialei Wang, and Peter I. Frazier. Multi-information source optimization, 2016.
URL wwtphtttt////ttti wwwp:aawwwwaaaaaaaaaaoaaaeooeeaa wwrgg/abs/16e 00389.
Team OLMo, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Shane Arora, Akshita Bhagia, Yuling Gu, Shengyi Huang, Matt Jordan, Nathan Lambert, Dustin Schwenk, Oyvind Tafjord, Taira Anderson, David Atkinson, Faeze Brahman, Christopher Clark, Pradeep Dasigi, Nouha Dziri, Michal Guerquin, Hamish Ivison, Pang Wei Koh, Jiacheng Liu, Saumya Malik, William Merrill, Lester James V. Miranda, Jacob Morrison, Tyler Murray, Crystal Nam, Valentina Pyatkin, Aman Rangapur, Michael Schmitz, Sam Skjonsberg, David Wadden, Christopher Wilhelm, Michael Wilson, Luke Zeutlemoyer, Ali Farhadi, Noah A. Smith, and Hannaneh Hajishirzi. 2 olmo furious, 2024. URL https::
/ /ttpvv..vvvvv./////////////.aav. org/ abs abss/501.... 00656.
Maurice Weber, Daniel Fu, Quentin Anthony, Yonatan Oren, Shane Adams, Anton Alexandrov, Xiaozhong Lyu, Huu Nguyen, Xiaozhe Yao, Virginia Adams, Ben Athiwaratkun, Rahul Chalamala, Kezhen Chen, Max Ryabinin, Tri Dao, Percy Liang, Christopher Ré, Irina Rish, and Ce Zhang. Redpajama: an open dataset for training large language models, 2024. URL https: ptatps://arxvv.. org/ hts/ab///a/ /t/aaabb//11... 2411. 12372.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi. and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. In Whirty-Fourth AAAI Conference on Artificial Intelligence, 2020.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv: 180.000033.... 2018.
Katharina Eggensperger, Philipp Müller, Neeratyoy Mallik, Matthias Feurer, René Sass, Aaron Klein, Noor Awad, Marius Lindauer, and Frank Hutter. Hpobench: A collection of reproducible multi-fidelity benchmark problems for hpo, 2022. URL httpp://aaaa
org/abst /2t///2199 2109.067166..111111111
Florian Pfisterer, Lennart Schneider, Julia Moosbauer, Martin Binder, and Bernd Bischl. Yahpo gym- an efficient mubti-objective multi-fideliin benchmark for hyperparameter optimization. In Isabelle Guyon, Marius Lindauer, Mihaela van der Schaar, Frank Hutter, and Roman Garnett, editors, Proceedings of the First International Conference on Automated Machine Learning, volume 188 of Proceedings of Machine Learning Research, pages 3/1-39, PMLR, 25-27 Jul 2022. URL https: : http////ececeee..e//////// mlr. pwwpss/iee/pppeeeeeeeeee http///////psttere2222 httll
Ciyou Zhu, Richard H. Byrd, Peihuang Lu, and Jorge Nocedal. Algorithm 778: L-bfgs-b: Fortran subroutines for large-scale bound-constrained optimization. ACM Trans. Math. Softw. 23(4): 550-560, December 1997. ISSN 0098-3500. doi: 10. 10.1145/279233.279236 URL https::///oo org/10. 100...114//29992 279236.
Eric Hans Lee, Valerio Perrone, Cédric Archambeau, and Matthias W. Seeger. Cost-aware bayesian optimization. CoRR, abs/2003.10823 2020. URL
10870.
13

---

Marius Lindauer, Katharina Eggensperger. Matthias Feurer, André Biedenkapp, Difan Deng. Carolin Benjamins, Tim Ruhkopf, René Sass, and Frank Hutter. Smac3: A versatile bayesian optimization package for hyperparameter optimization. Journal of Machine Learning Research, 23(54):1-9, 2022. http:// http://immmoorroppeaeeeeeeee/ http://p//..aaaoooo/////////////
.html.
Lisha Li, Kevin Jamieson, Giulia DeSalvo, Afshin Rostamizadeh, and Ameet Talwalkar. Hyperband: A novel bandit-based approach to hyperparameter optimization, 2018b. URL https://arxiv. org/abs/16033 166./06600...
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL attps://wwe wandb. com/. Software available from wandb.com.
Alon Albalak, Liangming Pan, Colin Raffel, and William Yang Wang. Efficient online data mixing for language model pre-training, 2023b. URL buttps://arte
abwgg//ab////..../ 02406.
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoc V.
Le, Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining, 2023. URL https: / //axxaixvv.ooo// /
Ce Ge, Zhijian Ma, Daoyuan Chen, Yaliang Li, and Bolin Ding. Bimix: A bivariate data mixing law for language model pretraining, 2025b. URL https: //t:///ara /arxivv.vvv. org/ ht///4bs// 14908.
Sachin Goyal, Pratyush Maini, Zachary C. Lipton, Aditi Raghunathan, and J. Zico Kolter. Scaling laws for data filtering data curation cannot be compute agnostic, 2024b. URL https://arxis: org/abs/24n 07177.
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, 08361.
2020. URL https: http///axvvvvv////////////a.aava www.riii...g/aaaa.
Margaret Li, Sneha Kudugunta, and Luke Zettlemoyer. (mis)fitting: survey of scaling laws, 2025.
URI.
2202 18969.
Tim Pearce and Jinyeop Song. Reconciling kaplan and chinchilla scaling laws, 2024. URL https: //arxiv. waaarii.ogg//ao ha///bb//22400 12907.
Rosie Zhao, Tian Qin, David Alvarez-Melis, Sham Kakade, and Naomi Saphra. Distributional scaling laws for emergent capabilities, 2025. URL /tttps:/aaxxx org/ abaaabs////00 17356.
Hiroaki Mikami, Kenji Fukumizu, Shogo Murai, Shuji Suzuki, Yuta Kikuchi, Taiji Suzuki, Shin ichi Maeda, and Kohei Hayashi. A scaling law for synthetic-to-reat reall transfer: How much is your pre-training effective?, 2021. URL httts:///rxivv org/ abs/2108. 11018.
Armen Aghajanyan, Lili Yu, Alexis Conneau, Wei-Ning Hsu, Karen Hambardzumyan, Susan Zhang, Stephen Roller, Naman Goyal, Omer Levy, and Luke Zettlemoyer. Scaling laws for generative mixed-modal language models, 2023. URL https: / //tpxv/xvvvv////////////////aava /tt///aavv....aaaaaaaaaaaaaa/aaa http://311.000 03728.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models, 2022b. URL https:// httpss///axxvvvia org/abs/2203. 15556.
Niklas Muennighoff, Alexander M. Rush, Boaz Barak, Teven Le Seao, Aleksandra Piktus, Nouamane Tazi, Sampo Pyysalo, Thomas Wolf, and Colin Raffel. Scaling data-constrained language models, 2023. URL httpst:t::/ htt.t:ttt///a//t// http://aaxvvvvoo/ee org/ htrgg/bbs/3300 2205.. 16264.
Yangjun Ruan, Chris J. Maddison, and Tatsunori Hashimoto. Observational scaling laws and the predictability of language model performance, 2024. URL settps///axxiv.rrg 22.///////./...... 10938.
James Bergstra, Rémi Bardenet, Yoshua Bengio, and Balázs Kégl. Algorithms for hyper-parameter optimization. In Proceedings of the 25th International Conference on Neural Information Process- ing Systems, NIPS' 11, page 2546 2554, Red Hook, NY, USA, 2011. Curran Associates Inc. ISBN
14

---

9781618395993.
Jian Wu, Saul Toscano-Palmerrn, Peter Frazier, and Andrew Gordon Wilson. Practical multi- fidelity bayesian optimization for hyperparameter tuning. CoRR, abs/1903.04703, 2019. URL
hstp::/aarxvv.rgge wttp.//////bbb/0 04703.
Stacie Bloom, Joshua Brumberg, Ian Fisk, Robert Harrison, Robert Hull, Melur Ramasubramanian, Krystyn Van Vliet, and Jeannette Wing. Empire ai: A new model for provisioning ai and hpe for academic research the public good. Practice and Experience in Advanced Research Computing 2025: The Power of Collaboration, PEARC '25, New York. NY. USA, 2025, Association for Computing Machinery. ISBN 97984007 000400777199,, doi: 10.114553708033333366600 URL https: //doi. org/10..... 1145/3700005 3736070.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR, abs/ Jbbs/142..990, 2014. uttpst //api. semantccscholaar
/ttp www.oolaarrrroorooooeeeee Tttg/CorpusID::660 6628106
15

---

Appendix
A Pretraining Runs Details
We use the OLMo 2 OLMo et al. (2024) package for training our language models. The model configurations are
Group
d model
n heads
n layers
Runs
20M
256
115
60M
512
8
71
150M
768
12
12
53
300M
1024
16
16
74
500M
1280
16
16
39
700M
1536
16
16
52
IB
2048
16
16
68
Table 4: Model Architecture Details by Group with Number of Runs
To decide the training mixtures of the experiments, for each experiment, we randomly sample from the probability simplex using the Dirichlet distribution of order (number of training datasets), where we use the parameters a; = 1, Vi [n]. Other training configurations (learning rate, momentum, etc.) are directly taken from OLMo's s configuration files (e.g. 700M) We study the compute optimal regime (Hoffmann et al., 2022a): for each 1B model run, we used 20B tokens in total for training. In the interest of collecting more runs, all other model scales trained on 10B tokens.
B Predictor Training
We train predictors using multilayer perceptrons (MLPs) consisting of three hidden layers with 64 hidden units each, ReLU activations, and dropout with a rate of 0. totaling approximately 5,000 parameters. The predictors accept three types of inputs: (1) the model size (number of parameters), (2) the number of training steps, and (3) the proportions each of the five dataset categories used during pretraining, namely (Wikipedia, Exccknxc, Github, ArXiv, and Book).
For each pretrained language model, the predictor outputs predictions for multiple metrics: the training loss, validation losses across Book, and held-out datasets CommonCrawl and C4), as well as downstream task accuracies on three evaluation tasks: HellaSwag, PIQA. and Arc Easy (Zellers et al., 2019; Bisk et al., 2020; Clark et al., 2018). Thus, each language model corresponds single row in our predictor's 's dataset, comprising 9 inputs described above and outputs spanning these 11 metries.
categories (Wikipedia, StackExchange, Github, ArXiv,
We train the predictors to maximize the coefficient of determination (R²), R
Σ (u 0z) 2. 9
where ys is the true metric value for data point i, 0 is the predictor's estimate, and is the mean of true values. Training of our predictor is conducted for epochs using a batch size of 64, an Adam optimizer with a learning rate of 0,001 and weight decay of 0.01, and data normalization (standard scaling) applied to both inputs and outputs.
C Bayesian Optimization Details
For MFMS-GP, the cost of evaluating a run at a particular model scale is taken from the number of FLOPS the corresponding model scale costs during the pretraining runs. The costs are scaled appropriately such that a unit cost corresponds FLOPS required to train 1B model for training step.
Additionally, since prohibitively expensive to optimize for each 19700 training steps, for multi-fidelig multi-scale GP, we limit the space of training steps to be Z {6000, 12000, 19700).
16

---

To initiate the hyperparameters of MFMS- GP, we randomly select 20 configurations to training step fit the kernel and mean functions' parameters. The GP hyperparameters are trained using the Adam optimizer (Kingma and Ba, 2014) with a 0.1 learning rate for 50 iterations.
As mentioned in Section 4.2, we optimize El within (m, tuple. For this optimization, initiate random probability weights drawn from the Dirichlet distribution and perform a gradient search for greater EI over the probability simplex.
Occasionally, the GP would be too certain of its posterior prediction such that the optimized Els are all small in magnitude. Therefore, when the optimal EI is below certain threshold, lower the length scales of the RBF kernels to encourage more exploration. The threshold set to be le and the length scales would be lowered to 95% of their original values.
As measure to encourage selecting higher cost evaluations later the optimization eycle, instead of using pl ())
[I()
we introduce an additional parameter a that controls the importance of cost,
c(m.s)
and pick the configuration that maximizes EI(A)
<(m,a) Initially 1, and it decays by 1% for every
step of the Bayesian optimization.
D Baselines
D.1 Baselines: Multi-Fidelity Bayes Opt
Algorithm 2 Hyperband with Random Forest, El
Input: Probability space A", Praining-step space Z, target model scale m", and cost function
c(m",-)
2: Initialize random forest surrogate model RF
3: Set initial design with Random Sampling
4: Initialize history H. 0
5: for each Hyperband iteration do
6: Split the total computation budget into brackets
7: for each bracket 8; do
8: Generate initial configurations Wi,..., wn at lowest fidelity = =
9: for each fidelity from 1 to =" do
10:
Evaluate configurations w, at fidelity z
11:
Store results in H.
12:
Fit RF on H.
13:
Select next Aveat using Expected Improvement Update H. with new evaluations
14:
15: end for
16:
end for
17:
end for
18:
Return best -t" (0) from configuration X = arg maxxeH (u())
As baseline for multi-fidelity Bayesian optimization, we use Hyperband implemented by SMAC: Sequential Model-Based Optimization for General Algorithm Configuration (Lindauer et al., 2022). This multi-fidelity Bayesian optimization uses a random forest as surrogate model, expected im- provement the acquisition function, and uses Hyperband (Li et al., 2018b), which is early stopping technique that focuses on efficiently evaluating multiple parameter configurations by pro- gressively eliminating poorly performing candidates, and exploring many combinations with fewer resources. Since the multi fideiiifi framework does not offer a straightforward way to incorporate the additional dimension of model scale, throughout the optimization, we fix the number of the model's parameters to the target model scale m
D.2 Baselines: Random Search
For baseline that does not consider utilizing the fidelity dimension, we consider random search. Random search selects data proportions that are uniformly drawn from our data proportion space. We then run it against the largest model size and training steps.
17

---

E Kernel Comparisons
The kernel function k(r,, x²) GP defines the covariance between different input points. In the main results, we use RBF kernel k(ur, w') = exp( d(w for the data proportions where d(w, w²) I w' 2 is the squared L distance between the two probabilities. In this appendix, we experiment with distance metries that may be more suitable for probabilities.
2
Hellaewag Accuracy
PIQA Aemuacy
ARC Easy Acouncy
0.72
0.000
0.70
0.55
0.17G
0.68
0.20
0.30
0.226
0.300
0.64
0.40
0.700
0.62
0.35
0.250
0.00
200
400
500
800
1000
1200
200
400
000
600
1000
1200
200
400
600
500
1000
1200
Number of Evaluations
Natsher of Ev Evaluations
Namaer of Evalustions
JSD
JSD full-acale
TV full-scale
RBF
full-scale
Figure 5: Comparing different Gaussian process kernels on maximizing accuracy in the downstream tasks.
Crand Entrogy
Entropy
Wiltpedis Eniray
Stack Tocehange Entragy
1.48
1.
27
0.32
Chas Enlmyy
26-
3.78
2.45
100
0.52
23
1.53
0.0
3.30
250
500
750
1000
1250
250
500
750
1000
1250
200
200
730
1000
1350
250
500
750
2000
1250
Nunzber of Evalaaliou
Namber Erabustions
Number of Evaluation
Namber of Evalaation
Gutbab Chos Ealropy
AcXiv Crow Eatropo
Book Croas Entropy
27
3.53
Oran Eaatoge
36
3SD fullosalee
0.65
2.45
25
TV
3.40
2.5
full scale
333
RBP
0.50
2.30
2.2
RBP full-acde
Numbee
Evalaations
Number Brabusttooe
Numbee
Evuluations
Figure 6: Comparing different Gaussian process kernels on minimizing the validation cross-entropy losses
Specifically, we consider the Total Variance (TV) distance and the Jensen-Shannon divergence (JSD), symmetric Kullback-Leible divergence, for the data proportions:
TV: : d(w, w') ||w w'll2
1 KL (w w) + KL (wo ())
1
ISD: : d(ur w²)
where KL(-) denotes the Kullback-Leibler divergence, and w witern
We found that on optimizing accuracy, the simple RBF kemels generally perform better. However, occasionally, JSD (on HcllaSwag of Figure. 5) or TV (on Common Crawl, Wikipedia, and Stack Exchange of Figure. 6) yicld better results.
18

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract claims introduction of new framework and speedup that comes from the formulation, which are both reflected in the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The final section describes limitations in the setting and ways to approach realistic use cases.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
19

---

Answer: [NA]
Justification:
The paper does not have theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: All the code used to produce the experimental results is provided.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
20

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: the code and data are provided.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
Mobllic/guides/Coeeetuee bosoodeSSummssssooooooccy for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
ctips.cc/pbblcc/eeeeccc laaiiccggudee... CoodeSummis ssionPolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The details are provided, either the main text in the appendices.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The results are accompanied by standard deviations.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
21

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The compute resources required are provided.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics attps://neurips. cc/ppblicc/chhicccuueeelleeee
Answer: [Yes]
Justification: There is no potential harmful impact of the research, nor human subjects involved in the experiments.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The paper studies a fundamental research topic and does not have direct application that can harmful.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
22

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper has no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The original creator/owners of the data and models have been properly credited.
Guidelines:
existing assets. The authors should cite the original paper that produced the code package or dataset.
The answer NA means that the paper does
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
23

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
The assets have been documented and anonymized.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/cood/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsoureing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (https : htt////eeerppp/eeeee (tttse//eearrpp..c//offfffeeeee //tp////eeeeeeeee///////////// for what should or should not be described.
25


Reviews:
[{"review_id": "k3aWt3osTf", "reviewer": "Reviewer_6aWf", "summary": "This paper addresses the problem of selecting an optimal training data distribution for large language models (LLMs). Given access to data from multiple sources (e.g., Wikipedia, Books, GitHub, Reddit), the goal is to determine the optimal mixture of these datasets to maximize model performance. To make this problem tractable in practice, the authors adopt a multi-fidelity Bayesian optimization framework.\n\nSpecifically, the authors treat the data mixture, model scale (i.e., number of parameters), and number of training steps as input variables to be jointly optimized. This allows them to adaptively balance training cost with the value of information gained about the performance of a full-scale model.\n\nThe underlying objective is modeled using a Gaussian Process with a kernel constructed as the product of three RBF kernels, corresponding to the three input dimensions. To capture expected positive correlations between model performance and both model scale and training steps, a linear mean function is employed. For decision-making, the authors use the Expected Improvement over Cost acquisition function to guide the selection of new configurations.\n\nOne particularly interesting aspect of this paper is the experimental validation demonstrating that \"low-fidelity\" samples can effectively predict full-fidelity performance. Specifically, the authors show that smaller models can reliably forecast the outcomes of larger models, and that earlier training steps provide useful insights into the results of later training steps by training a predictor using 472 language model pre-training runs with varying data compositions from the SlimPajama dataset.", "questions": "The paper *above* uses a different acquisition function, and you do mention it in the manuscript (I believe [1] corresponds to your reference [54]). However, my question is more about the overall setup. Is the multi-scale aspect of your approach truly different from prior work? I'm not sure I fully understand the distinction.\n\nTo me, the creative application and the empirical results presented are already strong enough to justify a solid contribution. That said, I remain unconvinced that the multi-scale component constitutes a novel contribution on its own. A clearer explanation or comparison with prior multi-fidelity or multi-information source frameworks would help strengthen this point. Alternatively, if the differences are minimal, consider avoiding the term framework — or feel free to clarify why my interpretation may be incorrect.\n\nI did browse the github a little bit but can you improve the documentation to offer the simulator stand out as a contribution? I think other researchers and practioners would appreciate to use the simulator as a benchmark as well.", "limitations": "yes", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "* Strengths* \n\nThis is a creative and compelling application of Bayesian optimization to select data mixtures for training large language models. The paper is well written, easy to follow, and the experiments are thoughtfully designed — the synthetic benchmark in particular stands out as a strong component of the validation.\n\n* Weaknesses* \n\nThat said, the paper would be stronger if the proposed methodology were demonstrated in a real-world application. As it stands, the experiments serve primarily to validate the approach in controlled settings, but do not yet showcase its effectiveness in a practical deployment scenario.\n\nThe use of Bayesian optimization in this context is certainly unique and well-motivated. However, it is not entirely clear that the “multi-scale” aspect of the method is novel. Specifically, how does this differ from prior work on multi-information source optimization, such as:\n\n> [1] M. Poloczek, J. Wang, P. I. Frazier. *Multi-Information Source Optimization*. NeurIPS, 2017.\n\nwhich also handles discrete fidelities in a Bayesian optimization setting?\n\nIf the authors are able to clarify this distinction, I believe this paper would represent a solid and meaningful contribution.", "was_revised": true, "final_justification": "I didn't change my score, I think this paper is an accept. I would bump my score if the authors convinced me that there is a real novelty on the Bayesian Optimization part OR if they present great results on an real application application; I think that would make this paper a solid accept. \n\nIn this version, I see many qualities of a paper that can be accepted: it shows a very clever application of Bayesian optimization for data-mixing for LLM training. That's certainty an interesting topic for several practitioners, and the surrogate model created is also of great interest for many researchers. The paper is very easy-to-read and the exposition is pretty good.", "rebuttal": "# General Response\n\nWe 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). \n\n\nReviewer 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: \n- (a) providing the simulator as a testbed for methodological development (echoing Reviewer 6aWf’s comment) \n- (b) proposing a sequential decision making framework that naturally describes data mixture optimization. \n\n(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]. \n\n(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.\n\n\nWe 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. \n\n[1] Zela et al., Surrogate NAS Benchmarks: Going Beyond the Limited Search Spaces of Tabular NAS Benchmarks, International Conference on Learning Representations, 2022\n\n[2] Eggensperger et al., Efficient benchmarking of hyperparameter optimizers via surrogates, Proceedings of the Twentynineth National Conference on Artificial Intelligence, 2015\n\n# Point-by-Point Responses to Reviewer 6aWf\nWe 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. \n\n>**how does this differ from prior work on multi-information source optimization… if the differences are minimal, consider avoiding the term framework**\n\nIn 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.\n\nWe 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. \n\nSpecifically, 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. \n\nThat 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.\n\n>**... offer the simulator stand out as a contribution** \n\nWe 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."}, {"review_id": "kySERI0d95", "reviewer": "Reviewer_H7gH", "summary": "This paper tackle with a critical and challenging problem of data mixture optimization for LLM pretraining. The authors propose a novel MFMS-BO method to apply bayesian optimization to predict the optimal data mixture composition from multiple data sets from previous training runs with various training steps and model scales. The empirical results on the SlimPajama datasets demonstrate the proposed MFMS BO method can estimate the optimal data mixture for larger-scale models from smaller-scale experiments, which achieves up to 3 times speedup above the random sampling baseline.", "questions": "1. Lack of comparisons to some critical baseline approaches. could you provide the comparison results to Regmix [1] and data mixing law [2], where only various model scales are considered?\n\n2. The benchmark is only limited to SlimPajama with 5 training domains. The scalability and effectiveness of the MFMS-BO method on larger number of domains has not been explored, e.g. The pile (22 domains) or hundreds of domains clustered from document embeddings.\n\n3. Can you provide a discussion on the computation complexity of the proposed method, and compare with traditional methods like DML and RegMix?\n\n4. How many simulation runs are required to obtain a relatively accurate estimation on the optimal data mixture? \n\n5. The current out-of-domain extrapolation results on the MFMS-BO method is limited, since CC and C4 are known as general web-crawled dataset containing diverse topics. Provide the results on more domain-specific datasets (e.g. freelaw, dm-mathematics in The Pile dataset) could help to improve the understanding on its OOD performance and strengthen the paper.\n\n\n[1] RegMix: Data Mixture as Regression for Language Model Pre-training.\n\n[2] Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance.", "limitations": "1. Lack of comparisons to some critical baseline approaches. For example, the Regmix [1] and data mixing law [2], where the optimal data mixture is predicted on experiments across various model scales, should be considered as baselines.\n\n2. The benchmark is only limited to SlimPajama with 5 training domains. The scalability and effectiveness of the MFMS-BO method on larger number of domains has not been explored? e.g. The pile (22 domains) or hundreds of domains clustered from document embeddings.\n\n3. The out-of-domain extrapolation results on the MFMS-BO method is limited, since CC and C4 are known as general web-crawled dataset containing diverse topics. Provide the results on more domain-specific datasets (e.g. freelaw, dm-mathematics in The Pile dataset) could help to improve the understanding on its OOD performance and strengthen the paper.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strength\n1. The paper tackles a critical and challenging problem of data mixture optimization in LLM pretraining.\n2. This paper introduce a critical consideration of data mixture optimization, the *fidelity* (length) of the small-scale training runs, which has been ignored by previous law-fitting based methods, such as RegMix [1] or Data Mixing Law [2].\n3. A novel probabilistic extrapolation framework framework is proposed based on bayesian optimization, which explicitly consider the uncertainty and fidelity of small-scale experiments and derive a more accurate estimation of the optimal data mixture in larger-scale pretraining experiments. The empirical results demonstrate a large speedup compared to random sampling baselines.\n\n### Weakness\n1. Lack of comparisons to some critical baseline approaches. For example, the Regmix [1] and data mixing law [2], where the optimal data mixture is predicted on experiments across various model scales, should be considered as baselines.\n\n2. The benchmark is only limited to SlimPajama with 5 training domains. The scalability and effectiveness of the MFMS-BO method on larger number of domains has not been explored? e.g. The pile (22 domains) or hundreds of domains clustered from document embeddings.\n\n3. The out-of-domain extrapolation results on the MFMS-BO method is limited, since CC and C4 are known as general web-crawled dataset containing diverse topics. Provide the results on more domain-specific datasets (e.g. freelaw, dm-mathematics in The Pile dataset) could help to improve the understanding on its OOD performance and strengthen the paper.\n\n[1] RegMix: Data Mixture as Regression for Language Model Pre-training.\n\n[2] Data Mixing Laws: Optimizing Data Mixtures by Predicting Language Modeling Performance.", "was_revised": true, "final_justification": "This paper tackles with data mixture optimization for LLM pretraining and proposes a novel MFMS-BO method to apply bayesian optimization to predict the optimal data mixture composition from multiple data sets from previous training runs with various training steps and model scales. The empirical results on the SlimPajama datasets demonstrate the proposed MFMS-BO method can estimate the optimal data mixture for larger-scale models from smaller-scale experiments, which achieves up to 3 times speedup above the random sampling baseline.\n\nDuring rebuttal, the authors provide additional results comparing to other data mixing baselines, e.g. RegMIX, which demonstrate slight performance gain over RegMix. I believe that the proposed MFMS-BO method can be an efficient and effective alternative to the traditional scaling law studies.", "rebuttal": "# General Response \n\nWe 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). \n\n\nReviewer 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: \n- (a) providing the simulator as a testbed for methodological development (echoing Reviewer 6aWf’s comment) \n- (b) proposing a sequential decision making framework that naturally describes data mixture optimization. \n\n(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]. \n\n(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.\n\n\nWe 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. \n\n[1] Zela et al., Surrogate NAS Benchmarks: Going Beyond the Limited Search Spaces of Tabular NAS Benchmarks, International Conference on Learning Representations, 2022\n\n[2] Eggensperger et al., Efficient benchmarking of hyperparameter optimizers via surrogates, Proceedings of the Twentynineth National Conference on Artificial Intelligence, 2015\n\n\n# Point-by-Point Responses to Reviewer H7gH\n\nWe thank the reviewer for your thoughtful comments and for appreciating the novelty of our approach and the significance of the problem.\n\n>**How many simulation runs are required to obtain a relatively accurate estimation on the optimal data mixture?** \n\nWe 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.\n\n>**... could you provide the comparison results to Regmix [1] and data mixing law [2]...**\n\nWe 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. \n\nComparison with Regmix:\n\n|                | Hyperband | RegMix (LightGBM, 512 models at 20M scale) | MFMS-GP full-scale |\n| -------------- | ---- | ------------------ | ------------------ |\n| Common Crawl   | 2.66 | 2.40               | **2.31**           |\n| C4             | 1.52 | 1.35               | 1.33               |\n| Wikipedia      | 1.06 | 0.97               | **0.88**               |\n| Stack Exchange | 1.90 | 0.83               | 0.84               |\n| Github         | 0.67 | 0.56               | 0.59               |\n| ArXiv          | 1.45 | 1.37               | 1.35               |\n| Book           | 2.67 | 2.41               | **2.29**           |\n\nBolded values indicate the best-performing method, with alternatives falling outside one standard deviation from 15 runs. \nNote 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.\n\nComparison with data mixing law: \nWe 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. \n\nAssuming 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. \n\n\n>**...scalability and effectiveness of the MFMS-BO method on larger number of domains has not been explored** \n\nIndeed 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. \n\n>**...discussion on the computation complexity of the proposed method**\n\nIn 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. \n\nThat 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). \n\n\n>**out-of-domain extrapolation results… could help to improve the understanding on its OOD performance** \n\nWe 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."}, {"review_id": "Q4d6XVvDd1", "reviewer": "Reviewer_bchD", "summary": "This paper proposes a multi-fidelity multi-scale Bayesian Optimization framework for optimizing data mixture compositions in LLM training. Instead of relying on intuition or costly trial-and-error approaches, the method treats data curation as a sequential decision-making problem that adaptively selects data mixtures, model scales, and training steps to balance cost and information gain. Using a simulator trained on 472 language model runs, the authors show their GP-based approach can leverage cheap evaluations from smaller models to guide decisions for larger ones, achieving 2.6x to 3.3x speedups compared to baseline methods while finding optimal data mixtures for models from 20M to 1B parameters.", "questions": "See weaknesses.", "limitations": "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.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strength**\n1. **Novel problem formulation:** First to treat data mixture optimization as a multi-fidelity multi-scale BO problem, providing a principled alternative to heuristic approaches.\n2. **Practical applicability: This paper addresses real-world computational constraints faced by practitioners training LLMs with limited budgets**.\n3. **Significant performance gains:** It achieves 2.6x to 3.3x speedups compared to baseline methods across multiple evaluation metrics, although it was achieved on a simulator benchmark.\n\n**Weaknesses**\n1. **Limited scale and scope:** The paper only tested on models up to 1B parameters and five data categories from SlimPajama, raising questions about generalizability to larger, production-scale models and diverse datasets. \n2. **Limited practical validation:** I know it might be a bit too much to ask the authors to perform BO on the full LLM training, but training an MLP on 472 records and using it as the simulator for validation has serious limitation. There is no demonstration on real-world deployment scenarios or comparison with industry-standard data curation practices, which can severely qualify the significance of the proposed claims.\n3. **Limited technical novelty:** It is kind of natural to treat the data mixing ratio as hyperparameters of the training framework and model this problem as BO. The current approach largely combines existing BO (standard GP), without introducing significant algorithmic innovations or theoretical contributions.", "was_revised": true, "final_justification": "The authors have addressed most of my concerns, especially the one about the pre-trained models.", "rebuttal": "# General Response \n\nWe 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). \n\n\nReviewer 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: \n- (a) providing the simulator as a testbed for methodological development (echoing Reviewer 6aWf’s comment) \n- (b) proposing a sequential decision making framework that naturally describes data mixture optimization. \n\n(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]. \n\n(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.\n\n\nWe 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. \n\n[1] Zela et al., Surrogate NAS Benchmarks: Going Beyond the Limited Search Spaces of Tabular NAS Benchmarks, International Conference on Learning Representations, 2022\n\n[2] Eggensperger et al., Efficient benchmarking of hyperparameter optimizers via surrogates, Proceedings of the Twentynineth National Conference on Artificial Intelligence, 2015\n\n# Point-by-Point Responses to Reviewer bchD\n\nWe 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.\n\n>**...raising questions about generalizability to larger, production-scale models and diverse datasets**\n\nTo 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.\nSimilarly, 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.\n>**might be a bit too much to … perform BO on the full LLM training ...simulator for validation has serious limitation…**  \n\nWe 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.\n\n[1] Zela et al., Surrogate NAS Benchmarks: Going Beyond the Limited Search Spaces of Tabular NAS Benchmarks, International Conference on Learning Representations, 2022\n\n[2] Eggensperger et al., Efficient benchmarking of hyperparameter optimizers via surrogates, Proceedings of the Twentynineth National Conference on Artificial Intelligence, 2015\n \n\n>**It is kind of natural to treat the data mixing ratio as hyperparameters… current approach largely combines existing BO…** \n\nWe 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. \n\n>**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.** \n\nWe 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."}, {"review_id": "U0ggCtf9xQ", "reviewer": "Reviewer_FwDJ", "summary": "The paper proposes an approach to data curation for training LLMs based on\nmulti-fidelity multi-scale BO. The authors describe their approach and evaluate\nit empirically, demonstrating speedups in training.", "questions": "See above.", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "The paper is well written and proposes an interesting idea that addresses one of\nthe main challenges in training LLMs. To the best of my knowledge, the idea is\nnovel and, while relatively simple, seems to work well in practice. The proposed\nmethod does not necessarily result in better predictive performance, but enables\nthe training process to achieve a particular level of performance faster, which\nis especially important in the context of large models such as LLMs.\n\nIt was unclear to me why the authors used MLPs as predictor models for\nestimating the performance of particular combinations of model size, dataset\nproportions etc. MLPs are not commonly used in BO, and later in the paper the\nauthors use the more common GPs as surrogate models. Did the authors evaluate\ndifferent approaches and settle on MLPs because of better performance? Please\nclarify.\n\nThe empirical evaluation relies exclusively on surrogate models instead of\nactual performance data as far as I can see. While the overall R^2 seems good,\nthe results for individual cases vary quite a bit (Table 3), especially for\nsmaller fractions of the overall data. The authors say that results were\naveraged over 3 random seeds, but do not give confidence intervals. A (small)\nadditional experiment on real data would help to convince the reader that the\npresented results hold true in practice.", "was_revised": true, "final_justification": "This paper presents a nice idea that works well in practice and should be accepted in my opinion.", "rebuttal": "# General Response \n\nWe 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). \n\n\nReviewer 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: \n- (a) providing the simulator as a testbed for methodological development (echoing Reviewer 6aWf’s comment) \n- (b) proposing a sequential decision making framework that naturally describes data mixture optimization. \n\n(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]. \n\n(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.\n\n\nWe 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. \n\n[1] Zela et al., Surrogate NAS Benchmarks: Going Beyond the Limited Search Spaces of Tabular NAS Benchmarks, International Conference on Learning Representations, 2022\n\n[2] Eggensperger et al., Efficient benchmarking of hyperparameter optimizers via surrogates, Proceedings of the Twentynineth National Conference on Artificial Intelligence, 2015\n\n# Point-by-Point Responses to Reviewer FwDJ\nWe 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. \n\n>**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** \n\nWe 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. \n\n>**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 …**\n\nWe 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."}]

Author rebuttal:

Meta-review: This paper presents a novel algorithm for constructing optimal data distributions for LLM training. The authors frame the problem of data distribution selection as a multi-fidelity, multi-scale Bayesian optimization, jointly optimizing data mixtures, model scale, and training steps in an adaptive manner. The proposed method is extensively evaluated at scale, demonstrating substantial empirical gains over random search baselines.

Most reviewers were impressed by the novelty, viewing the framework as an interesting and impactful application of Bayesian optimization.

While several reviewers raised concerns about the lack of real-world deployments and the focus on controlled experimental settings, the strengths of the work outweigh these limitations. There were also concerns about weak baselines, but I find these were adequately addressed during the rebuttal.

Overall, I align with the positive reviews regarding the novelty of the contribution, and I recommend acceptance.
Decision: Accept (poster)