MixMin: Finding Data Mixtures via Convex Minimization
Authors: Anvith Thudi
Evianne Rovers
Yangjun Ruan
Tristan Thrush
Chris J. Maddison
Venue: ICML 2025 poster (2025)
Keywords: Data Mixing
Bi-Level optimization
Convex Surrogates
Large Language Models
Foundation Models
Chemistry
Primary area: deep_learning->algorithms
TL;DR: We show that the bi-level optimization for data mixing reduces to a convex minimization as model classes become larger.

Abstract: Modern machine learning pipelines are increasingly combining and mixing data from diverse and disparate sources, e.g., pre-training large language models. Yet, finding the optimal data mixture is a challenging and open problem. We formalize this data mixing problem as a bi-level objective: the best mixture is the one that would lead to the best model for a downstream objective. Unfortunately, this objective is generally intractable. In this paper, we make the observation that the bi-level data mixing objective becomes convex as our model class becomes larger. We develop and study a gradient-based approach for optimizing this convex objective, which we call MixMin, and test it on language modeling and chemistry tasks. MixMin was the only method that uniformly improved the data mixture in all our experiments. With MixMin, we improved the data mixture using less than 0.2% additional compute for a pythia-$410M$ model trained on $8.2B$ tokens, resulting between 1-5% relative improvement to negative log likelihood on PIQA, ARC Easy, SciQ, and OpenWebMath. Crucially, we found that MixMin mixtures for smaller models improved training of larger models, suggesting that MixMin mixtures may be scale-invariant. When mixing bioassay data to train an XGBoost model, we saw improvements to average precision scores of $0.03-0.15$.

MixMin: Finding Data Mixtures via Convex Minimization
Anvith Thudi" 12 Evianne Rovers
Yangjun Ruan" Tristan Thrush S Chris J. Maddison 12
Abstract
Raffel et al., 2020; Achiam et al., 2023). These models are known to generalize to data-poor downstream tasks, core motivation behind scaling training data. However, as we increase the amount and diversity of data sources we can train on, a core challenge becomes choosing how to weigh these sources when training. The effectiveness of a model pre-trained on surrogate data is known to be directly related to how similar the surrogate (or "pretraining") data is to the downstream task (target distribution) (Ben-David et al., 2010; Isik al., 2024; Jain et al., 2024; Pouget et al., 2024).
Modern machine learning pipelines are increas- ingly combining and mixing data from diverse and disparate sources, e.g., pre-training large lan- guage models. Yet, finding the optimal data mix- ture is a challenging and open problem. We for- malize this data mixing problem as a bi-level ob- jective: the best mixture is the one that would lead to the best model for downstream objec- tive. Unfortunately, this objective is generally in- tractable. In this paper, we make the observation that the bi-level data mixing objective becomes convex as our model class becomes larger. We develop and study a gradient-based approach for optimizing this convex objective, which we call MixMin, and test it on language modeling and chemistry tasks. MixMin was the only method that uniformly improved the data mixture in all our experiments. With MixMin, we improved the data mixture using less than 0.2% additional com- pute for a pythia-410M model trained on 8.2 tokens, resulting between 1-5% relative improve- ment to negative log likelihood on PIQA, ARC Easy, SciQ, and OpenWebMarl found that MixMin mixtures for smaller mod- els improved training of larger models, suggest- ing that MixMin mixtures may be scale-invariant. When mixing bioassay data to train an XGBoost model, we saw improvements to average precision scores of 0.03 0.15.
Formally, finding the best mixture of data sources to train on for a downstream loss poses a bi-level optimization which is hard to optimize. We must find the mixture whose risk minimizer optimizes the downstream loss. This generally only admits expensive zero-order approaches, such as grid search (Dubey et al., 2024; Blakeney et al., 2024). Past work has considered learning to simulate the loss oracle for unseen mixtures (Liu et al., 2024), or leveraging large pool of existing models (Thrush et al., 2024). Other approaches consider solving data mixing objectives agnostic to down- stream loss, but are known to not consistently improve the baseline of training on all the data with their natural propor- tions (when evaluated on diverse downstream losses) (Xie et al., 2024; Jiang et al., 2024; Liu et al., 2024; Fan et al., 2023; Held et al., 2025).
Crucially, we
In this paper we show that the bi-level optimization for data mixing reduces to a convex minimization as the model classes become larger, which we can solve cheaply with gradient based methods. A key insight is that by restricting the loss functions to cross-entropy (CE) or mean-squared error (MSE), we gain additional structure which simplifies the data mixing objective. Practically, we show we can find the best mixture by first training a (cheap) proxy model for each source, learning the best mixture of their outputs for the target dataset (convex minimization), and then using this weighting to remix the data, see Figure 1.
1. Introduction
Recent progress ML has come from training on vast web-scale data (Dubey et al., 2024; Touvron et al., 2023:
'Department Computer Science University of Toronto, Toronto, Canada 'Vector Institute, Toronto, Canada 'Department of Chemistry. University of Toronto, Toronto, Canada 'Structural Genomics Consortium, Toronto, Canada Depparment of Computer Science, Stanford University, Palo alto, USA. Correspondence to: Anvith Thudi <anvith.thuce whu.th.hhuuieee @ mail.utorooto.ca>..
First, note that when our models are relatively close to Bayes optimal (the best possible functions), data mixing is equivalent to learning a linear model over the predictions of a fixed set of models (a convex objective). Learning this linear model requires comparatively little data to training modern large models, and is relatively cheap optimize. Our reduction specific to CE MSE, where we leverage
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

MixMin; Finding Data Mixtures via Convex Minimization
Source
Optimize
mixcures
Train cheap proxy
Target
with MixMin
Train on re-mixed data
Source
models
Source
Figure 1: Optimizing MixMin to find the mixture weights requires training few cheap models for each source and target dataset. Given the mixture weights we train a more expensive model using the mixture.
the fact the Bayes optimal model for a mixture is the mixture of the Bayes optimal models for each source. However, this reduction does not yet make data mixing amenable to first-order optimization. To evaluate convex objective's gradient need the Bayes optimal model for each source.
assays (e.g., distinct but predictive cytotoxicity assays). To summarize, our contributions are:
1. Observing data mixing reduces to single convex min- imization as model classes become more expressive
A key challenge is that the proxy models needed to evaluate this convex objective can be expensive to compute. We find that in practice, we can use cheaply computed mod- els and still have good data mixing performance. We call this approach to optimizing the convex objective MixMin. Empirically, we found the performance of MixMin did not significantly degrade between using proxy models computed with 100% the cost of full training run to 1%. In other cases, the MixMin mixture significantly improved over base- lines, but the ensemble of proxy models performed much worse than retraining, suggesting they were far from Bayes optimal but did not hinder MixMin performance.
2. Proposing MixMin which approximately solves the convex objective by using cheap proxy models and downstream data samples
3. Empirically showing MixMin improves past baselines for finding data mixtures across language modeling and chemistry tasks
4. Empirical and analytical evidence showing MixMin is robust to the use of weak/cheap proxy models
2. Preliminaries
We empirically compared data mixing with MixMin cur- rent baselines both language modeling and chemistry tasks. For language modeling we considered optimizing the mixture of source domains in SlimPajama (cer, 2023) for PIQA (Bisk et al., 2020), ARC Easy (Clark et al., 2018), SciQ (Welbl et al., 2017), or Open OehMWehhath (Paster al., 2023). all cases we found MixMin improved the base- lines for log likelihood, improving or maintaining the gap increased the scale of the models from Pythia-160M to pythia-410. (Biderman et al., 2023). Moreover, for pythia-410M the mixtures we used took 0.15% the com- pute of the training run to find and improved the negative log likelihood by 5%. For our chemistry experiments we considered mixing assay datasets PubChem (Beaini al., 2023; Kim al., 2016) to improve the performance of XGBoost (Chen & Guestrin, 2016) model several held-out assay datasets. In all cases, we observed MixMin improved over the natural distribution (the standard data mixture after filtering for this domain (Salem et al., 2020; Li et al., 2022; Ye et al., 2018)) as we increased the num- ber of surrogates to mix over, improving average precision scores between 0.03 0.15. We note an additional benefit of optimizing data mixtures for chemistry tasks is that the found MixMin weights could provide interpretability; we highlight patterns MixMin found in PCBA for predicting
We denote a hypothesis space by H. a model by f X O H. datapoint by (x, y) E X x V, and our loss function by L(o, O x - IR+ This paper considers the problem of mixing set of finite source distributions (dp E P where P is finite) to train on to do better on some downstream "target" distribution (dt). For example, mixing Wikipedia and arXiv data to better on benchmarks evaluating scientific knowledge.
Specifically, we are searching for the best weighting À A (where A is the simplex) of our sources to train a model on to perform well on our downstream distribution. Formally, using our loss function L as both our measure of downstream performance and training objective:
=J
DM H)
JXXY (Lf()))))ddt()) v)
Σ Ap I ())(t),
where fa(x) arg min
fEW
(Data Mixing)
In general min, (M H) a bi-level optimization that can hard to optimize: a standard approach is grid/random searching through mixture weights and solving the inner
2

---

MixMin: Finding Data Mixtures via Convex Minimization
Bigger Model Classes
to dp(x). MixMin also holds for conditional CE and MSE losses under no covariate shift: J [ (dp(a, y) and
Data Mixing
/ [lf(()) vllīdp(x, y). with dp(x) dp' (x) Vdp, dp' e P.
Target Loss
Theorem 3.1. Let the objective for Data Mixing be un- conditional CE. fllog(((x)) (d))(())))) or conditional CE or MSE with no covariate shift, (flog(f
MisMinn
y) and [1l((x) - yll2dp(z y), with p(x) = =='xx) Vp.p' p' E P. Suppose also H contains the Bayes optimal model for each mixture of the source distributions dp E P. Then
Training Mixture
Figure 2: The convex MixMin objective better approximates the Data Mixing objective as the model class becomes larger (and better approximates Bayes optimal).
X arg minxEAP (M H) iff
/ L (E>psfoo)),)) dt(x, (1)
minimization by fitting a (small) proxy model. In this paper we describe settings where the bi-level optimization col- lapses and we can use simple gradient based optimization to find the best mixture.
X arg min XEA xxxy
The proof is provided in Appendix A.1. This optimization is clearly convex, and can be handled with gradient based approaches unlike Data Mixing. However, running this optimization requires having the Bayes optimal f, and an estimator for the integral to compute the gradients. The main hurdle will be how to get reasonable approximations for the Bayes optimal fp. which we discuss in Section 4.
Data Filtering is not Data Mixing An often related but different approach to curating training data data filtering. Abstractly, data filtering is some map F that takes dataset D + D' where D' C D. Such approaches are motivated by training efficiency. e.g., coresets, and/or data quality, e.g., rejection sampling to be closer to a desired target distribu- tion. Data filtering can be composed with data mixing: one typically mixes amongst filtered data sources (Dubey et al., 2024). We leave studying such compositions to future work, but note some popular examples of data filtering include: perplexity correlations (Thrush et al., 2024), n-gram em- bedding similarity with a target distribution (Xie et al., 2023; Everaert & Potts, 2024), handcrafted classifiers (Xiao et al., 2024; Li et al., 2024a), and perplexity filtering (I.i et al., 2024a)
3.2. Not Well-Specified: 'Expressivity" is Enough
We now show that Data Mixing smoothly approaches Equa- tion I as the hypothesis space H contains models closer Bayes optimal, e.g., becomes larger. This follows from assuming some regularity in our loss function and the defi- nition of the data mixing objective.
Lemma 3.2. For L either CE or MSE, suppose J ((f())) v) is C-Lipschitz in f w.r. norm on the functions [|]())] Let
3. Data Mixing is Learning a Linear Model
XEA Jxxy x arg min [ L (EEnsulo).)) (ll..)).
We show that Data Mixing simplifies to a convex optimiza- tion problem as the model classes become larger. This intuition is presented in Figure 2.
Then, for any hypothessi space such that for all mixtures dpx E Conw(P), the dpx risk minimizer fx(H) in H. sat- isfies [I - Sall S e where Sx is the Bayes optimal model, we have the excess Data Mixing error is bounded by e:
3.1. Well-Specified: Data Mixing is just Risk Minimization
We show that, when our hypothesis space contains the Bayes optimal model fp for all dp E P, the best Data Mixing weights are also the best weighting of models trained on each source (a convex objective). The core idea is that for certain losses, cross-entropy (CE) and mean-squared error (MSE), the Bayes optimal model for mixture is the mixture of the Bayes optimal models' for the sources. However this is only true if there is no input distribution shift among the sources (p(x) are all the same), e.g., for generative tasks. To clarify, MixMin holds for the uncon- ditional CE loss, f [f (d))),((r) which matches f(x)
DMr(A" dd)) min (M, 2Ce XEAP
(2)
The proof is provided in Appendix A.2. The main intuition to take away from Lemma 3.2 is that as our model classes
'To ensure this for MSE one could consider only bounded output spaces, and for cross-entropy, bounded away from 0 and 'Effectively this means functions too far away will not have low error. Mathematically, this means there is lower-bound on the increase in loss from Bayes optimal in terms of the of the perturbation to the function (e.g., strong-convexity).
3

---

MixMin: Finding Data Mixtures via Convex Minimization
Algorithm 1 MixMin
Cost We only need to evaluate í2(x) Va E Dr once and can reuse the outputs for subsequent calls to the gradi- ent. Given this, the dominating cost with MixMin, with grid/random search based approaches to data mixing, is the number of models needed to train (i.e., fp). MixMin re- quires only as many models as sources, while the number of models grid/random search approaches require scales ex- ponentially (in the worst-case) with the number of sources. Note once the outputs and proxy models are computed, the per-iteration cost of MixMin is O([P) D:I), which is inde- pendent of model size.
Step size 7,, number of steps n, loss function C (either cross-entropy or E2), samples D. from the target distribution dt, and (cheap) models trained on each source fx(x)) v dp E P.
Require:
Initialize: Ap f [P]
- for all dp E P. and pre-compute
(fp()) E De,dpe P}
1: fori= 1, I,....n. do
2: f(xx) c Edpep pp5p())
3: 1t D. E (IDD)ED ((f((z).v)
4: g+ Vil
5:
Apt Apellup
Edpen Apo
for all dp e P
On the use of weak proxy models In our experiments, we found that proxy models Sp trained with very small frac- tions of total training compute achieved good results and additional proxy model compute did not significantly im- prove the quality of the final mixtures. This suggests that MixMin can perform well when the proxy models fp have a high loss and that the excess error of MixMin should be an- alyzed beyond the risk of the proxy models f.. Specifically, we empirically found that we could train with significantly less data (1%) and still have similar MixMin performance. Furthermore, we found the ensemble of the proxy mod- els could be far less accurate than the model retrained on the MixMin mixture weights, suggesting the proxy models were far from Bayes optimal 3 but still led to good mix- tures. We hope future work investigates and generalizes these phenomenon.
6:
end for
Return Axp dper
get larger and becomes closer to expressing the Bayes op- timal model, the Equation I weights become closer to the optimal Data Mixing weights. This is depicted in Figure 2
4. Mix
We now consider how to optimize the objective in Equa- tion 1. If we had a gradient oracle, we could optimize the objective using entropic descent (Duchi, 2018). However, the objective (and gradient) is typically intractable for sev- eral reasons which we alleviate with approximations.
Firstly, computing the integral for the objective (and gra- dient) is typically intractable. We can however compute an empirical risk given a dataset De ~ dp:. Secondly, we often do not know the Bayes optimal models fp for each dp E P. but instead can obtain "upproximations" fp. e.g., by training a model on a dataset Dr ~ dp E P. This gives us our approximate convex mixture objective, where we let Train(H. Dp) be some training function:
5. Related Work
Some past work has considered optimizing a pretraining data mixture for a specific downstream loss, ie.. task. Hwang & Whang (2021) train many small models with dif- ferent mixtures and learn to predict the task error of unseen mixtures, and RegMix (Liu et al., 2024) later used similar oracle to find better mixtures (the best mixture if the oracle is accurate). Similarly, Held et al. (2025) proposed optimiz- ing the mixture given estimates of downstream performance from different source datasets; they assumed data mixing was a linear model (with constraints) in the mixture weights and proposed mixtures by maximizing the performance of the linear model. We note Liu al. (2024) considered modeling data mixing linear objective, but found their non-linear models for predicting data mixing improved over linear models. Finally there are the normal grid-search ap- proaches (Blakeney et al., 2024), which random search is known to improve (Bergstra & Bengio, 2012).
X X(fihEE
Σ
C() ())(((())))))
= -arg min YEA D.
(a.))DD
pEP
where f = Train(H, Dp)
(MixMin)
Computing the gradient of the objective follows from chain rule, and hence we can run entropic descent to solve MixMin. Recall entropic descent is a variant of mirror descent which handles the constraint of being in the sim- plex. This gives us Algorithm 1, which we call MixMin for 'Mixtures by Minimization'.
Another line of work has considered data mixing objectives defined by just the sources and not downstream loss: e.g.. with the motivation of faster optimization or generalizing
Being close Bayes optimal would imply the ensemble was close to retraining, a contradiction.
4

---

MixMin: Finding Data Mixtures via Convex Minimization
better across the sources. However, when evaluating on downstream objective (as is the focus of our paper) methods here are known to not consistently improve over the natu- ral or balanced distribution of data (Fan et al., 2023; Jiang et al., 2024; A Albalak et al., 2023; Xie et al., 2024): see Held et al. (2025) for comparisons. Outside of language model- ing, we note the natural distribution of data is the standard data mixture used for transfer learning in chemistry (after filtering data) (Salem et al., 2020; Li et al., 2022; Ye al., 2018). Another line of work considered data mixing for dis- tributionally robust optimization (Thudi& & Maddison, 2024), which is distinct to downstream data mixing objective.
6.1. Pretraining Language Models
Data mixing is now a common step in pre-training LLMs (Dubey et al., 2024; Li et al., 2024b). We investi- gated how MixMin compared to other data mixing methods, its compute efficiency, and how weights found at smaller scale of models transfers to training larger models". 4.
6.1.1. EXPERIMENTAL SETUP
Datasets and Models We used the domains in SlimPa- jama (cer, 2023) as sources for pre-training. For the down- stream target, we considered several multiple-choice tasks that models in the 160M - 410M scale are known to do non-trivially at (Thrush et al., 2024), alongside another gen- erative task: SciQ (Welbl et al., 2017), PIQA (Bisk et al., 2020), Eas-Easy (Clark et al., 2018) and the first 10000 documents in Open MathWehtath (Paster et al., 2023) 5. For multiple-choice tasks we used MixMin to maximize the log probability of the correct question and answer (Q+A) sequence, i.c., treating these tasks new pretraining do- mains. Note that the optimal conditional of A given Q distribution is a conditional of the joint distribution of Q+A, so our objective subsumes the Q + A prediction problem. Alongside the objective loss, we also evaluated accuracy and the conditional cross-entropy of the pre-trained models for these tasks (i.e., multiple-choice performance). Open- WebMath is specialized pretraining dataset, and so the loss for MixMin and evaluation was the negative log probability of a sequence. Our experiments used the 160M and 410M pythia models (Biderman et al., 2023) trained for 3.2B and 8.2B tokens respectively (chinchilla optimal (Hoffmann et al., 2022)).
Given these past findings, we compare to RegMix and ran- dom search in our experiments. However, compared to MixMin, MegMix has significantly higher evaluation costs.
Our setting for Bi-Level optimization also presents sev- eral challenges for previous approaches, including those used for other hyperparameter optimizations. Firstly, the inner-opiimization is over an arbitrary function space (e.g., non-parametric models like XGBoost), and so lacks the para- metric gradients needed for many methods (Ji et al., 2021; Pedregosa, 2016). Alternative constraint based approaches through KKT suffer similar issues given the non-parametric space of the inner optimization (Shi et al., 2005). For many model classes, one could parameterize the inner optimiza- tion, but we then lose the necessary convexity for these methods (e.g.. Neural Networks).
6. Experiments
We compared MixMin to previous data mixing baselines for language modeling and chemistry tasks. For language mod- cling we further investigated the transferability of mixtures found using cheap proxy models to more expensive train- ing runs. For chemistry, we explored the performance of MixMin as we increased the pool of surrogate data sources. Experiments were run using A100 GPUs and AMD EPYC 7643 CPUs.
Data Filtering We included documents in SlimPajama with least 1024 tokens (using the Pythia tokenizer), and scraped documents for each domain until adding the next document would bring over 3.2B tokens for each domain.
Our experiments support the following list of (weak) condi- tions for when MixMin works:
Hyperparameter Tuning We took the largest batch size that fits on a single A100 gpu, which was 64 for the 160M- pythia model and 32 for the 410.M-pythia model for con- text length of 1024. For the 160.-ppyhhi model we in- creased the learning rate until training loss got worse on the natural distribution of domains in SlimPajama: we started from le - 4, 5e 4, le -3, 5e
1. The data sources do not have covariate shift (consider changing the objective to remove covariate shifts)
2. The downstream task is close to/contained within a mixture of our sources
3, le 2 and found loss got worse 1e 2 so chose 5e 3. For the -1yhha model we evaluated learning rates 5e - 3 and le 2 and found 5e - 3 was better. Other hyperparameters are the same as (Thrush et al., 2024). All hyperparemeters are fixed throughout the language modeling experiments.
3. Our final model class is very expressive
4. Our proxy model's learnt some signal for each source (e.g., even with very little data): consider iteratively using more compute for the proxy models until stability in the mixtures is reached
"p values from Welch's t-test are shown in Appendix B. 'We chose 10000 documents for computational efficiency.
5

---

MixMin; Finding Data Mixtures via Convex Minimization
1% of Triinin Cost
PIQA
0.1
0.1-
Nutual Gea Lass Medtod
Mehhod
0.-
-0.1-
-02-
Natural (ymx)
MixMin
MicMin
-0.3-
Uolanced
Random Search
Random Search
-0.4
RegMii
RegMix
PIQA (4.06)
ARC Easy (3.82)
SaiQ (3.26)
OpennMah (365))
138
S1
10
100%
Data Mixing Compate /Tsining Compute
Figure 3: MixMin consistently outperforms all baselines across the four target tasks using 1% of the final training run compute (Left). We report improvement the downstream generative loss of training the natural distribution (which is stated beside the task name): higher is better. Error bars indicate 95% confidence interval. Furthermore, we find that MixMin was robust to using less compute (Right), while RegMix and random search had their performance degrade with less compute.
PIOA
ARC Eagy
SeiQ
0:52
Nidttn
NhaMta
0/8
Nanaai
0.30
Natual
Nauni
0.48
0:46-
0.61
0.44
0.73
062
6.2
6:10
411
160M
415M
148M
AIOM
Figure 4: MixMin mixtures derived from small models continue to improve training for larger models as measured by target accuracy (in most cases). We report accuracy with the errors bars representing 95% confidence interval over 3 trials. MixMin weights were found using 1% the compute of the 160M model training run, which 0.15% the compute of the 410M training run.
Method Implementations We compared MixMin to the baselines of balancing the domains, the natural distribution of domains in SlimPajama, random search and RegMix (the latter known to improve or perform on par as previous data mixing methods). We split the target task into a random 80% training set, and 20% test set. The training set was used the MixMin optimization, and the evaluation of loss for random search and RegMix. Results are reported over 3 trials of random train-test splits.
compute we ran random search with the same number of proxy models (7). RegMix showed benefits by using many (but cheap) proxy models (Liu et al., 2024), and so we ablated the number of our cheapest proxy models to vary compute for it; this however increased its evaluation costs to the other baselines. All proxy models are trained with the same hyperparameters specified earlier, and use the 160M pythia model architecture. To achieve X% compute relative a full training of the 160M pythia model, tokens
trained each proxy model for 3,2(X/100) for MixMin and random search, and for RegMix train 7X proxy models where each proxy model trains for 3.2 (12(1) tokens (our lowest compute proxy models).
We ran MixMin using y 1.0 for 100 steps for all the experiments. For RegMix, we adapted the code available at https://github. com/sail-sg/ segmii//blob/main/eggrree.
regression pynb, changing the hard-coded natural distribution the SlimPajama natural distribution and inputting models' results. We implemented random search by sampling uniformly from the simplex and training proxy models on those mixture. Recall MixMin requires proxy model for each domain, and to normalize evaluation
1.2. RESULTS
Mix found good mixtures with 1% of train compute In Figure we show the generative loss results of MixMin, random search, and RegMix where the total compute for all the proxy models was 1 % relative to training the final
6

---

MixMin:
Finding Data Mixtures via Convex Minimization
Maxd
13
Mak 2004
Figure 5: MixMin mixtures derived from small models continue to improve training for larger models as measured by target generative loss (in all cases). We report generative loss with the errors hars representing a 95% confidence interval over 3 trials (lower is better). MixMin weights were found using 1 1% the compute of the 160M model training run, which is 0.15% the compute of the 410M training run.
160.-- model for 3.2B tokens (Chinchilla optimal). We compared the performance of these to a balanced distri- bution and the natural distribution. We found MixMin con- sistently improved generative loss over the baselines, We also found MixMin consistently matched or improved the predictive loss and accuracy over all the baseline methods (Figure 8 in Appendix B).
6.2. Chemistry
We explored how MixMin could improve predicting prop- crties of molecule libraries, in particular the endpoints of an assay (a predictive task). Many assays have few posi- tive samples and training on just that data can lead to only noisy performance. Past work has remedied this by transfer learning from a larger dataset, but had not considering opti- mizing the data mixture after filtering (Salem et al., 2020; Li et al., 2022; Sun et al., 2022; Ye et al., 2018). A common default mixture is the natural distribution, and here we show MixMin provided better mixtures than this, especially as the diversity of the source datasets grows. We found similar improvements over random search and RegMix. We hope MixMin's 's ability to scale helps future data curation efforts for predictive drug screening.
MixMin more cost effective than baselines In fig- ure 3 (and Figure 9 in Appendix B for other tasks) we compared the performance of MixMin, random search, and RegMix as we vary the the total compute for all the proxy models from 100%, 10%, 5%, 1% of the compute for train- ing the final 160--pyyhia model for 3.2B tokens (Chin- chilla optimal). We found MixMin was robust using less compute, where random search and RegMix benefitted from increasing compute (but do not meet MixMin perfor- mance even at the highest compute setting we tested). Note, we did not include RegMix at 100% compute given the significant evaluation overhead: 100 times the alternatives.
6.2.1. EXPERIMENTAL SETUP
Datasets and Models We worked with the PCBA dataset (Beaini et al., 2023) 6, which was a filtered scrape of PubChem (Kim et al., 2016). We used the first 10 bioassays as the target distributions, using subsets of all the bioassays (removing the target) sources. Our experiments used XG- Boost on ECFP fingerprints of SMILES (Rogers & Hahn, 2010) to fit each assay, known to be a strong baseline (Yang et al., 2019; Seidl et al., 2023) for this domain. We report AP performance (average precision score, a variant of AUPR) for our experiments.
MixMin mixtures for smaller models improve training of larger models Finally, we tested whether MixMin mix- tures derived from smaller models were useful for training larger models. We computed MixMin mixtures using pythia- 160M proxy models each trained with 0.14% the compute of full 3.2B token training run (for an overall cost of 1%). We then trained a pythia-410.M for 8.2B tokens those mixtures. Here we only compared to the balanced or natural mixture. as our previous experiments showed random search and RegMix found mixture comparable or worse than one of those mixtures for our tasks. Accuracy results are presented in Figure 4 with generative and predictive loss results pre- sented in Figure and Figure 12 (the latter in Appendix B). For accuracy and predictive loss found in most cases MixMin's 's performance improved over the baselines as we increased scale, despite not refitting the MixMin weights. Even in the cases where MixMin's ' accuracy and predictive loss did not strictly improve over baselines, the generative loss always improved with MixMin.
Data Filtering We trained models over the first 100,000 molecules in PCBA. We further removed any assay with no positive or negative instances among the first 100,000 molecules. For all assays (1328) the number after filtering was 1281, for the first 100 assays this left 86, and the first 10 were all kept. We performed data mixing among the filtered assays (excluding the target assay).
'The dataset can he found at https: / http///pllabbbii.oo io/ da dataseta/ /araa/graaaa graphium/pcba 1328-12-1 112228115644 1564kvv.
7

---

MixMin: Finding Data Mixtures via Convex Minimization
Micny Sumopiic Iatacds
Moony Strmpgii Dacods
Micnr 1279 Sumpppic Datants
Kaki Natsa
Masi Nevuil
MisMi Nakii
0.35-
03
01s
0.2
0.00
00
088
090
Figure 6: MixMin improved over using the natural distribution of data as we increased the number of surrogate assays. We report AP scores for the first 10 assays in PCBA, with error bars representing a 95% confidence interval over 3 trials.
Mixing 9 Sunogate Ditaoes
Mioing Surrginn Detisess
Mixing 1279 Sucugane Deissers
0.40
0.5
MidGa Retou
0.-
Mibiaa Gooekk
Mibia Retils
044
0.25-
4.31-
a.s
kan-
Figure 7: Retraining with the MixMin mixture performs better than using the MixMin weights to ensemble the proxy models. We report AP scores with error bars representing a 95% confidence interval over 3 trials. This implies the proxy models were far from Bayes optimal.
Method Implementation For every assay in PCBA, we used 64% 16%- 20% train-validation-testt split: original 80%- 20% train-test split, and further splitting the train set into a 20% validation set. We trained proxy models used for MixMin on just the train split, and trained baseline models on the train and validation splits. The baselines models represent the performance we get by just training on the target, as to normalize improvement from data mixing. Given a target assay, we used the proxy models on the other assays in PCBA to run MixMin, using the train set on the target task to fit the MixMin mixture weights (leaving the validation and test set unseen). Finally, given the MixMin mixture we retrained a new model on the train sets from all the other assays. For the natural distribution baseline, we trained on all the train sets from other assays as (no reweighting of sources). For all method we reported the mean test AP score improvement (from 3 trials) over the baseline of training the target distribution, with 95% confidence intervals.
(the final model trained on MixMin or natural mixtures) we fixed the n estimators 100 and max depth 6.
6.2.2. RESULTS
MixMin's 's advantage over the baselines increased with the number of sources As shown in Figure 6, we ob- served that MixMin performs mostly on par with the natural distribution when working with the first 10 assays and im- proved over natural as the number of assays to mix grew. This suggests MixMin was able to find relevant data sources amongst more diverse sets, and we note its absolute perfor- mance grew while the performance of the natural mixture decreased with the number of sources. We note in all cases both MixMin and training on the natural distribution im- proved over the baseline of training on the target assay, Similar improvements over random search and RegMix are presented in Figure 16 and Figure 17 respectively (in Ap- pendix B).
Retraining with MixMin Mixture performed better than ensembling proxy models An alternative approach to training on MixMin weights, suggested by the reduction for Bayes optimal models, is to ensemble the proxy models with the MixMin weights. We present a comparison between this and retraining with MixMin weights in Figure 7. Retraining
Hyperparameter' Tuning We selected hyperparameters for each proxy and baseline model by doing a 5-fold cross- validation on the train set. We grid search over all combi- nations of n estimators E [10,50 50, 100) and max depth E [4, 6, 8). For the models trained over all the surrogate assays
8

---

MixMin: Finding Data Mixtures via Convex Minimization
performed significantly better, suggesting the proxy models were too noisy to lead good ensembles. Specifically, this implies the proxy models were far from Bayes optimal, as the ensemble would have then also been near Bayes optimal. The fact we were able to improve over the natural distribution suggests MixMin did not need strong proxy models for these sources and targets.
Acknowledgements
Resources used in preparing this research were provided in part by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Insti- tute. We acknowledge the support of the Natural Sciences and Engineering Research Council of Canada (NSERC), RGPIN-2021-0344 Anvith Thudi is supported by a Vanier Fellowship from NSERC.
MixMin highlighted data sources with similar endpoint measure All the target assays tested cytotoxicity, and the majority of the top 4 data sources MixMin found for each assay also tested cytotoxicity (Table 2 Appendix B). More specific protein inhibition assays were also included, poten- tially providing insights into the compound's mechanism of action. However, an assay with no apparent connection to the target assay was also identified. We hope future work explores the full interpretability of MixMin weights for the bioassays we tested on.
Evianne Rovers acknowledges support given through the Matthieu Schapira lab the University of Toronto who re- ceives funding from NSERC [Grant RGPIN-2019-04416) CQDM (Quantum Leap-176), and MITACS accelerate (IT13051). The Structural Genomics Consortium is regis- tered charity (no: 1097737) that receives funds from Bayer AG, Boehringer Ingelheim, Bristol Myers Squibb, Genen- tech, Genome Canada through Ontario Genomics Institute [OGI-196],
Inno- vative Medicines Initiative 2 Joint Undertaking (EUBOPEN grant 875510], Janssen, Merck KGaA (aka EMD in Canada and US), Pfizer, and Takeda.
7. Conclusion
In this paper we formalized a bi-level data mixing objective for mixing data sources to minimize risk on a downstream objective. We showed that in the case of CE or MSE, this objective reduces to a convex minimization as our model classes become larger. We proposed using cheap proxy mod- els trained on the individual source distributions to optimize the convex objective, leading to data mixing method we call MixMin. Our experiments showed MixMin consistently outperformed previous baselines across language modeling and chemistry tasks, and was robust to using cheap proxy models. We hope future work explores explaining when and why MixMin can use cheap proxy model, or develops alternative empirical approaches to optimizing the convex objective for data mixing with large model classes.
We thank Chris Crebolder and the rest of the UofT Computer Science sysadmin team for their help in making many of the experiments in this paper possible. We would also like to thank Marta Skreta, Leo Cotta, Ayoub El Hanchi, William Held, and many others at the Vector Institute for discussions contributing to this paper.
References
SlimPajama:
627B token cleaned and version of RedPajama.
deduplicated
http::/wwww..ccerrrrr..eb/////pp
outeareeeeecaaccdddggllllleeeeee URL ht https: co/
2023.
datasets/cereee Sicaastts//erebrrrssSSiiaaaggai 627B.
Limitations Our reductions for data mixing were specific to CE and MSE, and required there be no covariate shift among the source distribution. While this covers generative tasks, we hope future work explores extending these ideas to other loss functions, and handles covariate shift. For example, our framework does not currently allow us to do data mixing for image classification, where covariate shift is common among sources. Also, MixMin need not be the best way to optimize the convex data mixing objective. Finally, an open problem is pushing the limit of how cheap the proxy model's 's can be without degradation performance.
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L.,, Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint S.XV::2303.0877,, 2023.
Albalak, A., Pan, L., Raffel, C., and Wang. W. Y. Efficient online data mixing for language model pre-training. In RO- FoMo: Robustness of Few- and Zero-shot Learning in Large Foundation Models, 2023.
Beaini, D., Huang, S., Cunha, J. A., Li, Z., Moisescu- Pareja, G., Dymov, O., Maddrell-
S., McLean, C. Wenkel, F., Müller, L.,, et al. Towards foundational models for molecular learning on large-scale multi-task datasets. arXiv preprint arXiv:2310.0429/ 2023.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Ben-David, S., Blitzer, J., Crammer, K., Kulesza, A., Pereira, F., and Vaughan, W. A theory of learning

---

MixMin: Finding Data Mixtures via Convex Minimization
from different domains. Machine learning, 79:151-175, 2010.
Hwang. S.-H. and Whang, S. E. Regmix: Data
mixing augmentation for regression. arXiv preprint arXiv:2106.03374 2021.
Bergstra, J. and Bengio, Y. Random search for hyper- parameter optimization. Journal of machine learning research, 13(2), 2012.
Isik, B., Ponomareva, N., Hazimeh, H., Paparas, D., Vas- silvitskii, S., and Koyejo, S. Scaling laws for down- stream task performance of large language models. arXiv preprint arXiv:2402.04177, 2024.
Biderman, S., Schoelkopf, H., Anthony, Q. G., Bradley, H., O' Brien, K., Hallahan, E., Khan, M. A., Purohit, S., Prashanth, U. S., Raff, E.,, al. Pythia: A suite for ana- lyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2430. PMLR, 2023.
Jain, A., Montanari, A., and Sasoglu, E. Scaling laws for learning with real and surrogate data. arXiv preprint arXiv:2402.04376 2024.
Ji, K., Yang, J., and Liang. Y. Bilevel optimization: Con- vergence analysis and enhanced design. International conference on machine learning, pp. 4882-4892. PMLR, 2021.
Bisk, Y., Zellers, R., Gao, J., Choi, Y., et al. Piqa: Reasoning about physical commonsense in natural language. In Pro- ceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Jiang, Y., Zhou, A., Feng, Z., Malladi, S., and Kolter, J. Z. Adaptive data optimization: Dynamic sample selection with scaling laws, arXiv preprint arXiv:2410.11820. 2024.
Blakeney, C., Paul, M., Larsen, B. W., Owen, S., and Fran- kle, J. Does your data spark joy? performance gains from domain upsampling at the end of training. arXiv preprint arXiv:2406.03-70 2024.
Kim, S., Thiessen, P. A., Bolton, E. E., Chen, J., Fu, G., Gin dulyte, A., Han, L., He, J., He, S., Shoemaker, B. A., et al. Pubchem substance and compound databases. Nucleic acids research,
Chen, T. and Guestrin, C. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd acm sigkdd inter- national conference on knowledge discovery and data mining, pp. 785-794, 2016.
2016.
Li, H., Zhao, X., Li, S., Wan, F., Zhao, D.. and Zeng, J. Improving molecular property prediction through a task similarity enhanced transfer learning strategy. Iscience, 25(10), 2022.
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try are, the ai2 reasoning challenge. ArXiv, abs/I 803.00333540 05450,,55,22 2018.
Li, J., Fang, A., Smymiss G., Ivgi, M., Jordan, M., Gadre, S., Bansal, H., Guha, E., Keh, S., Arora, K., Garg, S., Xin, R., Muennighoff, N., Heckel, R., Mercat, J., Chen, M., Gururangan, S., Wortsman, M., Albalak, A., Bitton, Y., Nezhurina, M., Abbas, A., Hsieh, C.-Y., Ghosh, D., Gardner, J., Kilian, M., Zhang, H., Shao, R., Pratt, S., Sanyal, S., Ilharco, G., Daras, G., Marathe, K., Gokaslan, A., Zhang, J.. Chandu, K., Nguyen, T.. Vasiljevic, I., Kakade, S., Song, S., Sanghavi, S., Faghri, F,, Oh, S., Zettlemoyer, L., Lo, K., El-Nouby, A., Pouransari, H., Toshev, A., Wang, S., Groeneveld, D., Soldaini, L., Koh, P. W., Jitsev. J., Kollar, T., Dimakis, A. G., Carmon, Y.. Dave, A., Schmidt, L., and Shankar, V. Datacomp-Im: In search of the next generation of training sets for language models. arXiv, 2024a.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024.
Duchi, J. C. Introductory lectures stochastic optimization. The mathematics of data, 25:99-186, 2018
Everaert, D. and Potts, C. Gio: Gradient information opti- mization for training dataset selection. ICLR, 2024.
Fan, S., Pagliardini, M., and Jaggi, M. Doge: Do- main reweighting with generalization estimation. arXiv preprint 100 15393, 2023.
Held, W., Paranjape, B., Koura, P. S., Lewis, M., Zhang, F, and Mihaylov, T. Optimizing pretraining data mixtures with llm-mmtinated utility. arXiv preprint arXiv:2501. 2025.
Li, J., Fang, A., Smymis, G., Ivgi. M., Jordan, M., Gadre, S., Bansal, H., Guha, E., Keh, S., Arora, K., et al. Datacomp- Im: In search of the next generation of training sets for lan- guage models. arXiv preprint arXiv:2406. 1174,, 2024b.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:200. 2022.
Liu, Q., Zheng, X., Muennighoff, N., Zeng, G., Dou, L., Pang, T., Jiang, J., and Lin, M. Regmix: Data mixture as regression for language model pre-training. arXiv preprint arXiv:2407.01492. 2024.
10

---

MixMin; Finding Data Mixtures via Convex Minimization
Paster, K., Santos, M. D., Azerbayev, Z., and Ba, J. Open- webmath: An open dataset of high-quality mathematical web text, 2023.
Welbl, J., Liu, N. F,, and Gardner, M. Crowdsourcing multiple choice science questions. In NUT@@ @MMNLP 2017.
Pedregosa, F. Hyperparameter optimization with approxi- mate gradient. In International conference machine learning, pp. 737-746. PMLR, 2016.
Xiao, S., Liu, Z.,, Zhang, P., Muennighoff, N., Lian, D., and Nie, J.-Y. C-pack: Packed resources for general chinese embeddings. arXiv, 2024.
Pouget, A., Beyer, L., Bugliarello, E., Wang, X., Steiner, A. P., Zhai, X., and Alabdulmohsin. I. No filter: Cul- tural and socioeconomic diversityin contrastive vision- language models. arXiv preprint arXiv:2405.13771 2024.
Xie, S. M., Santurkar, S., Ma. T., and Liang. P. Data se- lection for language models via importance resampling. NeurIPS, 2023.
Xie, S. M., Pham, H., Dong, X., Du, N., Liu, H., Lu, Y., Liang, P. S., Le, Q. V., Ma, T., and Yu, A. W. Doremi: Optimizing data mixtures speeds up language model pre- training. Advances Neural Information Processing Systems, 36, 2024.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21 (140):1-67, 2020.
Yang, K., Swanson, K., Jin, W., Coley, C., Eiden, P., Gao, H., Buzzan-Perezz A., Hopper, T., Kelley, B., Mathea, M., et al. Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling. 59(8):3370-3388) 2019.
Rogers, D. and Hahn, M. Extended-connectivity finger- prints. Journal of chemical information and modeling, 50 55:772-7554 2010.
Ye, Z., Yang, Y., Li, X., Cao, D., and Ouyang, D. An inte- grated transfer learning and multitask learning approach for pharmacokinetic parameter prediction. Molecular pharmaceutics, 16(2):533-525 --133,51,,,11 2018.
Salem, M., Khormali, A., Arshadi, A. K., Webb, J., and Yuan, J.-S. Transcreen: transfer learning on graph-based anti-cancer virtual screening model. Big Data and Cogni- tive Computing, 4(3):16, 2020.
Seidl. P., Vall, A., Hochreiter, S., and Klambauer, G. En- hancing activity prediction models drug discovery with the ability to understand human language. In In- ternational Conference on Machine Learning, pp. 30458- 30490. PMLR, 2023.
Shi, C., Lu, J., and Zhang, An extended kuhn-tucker approach for linear bilevel programming. Applied Mathe- matics and Computation, 162(1):51-56) 2005.
Sun, X., Zhu, Chen, B., You, H., and Xu, H. A feature transferring workflow between data- poor compounds in various tasks. Plos one, 17(3):e0266088, 2022.
Thrush, T., Potts, C., and Hashimoto, T. Improving pre- training data using perplexity correlations. arXiv preprint arXiv:2409.05816, 2024.
Thudi, A. and Maddison, C. J. Finding optimally robust data mixtures via concave maximization. arXiv preprint 2rXiv:2406.01447, 2024.
Touvron, H., Martin, L.., Stone, K., Albert, P., Almahairi, A., Babaci, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint arXiv:2307.09288. 2023.

---

MixMin; Finding Data Mixtures via Convex Minimization
A. Proofs
A.1. Theorem 3.1
Proof. Recall that the Bayes optimal model for conditional cross entropy is p(yla) and for conditional MSE is By In both cases, letting fx be the Bayes optimal for the mixture dp> Eper Apdp, and fp be the Bayes optimal for the individual sources, we have for cross-entropy and P.
Edpep {p5p(z)p(z)
Edger Appl (x)
and in particular, when there is no covariate shift amongst the sources we have fa = Edpep (p5z(x).
Note that the Bayes optimal for unconditional CE is just p(x), and so similarly fx(x) = Edpep Apfp(x).
With this formula for fx, we then have our Data Mixing objective reduces to just learning linear model over fz(2). Specifically plugging in fa Edpep (2))(2) into Data Mixing gives Equation 1.
A.2. Lemma 3.2
Proof. First note that, using the Lipschitz criterion with the condition the hypothesis space and the definition of data mixing, letting H" contain the Bayes optimal for all mixtures, we have AA E A that
DM (A, dt) DMH (A, dt)l s Ce.
Specifically, let fH, and , A be the minimizer of the mixture À in H and H respectively. By definition of H fu is Bayes optimal, and by the assumption in the lemma statement, ıI -ll I Now note IDM (A, di) (x, DM (A, dt)l = If (()H..(z), y)ddt) (x, y) - [(ff=.((ee (x), v)dt(z, y) < C'IlHHA- - fH- by the Lipschitz assumption.
Combining the two inequalities we have [ dt) - DM (A, dt) < Cc which gives the desired inequality.
This relation will be enough for the proof, and hence weaker criteria the lemma that ensure this relation will suffice.
Let Ā arg min DM (A, dt) be the optimal data mixture weights and recall X arg AF DM (A, dt) where 71* contains the Bayes optimal functions for all mixtures. Now note the left hand side of the inequality
DM*(A", di)
min DMH(A, di) = DM*(A" dt) - DMH (X", dt) + DM (X", DM+((M,(dd))) dt)
YEAP
The first difference is bounded by Ce by the earlier inequality for changing hypothesis spaces. Now bounding the second difference we have
(M DM dt) = DM (X*, DM (ī, dt) + DM (7, dt) DM dt)
S DM (2, di) - DMH(A, dt) s Ce
where the first inequality came from the definition X*, and the second from the inequality stated the beginning the proof for changing hypothesis spaces.
Hence conclude
DM di) min DM*(A,d1) S 2C
B. Extra Tables and Figures
12

---

MixMin; Finding Data Mixtures via Convex Minimization
Target
First
Second
Third
Fourth
ussaylD-1
assayID-89
assaylD-620
assayID-618
assayID-693
assaylD-3
assaylID-620
assaylD-5
assaylDD-618
assaylD-891
assayID-5
assayID-620
assayID-891
assayID-92
assayID-425
assaylID-7
assaylID-620
assaylD--52
assaylD-693
assaylD-618
assaylD-9
assayID-891
assayID-693
assayID-620
assayID-256
assayID-11
assayID-620
assaylD-953
assayID-692
assaylD-618
assaylD-13
assaylD-618
assaylD-710
assaylD-891
assaylD-451
assaylD-15
assayll-119
assaylID-20
assaylD--25
assaylD-758
assaylD-17
assayID-19
assayID-618
assayID-758
assayID-21
assayID-19
assayID-25
assayID-17
assayID-957
assayID-21
Table 1: The top 4 assays found by MixMin (over the 1328 assays in PCBA) for each target task. We list the PubChem
assaylD corresponding the assay in PCBA. Red indicated assays measuring cytotoxicity, green phenotypic screens, blue are specific protein binding/inhibition assays, and orange are other types of assays.
Target
First
Second
Third
Fourth
assaylID-1
0.036705
0.036184
0.035904
0.030910
as-ayID-3
0.05026778
0.04660426
0.0414932
0.03289703
as-aylD-5
0.05026778
0.04660426
0.0414932
0.03289703
assaylD-7
0.03170111
0.02936895
0.02839234
0.02707635
assayll>-9
0.03645385
0.02978397
0.02941164
0.02894659
assayID-11
ID-11
0.04995361
0,03905968
0.03367052
0,03212925
assay assayID-13
ID-13
0.05764348
0.04289928
0,03077811
0.02817639
assay assayID-15
ID-15
0.11996419
0.04302281
0.03967432
0.03457846
assay ID-17
0.22989839
0.04496526
0.04188494
0.03894247
assaylD-19
0.10650713
0.05704838
0.05019803
0.04435973
Table 2: The mixing weights for the top 4 assays found MixMin (over the 1328 assays in PCBA) for each target task. We list the PubChem assaylD corresponding to the target assay in PCBA.
1% of Training Cost
1% of Training Cost
Pred Loss Method PredLos
Aceurxy
0.05
Notaal
MisMin
0.04
Dalanced
-02
0.02
RegMiix
0.00
-0.4
-0.02
Natura
MixMin
0.04
0.6
Bolanced
-0.05
Randon Seaech
-08-
RegMix
-0.08
PIQA (4.01)
ARC Easy (6.15)
SciQ (6.84)
PSQA (0.02)
ARC Easy (0.15)
ScQQ (0.71)
Figure 8: In almost all cases MixMin improves or matches all baselines across the target tasks using 1% the final training compute (Left). We report improvement over the predictive loss (left) and accuracy (right) of training on the natural distribution, which is stated the x-axis with the task name.
13

---

MixMin: Finding Data Mixtures via Convex Minimization
Baseline
PIQA
ARC-EASY
SciQ
Open WebMath
Random Search 0.006
0.011
0.019
0.839
RegMix 0.007
0.004
0.003
0.051
Balanced 0.008
0.013
0.006
0.027
Natural 0.025
0.206
0.375
0.0003
Table 3: p-values for Welch's t-test with the null hypothesis being MixMin performs the same as the baselines for the 1 16-compute experiments presented in Figure (generative loss) over the 3 trials used.
Baseline
1%
5%
10%
100%
Random Search
0.006
0.0003
0.233
0.043
RegMix
0.007
0.011
0.023
NA
Balanced
0.008
0.0003
0.017
0.002
Natural 0.025
0.004
0.224
0.018
Table 4: p-values for Welch's 's t-test with the null hypothesis being MixMin performs the same the baselines PIQA for generative loss across compute budgets (as shown in Figure 3) over the 3 trials used.
Baseline
PIQA
ARC-EASY
SciQ
Open WebMath
Balanced
0.019
0.000
0.003
0.020
Natural
0.052
0.003
0.409
0.002
Table 5: p-values for a Welch's 's t-test with the null hypothesis being MixMin performs the same as the baselines on generative loss across tasks for the 410M models experiments shown in Figure 5.
OpeaMath
00-
-685-
Naecal Loo Mato
-800
-22-
-815
kodtt
Balaod
Ralaned
Ralaned
Rackee bnoch
Ravies Srock
Fimker Soarh
RecMta
Reghts
Reg
-035
1005
105
26
Miking Coogy Tebinn Coupee
Tai Theeigg Congee
Jun Kideg Cregue Thaaeg Orogge
Figure 9: MixMin consistently performed better than alternative methods on generative loss across compute budgets, and robust using less compute. We note both Random Search and RegMix tend improve with more compute. We report improvement over the generative loss of training the natural distribution.
PIQA
ARC Eayy
S:12
00-
0.2
084
-11-
Neaui
Netmi0 StaMte
Natara Mudiz
Moke
-03-
Regita
Aegitt
Figure 10: MixMin often performed better par alternative methods on predictive loss across compute budgets, and was robust to using less compute. We note both Random Search and RegMix tend to improve with more compute. We report improvement over the predictive loss of training on the natural distribution.
14

---

MixMin: Finding Data Mixtures via Convex Minimization
MQA
ARC-Blyy
Se1
0.05-
Nateel (ywc)
0:04
0.82
Balacod
-2.00
-0.00
2006
Dans
Figure 11: MixMin often performed better par to alternative methods on accuracy across compute budgets, and was robust to using less compute. We note both Random Search and RegMix tend to improve with more compute. We report improvement over the accuracy of training on the natural distribution.
M404
ABC Eay
Se0
4.1-
4.0-
6.4-
Molto 180511S
60-
Dutaved
Ninool
160M
100M
+3030
160M
400M
Figure 12: In most cases MixMin consistently improved predictive loss over the baselines as we scaled the models. We report predictive loss relative to pythia-160. 160M model trained with the natural distribution of tokens (lower in the figures is better). MixMin weights were found using % the compute of the 160M model training run, which is 0.15% the compute of the 410M training run.
MicMia 1% of Training Cait
Rundon Searc 16 of Training Cail
RegMEs 18 of Training Coal
0.5
7004
06
0.5
Ars Exp
0112
9:35
0117
4025
1.30
8.42
0080
0.90
U144
11:4
8090
M12
1117
4:15
0:50
4117
422
A11
S
00.5
0:42
000
0300
0814
0141
42
021
1000
Figure 13: The mixture weights found by MixMin, Random Search, and RegMix when using 1% the compute for 1ythia-160M training run on 3.2B tokens. Note the domains (X-axis) are "RedPajanaCCommoCCaa/// "RedPajamaC4", "RedPajamaGithub", "RedPajamaBook", "RedPajamaArxivv. "RedPajamaaWikipedia" "RedPajamaStncEExchangg re- spectively.
15

---

MixMin; Finding Data Mixtures via Convex Minimization
Evalazam
Eralenticn MicMin
Accumoy Evalzamir Misdir
HOK
633
3A
2.7
12
544
34
5D
Ope
8.16
79gA
Eay
30 PRgA
Seiq
PALA
MisMeidIM Treand
MoMe-diM Treand
KinMas 4I8M Truand
Figure 14: found the MixMin models for PIQA, Are Easy, and SciQ perform similarly across tasks, however were (almost) always best for their own task across metries. We report the cross performance of MixMin for different tasks the 410M parameter scale. The subfigures report generative loss, predictive loss, and accuracy respectively.
Gen Laos Bodtedion of Booelin
Poad Loss Beshseiin of Buoliin
Araracy Bvdbedire of Bearlts
0.44
Broaad for bog
30
555
840
076
balamd
Boolie LEA4 Toiked
Bandze (IM Troked
Ruriia <1Eb Trared
Figure 15: We report the performance of training on the natural and balanced mixture for different tasks at the 410M parameter scale. The subfigures report generative loss, predictive loss, and accuracy respectively. Comparing to Figure 14 we see both the natural and baseline mixture are matched or dominated in average performance across tasks by the MixMin weights for PIQA, Are Easy, and SeiQ.
Mixing
Mixing
Dassuts
05-
MisMaa
Rasicr Seuck
Ruddea Stuck
Ranien Seuth
3.4
64-
03
82
022
Figure 16: MixMin improved over using random search as increased the number of surrogate assays. We report AP scores for the first 10 assays in PCBA, with error bars representing a 95% confidence interval over 3 trials. We only report the first 3 assays for the 1279 sources setting.
16

---

MixMin; Finding Data Mixtures via Convex Minimization
Mising 9 Sumogan Dissets
Miciy Samprro Doowess
Mixiag 1279 Semoupsic Dilasets
66
Raghis
RegMis
Reghds
Ammy
Figure 17: MixMin improved over using RegMix as we increased the number of surrogate assays. We report AP scores for the first 10 assays PCBA, with error bars representing a 95% confidence interval over trials. We only report the first 3 assays for 1279 sources setting.
1% of Thining Cut Truds
1% of Toining Con: STrab
1% of ofTrmme Cre STra
Namm
Mobte
Lut Mothi Cos
01
0.00
Netui
-082
Figure 18: MixMin often performed better or on par to alternative methods on accuracy, generative loss, and predictive loss on PIQA for the 160M model using 1% compute for proxy models. Here we present results over 5 trials.
17


Reviews:
[{"review_id": "yJxoIeIECl", "reviewer": "Reviewer_KxeH", "summary": "The authors propose a novel method (called MixMin) for the problem of optimizing data mixtures for pre-training of large ML models in order to improve performance in downstream tasks. The MixMin method proposed by the authors solves this optimization problem with the following approach: First, MixMin trains a cheap/small “proxy model” on each of the individual data sources. Then, MixMin learns the best mixture of the “proxy model” outputs for the target dataset. This problem of finding the best mixture of “proxy model” outputs is a convex optimization problem, so MixMin solves this problem using simple entropic descent. Finally, MixMin outputs the optimized mixture of “proxy model” outputs as the optimal mixture of the corresponding data sources used to train each “proxy model”. The resulting mixture of data sources found by this MixMin method can then be used to train the larger, complex ML model that we want to pre-train and apply to the downstream task. The authors show that the larger the ML model of interest, the better the convex optimization problem solved by the MixMin method approximates the true data mixing objective. Thus, for models that are sufficiently large, solving the true data mixing optimization problem effectively becomes a simple convex optimization problem using the MixMin method. The authors demonstrate empirically that using MixMin to find optimal data mixtures improves performance of models on downstream tasks in language modeling and chemistry. Empirical results show that MixMin consistently outperforms baseline approaches for optimizing data mixtures, and that MixMin can find good data mixtures with significantly smaller computational cost compared to other approaches.\n\n## update after rebuttal\nAs I stated in my comment below, the authors answered all of my questions in their rebuttal, and I agree with the author's proposal of adding the provided list of \"weak conditions for when MixMin works\" to the paper. I therefore maintain that this work should be accepted for all of the reasons stated in my initial review.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and feedback! We elaborate on questions raised in the review below, and will update the draft with the typos pointed out by the reviewer.\n\n> Can the authors provide any additional insight into when they expect MixMin to work well in practice? In what scenarios (i.e. what types of data sources, downstream tasks, etc.) do the authors expect that using these cheap proxy models and applying the MixMin method will be successful? In what scenarios do the authors expect that MixMin would not be useful (i.e. data sources with covariate shift)?\n\nIf our proxy models are sufficiently close to Bayes optimal, then MixMin will find the near optimal mixture weights. This is if there is no covariate shift among our sources, and that we ultimately train another near Bayes optimal model on the mixture (i.e., big data and big model regime). Note, if the downstream task is not close to any mixture, then MixMin (or any data mixing) may not help performance on the downstream task.\n\nThe main question is when are the proxy models “practically” sufficiently close to Bayes optimal for the MixMin optimization. Our language modeling experiments showed that using 1% for the proxy models is comparable to using 100% the compute of a full training run. We believe this suggests a general approach to selecting the compute for the proxy models–we could iteratively increase the proxy model compute and run MixMin and see when the mixture becomes stable across scale. Ultimately, we did not need to do this as the cheapest settings we tested already performed better than the baselines.\n\n> It would be nice to have a concrete list of criteria for problems such that the authors think that MixMin would be useful.\n\nWe thank the reviewer for the suggestion of providing a list of criteria for using MixMin. Given our previous discussion we propose to add the  following list to the paper:\n\n“(Weak) Conditions for when MixMin works: \n\n1) The data sources do not have covariate shift (or we can change the learning problem to not have covariate shift)\n\n2) We believe the downstream task is close to/contained within a mixture of our sources\n\n3) Our final model class is very expressive (i.e., do not have significant architectural biases)\n\n4) Our proxy model learnt some signal for each source (e.g., even with very little data): if we have flexible compute, consider iteratively using more compute for the proxy models until some stability in the mixtures is reached”\n\n> Do the authors have any concrete ideas of how MixMin could be adapted in the future to handle data sources with covariate shift (i.e. image classification)? To be clear, I don’t think that concrete ideas for this are necessary to include in this paper, but I am intellectually curious.\n\nFor covariate shift, the primary bottleneck we hit was non-convexity. The formula for the Bayes optimal function of a mixture with covariate shift is a fraction with a linear model of the mixture weights in the numerator and denominator (see Appendix A.1 formula). When we naively ran entropic descent to optimize the mixture with this expression we hit instability, and couldn’t get convergence without carefully tuning the learning rate. Even then it was not clear to us if this was a global optima. In fact, in preliminary experiments directly optimizing the conditional p(A|Q) tasks for LLMs  (which was non-convex given covariate shifts in the question strings) gave worse mixtures to our approach of making the task generative, i.e., learning p(Q,A). \n\nAn analog of our approach, replacing p(y|x) tasks with p(y,x), could be done for image classification. People are currently training generative models to jointly learn the distribution of labels and images (multi-modal models). In this case MixMin (with no further modifications) can be used despite there being “covariate shifts” between the sources (i.e., image shifts).", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes, I checked the proofs in Appendix Section A and they are correct as far as I am aware.", "experimental_designs_or_analyses": "Yes, I checked all experimental design setups and experimental results provided in the paper. All are valid as far as I am aware.", "relation_to_broader_scientific_literature": "Finding an optimal data mixture for machine learning pipelines is a very relevant problem to the broader ML research community because improving the data mixture used to train a model can significantly improve performance for downstream tasks. This problem is a particularly relevant problem for researchers who work on training and applying LLMs. The author’s proposed MixMin method for solving this problem is therefore relevant to the broader research community, especially since the authors provide convincing empirical evidence that MixMin consistently outperforms baseline approaches.", "essential_references_not_discussed": "There are no essential references missing as far as I am aware.", "other_strengths_and_weaknesses": "Strengths:\n\nStrength 1: Relevancy of problem setting: The problem of finding an optimal mixture of data to pre-train models efficiently is clearly a relevant problem that can improve performance of models in downstream tasks across a wide variety of applications and ML pipelines. \n\nStrength 2: Empirical results: The empirical results in the paper provide convincing evidence that the proposed MixMin method consistently outperforms other methods for finding optimal data mixtures across a variety of downstream tasks. This is further strengthened by the inclusion of performance on relevant, real-world chemistry tasks, and by the additional results demonstrating that MixMin can achieve good mixtures with smaller computation cost. \n\nStrength 3: Quality of Writing: The writing is very clear and easy to understand. \n\nWeaknesses: \n\nA minor weakness of this paper is that the novel methodological contribution is on the smaller side because the proposed MixMin method is a fairly simple approach to optimizing data mixtures by training proxy models and then optimizing a mixture with entropic descent. That being said, this is a very minor weakness because the authors show that the MixMin method works very well empirically, consistently outperforming other approaches and cutting computation cost. I therefore still recommend that this paper be accepted.", "other_comments_or_suggestions": "Typos: \n\n1: Line 125/126: \"filering\" should be \"filtering\" \n\n2: Line 209/210: \"resuse\" should be \"reuse\"\n\n3: Line 572/573: \"left hand sign\" should probably be \"left hand side\"?", "questions_for_authors": "Question 1: \nCan the authors provide any additional insight into when they expect MixMin to work well in practice? In what scenarios (i.e. what types of data sources, downstream tasks, etc.) do the authors expect that using these cheap proxy models and applying the MixMin method will be successful? In what scenarios do the authors expect that MixMin would not be useful (i.e. data sources with covariate shift)? It would be nice to have a concrete list of criteria for problems such that the authors think that MixMin would be useful. Such a list would allow practitioners to more easily determine if MixMin is applicable to their particular problem setting, so I think adding something like this would strengthen the paper. \n\nQuestion 2: \nDo the authors have any concrete ideas of how MixMin could be adapted in the future to handle data sources with covariate shift (i.e. image classification)? To be clear, I don’t think that concrete ideas for this are necessary to include in this paper, but I am intellectually curious.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qr4kJ19Pdr", "reviewer": "Reviewer_YnxL", "summary": "This paper proposes MixMin, a simple but effective method for solving the data mixture coefficients in large language model pretraining. The authors identify that the bi-level optimization objective for solving data mixture is intractable. But luckily, such an objective tends to be convex when model classes become larger. Thus, such an objective can be effectively solved using a gradient-based approach. Some theoretical justifications are provided along with experimental evaluations for LLM pretraining experiments to justify the effectiveness of the proposed MixMin method.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and feedback. We discuss specific questions below.\n\n> Only models under the parameter size of 1B have been evaluated. It's not clear if the method scales well to larger model sizes, such as around 7B.\n\nWe agree it would be nice to extend our analysis to larger scale experiments, but at the moment, we do not have the computational resources to run those experiments. We did find MixMin’s relative performance to be stable over the two scales that we tested (160M, 410M), and also in non LLM domains (chemistry). Although it's far from conclusive, it suggests that the benefits of MixMin may continue to hold.\n\n>It seems that the experiments directly use domains contained in SlimPajama. I wonder if the method also works for data domains like web crawl data, code, books, arXiv/Wikipedia, etc., which are more common in LLM pretraining practices.\n\nSlimPajama domains already contain the mentioned categories. The seven domains in SlimPajama are: CommonCrawl, C4, Github, Books, arXiv, Wikepedia, and Stackexchange. We also tested MixMin on sources from a completely different domain (chemistry), with varying combinations of sets, and MixMin continued to perform well.\n\n> It seems that the major computational overhead comes from training small LLMs on sub-data domains. I wonder how large the small-scale LLMs need to be and what can be done when there are many data domains.\n\nOn the question of many domains, in our chemistry experiments we tested MixMin across many sets of sources, ranging from 9 to 1280 sources. We found MixMin improved as the number of sources grew.\n\nIn terms of complexity, we found MixMin was more efficient than past approaches. MixMin still needs to train a proxy model for each source which scales with the number of sources, but this is true of past approaches. Importantly, we found the proxy models for MixMin could be obtained cheaply by using smaller models and less data (see our LLM experiments where we used <1% of the cost of the full training run to compute all the proxy models).\n\nFurther experiments could explore pushing the cost of data mixing smaller than what we tested (<1% of the full training run), perhaps to find when MixMin would break. We will acknowledge this open problem in a revised draft.", "claims_and_evidence": "The claim of MixMin is that it improves model performance (in terms of better benchmark evaluation scores) with marginal additional computing overhead. This claim is well supported by the empirical evaluations.", "methods_and_evaluation_criteria": "Both the proposed method and evaluation criteria make sense and follow the common practice of the LLM-related research field.", "theoretical_claims": "I checked the proofs of Theorem 3.1 and Lemma 3.2; both seem to make sense.", "experimental_designs_or_analyses": "The experimental design and analyses make sense and are pretty thorough.", "relation_to_broader_scientific_literature": "The proposed method can potentially improve the data curation and training recipe of large language models or foundation models in general.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n- The paper is well-written and well-motivated.\n- Optimizing data mixture for LLM pretraining is an impactful research direction.\n- The proposed method is intuitive and easy to understand.\n- Both theoretical and empirical evaluations have been provided to justify the performance of the proposed method.\n\nWeaknesses:\n- Only models under the parameter size of 1B have been evaluated. It's not clear if the method scales well to larger model sizes, such as around 7B.\n- It seems that the experiments directly use domains contained in SlimPajama. I wonder if the method also works for data domains like web crawl data, code, books, arXiv/Wikipedia, etc., which are more common in LLM pretraining practices.\n- It seems that the major computational overhead comes from training small LLMs on sub-data domains. I wonder how large the small-scale LLMs need to be and what can be done when there are many data domains.", "other_comments_or_suggestions": "Please see the ones in \"Strengths And Weaknesses\".", "questions_for_authors": "Please see the ones in \"Strengths And Weaknesses\".", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "w26puR0OdR", "reviewer": "Reviewer_nsAa", "summary": "This submission addresses the optimization of data source mixtures, formulating it as a bi-level optimization problem. The key result is Theorem 3.1, which states that under certain conditions (cross-entropy or mean squared error loss, hypothesis class contains Bayes optimal models), the optimal mixture weights are the same as those for optimally mixing the outputs of Bayes optimal models for each data source. The authors then propose the MixMin method as a practical realization of the key result, where less expensive proxy models are substituted for Bayes optimal models. Experiments are conducted on language model pre-training for better downstream task performance and on transfer learning for chemical assay prediction. MixMin's data mixtures achieve larger gains than the baselines considered in these experiments.\n\n### Update after rebuttal\n\nThanks to the authors for their follow-up reply. The verbatim changes to the theorems and proofs are good and I approve them.\n\nI am still uneasy about reporting confidence intervals or performing statistical tests on only 3 experiment trials. I recommend that the authors take the time before the camera-ready submission to increase the number of trials to at least 5, for all experiments, and perform Welch's t-tests as in the rebuttal and/or recompute confidence intervals.\n\nOverall, I am placing not insignificant trust in the authors to carry out the revisions (including the clarifications on experiments and bi-level optimization mentioned previously) and hope that they will do so for the camera-ready version.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and feedback. We discuss the main questions below, and will incorporate the other suggestions to our revised draft. \n\n> For the proof of Lemma 3.2, it is not clear how the first displayed inequality is obtained…\n\nWe describe the derivation below, and will add this to our revised draft. \n\nLet $f_H$ and $f_{H^*}$ be the minimizer of the mixture $w$ in $H$ and $H^*$ respectively. By definition of $H*$, $f_{H^*}$ is Bayes optimal, and by the assumption in the lemma statement, $||f_H - f_{H^*}|| \\leq \\epsilon$.  Now note  $|DM_{H}(w)- DM_{H^*}(w)| = | Loss(f_H) - Loss(f_{H^*})| \\leq C ||f_H - f_{H^*}||$ by the Lipschitz assumption. Combining the two inequalities we have $|DM_{H}(w) - DM_{H^*}(w)| \\leq C \\epsilon$, which gives the desired inequality.\n\n> I do not understand the claim in lines 157-159… \n\nWe mean MixMin holds for the unconditional CE loss, $\\int -log(f(x)) dp(x)$, which matches $f(x)$ to $dp(x)$. Mixmin also holds for conditional CE and MSE losses under no covariate shift: $\\int -log(f_{y}(x)) dp(x,y)$ and $\\int ||f(x) -y||_2^2 dp(x,y)$, with $dp_i(x) = dp_j(x)$ for all sources. We will add these precise statements to our revised draft.\n\n> I think 3 train-test splits is too few…\n\nWe evaluated MixMin over many downstream tasks, and sources (with 3 trials for each setting). Over these hundreds of experiments we saw MixMin consistently improved over the baselines.  As a sanity check, we also increased the number of trials for our PIQA experiments with 160M models to 5 trials, and found the results were consistent. See the figures here: https://anonymous.4open.science/r/MixMin_LLM_extra_trials-54C6\n\n> Why choose the loss to be the (negative) log probability of the question and correct answer… \n\nThe A|Q task is a conditional CE task that can have covariate shifts (different probability of Qs among sources): MixMin does not apply. Hence we chose to learn A+Q which is an unconditional CE task and so MixMin applies. See our response to Reviewer KxeH for more discussion on the issue of non-convexity for covariate shift data mixing.\n\n> I do not understand the claim that the RegMix baseline has higher computational cost…\n\nThe difference is the evaluation cost. RegMix requires evaluating 7X models (where X can be 10 in our experiments) but MixMin requires evaluating just 7.\n\n> What is meant by \"baseline models\"? Why are they trained on the train + validation splits, while MixMin's proxy models are trained only on the train splits?..\n\nWe meant models (XGBoost) trained solely on the downstream task data, using the train + validation splits for that task. MixMin models are trained on the source datasets\n\nWe left the validation sets untouched to allow for unforeseen future experiments that require a validation set (e.g., additional hyperparameter tuning on the MixMin mixture). In this paper the default XGBoost hyperparameters on the MixMin mixture already worked better than the baseline model (with hyperparameters selected from cross validation), so they were unused.\n\n> Does the training data for the final model include the train split (or train + validation) for the target assay as well?\n\nNo, it does not include the target assay. The train split of the target was only used to tune the MixMin weights. \n\n> For the chemical assay experiment, why is the natural distribution the only baseline considered?...\n\nOur language modeling experiments already showed the other baselines often performed worse than natural. We have now tested RegMix and Random Search on the chemistry experiments. We observed the same trends as natural; the baselines become worse while MixMin becomes better with increasing sources (only tested 3 targets for 1280 sources due to compute). The figures are here: https://anonymous.4open.science/r/MixMin_Chem_More_Baselines-286C/\n\n> To what extent are gradient-based methods for bi-level optimization applicable?\n\nThe methods from this literature were not applicable to our problem, so we did not discuss them. To clarify this we plan to include the following paragraph in our revised draft:\n\n“Our setting for Bi-Level optimization presents several challenges for previous approaches, including those used for other hyperparameter optimizations. Firstly, the inner-optimization is over an arbitrary function space (e.g., non-parametric models like XGBoost), and so lacks the parametric gradients needed for many methods [1,2]. Alternative constraint based approaches through KKT suffer similar issues given the non-parametric space of the inner optimization [3]. For many model classes, one could parameterize the inner optimization, but we then lose the necessary convexity for these methods (e.g., Neural Networks).”\n\n[1] K. Ji et al “Bilevel optimization: Convergence analysis and enhanced design”\n\n[2] F. Pedregosa, “Hyperparameter optimization with approximate gradient”\n\n[3] C. Shi et al “An extended kuhn–tucker approach for linear bilevel programming”", "claims_and_evidence": "On the theoretical side, I have issues with the proof of Lemma 3.2 and the applicability of the results to generative models. On the experimental side, I think that 3 train-test splits is too few and have several questions about the experimental design. Please see the respective sections below.", "methods_and_evaluation_criteria": "The proposed MixMin method makes sense for data mixture optimization thanks to the connection given by Theorem 3.1. The two experimental evaluations, data source optimization for language model pre-training and transfer learning for chemical prediction, are also sensible.", "theoretical_claims": "I checked the proofs of Theorem 3.1 and Lemma 3.2. I have two major issues and three minor ones.\n\nMajor issues:\n1. For the proof of Lemma 3.2, it is not clear how the first displayed inequality is obtained, which is key for the remainder of the proof. More steps should be shown, and perhaps the authors have to be more precise about which norm is used for functions $f$.\n1. Although this is not a formal theoretical result, I do not understand the claim in lines 157-159 that \"distributions for generative tasks are only over 'labels' $p(y)$.\" This statement should be explained as it supports the application of MixMin to language modeling.\n\nMinor issues:\n1. \"No input distribution shift/covariate shift\" should be formally defined in Section 3.1 as it is not clear without reading the proof of Theorem 3.1.\n1. Footnote 2: I do not see how the claim that \"functions too far away will not have low error\" follows from the near-Bayes optimality assumed in Lemma 3.2.\n1. Lines 572-573: What does \"left hand sign of the inequality\" mean?", "experimental_designs_or_analyses": "I reviewed all of the experiments section (Section 6) and think the following issues are important:\n1. I think 3 train-test splits is too few. It does not seem meaningful to show 95% confidence intervals from three splits. How are these computed?\n1. For the language model multiple-choice tasks, why choose the loss to be the (negative) log probability of the question and correct answer (Q + A) sequence, as opposed to the probability of A given Q? It seems that evaluating the model on also generating the question Q makes the task harder than necessary.\n1. I do not understand the claim that the RegMix baseline has higher computational cost. As I understand it, for RegMix, the total number of tokens for training proxy models is (3.2B / 700 tokens/model) * (7X models) = 3.2B * (X / 100), which is the same as the (3.2B X / 700 tokens/model)  * (7 models) for MixMin and random search.\n1. For the chemical assay datasets, I do not understand the choices of training data that were made.\n    1. What is meant by \"baseline models\"? Why are they trained on the train + validation splits, while MixMin's proxy models are trained only on the train splits? What prevents MixMin's proxy models from using train + validation also?\n    1. Given the MixMin mixture, does the training data for the final model include the train split (or train + validation) for the target assay as well? I have the same question as above (train only vs. train + validation) for all assays included in this final training.\n1. For the chemical assay experiment, why is the natural distribution the only baseline considered? Why not include the same baselines from before: balanced distribution, random search, RegMix?", "relation_to_broader_scientific_literature": "I question the claim that bi-level optimization \"generally only admits expensive zero-order approaches\" (lines 19-21). A quick search yields survey/tutorial papers such as [Zhang et al. (2023)](https://arxiv.org/abs/2308.00788) that discuss gradient-based methods, based on implicit gradients or unrolling of the lower-level optimization. I think it is a significant omission to not at least discuss such methods. Furthermore, while gradient-based methods may make assumptions that do not hold for the data mixing problem here, it may still be possible to run the algorithms and compare results (particularly on the chemical assay datasets that I presume involve less computation than language modeling). Doing so would better contextualize MixMin's performance in terms of mixture quality and/or computational cost.", "essential_references_not_discussed": "Gradient-based methods for bi-level optimization (please see the previous comment).", "other_strengths_and_weaknesses": "I think the key result of Theorem 3.1 is also the main strength of the work, showing that under some conditions, it may be possible to avoid the bi-level optimization inherent to data mixture optimization and instead optimize the mixture of outputs of source-specific models.", "other_comments_or_suggestions": "Minor comments and questions:\n- Line 84, right column: Should \"a set of finite source distributions\" be \"a finite set of source distributions,\" i.e., \"finite\" applies to the set, not the distributions?\n- Lines 85-86, right column: Related to the previous comment, it is unclear at this point what the notation $dp$, $dt$ means.\n- Data Filtering is not Data Mixing: Is data filtering (at the data source level) not a special case of data mixing where some mixture weights are set to zero and the remaining non-zero weights are equal? Is the distinction being drawn here more about data sample level vs. source level? \n- Lines 171-172, right column: What is meant by \"the excess error of MixMin should be analyzed beyond the risk of the proxy models $\\hat{f}_p$?\n- Lines 325-326, left column: Does \"predictive loss\" mean cross-entropy with respect to the answer choices? This could be clarified.\n- Lines 335-336, right column: Does \"trained models over the first 100,000 molecules in PCBA\" mean the first 100,000 molecules in each assay?\n- I would appreciate a little more information about the chemical assay datasets. What is the data type of the target variable? Is the data tabular or otherwise structured?\n- \n- Figures 6 and 7: I suggest labelling the x-axis as \"target array\" for clarity.\n- Lines 432-433, left column: Should Table 2 be Table 1?", "questions_for_authors": "Please see:\n- Points 2-5 under Experimental Designs or Analyses\n- Major issues 1, 2 under Theoretical Claims (although the latter may require more involved responses)\n\nAlso from \"Relation to Broader Scientific Literature\": \n- To what extent are gradient-based methods for bi-level optimization applicable?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a method for optimizing the proportions of data that should be used from multiple sources in training large machine learning models. The idea is to train inexpensive proxy models on each data source individually, and then leverage the performance of these proxy models on downstream tasks to reason about the value of the individual data sources.

The reviewers highlight the empirical results and relevant topic tackled by the paper, and appear to have been satisfied with the author feedback responding to any questions, which largely amounted to minor requests for additional clarity.
Decision: Accept (poster)