Aioli: A Unified Optimization Framework for Language Model Data Mixing
Authors: Mayee F Chen
Michael Y. Hu
Nicholas Lourie
Kyunghyun Cho
Christopher Re
Venue: ICLR 2025 Poster (2025)
Keywords: data mixing
language models
data curation
Primary area: foundation or frontier models, including LLMs

Abstract: Language model performance depends on identifying the optimal mixture of data groups to train on (e.g., law, code, math). Prior work has proposed a diverse set of methods to efficiently learn mixture proportions, ranging from fitting regression models over training runs to dynamically updating proportions throughout training. Surprisingly, we find that no existing method consistently outperforms a simple stratified sampling baseline in terms of average test perplexity. To understand this inconsistency, we unify existing methods into a standard framework, showing they are equivalent to solving a common optimization problem: minimize average loss subject to a method-specific mixing law---an implicit assumption on the relationship between loss and mixture proportions. This framework suggests that measuring the fidelity of a method's mixing law can offer insights into its performance. Empirically, we find that existing methods set their mixing law parameters inaccurately, resulting in the inconsistent mixing performance we observe. Using this insight, we derive a new online method named Aioli, which directly estimates the mixing law parameters throughout training and uses them to dynamically adjust proportions. Empirically, Aioli outperforms stratified sampling on 6 out of 6 datasets by an average of 0.27 test perplexity points, whereas existing methods fail to consistently beat stratified sampling, doing up to 6.9 points worse. Moreover, in a practical setting where proportions are learned on shorter runs due to computational constraints, Aioli can dynamically adjust these proportions over the full training run, consistently improving performance over existing methods by up to 12.012 test perplexity points.

Published as a conference paper at ICLR 2025
AIOLI: A UNIFIED OPTIMIZATION FRAMEWORK FOR LANGUAGE MODEL DATA MIXING
Mayee F. Michael Y. Hu Nicholas Lourie,, Kyunghyun Cho Christopher Ré'
1 Computer Science Department, Stanford University; 2 Center for Data Science, NYU; 3 Computer Science Department, NYU; Prescient Design, Genentech
ABSTRACT
Language model performance depends on identifying the optimal mixture of data groups to train on (e.g., law, code, math). Prior work has proposed a diverse set of methods to efficiently learn mixture proportions, ranging from fitting regression models over train- ing runs to dynamically updating proportions throughout training. Surprisingly, that no existing method consistently outperforms a simple stratified sampling baseline in terms of average test perplexity. To understand this inconsistency, methods into a standard framework, showing they are equivalent to solving a common optimization problem: minimize average loss subject to a method-specific mixing law- an implicit assumption on the relationship between loss and mixture proportions. This framework suggests that measuring the fidelity of method's mixing law can offer insights into its performance. Empirically, we find that existing methods set their mixing law parameters inaccurately, resulting in the inconsistent mixing performance we ob- serve. Using this insight, we derive a new online method named ATOLI, which directly estimates the mixing law parameters throughout training and uses them to dynamically adjust proportions. Empirically, AIOLI outperforms stratified sampling 6 out of 6 datasets by an average of 0.27 test perplexity points, whereas existing methods fail to con- sistently beat stratified sampling, doing up to 6.9 points worse. Morcover, in a practical setting where proportions are learned on shorter runs due to computational constraints, A can dynamically adjust these proportions over the full training run, consistently improving performance over existing methods by up to 12.012 test perplexity points.
find
unify existing
INTRODUCTION
important to determine what data to train on for language model (LM) to acquire a range of capabilities, from generating code to understanding scientific literature and conversing with users (Albalak et al., 2024; Longpre et al., 2024; Li et al., 2024). To achieve this, practitioners mix data from various groups (such code files, scientific papers, and chat logs) in specific proportions compose overall training dataset-a procedure known as data mixing. Identifying the optimal mixture proportions is critical LLM performance. However, brute-force trial-and-erro search over the proportions computationally expensive, requiring many training runs.
Recent work introduces two types of data mixing algorithms that learn mixture proportions: offline and online methods. Offline methods conduct multiple training runs with varying proportions, fit regression model to predict performance, and use this model to determine the optimal static mixture (Ye al., 2024; Liu et al., 2024). Online methods adjust the mixture proportions dynamically throughout training using information from the model, such as its loss and gradients (Chen et al., 2023; Fan et al., 2024; Xie et al., 2023a; Albalak et al., 2023). All mixing methods require at least one training run to learn the proportions but are more efficient than a brute-force search.
Given the wide range of methods available, it is important to determine which ones are effective. However, when we evaluated existing methods, we found that no method consistently oupprfoormed stratified sampling-- simple baseline that uniformly mixes groups and requires extra training acssssross all sets of data groups in terms of average test perplexity (Table 2). This surprising outcome suggests that all existing methods suffer from some common weaknesses. To make progress in data mixing, we identify three objectives: 1) improve our understanding of the underlying assumptions of existing methods, 2) assess the fidelity of these assumptions in practice to better understand performance, and 3) apply insights to develop principled new data mixing methods.
"qqual contribution. Contact: mf woffche@@seee @st.nfordd.... edu, mi chael .h@@nyu.edu

---

Published as a conference paper at ICLR 2025
Unified Framework: Linear Mixing Optimization
2
Analyzing fidelity of existing methods
3
Our method: Aioli
m - Σ
A "(ttted)
Ajoxistingg)
shiiee E aaepp)
Estimate AAioi
ADDRRMI
MDDaE
Aski-it
Accea ring stopo
Ascietig
Proporson son
Properion BlackExchange
Liata mitetumn
modelss
Figure 1: Left: existing methods can be expressed in a unified optimization framework, in which they im- plicitly assume a linear or log- log-linear loss- proppoooootion relationship. Center: the (log)-linear parameterizations are well-specified, but existing methods set their parameters incorrectly. Right: AIOLI, an online mixing method that more accurately estimates the parameters that capture the true loss-proportion relationship.
In this paper, we improve our understanding of data mixing methods by showing that many existing methods can be expressed in unified optimization framework, which we call Linear Mixing Optimization (LMO) (Section 3). These methods are equivalent to solving an optimization problem that sets proportions minimize the average loss per data group, subject to an implicit methoddddependen mixing law-an an assumption relating loss per group and mixture proportions. We find that all current mixing laws share the same parameterization: for training round t from 1 T,
L (r' (A''''
where pl E Ath (the simplex) are mixing proportions over m given data groups at time t, L (p'): Am + (R+)" are the losses per group at the next timestep. A' E Rmx is a param- eter matrix, σ = Id or exp, and lin means equal up to linear transformation. Existing offline methods assume static (T =-=1 1) log-linear parameterization of the mixing law, while online methods assume linear dynamic mixing law. All methods set the parameters of their mixing laws differently (Table 1), and offline methods solve the optimization problem directly while online methods solve greedily using exponentiated gradient descent. Our framework reveals the underlying assumptions of each method in terms of the mixing law's parameterization, the values of the parameters, and how the optimization problem is solved. Furthermore, the fidelity of the mixing law and solving strategy dictates the optimality of the method, providing us with a new tool for understanding data mixing methods.
Applying the LMO framework, we test the fidelity of existing methods' assumptions, examining if they hold in practice (Section 4). Both the log-linear static and linear dynamic parameterizations capture the true loss-proportion relationship across datasets, achieving an average of 0.0005 MSE and 0.969 R?2 We then show that although existing mixing laws are well-specified, methods can set their parameters (A') inaccurately, causing poor performance. We compare each method's parameters to the optimal parameters, which we approximate by fitting the mixing laws to training runs. We find that the method's 's parameters differ significantly from the optimal parameters, and the extent of these deviations is correlated with method performance relative to stratified sampling (Figure 3), helping explain our initial observations. Finally, validate the assumptions used in solving the optimization problem, finding that the greedy approximation in online methods is reasonable proxy for the full objective. Our analysis shows that existing methods' parameterizaaions and solving strategies are of high fidelity, but their parameters are not.
To validate these insights, we develop AIOLI, simple new online data mixing method derived from the LMO framework (Section 5). Unlike existing online methods, AIOLI directly estimates the parameters A' from the current training run by fitting the mixing law on the history of losses and dynamic mixture propor- tions so far. AIOLI is thus able to dynamically adjust proportions without requiring any extra training runs.
We evaluate AtOLI in two settings by training 160M models on various combinations of data sources from SlimPajama (Soboleva et al., 2023) (Section 6). First, we compare AIOLI to existing data mixing methods and find that AIOLI consistently outperforms stratified sampling on all 6 datasets, by an average of 0.274 and up to 0.439 points in test perplexity. On the other hand, existing data mixing methods do worse than stratified on at least one dataset by up to 6.9 perplexity points, despite using extra training runs, As ex- pect, the parameters of AIOLI are also more consistently close to the optimal parameters (Figure 2). Second, we consider a scenario with limited additional computational resources, in which practitioners camnot ex- periments for learning mixture proportions for the full training duration. In this setting, mixture proportions leamed shorter run may not perform well on the longer final run. We find that using ATOLI to dynam-
2

---

Published as a conference paper at ICLR 2025
ically adjust these learned proportions throughout the final training can improve performance by an average 1.202 perplexity points in 28 out of 30 cases, compared to using the learned proportions directly.
PROBLEM SETUP
We formalize the data mixing problem and establish notation. In data mixing, we have mn data groups of text, such as GitHub, BooksCorpus, and arXiv. We are given train, validation, and test sets for each data group which we denote as : Diss for the ith group. Define Dimin
{ trin tra.n...... D } and
similarly define Daal and Dies.
Data & Mixing, During training, we show the model total of N examples from Duain over S training steps. To express how data proportions can change throughout training, we divide training into T equal rounds. Each round uses mixture proportion from the probability simplex: p' = Wi .i.....PPm) E A Static mixtures use only single round (T [I'=1] 1): p= (p'), while dynamic mixtures use several (T 1): p- (p'....''))
Model & Loss. Let (p.t) refer to the language model, f. at the beginning of round t where the model has been trained on data sampled using mixture proportions
.''p''' so far. Given model compute its Ioss on each group using the training data, Luain (f)= (Luain, ( (fn..,,(in m (f)), and similarly with the validation, (f). and test data, [f)a /) In this notation, the loss at the end of training be expressed as (f,,T +))++ When the f being referred to is obvious, we simply write Li) (p), and for static mixtures drop the superscript: (p)
Data Mixing Problem. Given a set data groups, an LM| f to train for S steps with N samples, and static or dynamic proportions), aim to determine the p that T+
T rounds of training (i.e., whether minimizes the total test loss across groups: minimize m E1
PEATX"
This objective aims to produce a trained model that does well on many data groups, which can serve proxy for downstream performance. However, without assuming additional structure on LT+ (p). this problem can only be solved with a brute-force search over p. which requires training many different models. In the next section, LMO framework imposes constraint T.+(pp) (p) that allows many existing methods to be expressed as approaches to solving this problem.
A UNIFIED OPTIMIZATION FRAMEWORK FOR DATA MIXING
We introduce the LMO framework by stating the general optimization problem (Section 3.1). Then, we show how this framework can express several existing methods (Section 3.2, 3.3), with summary of insights regarding these methods in Section 3.3.3.
3.1 LINEAR MIXING OPTIMIZATION (LMO) FRAMEWORK
The LMO framework consists of an optimization problem that is equivalent to the data mixing minimization problem (Section 2), subject to an additional constraint:
minimize
Pxm
T+1 val.á (p)
s.t. At.L (p) = -
A.yp)) Vie [m),te [
(2)
for some A Ar,bfcc and σ. A Rm is a matrix that encodes cross- cross-group group interactions, where A intuitively describes how much training on group j at a impacts group i's loss. 6f,c² E are group-specific parameters. :R - IR is either the identity function (Id) the exponential function (exp). We refer to the constraint in (2) as a mixing law that specifies the assumed relationship between loss and proportions.
There three components of this problem that need to be specified to yield a way to set p: a) the parameterization of the mixing law (T. σ), b) the values of the parameters (A', bf ,c'), and c) how to solve the problem. We express existing methods in LMO by specifying these components.

---

Published as a conference paper at ICLR 2025
Method
1) Mixing Law Purameteriiationn
2) Parumeters
3) Solver
DMI.
(p) m:; 1oxxtthxpppp22 Agp5))
Fit from 2 m-+t training runs
Direct
Skill-I
Fai (p)= -p) -
(p) (p) - - 1 4gp
A ( (T-al (1,) Fig (1,1/
FGD
DoReMi
A miin L
Lonin ((mn).0) Ai = (9 Lig,, (p). ;
BGD
DOGE
(p)= Lin, (p)-b'' -b' 612 Alpp
BGD
AJOLI
Le,, ------1-- Alpp]
Füt from history of Lost and
BGD
Table I: Summary of how existing methods and ATOLI are expressed in the LMO framework (1).
3.2 PRELIMINARIES FOR UNIFYING METHODS
We discuss preliminaries before presenting existing methods and explaining how they can be expressed in the LMO framework. First, we formally define what means for method to be expressed in the LMO framework. Then, we present result that allows us to convert between linear dynamic mixing laws and a way set p. which we will use to express online methods in our framework in Section 3.3.
Definition 1.
We say that a data mixing method can be expressed in the LMO framework if its exact algorithm how it sets proportions p and trains model f in terms of p can be equivalently constructed by specifying a mixing law and way of solving the LMO optimization problem.
This definition allows us cast existing methods way of solving the LMO optimization problem based on how they set p and train according to p. even if the methods themselves are not originally designed to minimize average test loss.
Converting mixing laws into update rules. When T> 1, natural way to solve the LMO optimization problem is via exponentiated gradient descent (EGD) (Kivinen Warmuth, 1997; Arora et al., 2012), which updates y' greedily while ensuring that it remains on the probability simplex. The following lemma presents the EGD update rule for the LMO optimization problem when o =
Lemma The FGD update rule for (1) subject to I. (p)
=
DE Aort (iE(m]
[+1 1 Pj
z ppppp
M B kjE [m]]
(3)
where 1) O is the step size and Z is normalizing constant such that pj ++1 E
This lemma shows how to adjust p' dynamically solve the LMO optimization problem. Notably, this update rule is defined in terms of the mixing law parameters, A' and b This gives us way to convert between how a method sets p and the implicit assumption it makes in the mixing law.
3.3 UNIFYING EXISTING METHODS
We discuss four existing data mixing methods and express them as specific instances of the LMO framework. A summary of our insights is provided in Section 3.3.3 and Table 1. In Appendix B. we comment on how several other online and offline data mixing methods are related to our framework, and all proofs for this section are in Appendix B.2.
3.3.1 OFFLINE METHODS
Data Mixing Laws (DML). Ye et al. (2024) propose an offline method using a static mixing law (T'= 1): Lva.: = b, A.pp;) for iE [m], with A,b,c learned by sweeping training runs over static proportions (2 m + runs to avoid being underdetermined). They select the proportion that minimizes the predicted validation loss. This law can be derived from (2) with o= exp, showing that LMO with a) log-linear static mixing law, b) fitted parameters, and c) direct computation of can express DML.
3.3.2 ONLINE METHODS
We provide colloquial description and an algorithmic description of the following three online methods. Then, in Theorem we demonstrate how they all are expressed in LMO using a linear dynamic mixing law, the EGD update rule, and method-specific mixing law parameters.
4

---

Published as a conference paper at ICLR 2025
Skill-It. Chen et al. (2023) is online method motivated by curriculum learning that dynamically adjusts mixture proportions. Data group interactions are expressed in a "skills graph," where each edge denotes how much the loss one group changes when trained on another. The skills graph is learned in advance using m training runs and then used to update proportions p¹ throughout training.
(1,)) (1, (1₂) indicating the relative decrease loss on group i when training a model on group j only. This is used in the Skill-It update rule, P2 (xfe m AAS T : for all j [m]] and learning rate 17 O. This rule determines + which is then used sample Dimin for training in the next round.
Concretely, the skills graph matrix A has entries A - - (1;) Lat,
DoReMi. Xie et al. (2023a) is an online method that applies ideas from distributionally robust optimization to data mixing, where the training objective minimizes the worst-group excess loss over a model trained with stratified sampling. p' is updated dynamically to minimize this excess loss and then averaged for the final run. DoReMi requires two additional runs Ieam a static p.
Concretely, let fer = f(Unif(m)) T"' + 1) denote a "reference model" that is first trained using stratified sampling. Then, a "proxy model" uses dynamic proportions according to the update rule P2 x pj expx(pmmaxx Luinn (p) (p) (p)-Lnaing ( (fr)))00)) for all j E [m]] and step size 17 > 0. This p² is used to weight the training objective, such that the proxy model is updated to minimize Em=,, at the next timestep. The averaged static proportions + - are then used in the final run.
1+1 LLmin, (f)
DoGE, Fan et al. (2024) is an online method that solves a bi-level optimization problem in which p' is updated minimize the average training loss at each step. By using first-order Taylor approximation of the training loss, p' is updated using the gradient of each data group. The dynamic proportions are then averaged for the final run. DoGE requires one additional run to learn static p.
Concretely, a proxy model is trained using p1+ p) exp((((V Lumin.3 (ft), V V.vva. (f))), and f is updated to minimize the training loss weighted by p', similar to DoReMi. The averaged static proportions +27 p' are used in the final run.
Framework expression. All three online methods use an update rule yj ++1 x pxp((()) which is similar to (3). This provides intuition for our main theorem, which expresses these methods in LMO.
Theorem 1. Define the following parameters for each method:
A e @Rmxm where A.g A =T: .s (p) val,i (1,) (TTT+111)
1 nad, (1)))) ( for all i.j [m).
AU.RRM ERmxm where A mmi T.main (p) (p) - Temain, i (m) and A,D = for i
ALDOGE E Rmx" where Atj.uGE = =(V (p), Linin for all e.je(mn [m]]
Instantiating the IMO framework (1) with a) linear dynamic mixing law T Twal( ( b' ΣT= Aiyp' b) parameters A A.SH.--DRRMDGGG and e) EGD to solve for p allows for us to express Skill-It, DoReMi, and DoGE, respectively.
3.3.3 SUMMARY OF LMO FRAMEWORK INSIGHTS
Table summarizes how existing methods are expressed in the LMO framework. LMO reveals the assump- tions each method makes through how the components of the framework are specified. First, all mixing laws are either linear or log-linear. Second, the mixing laws differ in the values of the parameters used. For example, Skill-It's A' the current loss times a static skills graph matrix, while DoReMi's A' is diagonal. Third, offline mixing methods solve for p directly while online mixing methods use EGD, which uses greedy approximation. If the mixing law and solving strategy assumptions hold true in practice then the method yields optimal mixture proportions. In the next section, study the fidelity of these assumptions.
4 ANALYZING FIDELITY OF EXISTING METHODS WITH THE LMO FRAMEWORK
We examine the fidelity of the assumptions made by existing methods in terms of the three components of the LMO framework: a) the mixing law parameterization, b) values of the mixing law parameters, and c) how to solve the optimization problem for p. After providing experiment details (Section 4. I). we discuss these three components in order (Section 4.2-4.4).
5

---

Published as a conference paper at ICLR 2025
Log rinelizeer static mixing law on Atnav/StaccExchanng
Linear dynamic mixing law ArxsrSttacdxxehaange
1ig maadeechaan
Log 0.e0 20 um
0.2
0. Propection of STRRY
3.0
64 2.0
0.2
0.4
0.0
022
0.2
Proportion stackexchange
Proportion
Proportion of starbexcchange
Figure 2: Left: P. vs log ( fitted linear dynamic mixing law. Colors represent random seeds (left) and initial yP EP (right, blue 0.7,0.3). Both laws fit the true loss- loss-proportion relationship well.
-c) (p)
with fitted static log-linear mixing law. Right: pt vs Lvat, with
4.1 EXPERIMENT DETAILS
Data settings. We use sampled version of SlimPajama (Soboleva et al., 2023; Yoon, 2023), a pre-processed version of the RedPajama pretraining dataset (Together.ai, 2023). SlimPajama consists of 7 data groups: ArXiv, Books, CommonCrawl, C4 (Raffel et al., 2019), Github, StackExchange, and Wikipedia. To develop a fine-grained understanding of data mixing, we create 6 settings by extracting combinations of these groups. We study three settings with m = 2: Arxiv/Stackexchange. Github/C4, and Book/StackExchangge We study two settings with m m=3: 3: DoriwBooo//ttckExxhhaeee CaaeoaCCraa//GGthb//WWiiiipeeeee Finally, we study mixing over the full SlimPajama dataset with m=
and
Models. We train 160M parameter GPT-style decoder-only LLMs with batch size 8 and context length 2048. For m m=2,3, 2,3, we train for 5K steps, and for m= 7, we train for 40K steps.
Training sweeps. To assess the true loss- puss-propottion relationship and compare it the assumptions made by existing methods, we conduct training sweeps over different mixture proportions, denoted as P. For m 2, we set P [[0.1,0.9] [0.2,0.8).... (.000.0..).0.22008888 009 1...0..... For = -3 and 7, we set P equal to 10 p's and 40 p's drawn from the Dirichlet distribution with a = 1.0 and 1.5, respectively.
4.2 MIXING LAW PARAMETERIZATION
We examine whether existing methods' mixing law parameterizattons- log-lnnar static and linear dynamic dynamic-ccpturre the true loss-proportion relationship. By empirically fitting them pos-prrooottion pairs, we find that both parameterizations are indeed well-specified. Full results for both mixing laws in Table 5 in Appendix C.I. We discuss the generality of these parameterizations across training scales and other datasets, as well as higher-order parameterizations,, in Appendix C.1.1.
Setup. For the log-linear static mixing law, we study if there exists b, such that (pa,; can be expressed as c b1exxp(( -Agp;) for all E [m]. We fit the parameters using full training runs on P. For the linear dynamic mixing law, we study if there exists A' such that L (p) can be expressed as (p) En= for all E [m] (b' is absorbed into A¹). To fit A", we select a timestep t and train on static proportion p" eP for all p² .'..p' until time t, and at t++1 sweep the values of p/+1 E eP.
Results. On average across our 6 data settings, the mean squared error (MSE) of the fitted log-linear static mixing law is 8.9 x 10-4, and the R²2 coefficient of determination is 0.991. The average MSE of the fitted linear dynamic mixing law 1.0 x 10-4 and the R22 0.947. See Figure for examples. Since both parameterizations have high R and low MSE, conclude that they capture the true loss propprronntion relationship well and are high fidelity.
4.3 VALUES OF MIXING LAW PARAMETERS
As shown Table I, each method sets the parameters of its mixing law differently. We study how close the method-specific parameters are to the optimal parameters that are obtained when fitting the method's mixing law to the true loss oossprroooniion relationship, and if these parameter disparities are reflected in method performance. We find that existing methods' differences in mixing law parameters are largely responsible for their performance. We omit studying DML since its parameters are fitted from full training runs and hence differ from the optimal in estimation emor only.

---

Published as a conference paper at ICLR 2025
Setup. For Skill-It, DoReMi and DoGE, we select a step t and obtain the method-specific A'. We then sweep P for the next round 1+ This sweep is used to approximate an optimal A that captures the true loss-mixture relationship, _ (p) = (p) AA"pp,, as well as fit b' E R used for scaling A' (details in Appendix C.2). We study the relationship between A b' A and A and how is related the performance of the method.
To express similarity between A' and A in way that is reflected in performance, we observe that from Lemma 1, p" is updated using the column sum of At, 1 A'. Moreover, the magnitude of A' is not critical to performance since the step size 7) can always be tumed to control this. Therefore, we compare the vectors a' = 1 A || Af and at* = 1 A /|| 11", A Finally, we note that the order of the elements of a' determines the update direction from y' to p*+1 Lemma Therefore, that is an average of cosine similarity and the Spearman rank correlation, sim( A', A 0.5cossim (ā', ,a')) + 0.5Spearman(a* (a), This metrie is bounded between -1 and 1, where 1 indicates a' =a and -1 indicates a'=-a"+
Stratified loss method loss
0.00
-0.05
Skill-It
0.10
DoReMi
DoGE
Aioli (ours)
propose similarity score
-0.15
0.0
0.4
0.8
slm(A',
Figure 3: Improvement over stratified sampling versus optimality of A'. Each dot represents method applied to dataset. The red region shows that existing methods are worse than strat- ified on least dataset. The vertical dashed line serves as visual aid.
Results. In Figure 3, we plot each method''s sim( , versus each method's improvement over the stratified sampling baseline, which sets Pą 1/m for all iE (m]. for each dataset in the 2,3 data settings. We find that no existing online method works well across all datasets (also see Table 2). and that our metric and loss improvement have a moderate positive correlation (R² 0.4991. This suggests that A accuracy critical to the performance of online methods, and that existing methods' A' are not consistently accurate across the datasets. In Appendix C.2.1, give more details on the structure of A providing intuition for why existing methods parameters cannot express
4.4 SOLVING STRATEGY
We study the assumptions made how existing methods solve the LMO optimization problem. We find that the greedy approximation used by EGD, minimize, Ei= L (p), does not significantly compromise performance compared to full optimization of dynamic proportions, which has an exponentially large solution space. particular, we study if greedily selecting p' from P each yields the optimal dynamic proportions in PT and find that this holds in 2 out of 3 data settings (Table 8). This suggests that the greedy approximation can simplify optimization without substantial performance loss. We also comment on other possible solving strategies in Appendix C.3.
5 AIOLI: A METHOD FOR IMPROVED DATA MIXING
To validate our insights from Section 4, we develop AIOLI, an online method derived from the LMO framework. We have three takeaways from section 4:
a) A linear dynamic mixing law, L (p) = (p) E/ Atyps for all E [m]] can capture the loss propprronn relationship with high fidelity (Section 4.2).
b) Existing online methods often set the parameters A' to be very different from true 4 (Section 4.3).
c) Exponentiated gradient descent can recover nptr-optimal performance while simplifying the optimiza- tion problem, avoiding an exponential solution space (Section 4.4).
We thus directly specify the linear dynamic mixing law parameterization and EGD as two out of three LMO components of AIOLI since we found that their assumptions generally hold in practice. According Lemma the update rule given these two components is p x p/J rf)exp(r _ Af,) (D' is absorbed into A'). Thus, our primary mandate in creating AIOLI is to construct and utilize an A' that is an accurate estimate of the true A in the linear dynamic mixing law, which existing online methods fail to achieve.
7

---

Published as a conference paper at ICLR 2025
Estimating A To build intuition, first consider a high-cost naive approach. For each round, we could conduct a training sweep of m different proportions p ,....p'." and observe each resulting change in loss. We could then solve a system of m equations for each i: - (r-, - obtaining vectors A,.....A.. However, this approach effectively requires m extra training runs, AIOLI similarly solves a system of equations, but computes loss changes per sweep mixture without requiring extra training. First, it allocates 8 fraction of the training round for learning A'. Second, it partitions this 5 into K = mn intervals and trains according to an interleaved order on p ......... After training on each p²), we record the resulting change in validation losses, and we average over all of pt intervals. Intuitively, the interleaving ensures that the model is trained on each p²u for several intervals throughout 8, which can approximate if we were to train on p² for the entire ò (which approximates the entire round). This procedure outlined in LEARNPARAMS (Alg. 2). with more details Appendix and Figure 7.
for s E [mm]
AIOLL First, set pD to be uniform. In each round, estimate A' using LEARNPARAMS and then normalize the entries of A', producing A'. Otherwise, A' decreases along with loss over time, resulting in the first few p' updates being much larger in magnitude than others. Then, we update the proportions using A, as in Lemma 1, and train for the remainder of that round using p/.-
P) "xpp/"exp(11 exp(1
Finally, we design AIOLI so that it can also be used to improve other data mixing methods, which study in Section 6.2. Mixture proportions can be updated using AIOLI either from the start of training or from the middle of run. In the latter case, denote initial static mixture piu E " and initial number of steps Sma If Sma is nonzero, AIOLI trains according to Pmnit for the first Smn steps before beginning to update the mixture proportions. AIOLI is presented Algorithm
Algorithm AIOLI
1: Input: data Duain Deat, model fi, Initial steps Smit initial proportions piuit E Am, rounds over S - Smit remaining steps, 5 fraction per round for learning parameters, learning rate 7, one-hot smoothing factor E.
2: If Smit O, train f piit for Sinit steps.
3: Set petp" -nnif(mm))
4: for 1o..= do
5: Set f + LEARNPARAMS for all [m].
Arain:Delldd, .e) (Alg. 2), and normalize A to get A
6: pjppp 'exp(n) m At "j
7: Train model fe with &- (1 5) steps from mixture p' over Dusin- Obtain updated ft+1.
8: end for
Algorithm LEARNPARAMS
Input: Deat, 8, model f², number of sweeps k, one-hot smoothing factor e.
2: Split the fraction of a training round o into K intervals, where K =mk.
3: Set 3=0mmm
4: Define
for E [m], and define P P= p......mm E Amx 5: Randomly shuffle instances of each ie [m]] to create an order I E [m]]
6: for 1= do
7: i= Train model on mixture ydd of Duain for one interval, obtain f
8: for io [m] do
9: Update Sij B + Lvva, (fttrr((24--166// Loa ( with loss difference on D
10: end for
11: end for
12: Update 3 K.
13: Set A P-1 BB for each i (m].
14: Return A' E mm fo+6
6 EXPERIMENTAL RESULTS
We evaluate all methods in the LMO framework, including AIOLI, in two settings. First, we consider an unrestricted additional training budget setting to assess how AIOLI compares to other methods in their original form, since each method uses different number of extra training runs leam proportions
8

---

Published as a conference paper at ICLR 2025
Table 2: Difference in average test perplexity compared to stratified sampling in the unrestricted setting, where all methods can use < 10 extra runs to learn Negative values (green) improvement. A=Arxiv, B=Books, GH-GitHub, SE=SStckE W=Wikipedia.
Method
ASE
GH/C
BSB
A/BISE
CCYGHI
SlimPajamu
stratified
extra runs
Stratfiidd
16.532
35.991
17.192
35.114
41.583
26.426
GS
-0399
-0.407
-0.645
-0.247
0.298
0.490
10
DML
-0.241
-0.110
-0,644
-0.599
0.242
1.611
10
Skill-I
-0.326
0.551
-0.728
-0.568
-0.195
-0.184
DoReMi
0.307
5.303
-02217
0.703
DoGE
0.419
0.184
0.678
1.843
0.604
0.949
AJOLI
0.205
-0.340
0-439
-0.226
-0.196
-0.240
(Section 6.1). Second, we consider a restricted training budget setting to assess if AIOLI can enhance existing methods in practical, budget-constrained conditions, where existing methods have less than full training run to learn mixing proportions (Section 6.2). Hyperparameters and experimental details, including proportion trajectories are available in Appendix E. Downstream evaluation, ablations, experiments on larger models, and results adapting AIOLI to an out-of-domain setting are in Appendix F.
Data settings and models. We use the same data settings and models as in Section 4.1, where we train for S S=5K 5K steps for m 2,3-group settings and 40K steps for the full SlimPajama.
Baselines and evaluation. We consider three online methods (Skill-It, DoGE, DoReMi) and one offline method (DML). We also consider grid search (GS), which sweeps training runs and selects with the lowest average validation loss, and stratified sampling, which sets p. y2 for all EE [mm]. For each method, we report the average test perplexity per group of the trained model. This metric is considered proxy for downstream performance (Fan et al., 2024) and also represents the objective in the data mixing problem.
6.1 UNRESTRICTED SETTING
Setup. We allow methods up to 10S additional training steps leam the mixture proportions. Approaches like grid search and DML can use the entire budget (searching and fitting over 10 full runs), while Skill-It, DoReMi, and DoGE use mS, 25, and S extra training steps, respectively (see Section 3.3). Stratified sampling and AIOLI use no extra training steps. We evaluate ATOLI with Smm=0.
Results. Table 2, we find that AIOLI robustly outperforms stratified sampling in all data settings by an average of 0.274 perplexity points, while all other methods do worse than stratified sampling at least I set of data groups by up to 6.9 points. The performance of AIOLI and other online methods additionally reflected in Figure in which find that AIOLI'S A' similarity with A is correlated with performance. While AIOLI'S parameter similarity is not always the highest, we note that its lowest similarity score is much higher than that of other methods, providing evidence that ATOLI'S parameter estimation procedure more consistently accurate than that of other methods. Lastly, regarding offline methods, we hypothesize that their poor performance on settings with larger m is due the training budget being limited to 10.,, and that increasing this budget would eventually allow them to perform well.
6.2 RESTRICTED SETTING
Motivation. We introduce the restricted setting because practitioners may not have the resources or desire to complete multiple full training runs, especially as recent LLMs are trained for longer and on more data (Muennighoff et al., 2023). As result, practitioners may only use data mixing methods on shortened runs, producing learned proportions that may be suboptimal on the full run. We study if AIOLI is able to improve performance by dynamically adjusting previously leamed proportions throughout the full training run.
Setup. We allow all existing methods up to 0.5. additional training steps to learn the mixture proportions. This requires methods to learn method over shorter runs of Smathod steps each. For instance, grid search will conduct 10 runs of length S/20 (see Table 9). We evaluate each method by using p meliial learned from shorter runs to train the model the full run of S steps. We use AIOLI to dynamically adjust each p method throughout the full run. That is, for each existing method, we run AIOLI with p p methood and Smn = Smehhd, referring to this AIOLI +method.
Results. In Table 3, we find that adding AIOLI to any existing method that learss proportions over shorter nuns improves average test perplexity per group in 28 out of 30 settings, by average of 1.202 and

---

Published as a conference paper at ICLR 2025
Table 3: Average test perplexity in the restricted setting, where each method learns p on shortened runs, and
AIOLI +method dynamically adjusts throughout training, green=AIOLI +method outperforms method.
Method
ArxivSE
GH/C4
Books/SE
Arxiv/Book/SEE
CC/GH/Wiki
SlamPajama
GS
16.373
36,315
47.063
35.174
42.767
27.741
A GS
16.388
35,925
46.667
34.705
41:378
25.654
DMI.
16.659
36.658
46.846
34.585
42.731
37.696
AIOLI + DML
16.277
35.856
46.710
34.529
41,595
27.654
Skill-it AIOLI + Skill-it
16.246
37.255
46.667
34.539
42.069
20.734
16.261
36.153
46.586
34.563
41.732
26.073
DoReMi AtOL1 + DoReMi
16.322
37.812
46,489
34.934
42.738
28. 762
16.347
35.626
46.163
34.770
41.800
26.587
DoGF
16.853
35.795
46.743
35.775
41.790
32.301
AIOLI + DoGE
16,473
35.632
46.145
34.771
41.378
26.073
maximum of 12.012 points. Furthermore, A1OLI can help methods that initially underperform stratified sampling surpass it, such DoGE across all settings. In some settings, such as Books/StackExchange. AIOLI improves methods that already outperform stratified sampling. This shows that AIOLI can enhance wide variety of static proportions, regardless of their initial performance. For the two settings where AIOLI underperforms the base method, the base method already outperforms stratified, and adding AIOLI maintains this trend, worsening perplexity by at most 0.025 points.
7 RELATED WORK
Data mixing. Albalak et al. (2023) frames online data mixing as a multi-armed bandit problem. Recent works have also studied how to generalize data mixes from smaller to larger models (Kang et al., 2024; Ge et al., 2024; Liu et al., 2024; Allal et al., 2025). Thrush et al. (2024) optimizes model performance on downstream tasks by constructing like data interactions matrix using pre-trained model perplexities. Jiang et al. (2025) mixes data online using each data group's loss, effectively using a diagonal A' matrix.
Curriculum Learning. Bengio et al. (2009) initially introduced curriculum learning as training models over samples from easiest to hardest. While early work focused on manually designed curricula, later work emphasizes model-driven ones (Hacohen & Weinshall, 2019; Varshney al., 2022; Fan & Jaggi, 2023; Mindermann et al., 2022). Curricula can encourage skills-based generalization (Huang et al., 2024), or emphasize high quality data improve downstream task performance (Blakeney et al., 2024).
Data Selection. A common way to curate datasets besides mixing is to select data at the per-sample level (Albalak et al., 2024). Techniques here can be broadly classified data filtering, data matching, and data condensation. In data filtering, low-quality samples are removed using simple heuristics, such as GitHub file lengths (Together.ai, 2023; Touvron al., 2023), or via deduplication (Abbas et al., 2023; Tirumala et al., 2023; Lee et al., 2022). In data matching, samples that are most similar to reference dataset are selected. Similarity can defined in terms of embeddings Xie et al. (2023c), gradients (Xia al., 2024; Engstrom et al., 2024), or directly using machine learning models to score samples (Brown et al., 2020; Grave et al., 2018; Moore & Lewis, 2010). Lastly, data condensation aims to identify a subset of samples that captures the full training dataset's properties. Selection mechanisms include using gradients, model predictions, and embedding distances (Toneva et al., 2019; Paul et al., 2021; Sorscher et al., 2023).
8 DISCUSSION
We introduce the LMO framework, which unifies existing data mixing methods by viewing them as solutions to a common optimization problem involving an implicit method-dependent mixing law. Using this framework, we find that existing methods perform poorly on some datasets due to inaccurate mixing law parameters. This insight inspires AIOLI, whose performance gains are rooted in its ability to estimate parameters A' of the linear dynamic mixing law throughout training.
Limitations and Future Work ATOLI incurs extra inference cost via the repeated evaluations LEARN- PARAMS (Alg. 2). This can be reduced by computing Lval over subset of Deat. and by using each AA' for longer (decreasing 7). Another direction is understanding the role of data group partitions. For example, C4 is subset of CommonCrawl, and it is unclear if disjoint groups could improve performance.
The LMO framework itself is an invitation for future work. It shows that data mixing methods can be improved and analyzed by studying their assumptions how models leam from data. By exposing such assumptions, LMO identifies key axes for improvement (mixing law parameterization, parameter estimation, and how to solve for p), which we hope will inspire new principled data mixing methods.
10

---

Published as a conference paper at ICLR 2025
9 ACKNOWLEDGMENTS
We thank Sabri Eyuboglu, Neel Guha, Ben Viggiano, Dan Biderman, Dan Fu, Michael Wornow, Jon Saad-Falcon, Alyssa Unell, Owen Dugan, Jerry Liu, and Gautam Machiraju for their feedback. We thank NYU HPC and Stanford NLP for providing compute and research support. This research project has benefited from the Microsoft Accelerate Foundation Models Research (AFMR) grant program.
We gratefully acknowledge the support of NIH under No. U54FB0000405 (Mobilize), NSF under Nos. CCF2247015 (Hardware-Awaree, CCF1 763315 (Beyond Sparsity), CCF1563078 (Volume Velocity), and 1937301 (RTML); US DEVCOM ARL under Nos. 191 (Long-contex() and W91 INF 21-2-0251 (Interactive Human-AI Teaming); ONR under Nos. N0001423 12633 (Deep Signal Processing); Stanford HAI under No. 247183; NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASE, Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), the Samsung Advanced Institute Technology (under the project Next Generation Deep Learning: From Pattem Recognition to AI), the NSF Graduate Research Fellowship (MYH), and members of the Stanford DAWN project: Meta, Google, and VMWare. The U.S. Govermment is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government.
REFERENCES
Amro Abbas, Kushal Tirumala, Dániel Simig, Surya Ganguli, and Ari S Morcos. Semdedup: Data-efficient learning web-scale through semantic deduplication. arXiv preprint 23Xh::233300554, 2023.
Alon Albalak, Liangming Pan, Colin Raffel, and William Yang Wang. Efficient online data mixing for language model pre- pre-training. 2023. URL https: : //t////aaii aawaiaiivvv..a//aaaaaaaaa wwwpxiv.orrroooooo abs/ C2406.
Alon Albalak, Yanai Elazar, Sang Michael Xie, Shayne Longpre, Nathan Lambert, Xinyi Wang. Niklas Muennighoff, Bairu Hou, Liangming Pan, Haewon Jeong, Colin Raffel, Shiyu Chang, Tatsunori Hashimoto, and William Yang Wang. A survey on data selection for language models. arXiv preprint arXhv:240.118827 2024. https: /aaxiviivvv../// org/ abs 222022422222 16827.
Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Gabriel Martín Blizquez, Guilherme Penedo, Lewis Tunstall, Andrés Marafioti, Hynek Kydlfček, Agustín Piqueres Lajarín, Vaibhav Srivastav, Joshua Lochner, Caleb Fahlgren, Xuan-Son Nguyen, Clémentine Fourrier, Ben Burtenshaw, Hugo Larcher, Haojun Zhao, Cyril Zakka, Mathieu Morlon, Colin Raffel, Leandro von Werra, and Thomas Wolf. Smollm2: When smol goes big data-centric training of small language model, 2025. URL hittps://arxt wwwp:vvvorrrae.oooo alos/ 2502. 02737.
Sanjeev Arora, Elad Hazan, and Satyen Kale. The multiplicative weights update method: a meta- algorihgo and applications. Theory computing, 8(1):121-164, 2012.
Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In Proceedings of the 26th Annual International Conference on Machine Learning, ICML '09, pp. 41-48, New York, NY, USA, 2009. Association for Computing Machinery. ISBN 9781605585161. 11144//155334.15533385
doi: 1553380.
URL https: doi http://::///o..../...o1111 wtt.////...1111111111 /1/
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradlcy, Kyle O' O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piga: Reasoning about physical commonsense in natural language. In AAAI Conference on Artificial Intelligence, 2019. URL Io:pus 200.:ppui::::020999:::9
https: semant. wet.chcclossss aar.arroooo orgg
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. In First Conference on Language Modeling, 2024. URL https: //tp:///op//e//eeee net. forum?idee d=w.wwwwttweDFF
11

---

Published as a conference paper at ICLR 2025
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter; Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Bemer, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL https: /tt///axxvvvv/ .rg/ abo5/005. 14165.
Luis F. Chaparro and Aydin Akan. Chapter sampling theory. In Luis F. Chaparro and Aydin Akan (eds.), Signals and Systems Using MATLAB (Third Edition), pp. 449-485, Aca- demic Press, third edition edition, 2019. ISBN 978-0-12-814204-2. doi: https///oo..or//00 URL https:: htttps://ww.....t sttps///ieciiieeerrrreee.e.. com/ htmm/science articleepi1//s9700222 0TAT778000224422000009999
001669978-0112-844444-20000000
Angelica Chen, Sadhika Malladi, Lily H. Zhang, Xinyi Chen, Qiuyi Zhang, Rajesh Ranganath, and Kyunghyun Cho. Preference learning algorithms do not learn preference rankings, 2024. URL
Mayee Chen, Nicholas Roberts, Kush Bhatia, Jue WANG, Ce Zhang, Frederic Sala, and Christo- pher Ré. Skill-it! a data-driven skills framework for understanding and training language mod- cls. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (cds.), Advances in Neural Information Processing Systems, volume 36, pp. 36000-36040. Curran Associates, Inc., 2023. URL https: httt:::///oocccdeeeeeeeeee neurips. wttppppppppppppppppppppppppppe.e httpppppereeeeppppeeereeeeeeeeee /appp/////aaleeeapppp/////////// https/pppp//2///////lm
pdf.
Christopher Clark, Kenton Lee, Ming- Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes'no questions. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 2924-2936, Minneapolis, Minnesota, June 2019, Association for Computational Linguistics. doi: 1863/v IN19-1300 URL https::tt//. : /ac//aaaaanttt /at/a/aaaanntoooooooooooo
/N19 1300.
Peter Clark, Isaac Cowbey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the A12 reasoning challenge. CoRR, abs/1803.054577 2018. URL http: / wttpiv.ivvv.///av///ee//aaaaa org/ abs wtaaa//ss.0 1803. 05457.
Tri Dao, Dan Fu, Stefano Ernon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35: 16344,16359, 2022.
Tobias Domhan, Jost Tobias Springenberg, and Frank Hutter. Speeding up automatic hyperparameter optimization of deep neural networks by extrapolation of learning curves. Twenty-fourth international joint conference artificial intelligence, 2015.
Logan Engstrom, Axel Feldmann, and Aleksander Madry. Dsdm: Model-aware dataset selection with
datamodels. Forty-first International Conference on Machine Learning, 2024. URL https:
//openreview.ce
net.ewwwweeeeeeeeeeeeeeeeeeee.ee net wermtorut/tteeetoe furum? hid.ciiHkeee..........
Simin Fan and Martin Jaggi. Irreducible curriculum for language model pretraining. arXiv preprint
arXiv:2310. 15389, 2023.
Simin Fan, Matro Pagliardini, and Martin Jaggi. Doge: Domain reweighting with generalization estimation, 2024. URL https: aaaxiv.iiiii///ii//aaaaaaaaiavva org/ aaaaab/2333 18301.11.3....
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf. Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang. Kevin Wang. and Andy Zou. A framework for few-shot language model evaluation, 07 2024. URL https::/ / zeeoooooo.oo org/ records / / 608602.
Ce Ge, Zhijian Ma, Daoyuan Chen, Yaliang Li. and Bolin Ding. Data mixing made efficient: A bivariate scaling law for language model pretraining, 2024. URL https: / arxiv. org/ abs/ 2405. 14908.
12

---

Published as a conference paper at ICLR 2025
Edouard Grave, Piotr Bojanowski, Prakhar Gupta, Armand Joulin, and Tomas Mikolov. Learning word vectors for 157 languages. In Nicoletta Calzolari, Khalid Choukri, Christopher Cieri, Thienry Declerck, Sara Goggi, Koiti Hasida, Hitoshi Isahara, Bente Maegaard, Joseph Mariani, Hélène Mazo, Asuncion Moreno, Jan Odijk, Stelios Piperidis, and Takenobu Tokunaga (eds.), Proceedings of the Eleventh International Conference Language Resources and Evaluation (L.RCC 2018), Miyazaki, Japan, May 2018. European Language Resources Association (ELRA). URL https: hattts://aclanthhott hww.nnhooooyy org/ L18 1550.
Guy Hacohen and Daphna Weinshall. On the power of curriculum learning in training deep networks. In International Conference on Machine Larningg 2019. URL https: /api semant cscho lar. ar. org/ Corpus :nteemI:::120000 102350936.
Roger A. Horn and Charles R. Johnson. Matrix Analysis. Cambridge University Press, 1985.
Daniel Hsu, Sham M. Kakade, and Tong Zhang. Random design analysis of ridge regression. In Shie Mannor, Nathan Srebro, and Robert C. Williamson (eds.). Proceedings of the 25th Annual Conference on Learning Theory, volume 23 of Proceedings of Machine Learning Research, pp. 9.1-9.24, Edinburgh, Scotland, 25-27 Jun 2012. PMLR. URL https: / /ptp/eeeecceesss.//oe/e/pc/..sse mlr www.r.pess//ee hatpsss///////uuuu html.
Yuncheng Huang, Qianyu He, Yipei Xu, Jiaqing Liang, and Yanghua Xiao. Laying the foundation first? investigating the generalization from atomic skills to complex reasoning tasks, 2024. URL https:/ /atxiiiiiv.////////a//aaaa org/ altp//b///// /203. 09479.
Yiding Jiang, Allan Zhou, Zhili Feng, Sadhika Malladi, and J Zico Kolter. Adaptive data optimization: Dynamic sample selection with scaling laws. In The Thirteenth International Conference on Learning Representations, 2025. URL https: //ppeneeeeieweeccooo//e/o
Sham Kakade. Lecture 22: Exponentiated gradient descent. https: /wweesooeeeeeeeeeeeeeeeeeeeeeeee cs. washington. edu/ sae//sshane wwtha///ooureeeeeeeeeee wttp://tteeeeeeeeeeeeeeaaae wtwp://eeeeeeeteeeeeee222222 pddf pdf, n.d. Accessed: September 29, 2024.
Feiyang Kang, Yifan Sun, Bingbing Wen, Si Chen, Dawn Song, Rafid Mahmood, and Ruoxi Jia. Autoscale: Automatic prediction of optimall data composition for training Ilms, 2024. URL https: aaxiivoorg abs/ 0/bs/2407 2407 20117.
Jyrki Kivinen and Manfred K. Warmuth. Exponentiated gradient versus gradient descent for linear predictors. Information and Computation, 132(1):1-63, 1997. ISSN 0890-5401. doi: https://doi.or// 10. 10.10006/ico..1966.. 1996.2612. URL pettps://ww/ wwweeceeiieeett....c..c.ooce com/ sttp:scccccccecceee/ee/eeeecccee article/ pi1//009054419622112
Katherine Lee, Daphne Ippolito, Andrew Nystrom, Chiyuan Zhang, Douglas Eck, Chris Callison- Burch,, and Nicholas Carlini. Deduplicating training data makes language models better, 2022. URL ht https:
//arxiv. org/ / abs/ 2017. 7. 06499.
Mosh Levy, Alon Jacoby, and Yoav Goldberg. Same task, more tokens: the impact of input length on the reasoning performance of large language models. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (cds.), Proceedings of the 62nd Annual Meeting the Association for Computational Linguistics (Volume 1: Long Papers), pp. 15339-15333 Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10.18653/v122244 aal-- URL ht /acl anthcllnnhoooggg . wwrg//00000 2024. acl- 818.
Jeffrey Li, Alex Fang, Georgios Smymis, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash 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 Abbas, Cheng- Yu Hsieh, Dhruba Ghosh, Josh Gardner, Maciej Kilian, Hanlin Zhang, Rulin Shao, Sarah Pratt, Sunny Sanyal, Gabriel Ilharco, Giannis Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevic, Sham Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri, Sewoong Oh, Luke Zettlemoyer, Kyle Lo, Alaaeldin El-Nouby, Hadi Pouransari, Alexander Toshev, Stephanie Wang. Dirk Groeneveld, Luca Soldaini, Pang Wei Koh, Jenia Jitsev, Thomas Kollar, Alexandros G. Dimakis, Yair Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. Datacomp-Ims Im: In search of the next generation of training sets for language models, 2024. URL https: : / / arxii org/ abs/ 2406 11794.
13

---

Published as a conference paper at ICLR 2025
Lisha Li, Kevin Jamieson, Giulia DeSalvo, Afshin Rostamizadeh, and Ameet Talwalkar, Hyperband: A novel sundit-based approach to hyperparameter optimization. Journal of Machine Learning Research, 18(185)1152, 2018.
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better downstream: Implicit bias matters for language models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardd, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 1188-22214, PMLR, 23-29 Jul 2023. URL https::/p//........ /ptpceeroeeeeeeeeee ngs 1gg.mm.. m.r..press/ pearess/v000 v2.02/ 1iu23ao. html.
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, 2024. URL
http://axvvvvoo
aaaa////0000 014
Shayne Longpre, Gregory Yauney, Emily Reif, Katherine Lee, Adam Roberts, Barret Zoph, Denny Zhou, Jason Wei, Kevin Robinson, David Mimno, and Daphne Ippolito. A pretrainer's s guide to training data: Measuring the effects of data age, domain coverage, quality, & toxicity. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.). Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume I: Long Papers), pp. 2245-3276 Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 10. 10.11663//1/2024..o naacl-4naacllooog 179, URL htt.ttss://aclaneeeeee http///aaaaaattloacceeee org/ 2024 maacl-long 179.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? new dataset for open book question answering. In Ellen Riloff, David Chiang, Julia Hockenmaier, and Jun' Tsujii (eds.), Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 2381-2399 Brussels, Belgium, October-November 2018. Association for Computational Linguistics. doi: 10. 18653/v1/D18 1260. URL https:
1260.
Sören Mindermann, Jan M Brauner, Muhammed T Razzak, Mrinank Sharma, Andreas Kirsch, Winnie Xu, Benedikt Höltgen, Aidan N Gomez, Adrien Morisot, Sebastian Farquhar, al. Prioritized training on points that are learnable, worth learning, and not yet learnt. In International Conference on Machine Learning, pp. 15630-15649, PMLR, 2022.
D.C. Montgomery, E.A. Peck, and G.G. Vining. Introduction to Linear Regression Analysis. Wiley Series in Probability and Statistics. Wiley, 2021. ISBN 9781 119578777
Robert C. Moore and William Lewis. Intelligent selection of language model training data. Jan Hajič, Sandra Carbery. Stephen Clark, and Joakim Nivre (eds.), Proceedings of the ACL 2010 Conference Short Papers, pp. 220-224, Uppsala, Sweden, July 2010. Association for Computational Linguistics. URL https: / ac aellaonnhhoooooooo wow..ooooooooooooooooooooooooo.. org /Pl///100 P10- 0-2011.
Niklas Muennighoff, Alexander M Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin Raffel. Scaling data-constrained language models. In Thirty- seventh Conference on Neural Information Processing Systems, 2023. URL https: //tp////eee......e////e/ net/ dow.mmidddisoiio.o..ooooii jadd==jsBBrrree
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bemardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández, The LAMBADA dataset: Word prediction requiring a broad discourse context. CoRR, ab6/1606.000311, 2016. URL http://arxiv. org/abs/ 6106. 06031.
Mansheej Paul, Surya Ganguli, and Gintare Karolina Dziugaite. Deep learning on a data diet: Finding important examples early in training. In M. Ranzato, A. Beygelzimer, Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 20596 20607. Curran Associates, Inc., 2021. URL https: httpp://ppooeeeanggss... neurips www..ps.cc/pppp......ee files/ httpp/ppprr/222/
1 6 (eweefeeeeeeeeeeeee2222 1551 -aw.e55--pprr.... pdf.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yangi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv e-prints, 2019.
14

---

Published as a conference paper at ICLR 2025
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Proceedings of the AAAI Conference on Artificial Intelligence, (4(05):8732-8744) Apr. 2020. doi: w00/////....34....... Aatpp//AAA/// art www.cleeciiieeeeeee 6399.
URL https://:// aaai org/
R. J. Serfling. Probability Inequalities for the Sum in Sampling without Replacement. The Annals of Statisticx, 2(1):39 - 48, 1974. doi: / URL hwwp:: httttt/////daii http:///d..oogogooo/o////o 10. 1214/
aos/117634263
Daria Soboleva, Faisal Al-Khateeb, Robert Myers,
Jacob R Steeves,
Joel Hes- token cleaned and dedu-
tness, and Nolan Dey. plicated version of RedPajama. s1 impa jama-a 2023. URL https: htw.h.t.ss/hhhggnnngeeee
SlimPajama:
A 627B
https: http:///wwwweebbb.eeeeee.ee.ee net/blgg/
jama,
Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari S. Morcos. Beyond neural scaling laws: beating power law scaling via data pruning, 2023. httptttt////axi...//////a abs/ 2206. 14486.
Kevin Swersky, Jasper Snock, and Ryan Prescott Adams. Frecze-thaw bayesian optimization. arXivv preprint vrXiv:1406.3888 2014.
Yi Tay, Mostafa Dehghani, Samira Abnar, Hyung Chung, William Fedus, Jinfeng Rao, Sharan Narang. Vinh Tran, Dani Yogatama, and Donald Metzler. Scaling laws vs model architectures: How does inductive bias influence scaling? In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Findings of the Association for Computcational Linguisticx: EMNLP 2023, pp. 12342-12233 -2364, Singapore, December 2023. Association for Computational Linguistics. doi: 10. 18653/v 1/2023 findings-emmlp.8255 URL https: http://aaaaatttta /at/aaaaaannoooooooooooao org/2023. findings emnlp. 825.
Tristan Thrush, Christopher Potts, and Tatsunori Hashimoto. Improving pretraining data using perplexity correlations, 2024. URL https: /ttp/vrxvvvvv///.e//.. abs 2409. 05816.
Kushal Tinumala, Daniel Simig, Armen Aghajanyan, and Ari Morcos. D4: Improv- ing Ilm pretraining via document de-duplication and diversification. In A. Oh, T. Nau- mann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neu- ral Information Processing Systems, volume 36, pp. 53983-53995. Curran Associates, Inc., http://eeeeriss.spppppe__ii_i_i and Benchmarks pdf.
2023. URL https: htttt:::s//pooceenngggssssee nwwrips. cctpsper.ppppppeeeeeeeeeeeee faterrrfills/eeee www.eeeepppeeee petpeeepeee/22///////ee/ 1xf8:0077aff22267777777776022220
Togetherai. Redpajama: an open dataset for training large language models, October 2023. URL
http::///...cctccccooco/ttteeeep http://///tceeeeee//ooeepeeeeeee hatpar/cccoomttteeeeaaaaaaaaaaee jama-Daaa. https: https://github.coo
Mariya Toneva, Alessandro Sordoni, Remi Tachet des Combes, Adam Trischler, Yoshua Bengio, and Geoffrey J. Gordon. An empirical study of example forgetting during deep neural network learning. In International Conference on Learning Representations, 2019. URL https: /wwp//wweee..eeee///////.. net/
Forum?id=Bal xm30ckm.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023.
URL https: //atii arxivv. org/ abss/ 2302. 13971.
Neeraj Varshney, Swaroop Mishra, and Chitta Baral. Let the model decide its curriculum for multitask learning In Colin Cherry, Angela Fan, George Foster, Gholamreza (Reza) Haffari, Shahram Khadivi, Nanyun (Violet) Peng, Xiang Ren, Ehsan Shareghi, and Swabha Swayamdipta (eds.), Proceedings of the Thind Workshop on Deep Learning for Low-Resource Natural Language Processing, pp. 117-125, Hybrid, July 2022. Association for Computational Linguistics. doi: 10.186653/1//022.depppo 1.13. URL ht tttpps:s : / awwaanttooooooooo org/ 2022. deeplo-1.13.
Mengzhou Xia, Mikel Artetxe, Chunting Zhou, Xi Victoria Lin, Ramakanth Pasunuru, Danqi Chen, Luke Zettlemoyer, and Veselin Stoyanov. Training trajectories of language models across scales. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Proceedings of the 61 Annual Meeting of
15

---

Published as a conference paper at ICLR 2025
the Association for Computtational Linguistics (Volume 1: Long Papers), pp. 13711-13738, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 2o..8666//1/22223.ll1......./ https::
wwwp:/hhhoooooooo///eeee2.e. acl-long. 767.
Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. Less: Selecting influential data for targeted instruction tuning, 2024.
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. In Thirty-seventh Conference Neural Information Processing Systems, 2023a. URL ptttps://openreeeee net forum? 1ineedddlllBsa d-lXuByUeHhd.
Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models importance resampling, 2023b.
Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models via importance resampling, 2023c. URL https: //tp///aaavv.//aa////a/eaaeaawia wwtp:vv..rgaab///222222 /ttg:/aabb2000.01.....
Jiasheng Ye, Peiju I 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.org www.ooggg//b//.o..ooooo..../..b htbp/2403.1655
Dongkeun Yoon. Slimpajama-6b.
https: //rpss:/hae S1 imPa jama- - 6R, 2023. Accessed: September 24, 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can a machine really finish your sentence? In Anna Korhonen, David Traum, and Lluís Märquez (eds.). Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4791-4800 Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/119 1472. URL https: //acl anthoooooooo oog.. org/p115 9-1..
16

---

Published as a conference paper at ICLR 2025
APPENDIX
In Appendix A, we provide a glossary of notation used in the paper. In Appendix B, we discuss how additional data mixing methods related to the LMO framework and provide proofs that existing methods can be expressed in our framework. In Appendix C, we provide additional results on our analysis of existing data mixing methods. In Appendix E provide additional details for our results in Section 6, and in Appendix we provide additional results, including downstream evaluation and ablations.
A NOTATION
The glossary is given in Table 4 below.
Symbol
Used for
m
The number of data groups. Examples of data groups include pre-training domain or instruction-tuning task.
Dusineathes
Training, validation, and test datasets comprised m groups, where Diy group i's training/validenn validationntest data.
N
Total number of samples from Duain train Number of steps to train for (i.c., S= N batch size).
T
Number of rounds to divide training into, where each round is _ steps.
P
Mixture proportions where p = pí, .p... E A probability distribution.
(=(p") for T= (static) and p= (p',p.p. ---pp" for 7 (dynamic),
A language model (can be either pre-trained initialized from scratch).
f(p.))
The model at beginning round after being trained on p² .....p"" so
Tnnanivallee (f)
Tumin (f)== (ff)....TTunnn... (f)) is the vector of f training losses over each data group; similarly defined for validation and test losses.
(p)
Shorthand for (fp.t)) When dealing with static mixtures, also (p)(p) Parameter matrix A mm used in mixing laws (2), capturing cross-group interactions.
A
See Table for instantiations.
Ved
croup-specifie parameters ER" used mixing laws Note that value of é' does not impact the LMO framework, and neither does 6' when all of are equal. Fither o: + R= exp.
z
Used for normalization proportion update rule. Step size > 0 used in proportion update rule.
7
P
The set of mixture proportions that comprises training sweep.
A
Approximately optimal A' for the linear dynamic mixing law, ohtained by fitting Lvall++(pp) Lval' (p) -A' p over training sweeps.
A'
Method -seiodd- A =b" A", where A obtained directly from the method b' E R is learned from training sweeps.
(im(A",A*)
Similarity between method method-specific and optimal A', defined as an average cosine similarity and Spearman rank corrlattin over A''s normalized column sums.
one-hot smoothing factor used define pr* =(1 e)1, + + +eUnif(m))) smoothed one-hot distributions we use to leamn A' AIOLI.
5
The fraction per round dedicated to learning A' in ATOLI
Number of sweeps per group to average A estimates over in AIOLI.
plu
Initial mixture p! E A that AIOLI can dynamically adjust.
Smnt
Number of steps to train according to p
Table 4: Glossary variables and symbols used this paper.
B LMO FRAMEWORK DETAILS
B. ADDITIONAL EXISTING METHOODS
We comment on two other popular data mixing methods, Online Data Mixing (ODM) (Albalak et al., 2023) and RegMix (Liu al., 2024).
In ODM (Albalak et al., 2023), data mixing is framed as a multi-armed bandit problem, where each am is data group that batch is trained on, and the reward function is defined in terms of the training
17

---

Published as a conference paper at ICLR 2025
loss of each group. ODM uses the EXP3 algorithm to explore training on different data groups. p'. which is used to determine which group the entire training batch is comprised of, is updated according explec-11tt)
to p = (1 11-1111 me{) En explec-112)) Eg. Ee is an exploration rate, and the reward function Rt RR + (1 a) limm,
the jth group is selected at time i; otherwise, Rj - R While the exploration and the smoothing of p' and It make this method not directly expressible in our framework, we note that the update rule can be loosely interpreted as allocating larger proportions groups that have high loss. This update rule does not consider cross group interactions and is thus similar to DoReMi's update rule, which utilizes a diagonal A'' defined in terms of current loss.
RegMix (Liu et al., 2024) conducts many training runs on smaller models at shorter scales. Similar to DML (Ye et al., 2024), a regression model is fit to these runs and used to predict mixture proportions larger model. They consider using a linear regression model, i.e., the mixing law (p) ( -jj A:ppj. ; but find that the R is relatively low (0.87). Instead, their main approach uses LightGBM, a tree-based gradient boosting approach, i.e., using an ensemble of non- non-linear decision trees as mixing law. We note that ATOLI could be used in conjunction with RegMix in their settings, exciting direction for future work.
for longer run
B.2 PROOFS FOR SECTION 3.3
B.2.1 BACKGROUND ON EXPONENTIATED GRADIENT DESCENT
We provide background on exponentiated gradient descent (EGD) taken from Kakade (n.d.). In EGD, we have a sequence of decisions w ,..., w where w' = wit E A We also have a sequence of cost functions c² .....c":: A - R. To minimize the total cost _t=1 (w"), the EGD update rule
sets w Uni, and updates according to w w+11_ wsexp(- n cw"") Ze ensures that wt+1 ".. de' (w(), EGD is known to have certain regret guarantees on the Bhe value of costs incurred by playing w w'...,"" versus always playing the best fixed point in hindsight:
Ze
y is step size, and (w") denotes
El (w") -nmfi)-inffee -n df
We now are ready to prove Lemma Lemma 1. The EGD update rule for (1) subject to (P) (p)
D A AiEs
1
m
pf
z: expexp
bj E [m).
(3)
where 77 0 the step size and Z' a normalizing constant such that p E A..
Proof. The cost function at each timestep in our setting is Σm 1 (p), and the decision we make is
p¹. The mixing law constraint in (2) with a 1+1 (p) o 1 for all i [m]] so our objective (I) can be written as
m
m
The gradient of this expression with respect to pf for j [m] is - Plugging this into the EGD update rule, we obtain the update p ž, p.epppp - B.,4t)).
B.2.2 PROOF OF THEOREM 1
To prove Theorem 1, write out individual propositions 1, 2, 3 for expressing each online method in the LMO framework.
By our definition of what means to express method in LMO, must consider how each method 1) trains f and 2) sets p'. We must see if this procedure can be replicated by solving some specification of the LMO optimization problem in data mixing setup.
18

---

Published as a conference paper at ICLR 2025
Critically, note that this definition of "expression" does not claim that the optimization problems proposed in existing methods are exactly the same as the LMO optimization problem. Instead, we are stating that the training procedures used in their methods can be equivalently viewed as a way of solving the LMO optimization problem subject to certain assumptions on the loss-proportion relationship.
Proposition (Skill-It Derivation). Using a) a linear dynamic parameterization (p)) i B' B'E P. b) parameters Ai, = , ( (1z) La / (1,), and c) exponentiated gradient descent (EGD) to solve for p. the LMO framework (I) can express Skill-Ir.
Proof. The Skill-It algorithm sets p' in each round and then samples from Durin according to p to train I for a round. This training procedure is directly specified in our data mixing problem setup (Section 2). Therefore, we simply need to show that the Skill-It update rule can be converted into a linear dynamic mixing law. By comparing Lemma and the Skill-It update rule pj = ž. 12jexp(1E A L we can match A in the lemma with A.5 in Skill-It, and we can match by in the lemma with L (p). Therefore, Lemma tells us that using (p) = ce -b ""j" with exponentiated gradient descent recovers Skill-It (since the b' and oa can be dropped and are only used for scaling A').
(p)) P) in the LMO framework
Using the definition of A we can rewrite the mixing law as (p) -EE A 1 = ,u, (p) (L (1₂) L.a, (1;)) Laa (1)). Lastly, note that can replace of with any other value, including a (P). due to the fact that p' has m degrees of freedom (see Lemma 2).
Skiilll- P where
We note that Chen et al. (2023) explicitly specify their mixing law in equation 2 of their paper, along with the same objective function as ours in the LMO framework.
Proposition 2 (DoReMi Derivation). Using a) a linear dynamic parameterization L (p) L' vol. (p) D' "'E m Agppy b) parameters Aly- mmin Li (p) (p) [ (fre))00} for i-j and A,y - otherwise, and <) EGD to solve for p. the LMO framework (1) can express DoReMi's 's proxy model.
Proof. When training the proxy model for DoReMi, p' is set in each round, and then f is updated to minimize E _i=p; Lamin. ) Using Lemma 3, we establish that DoReMi's weighted training objective at each timestep is equal in expectation to the objective of training on data sampled from p², which is what our problem setup focuses on. Having established that the training procedure the same in expectation, we now need to show that the DoReMi p' update rule can be converted into linear dynamic mixing law. By comparing Lemma and the DoReMi update rule pf+ x px]eexx(maaxx ['iiain. (p) T (0)), we can match A'y in the lemma with 0 for jzj. and A with max Linin (p) Luain (fre)).). Therefore, Lemma 1 tells us that using L = - - Ai,ps with AG, max Limin (p) Luiin (fret) ).}} can express the DoReMi proxy model training. We include b' to allow for scaling A', but since this does not impact the optimal p. it is not in the update rule. Lastly, applying Lemma 2 lets us write the mixing law as T - Tha, (
We comment on the fact that DoReMi's proxy model is trained with DRO (distributionally robust optimization) min-max objective, namely, minimize, maximize,_.=n1 [ rain (f). This objective, which differs from our data mixing objective, yields the y' gradient ascent and f² gradient descent updates. However, we are still able to express this training procedure in the LMO framework, since our claim is: if we assume that the - Li (p) -b'E E>= L and p', then training according to the DoReMi proxy run should not only guide ff and p to optimize the DRO objective, but also to optimize the average validation loss per group.
1.DRM A4, P) mixing law captures the relationship between
Proposition (DoGE Derivation). Using a) a linear dynamic parameterization L (p) = Lini, (p) b' "'> At Alppy, b) parameters At.- = (VVL (p), Linin (p)) for all i.j E [m]] and c) EGD to solve for the LMO framework (1) can express DoGE's proxy model.
Proof. When training the proxy model for DoGE, p' is set in each round, and then f is updated to minimize m
Em 1 Immin,(() Using Lemma 3, we establish that DoGE's weighted training objective at each timestep
19

---

Published as a conference paper at ICLR 2025
is equal in expectation to the objective of training data sampled from p'. Next, show that the DoGE update rule can be converted into a linear dynamic mixing law. By comparing Lemma and the DoGE that A in the Lemma
update rule p pfppppp exp(n) (v Lumin
[)). we
be
matched with (V Lunin ( (f"), (f))) Therefore, using the mixing law L = c: - b4 b Aipp) with A =4,= -V (J*)), =J")) allows LMO to express DoGE proxy model training. Again, D is included for scaling but does not impact optimization, and by applying Lemma 2, replace of with (p).
Lemma 2. Let 1 (p) & -ij Aigh? for some et and A Then, there exists an Big such that (p) (p)) Bijji
Proof. Since pl mh we can write the probability p', as 1 11- written as
Then, the first equation can be mm 1j.
m-1
1.41 (p)
Alspl
A
1
Σ
(5)
j=1
==1
= c
(A1
Amp lnn -Ah
m--1
- - AAmmD' -c
j=1
m-1
Lp(pp) Σ (Ai,- - -A + + - - - (A - Ens
m-1
m-1
= -p(pp)
Σ (A1y-c4 - - -4 - Σ p.).
j=1
j=1
Let Bi P A - + (p) for all j jmm.]. Then, this equation becomes
m-
m-
4, =
Σ
By
B (1
Σ
ps)
(6)
5=1
5=1
= =T((p)
B
B
Lemma 3. Let L (f.p) be the total training loss of f on a batch of size B sampled from D Dran according to pE Am, and let (J.p() the total training loss on samples from group i that batch. Then, the average loss over a uniformll sampled batch weighted by y is equal in expectation to the average loss per group over batch sampled according p²:
(f(fp')) m
IE
P. P.. (F:
=1E
Proof. Let each group consist of samples from the distribution P,. and let Lusin. i(f) = Epw., [((f.x))] be the population-level loss on group i, where {(fxx) f's loss on sample x.
If a batch is uniformly sampled, each group has B/m samples. We can then write L ( (f,,li(()))) B/m ((fart)) where of, is the kth sample of group i. Then, Eken
/m
v T sf))
E
Σ U.
=F
(
m
20

---

Published as a conference paper at ICLR 2025
Next, if batch is sampled according p'. then group i has Bpf samples in the batch. We can then write -= - >
(f Then,
(f
B
((faat)
m
B
E
E
Lariin is))
m
m
m
This hence establishes the equivalence in expectation between weighted training objective and training on data sampled according to p.
C ANALYSIS DETAILS
C.I MIXING LAW PARAMETERIZATION
Table 5:
Comparison of log-linear static and linear dynamic mixing law parameterizations across different data settings with MSE and R2 metries. Both log-linear and linear dynamic mixing laws fit the relationship between mixing proportions and losses well.
Arxiv/SE
GH/C4
Books/SE
Parameterization
MSE
R²
MSE
R²
MSE
R²
Log-linear static
2c-4
0.990
5c-4
0.989
6c-4 4e-5
0.987
Linear dynamic
2c-4
0.936
le-4
0.948
0.926
ArxivBookssSEE
CC/GH/Wiki
SlimPajama R²
MSE
R²
MSE
R.a
MSE
Log-linear static
6e-4
0.991
0.001
0.989
0.002
0.997
Linear dynamni
6e-5
0.957
lc-4
0.975
5e-6
0.938
We describe how we performed the linear and log-linear parameterization experiments.
For the log- linear static parameterizations, we train our model p P sweeps and fit the parameters using code provided in Ye al. (2024) (i.e., using PyTorch and L-BFGS minimize the Huber loss of the mixing law). We do this over 5 random seeds for k: = 2,3 and over seeds for the full SlimPajama.
For the linear dynamic parameterizations, for 2=3 we train the model for 2000 steps according to some pº EP, and then sweep over P for the next 100 steps. We do this for one random seed, performing |P|2 total runs. For the full SlimPajama setting, we train the model for 10000 steps using stratified sampling, and then sweep over P for the next 5000 steps. We fit the parameters using Pytorch and L-BFGS.
C.I.1 ADDITIONAL PARAMETERIZATUN EXPERIMENTS
Parameterization across checkpoints. We investigate whether the log-linear static and linear dynamic mixing laws remain well-specified in later stages of training and on other datasets. To do so, we take various Pythia 160M checkpoints (Biderman et al., 2023), sweep mixing proportions, and fit the linear dynamic and log-linear static mixing laws. We train for 2000 steps according to the learning rates and leaming rate scheduler reported in (Biderman al., 2023). We fit the static mixing law on full runs of 2000 steps, and the linear dynamic mixing law at t= 500, after which we training sweep over the next 500 steps. In Tables and 7, we find that the strong fit for log-linear static mixing laws continues hold during pre-training at checkpoint 72K (roughly halfway through training Pythia-160M, and after pre-training, with an average R² of 0.982 and 0.991. respectively. However, the linear dynamic mixing law's R² coefficient is lower, averaging 0,815 at checkpoint 72K and 0,830 at the end of pre-training. It thus may be interesting further study if the dynamics of the loss- pos-proportion relationship evolve in structured way throughout training, or if these results are due to more noise in how models leam at later stages of training.
Parameterization across other sets of data groups. In Figure 4, we identify an example set of data groups that exhibits non-linear relationship between loss and proportion: Books/C4 from SlimPajama.
21

---

Published as a conference paper at ICLR 2025
Noar step Loss on book
4.20
Nott-blp Loss
4.40
183
10-1
1.35
10-1
1.30
4,2
044
Lo0
20-1
Progortion
Progortion
022 6.5 06 64 10 050
Proportion book
Propurtion
044 6-6 2.
yoarr
Figure 4: Top: Log-linear static mixing law fit on Books/C4 across random seeds. Bottom: Linear dynamic mixing law fit on Books/C4 on random seed. Each color a different initial mixture p E trained for 2000 steps, and the fitting sweeps done over 100 additional steps.
For these two data groups, see that as the proportion of Books increases while C4 decreases, the loss on Books starts increasing past a certain p, suggesting quite counterintuitively that performance on Books optimized by allocating some proportion to C4. In this case, neither log- oog-linear static or linear dynamic mixing laws have good fit to the proportion-loss relationship, as neither can represent the non-linearity. In particular, the average MSE and Rº for the log-linear static mixing law is 0.003 and 0.558, respectively, and the average MSE and R2 for the linear dynamic mixing law is 0.0002 and 0.721.
Fortunately, because these nonlinearities exist the boundary of the simplex and tend to incur high loss, they tend to have little impact on the optimization of p. which strives to minimize the average loss. For instance, we found that the optimal proportion according Ye al. (2024)'s log-linear static mixing law on one random seed was (0.176.0.88 0.844], and the true optimal from grid search was [0.2,0.8]. However, it is important to further investigate this non- con-linear phenomenon on additional data groups and training regimes, which we defer to future work.
Table 6: Comparison of log- binllinear static and linear dynamic mixing law parameterizations when training from the 72K Pythia-1666 160M checkpoint.
Arxiv/SE MSE
GH/C4
Books/SE
Parameterization
R²
MSE
R²
MSE
R²
Log-linear static
2e4
0.975
7e-5
0,992
2e-4
0.981
Linear dynamic
4e-4
0.834
7e-4
0.815
6e-4
0.796
Table 7: Comparison of log- log-linear static and linear dynamic mixing law parameterizations when training from the pre-trained Pythia-160M)
Arxiv/SE MSE
GH/C4
Books/SE
Parameterization
R2
MSE
R²
MSE
R2
Log- linear static
3
0.994
4e-6
0.992
6e-6
0.986
Linear dynamic
5e-5
0.896
8e-5
0.824
lc-4
0.769
Checking for interactions among groups. It is natural to ask whether linear mixing law is sufficient to model how mixing proportions affect the loss. In linear regression, such assumptions are often evaluated using visual diagnostics called residual plots (Montgomery et al., 2021). Residual plots graph the prediction error from each data point (the residuals) in order to reveal different kinds of structure. For example, it is common to plot the residual against the predicted value to check for nonlincarity.
Figure shows several such residual plots for the dynamic mixing law experiments with 3 domains (Arxiv, Books, and Stackexchange). The figure checks for interactions when predicting Arxiv's 's loss. The corresponding plots for the other domains look similar.
The top visualizes the residuals inside the simplex. If strong interactions present, then they would cause clustered patterns in the residuals- rssiduals-regions where the linear model consistently gives predictions that too low too high. Strong patterns do not apparent.
22

---

Published as a conference paper at ICLR 2025
Anat Interaction Elagnostics
3.094
0.0973
1439
0.1653
0,2913
0.4352
0,3000
0.8577
0.893
12
81
10
-2.-
02
01
PAPy
PP
PA
9
AP
Nescual
-25
0.0
Figure 5: Residuals plots check for interactions in the dynamic mixing law experiments with domains (Arxiv, Books, and StackExchange). The target loss is Arxiv. Columns correspond to different initial mixing proportions. Data points show the (extemally studentized) residuals of different mixing proportions after fitting the linear mixing law. Top row: Each point in the simplex corresponds a different mixture of the 3 domains, with color giving the residual's 's value at that point (red is positive, blue is negative). Bottom 3 rows: each row shows the residual plotted against a different interaction term: P1 Pz. P1 P.P3. and P2P3. Dotted gray lines show upper and lower 99% confidence limits for the residuals, assuming the linear regression assumptions hold.
The bottom three rows plot the residuals against different interaction terms. A consistent trend in the residuals above or below zero would suggest the term captures a meaningful interaction. The scatter plots show no consistent trend. The first three charts on the bottom row hint that small interaction could be present in those cases, however, it is difficult to say without larger samples. Considering the linear model'ss excellent fit and high R², if such an interaction is present then it is likely small.
To summarize: the linear model seems sufficient. While we can not rule out the possibility of small interactions, the diagnostics do not reveal any major departures from linearity that might compel us to use more complex model.
C.2 VALUES OF MIXING LAW PARAMETERS
We explain how to compare method-specific A to an approximation of the true A First, after performing method -speciffc initialization, such as training reference models, we run each online method (Skill-It. DoReMi's proxy model DoGE's proxy model, Skill-it, and AIOLI) for t steps. For Skill-It, DoReMi, and DoGE, we use the unrestricted setting configuration of hyperparameters presented Section E. For AIOLI, we analyze the parameters of ATOLI +GS from the restricted setting, since we found that this had less noisy fluctuation in the weights than the unrestricted setting. For m 2, we set / = 1000 for Skill-It and -= 500 for DOGE, DoReMi, and AIOLI since Skill-It is updated less frequently. For m m= we set t- 1000 for DoGE, DoReMi and Skill-It, and t= 1500 for AIOLL We then checkpoint the language model and the method's A', For DoGE and DoReMi, we compute a smoothed A' = Too 100 A 40-100++ because each A' is computed at the batch level, and can thus be noisy. For AIOLI, we also smooth the A' by averaging the previous timestep parameters.
To approximate * we then run a training sweep of p' over P for 100 steps on the checkpoint. We use
this training sweep fit A from the dynamic mixing law v (p) T (p) - Alppy.
23

---

We also have similar findings on Published as a conference paper at ICLR 2025
Before compare parameters, scale A' by some D' where - Lp) (p) for all iE [m]] This allowed since bf does not influence the optimal p and does not need to be the update rule. fit single 0' across each group's mixing law and set A' = A', We can then compare A' and A using the metric sim(A', = 0.5cosiimm a +0.5pperrman(( ,a'a,, which we proposed in Section 4.3.
C.2.1 PROPERTIES OF A
We discuss some properties A finding that 1) A' can vary significantly across time, and 2) A' needs to be modeled as full matrix. To do this, for each initial mixture p P, train for fort= 20000 steps and then sweep over P for the next 100 steps. We repeat this setup for t= 4000 to obtain 42000~ and A 40000. We do this experiment for Arxiv/Stackexchange and Github/C4.
Extent of time variation of A', We find that the column sums of A' can change order over time, meaning that the pf "changes direction" in terms of which group has the largest proportion. In particular, for [0.5.0.5 and Github/C4, we have that
2000-~
0.148
0.011
44000
0.015
0.001
(10)
0..03
0.087
0.001
0.015
The column sums are 1 -0000~ [0.135,0.098] and 1 A - =0.006,0000017 showing that the ordering of proportions of the groups changes. This suggests that the optimal p can change significantly across time, prioritizing Github initially and later C4, which is also reflected for Github/C4 in the greedy row of Table 8.
However, for Arxiv/Stuckexchang. the column sums of 2000 and 44000+ never change in terms of the ordering of proportions of the data groups, across all p" e As a result, the optimal p' never changes direction. This suggests that how much A' varies in ordering a result, methods like Skill-It, which use a time-invariant SS multiplied by validation loss, may not be able match the observe in Github/C4.
time depends on the data groups. As
f if the groups' validation losses do change ranking across time, which we
Modeling * as a full diagonal matrix. We find that modeling the off-diagonal entries of A is important. For each sweep, we fit both 4 as described above and a diagonal matrix AT, We compare if the column sums of * and A differ in the order of elements.
We find that for Arxiv/StackExchanec pd =0.4, and both t 2000 and t= 4000, setting p based on the full matrix would put larger proportion on StackExchange, while setting p based on the diagonal matrix would put a larger weight on ArXiv. In particular, the full and diagonal matrices for t 2000 are
T..29
0.058
00.244
2000
A2000
(11)
0.025
0.224
0.238
The second column sum is larger for 2000 and smaller for A2000. Github/C4; for pd 00.6 and t= 2000, we have
20000-
0.119
0.027 0.104
A
2000- = 035 O
(12)
0.010
0.098
Using the diagonal matrix for Github/C4 would result in prioritizing training Github, though the full matrix suggests that C4 should be prioritized. Therefore, it important to model 4 as a full matrix. As result, methods like DoReMi, which use diagonal A", can perform suboptimally.
C.3 SOLVING STRATEGY
We present our results on examining the assumptions made in how existing methods solve the LMO optimization problem. All online methods use exponentiated gradient descent, which updates p' using the gradient at the current timestep. This involves a greedy approximation of the objective function. We study if the greedy approximation yields p is close the true optimal p.
24

---

Published as a conference paper at ICLR 2025
Nrat www.koethhhr.rr
2.0
0.8
0.6 08 Propoction stackexchangg
Propertion
priee
Figure 6: The linear dynamic parameterization results from Figure 2 (right), with p' p0= and [1,0] also plotted. We see that the linear dynamics are misspecified pt -0 for both
For m==22 data settings, we take our 5000 steps and split it into T rounds. We perform brute-force sweep at each round over P. which sweeps pl -...0........0.. In total over one random seed, we conduct 81 training runs for each of Arxiv/Stackexchangee Github/C4, and Books/Stackexchange.
We determine the gpeexxappeo p by selecting the best p². Then, conditioning on this p², we select the best p². We report what the greedy and its performance is in the first of Table 8, and report the optimal and its performance in the second row. Note that this protocol does not depend on the mixing law or method for setting p.
We find that for Arxiv/StackExxcangg and Books/StacEExchange, the greedy proportions and the optimal proportions are identical. However, for Github/C4, the greedy approximation fails to recover the optimal proportions. Therefore, the greedy approximation recovers the optimal dynamic proportions in 2 out of 3 cases.
Table 8: Comparison of the greedily selected p' y² versus the optimal p', p² for a T -= 2 rounds data mixing problem. On 2 out of 3 datasets, the greedily selected proportions match the optimal proportions.
Arxiv/EE
GH/C4
Books/SE
Solving
pl.pi
Avg test PPL
pl.pf
Avg test PPL.
Pl.
Avg test PPL
Greedy
0.4,0.4
16.039
0.6,0.4
36.525
0.3,0.6
45.513
Optimal
0.4,0.4
16.039
0.3,0.6
34.709
0.3,0.6
45.513
Beyond exponentiated gradient descent, one may wonder if exactly solving the greedy objective could suffice. For the linear dynamic mixing law Li+1(p) = L'(p) A'p', the optimal p is 1,, where j= argmax 1 However, we find in Figure that the loss-proportion relationship can be nonlinear the edge of the simplex where 11, Exponentiated gradient descent, which uses entropy regularization, is hence able to implicitly avoid extreme p where the linear mixing law is misspecified and thus is practical technique for LMO.
D ADDITIONAL ALGORITHMIC DETAILS
In AIOLI, LEARNPARAMS is used in each round learn A'. Then, A' is used to compute pf, which is used for training during the round. We provide a derivation of LEARNPARAMS by first presenting a naive, high-cost method for estimating A' (Appendix D.I.. This involves checkpointing the model at each round, running training sweep over the round and observing the changes in validation losses, and fitting A' to these changes. Then, we layer on two modifications that compute slightly different loss changes, helping lower the cost of estimation. First, we shorten the training sweep to be only over a fraction of the round, and use these shortened changes validation losses to fit A' (Appendix D.2). Second, we simulate simultancous training sweep by partitioning the ö fraction of the round into many small parts, interleaving the different sweep mixtures at a fine granularity and averaging the loss changes for each sweep mixture (Appendix D.3). This idea, with similarity concepts like time-division multiplexing in signal processing (Chaparro & Akan, 2019), enables AIOLI to require no extra training while trading off accuracy the estimate. We provide sketch of our derivation in Figure 7.
Finally, in Appendix D.4 we convert our first principles description of why LEARNPARAMS works above into theoretical statement, bounding the difference between A' leamndd from LEARNPARAMS and A
25

---

Published as a conference paper at ICLR 2025
Naive approach with full training sweeps
DR1
p11
p
p
D1.
2. Shorten training sweeps
D11
p²
3. Aloli: interleave training sweeps
LearnParams
= 3.
1= 2,
p
Figure Derivation of AIOLI. Top: a naive high-cost approach where training sweeps are conducted to fit A' at each round (Appendix D.I). Middle: a modification that shortens the training sweeps used learn A' (Appendix D.2). Bottom: final modification that interleaves the sweep mixtures high frequency (large k) in one single run, enabling AIOLI's LEARNPARAMS to require no additional training (Appendix D.3).
D.I NAIVE TRAINING SWEEP APPROACH
This approach is depicted Figure (top). By conducting a training sweep over round we can use linear system of equations to estimate A from the linear dynamic mixing law L.d. (p) L (p) - -i= Let p'1 ./......p'' E " comprise training sweep the duration of round t. First, checkpoint the model f', and for simplicity denote f' validation loss on group as L For each y² we train ft for the entire round using pfJ, We then record how much the validation loss on each group changes, L Lai (n²") for all iE [m]. By the end this procedure on each p!u, we have the following system of equations for each i [m]]
Σ
Appp =
val, L. (p
I1+1
(13)
7=1
Alyp
1.2
L'utt
L val. (2)
1=1
Arp tm
That
I.+1 (pt.ma
This is system of linear equations with unknowns: A Aim. We can write it in matrix form as:
Pj
1.1
Pm
A
D (1
2
A
The I 1+1
vall 1 (2)
P1
P2
pm
(14)
I.
1.m
p
L
L - (ot
P2
26

---

Published as a conference paper at ICLR 2025
Let Pe "eRmxm be the leftmost matrix and B, E " be the vector on the right hand side. Then, we can write A We solve this system for each ie [m]] to obtain A'.
The advantage of this method is that it directly estimates the optimal A that is used in the mixing law. However, it requires sweeps per round, because the key quantity we must observe to learn A' is L' (p): the change in loss after training through the entire round As a result, this approach requires m extra full training runs to learn A'. Below, we will describe how we can compute cheaper alternatives to I " -
D.2 MODIFICATION 1: SHORTENING TRAINING SWEEPS
This modification is depicted in Figure 7 (middle). A simple way to reduce the number of extra training runs needed estimate A' is to train each mixture p for less than a round. Let denote the fraction of the round use for the training sweep. Then, system of equations in 14 uses L L - (p'");; we simply record the loss difference over å of the round rather than the entire round, and use this to solve for A'. Now, this approach effectively requires mő extra training runs; however, this cost is still linear in the number of data groups. Moreover, there is some inaccuracy incurred by using 5 of a round to approximate the entire round.
D.3 MODIFICATION 2: "INTERLEAVING" TRAINING SWEEPS
This modification is depicted in Figure 7 (bottom). Our final modification to derive LEARNPARAMS is convert the training wweep-where we checkpoint the model and execute separate runs for ó of round-into one round without requiring any checkpointing or rolling back of training. Our intuition is that if we interleave different mixtures sequentially at high frequency, mixtures simultaneously. This is similar to a concept in signal processing called time-division multiplexing, in which two or more signals bit streams are transferred appearing simultaneously as sub-channels in one communication channel, but are physically taking tums on the channel'.
simulate executing these
Formally, we break down the 85/T steps allocated for learning A' into K intervals, where Kmmk and k is the number of sweeps per mixture. We construct an interleaved order of "'.."".''''"" over these K intervals, and we denote their index order as I (m] K Let I, denote the mixture at the rth position in I. We can denote the model at the end of each interval as 5/ 25/K Bo/.K..... 5. During the rth interval, we train on one p/-T- and observe the change in loss, (f (ft (r-)6/(11//))------ Lvat. (f 66/KK) K)) for each validation group i. Let T, (T jj} be all the intervals where pn is assigned. We approximate Li L (y/-) with + >EET, (f 1)6/K) ( These approximated loss differences are then used to recover A' from the system of linear equations.
Lastly, note that the choice of is controls the interleaving frequency and the bias of the estimated A Suppose that ==1. This means that each mixture is only assigned to one interval, and this could be at the beginning, middle, or end of the 8 round. Then, the change in loss a poor approximation quantity T. - - trained on in the beginning, middle, and end of the 5 round, allowing for a less time-biased estimate of the loss change.
the original due to dependence on time. However, as we increase k, the mixture pd will be
( -
With this modification, LEARNPARAMS now requires no extra training. However, there are still some performance tradeoffs. First, in order to save compute, our estimate of A' via the shortened interleaved sweeps is less accurate than the naive approach. Second, without rolling back training, ATOLI has both an "explore" and "cxploit" phase, where the former Iearns A' over of the round and the latter uses A' to set o and mix data accordingly for the remainder of the round. If 8 is large. the estimate of A' may be relatively more accurate. However, training for longer on the sweep mixtures p' ...p'' may be suboptimal for the performance of the model. Moreover, the training duration that utilizes the p' that is updated using the more accurate A' is now shortened. Therefore, adjusting 8 is key to ensuring that A' accurate and the model performs well.
mm. wik http:///wwwaaaaaaaoa/o//////////
27

---

Published as a conference paper at ICLR 2025
D.4 THEORETICAL RESULTS
We now provide bound on the difference between A the optimal mixing law parameter, and A A the parameters that LEARNPARAMS estimates.
Our proof approach follows similar break down sections D.I, D.2, and D.3. Let A be the parameters learned from the naive approach in Appendix D.I. where we fit the parameter from full training sweeps over the round. Note that the difference between A' and A in this case is due to finite-sample regression (i.e., over the runs rather than the entire population of mixture proportions). Let A be the parameters leamed from the shortened approach in Appendix D.2, where we fit the parameter from shortened training sweeps. The difference between A' and Å comes from variation how the losses when trained on each mixture evolve from + to t++1. 1. Finally, let A AtAioli be the parameters learned from LEARNPARAMS: which interleaves different proportions throughout 8. The difference between A and A comes from how frequent the interleaving of mixtures is.
Our goal is to find an upper bound on || -ur=-nAAiolllly since differences in magnitudes can be offset by adjusting Using the triangle inequality, we now have the following decomposition:
|| p. (Aiil
- | - 12 11 -v
11 | μ.A tAioli
||
(15)
OLS
interieaving
Below, we state a set of conditions we use to bound all three quantities. Then, we present our main result, Theorem 2. before providing our proof.
Setup for OLS error. We frame the error from A being learned on finite samples as the "random design" prediction error of ordinary least squares regression (Hsu et al., 2012). We specify the data model (covariates and response variable) and use it to formally construct * and A
Let the covariate p E A separate regression problems. Note here that we have dropped the remainder of the dynamic p in the loss notation in order to model that only the mixture at time I, is the covariate. For simplicity, we consider single regression problem over the ith validation group.
random vector. We have m response variables, L - Li (p). and m
We are given samples, (p/ I algorithm, we m)) We make the following assumptions that define A and our data-generating model.
(p.."))......""""""""" (yl -I T.
())))) to estimate each A (in our
Condition 1 (Ordinary least squares data model). Suppose the following hold true regarding the true data model of p and (p.
1. p drawn from centered Dirichlet distribution. That is, p~ Do-Drichlee(((s......l.. where a2 =0 for all jE (m).
2. The true loss propporonn relationship follows the form,
L - L (p) > AA55 p
(16)
==1
for some At ER", where EOLS is zero-mean noise term. That is, assume approximation error in our dane-gene model.
3. The noise term EOLS is Subgaussian, meaning that there exists dols 2 0 such that almost surely.
[((exp(nsous) S exp(1] / ER.
(17)
We assume that p is drawn from the Dirichlet distribution with a 1, since exponential gradient descent uses entropy regularization term that pushes proportions away from the edges of the simplex. We assume that there is no approximation eror in the linear mixing law, since we have established that it is of high fidelity.
28

---

Published as a conference paper at ICLR 2025
Now, formally define A Let E denote the covariance E [ppp" Let Σ be the sample-level estimate of the covariance, and let BB] express sample-level expectation on our size n dataset. Then, our ordinary least squares estimator is
= LeC L (
(18)
A, is an unbiased estimate of AT Moreover, for m, this estimator is equivalent to setting A, P Lat, in LEARNPARAMS.
L L.nt over the linear system of m equations, where P is defined
Setup for error from d. We can express A as the solution to a linear system of equations where the response variable computed over shortened runs. First, we abuse notation from LEARNPARAMS and define 3, ER" to be the response variable vector over our samples, where Biy L 1-11 (p*.)).
Then, we recall that A = 1--P-1B. Br, where satisfies Pij for izj and P.: = 1--++- E.
Next, define 83 ER", where 1 = L'ali L (1 over of a round. We have that
the loss difference in group when training on p²
705 = p-1 18g
(19)
03 I' - (pt VjE
(20)
The next condition characterizes the relationship between B and Bs, connecting the full round's loss difference the shortened loss difference.
Condition 2 (Linear + noise time model). There exists (0.1) and ex (00 such that
=c [iE[m].
(21)
Intuitively, this means that when shortening training run from one round to of round, the loss change scaled by a linear factor and then adjusted by some Gaussian noise. Eş captures how much variance there in the shortened training runs relative the full training runs. Note that as 8 approaches c6 approaches 1 and Es decreases in variance.
Setup for error from interleaving. We define AtAloli now. Recall that we have taken ó of the round and split into K equal intervals. During each interval, we train on some mixture pd Let the order of interleaved indices be captured by I E [mi K I, is the index of the mixture that is trained on in the 7th interval, and 7, =I=}] are the intervals for which the model is trained using ydJ I is constructed such that 17,1=k =k for all j [n]. Moreover, we assume that each set of indices for is selected with equal probability (i.e., Pr((T)=1/((A
Due to the dynamic nature of this training procedure, we abuse loss notation to capture what the model has already been trained
far. Let L (pl[-)) denote that the model up to interval was be the response variable vector over n samples, where each Big equal to (T - t-+5/K (fr=((Irr-))). Here, have annotated the losses to show that the loss differences are always computed before and after training an interval using pdd, even though the remaining intervals are trained on other mixtures. Using this definition of 8 we have
trained on [:]. Let 3 E IR
!ETE,, (I--((1- (r-1))6K
AAinli
p-1 33,
(22)
8/
1
++ 1)6/K (p. I[I L
4t76/K (p' (I.r))))))))
(23)
k
TET
29

---

Published as a conference paper at ICLR 2025
Moreover, write L in terms of interval-level loss differences:
Bij Lill ( (Liss, - k (ff + (Lial (+5/K (ph Viall
2426/K
(1/-01191))
(24)
+ vall
118/K (1/0 L (p' (rb)))
(25)
We make the following assumption relating Sij and Pij Condition 3 (Independent loss differences). Fix a j [nl. For any interval T TTy
(1)( (p' I I1:--2 - / (p' I1:--1) Ludl (+(r-1)6/1111 volii
(p-)
L vat,
1478/K
val.i
Intuitively, this means that the loss difference before and after train on pyd at interval 7 is independent of what we trained on before this point; this quantity only depends on pk1 and т.
Main result. We are now ready to present our main theorem that bounds || - j.rrAioli I1. We provide few extra definitions, well reminders of key terms.
Recall that m is the number of domains and n is the number of samples used to fit the training sweep (we n= m in practice). or is the parameter of the Dirichlet distribution for generating mixtures. dors is the Subgaussian parameter for the noise in the mixing law. is the proportion of the round allocated for leaming A'. of and сj are the noise and slope in the linear model, respectively, for how loss change over d of round is expressed in terms of loss change over the entire round. is the one-hot vector smoothing factor used LEARNPARAMS. K is the total number of intervals in 5, and & the number of intervals per mixture, where kn= K. Let a:y = max, L. (476/K (p'd) (p²') and let by = mmnn val, r (aj -bi--
/ (y'- (p")). Define
Theorem Suppose that Conditions 1, 2, and hold. Let satisfy t>> max llgm)) and assume
n 6(lggm++t))
m[ma+1]" There exists some p=> O such that the estimation error of LEARNPARAMS can be bounded
as follows:
||A -uA Mholi 1a Err Errz w.p. 21 11----mmppp----- exp( mt (-mp((1))) nmtt (26)
4
where there exists some finte e>> and
Errols= cm +m m/m//l++/ 21) m/a+ I
(27)
(1 (1-5)033m3/2 (1++t)
Errg=
(28)
(1 (1-E)c}
Errz= rK
nmmln(1((((1 k-1
(29)
Ca Ca(1 (1-e)
2k
We briefly discuss each of the three terms:
Emoos: this describes the enror incurred when learning A' finite number of samples. We see that there are several quantities controlling this error. As the number of samples increases, this error goes O. Similarly, if the linear dynamic mixing law is not very noisy, then Lo s small and the eттor will be negligible. Lastly, note the dependence on the key parameter dictating how sharp the Dirichlet distribution is.
Erry: in this tern, we can see that if 8 close to 1, the full round, then the error is negligible. Moreover, if the shortened rounds
noisy and of is low, then the error is also small. Lastly, observe the 1 -- in the denominator. Intuitively, this means that if the pf) ...pp' are spread out, things are easier leam.
30

---

Published as a conference paper at ICLR 2025
this tern is controlled by the relationship between k how many intervals you out of K to train on each phu, and K, the total number of intervals. If k K, then error is roughly constant. However, if & is small, the emror from interleaving be high because only using fraction of ő for estimating loss differences.
Emy:
D.5 PROOF OF THEOREM
We break down our proof into three lemmas that bound ||A A9112, 1|A v226112 and / (tiiiilli 2
D.5.1 OLS ERROR BOUND
First, restate Theorem 11 from Hsu et al. (2012), which apply obtain our bound on -| - Theorem (Hsu et al. (2012)). Denote
||||| to be Vx Σx. Assume that Condition holds. Let satisfy |X |X~1/"pla S Po. Assume that the number of samples satisfies
(02.6 logm). Let po satisfy 226gm((ogmm+)) Then, with probability at least -3e we have
vmm
2 1-6g
om + +2/mtt +2)
||A
A
(30)
7
where 6, =
4pfm((oomm+))
2ph (logm+++))
3n
To adapt this theorem to our setting, we must compute Po and 8g. and convert the bound into result on the 2-norm of the matrix difference.
Lemma 4 (OLS random design bound). Let sctissfy t> 6(logm++()
Assume Condition I holds, and the number of samples n satisfies n 2 mlma+1]' Then, with probability at least 1-3me we have
2 fasmmmm(m++2ym 2/ +
|| A
(31)
1-
where 55 5,=2(p
logm-+t nm (maa+1
logm-+t [am[mm+++]]]
Proof.
We first compute Po and 8g. From the triangle inequality and definition of the norm, have
that
E mm
IPlll,, Using Lemma and the fact that ||p||2 S 1, we can set Po v/m
mntma 11 Vm
Then, have that
(1)++1)
6, 6.=2(1 (
logm+t logm++
mm(ma+1) 3mm(ma+1))
(32)
To convert the bound in (30) into a bound on || - At- we note that ||A; A
|1(A
/ 2/2
|2|1
1/2
-x> Then, using Lemma 7, we
have
2 d m mm(m++1)
|| <
(33)
1-6,
n
Then, use || - n 1a ||'
A
Σm
A
A / -1/1122 Applying union
bound, we have
L (m m2
mtt - 21) ) ma+II
||
A
(*1122
<
w.p. 1-me
(34)
1 5
31

---

Published as a conference paper at ICLR 2025
D.5.2 ERROR FROM Ó BOUND
Lemma Define & [0,1] and assume that Condition holds. There exists such that
m
-v'
122
|s0/co//111--1100
m
(35)
Proof.
We first simplify the left hand side and plug in the construction from Condition 2:
|| 11
125
A
12
Σ 1 -3 1 112
18 -
(36)
Setting / equal to we have
m
- 112
12
Cs 2 vm
(37)
From Lemma 8, we have that |P~1|p ||2 under our construction in LEARNPARAMS T Moreover, note that [|ea]l is equal to (1=5)2oūx2. where x is the Chi-squared distribution with m degrees of freedom. Using Chernoff bound, we have that Pr[[ee|]? < (1 - (5))))m(((11 < - exp( 2 (1+ 147
for
some 1>00. 0. Putting everything together, we have
|| \ (1
(1-6)03mmmm(1++ w.p.1 - -epp ( (- 11 [+In 1+t
--
(38)
(1-c)e}
D.5.3 ERROR FROM INTERLEAVING BOUND
Lemma 6. Assume that Condition 3 holds. Let 0ij = max I I
T,46//K (phd)phh)) and let bij
1tr6/K (ptt)pp')) Define m bug. Then,
min, vol.i
|| - Miool ||
rK cs(1 e))
amin(1/t)-(( k-1
w.p. 2 1-mmm
(39)
2k
Proof.
We write | - (Ailll as
m
| -Je Su||P- 112 12,
1135
11
371
Σ
1.89
H v
3
(40)
1
We write 83 = =115 + 02j + ... + Økij where Ezij
L
11-11 (pt- (ptt) L L
Similarly, we write 3 = x1y + Jzzz
TKIZ, where 2tpp = L
(+(T-1)6// (
L 1475/K (I) From Condition 3, we have that Fzlz = Zrjz for any E ET,. This means that
83
1 Σ
Ttli
(41)
K
TET
32

---

Published as a conference paper at ICLR 2025
This gives us an elegant way of comparing
3 and ( in expectation, we have that IE [5 Then,
K
|| - (39 50 K, then we get
(Z "ズ - If we set / to be equal to =
14 3112 K2 n1
Σ xTイ )² TET;
11
(42)
Applying Theorem 5 and a union bound, this becomes
| 3 n
(1(1))) A )r
w.p. 10.21-1 21-nt nt
(43)
2k
where r the range of m Finally, we plug this into the original bound and use == o6 again.
and a union bound
rk
nmaln(1/0)-L
A-1)
|1: - (
w.p. 11-21-100 nmm.
(44)
c6(1 (e
2k
D.5.4 AUXILIARY LEMMAS
Theorem 4 (Horn & Johnson (1985)). matrix A E Rmxm has eigenvalues X, Am and x is the eigenvector such that Ax= Xx, for any V R"" A+ has eigenvalues X+ a...
Lemma 7. For p Dirichlet(a where a; a for all i E [m], the L2 norm of the covariance matrix -E [pp is
E m(ma+1)
(45)
Proof. First, the Dirichlet covariance matrix with constant ogy has the following form:
i
Tn (ma++)
(46)
m-
i=j
(ma++)
We can write this covariance matrix as sum of a rank one matrix and diagonal matrix:
Σ diag m(ma++))
m²(ma+11) 11
Note that the diagonal matrix has all eigenvalues equal to mimo+1]] and every vector in I is an eigenvector.
Next, we apply Theorem 2.4.10.1 from Horn & Johnson (1985), which restated in Theorem 4. We A == diag( mm ), -x=1= v= m=(ma++1)" 1. Then, the eigenvalues of Σ are mima++) + mm(ma+1)
0 (multiplicity 1) and m[ma+1) (multiplicity m m-1). 1). Therefore, we have that | m(mma+1)
Lemma 8, The matrix PE
PER mm is defined in LEARNPARAMS as:
i=j
P.
(47)
8
i+j*
where is one-hot correction factor. Then,
1-11
112
(48)
33

---

Published as a conference paper at ICLR 2025
Proof. First, observe that we write P as the sum of rank one and diagonal matrix.
P 111 diag(1 66)
(49)
m
The diagonal matrix has all eigenvalues equal to Next, we apply Theorem 2.4.10.1 from Horn & Johnson (1985), restated in Theorem 4, the sum of the diagonal and rank-one matrix. The eigenvalues are 1 1 ----+==11 (multiplicity 1) and 1 1-s (multiplicity n n-1). The cigenvalues of p-1 are then and and T-c the largest eigenvalue.
Theorem 5 (Serfling (1974)), Let 31 .I..KKKK a finite list of values. We draw a sample of size without replacement, each with equal probability denote this as XI.....N. Define minar and b= max, Then,
h
K
2kt
Pr
X
Σ
FK
28
Scxp
(50)
K
(1-
This also be written
K
ab)
S
w.p. 2 1-
(51)
k
K
2k
1=1
E EXPERIMENTAL DETAILS
E.
DATA
To obtain Yoon, 2023) in half.
shuffle and split the validation from SlimPajama-61 (Soboleva et al., 2023;
To perform training sweeps and emulate grid searches in static settings for m 3,7, we oversampled from the Dirichlet with a= by 4x the number of points and then hierarchically merged closest points into a centroid until we obtained x points. For example, obtain 10 points in the 7-dimensional simplex for SlimPajamaafull. we would sample 40 points in the simplex and hierarchically merge closest points until 10 points remain. This to ensure that near-duplicate p's are not included the sweep. This procedure used in Grid Search (GS) and DML. in Section 6 and in our analysis in Section 4
E.I.I TRAINING
Here, we discuss the training setups for the restricted and unrestricted settings. For the m =0: 2,3 settings, we train 160M model using Pythia- 160M's configuration for S = 5000 steps and results averaged over random seeds. For m m=7, we train a 160M model using Pythia-160M's configuration for S - 40000 steps results are averaged over random seeds. All settings use FlashAttention (Dao al., 2022), batch size of context size of 2048, and cosine learning rate decay from starting learning rate of 5e-5 to le-5 with 500 steps of learning rate warmup.
For the m 2,3 settings, experiments were run on a NVIDIA RTX 6000 Ada Generation GPU, For the m= setting, experiments were run on a NVIDIA 100 80 GB GPU.
Restricted versus unrestricted. Both the restricted and unrestricted settings share the same length of the
final training runs (5000 and 40000 steps, as above). The unrestricted setting gives methods up to 10 training runs to initialize mixing algorithm parameters, or 10S steps, while the restricted setting give 0.58 steps. See Table for training budget allocations in each setting. AIOLI and stratified sampling do not use extra training runs.
E.2 DATA MIXING METHODS
AIOLI-Specffic hype hyperparameters In the unrestricted setting, we found sometimes helpful to use an exponential moving average with proportion 7 over A' for AIOLI Formally, the standard p update rule in
34

---

Published as a conference paper at ICLR 2025
Table 9: Training budget allocations for restricted and unrestricted settings.
Setting
772
Method
Runs within training budget
Unrestricted
2
DML
10 runs, 5000 steps
Skill-it
2 runs, 5000 steps
DoReMi
2 runs, 5000 steps
DoGE
run, 5000 steps
3
DML
10 runs, 5000 steps 3 runs, 5000 steps
Skill-it
DoReMi
2 runs, 5000 steps
DoGE
run, 5000 steps
7
DML
10 runs, 40000 steps 7 runs, 40000 steps runs, 40000 steps 1 run, 40000 steps
Skill-it
DoReMi
DoGE
Restricted
2
DML
10 runs, 250 steps 2 runs, 1250 steps
Skill-it
DoReMi
2 runs, 1250 steps run, 2500 steps
DoGE
3
DML
10 runs, 250 steps 3 runs, 833 steps
Skill-it
DoReMi
2 runs, 1250 steps rum, 2500 steps
DoGE
7 DML
runs, 2000 steps 7 runs, 2814 steps runs, 10000 steps 1 run, 20000 steps
Skill-it
DoReMi
DoGE
Algorithm can be unrolled as pf+ ExpfexptnE==1""> which places equal weight on every Alj- To incorporate the EMA, we define A = A and ALm. = (I- We then use the update rule 1/51 x Alm) This allows AIOLI to gradually decay the contributions of A", such that the value of p' is less dependent earlier proportions in the training.
We summarize the hyperparameters used in AIOLI, providing their default values as well as guidelines for how to set them. Refer to Algorithm and 2 to see how they are used:
Number of rounds T: we set this to 20 in all experiments. Larger means more frequent updates to the mixture proportions.
Sweeps A: we set this to be 4 for m m22,33 2,3 and for the full SlimPajama experiments. We did not adjust this hyperparameter otherwise. Intuitively, larger & will give more accurate A', because this means that each yf* will be trained on more frequently throughout the 5 proportion of the round; however, this will also result in less of the round being allocated exploiting A'' via using p².
one-hot smoothing factor: we set this to be 0.75 in all experiments. In general, must be set between 0 and 1, where results in the training sweep using one-hot mixture proportions to leam A', which means that each batch only consists of one data group and can result in poor learning dynamics. E 1, on the other hand, means that our training sweep would only consist uniform proportions.
EGD step size 1f: we sweep (0.1, 0.2, 0.3, 0.5). with higher 7 resulting in greater magnitude of the proportion update.
Proportion of round 8 dedicated to learning A We use 0.128, 0. 0.288, 0.007 for m 2,3,7 respectively, Intuitively, a larger 8 will give more accurate A' because the parameter is leamedd on more data, but this will also result in less of the round being allocated to exploiting A' via using y².
EMA parameter 7: we sweep None, 0.1.0.5. Intuitively, None means that the p' update is equally dependent on all previous p' while a small y == 0 means that the p' update is only a function the current A
35

---

Published as a conference paper at ICLR 2025
For the last three hyperparameters, n,,d,, used different values of them in different experiments. Tables 10, 11, 12, 13, 14, and 15 list exact values for the unrestricted and restricted settings for m 100=2,3 2,3,7. In addition, Appendix F.3 provides results on hyperparameter sensitivity for 17.6, and y.
Table 10: Unrestricted hyperparameter values for each data mixing algorithm for experiments where (corresponding to Table 2 results).
Data groups
Hyperparameter
Value
arXiv/SE
proportion of round
0.128
EGD learning rate 7)
0.2
EMA parameter
0.1
GitHub/C4
proportion of round 8
0.128
EGD learning rate 7)
0.3
EMA parameter
0.5
Books/SE
proportion of round
0.128
EGD learning rate
0.1
EMA parameter ?
None
Table 11: Restricted hyperparameter values for each data mixing algorithm for experiments where m= (corresponding to Table 3 results).
Data groups
Hyperparameter
Value
arXiv/SE
proportion of round
0.128
EGD learning rate 7)
0.2
EMA parameter y
None
GitHub/C4
proportion of round
0.128
EGD learning rate r)
0.2
EMA parameter
None
Books/SE
proportion of round
0.128
EGD learning rate 7)
0.2
EMA parameter
None
Table 12: Unrestricted hyperparameter values for each data mixing algorithm for experiments where m==3 3 (corresponding to Table 2 results).
Data groups
Hyperparameter
Value
arXiv/Books/SE
proportion of round
0.288
EGD learning rate v
0.5
EMA parameter ?
None
ConmonCraw/CGithuh//ii
proportion of round
0.288
EGD learning rate y)
0.3
EMA parameter
0.5
Baseline hy hyperparameters. We consulted the original papers and implementations determine how to set the hyperparameters for each baseline, ensuring that the updated proportions were changing significantly not oscillating under these configurations.
Skill-It: the hyperparameters are the number of rounds 7, the EGD leamning rate 7), and the multiplicative weights window w. Our default configuration was 10, = and 3 However, we made two exceptions in the unrestricted setting after conducting a sweep over 7' (5,10) and e (0..0....0.,,0.8):: we used 7) =
for GitHub/C4, we used T = 5 and 7- 0011, and for Books/StackExchang
DoReMi: the hyperparameters are the EGD learning rate and a smoothing factor e (0 = no smoothing). For all experiments, set 7) 0001 and 1.-3.
36

---

Published as a conference paper at ICLR 2025
Table 13: Restricted hyperparameter values for each data mixing algorithm for experiments where m=3 (corresponding to Table 3 results).
Data groups
Hyperparameter
Value
arXiv/Books/SE
proportion of round 6
0.288
EGD learning rate ?)
0.2
EMA parameter Y
None
ConmonCraw/Gitttu/WWit
proportion of round
0.288
EGD learning rate 7)
0.2
EMA parameter ?
None
Table 14: Unrestricted hyperparameter values for each data mixing algorithm for experiments where
m= 7 (comesponding to Table 2 results).
Data groups
Hyperparameter
Value
SlimPajama, full
proportion of round 5
0.07
EGD learning rate r)
0.2
EMA parameter y
0.1
Table 15: Restricted hyperparameter values for each data mixing algorithm for experiments where m (corresponding to Table 3 results).
Data groups
Hyperparameter
Value
SlimPajama, full
proportion of round
0.07
EGD learning rate 7
0.2
EMA parameter y
0.1
the hyperparameters are the EGD learning rate 1). the smoothing factor e, and the proportion of the training batch that is allocated for the validation dataset r; this is needed to compute the gradient dot-product each step. We use = for all experiments. For m= we 0.25 and for m=3.77 use 0.011. For Github/C4,
DoGE:
set 0.5. For all experiments besides Github/C4 and SlimPajama, we use = and for SlimPajama we used 0.1 and r 0.03 for unrestricted and restricted settings, respectively.
Weight trajectories. In Table 16, we provide the mixture proportions for each method (averaged across training steps) for each dataset on one random seed. In Figure 8, we provide all of AIOLI's proportion trajectories throughout training in both the unrestricted and restricted settings on one random seed for the m settings. In Figure 9 and Figure 10, we provide AIOL.'S trajectories in the unrestricted and restricted settings on one random seed for Ariw/Books/Saackxchaangg and SonnmonrrwW/Gittu/WWeen Wikklediaa respectively. All of our trajectories demonstrate that AIOLI can significantly adjust proportions over time, and that conditioning on different initial proportions can drastically change the behavior AIOLI?
F ADDITIONAL EXPERIMENTS
F.1 DOWNSTREAM TASKS
We find that lower perplexity is positively correlated with worse performance on downstream tasks. We evaluated all models trained SlimPajama on ARC-Challenge, ARC-Fasy (Clark et al., 2018), BoolQ (Clark et al., 2019), HellaSwag (Zellers et al., 2019), LAMBADA (Paperno et al., 2016), OpenBookQA (Mihaylov al., 2018), PIQA Bisk al. (2019), and WinoGrande (Sakaguchi et al., 2020) using the
Note that for the restricted setting, AIOLI's trajectory consists of using the base method for certain amount of steps, and then roughly reverting to the unifomm distribution before adjusting the proportions. This is expected hehavior, since initial proportions pd uniform Algorithm this avoids "hiased" proportion update.
37

---

Published as a conference paper at ICLR 2025
Table 16: Average proportions over the entire training trajectory for the unrestricted setting, on one random seed.
Data groups
Method
Average Proportions
arXiv/SE
Grid search
[0.4, 0.6]
DML
[0.404, 0.596]
Skill-it
[0,437, 0.563]
DoReMi
[0.37, 0.63]
DoGE
[0.624, 0.376]
AIOLI
[0.507, 0.493]
GitHub/C4
Grid search
[0.3, 0.7]
DML
[0.46, 0.54]
Skill-it
[0.583, 0.417]
DoReMi
[0.858, 0.142]
DoGE
[0.352, 0.648]
ATOLI
[0.505, 0.495]
Books/SE
Grid search
[0.3, 0.7]
DML
[0.381, 0.619]
Skill-it
[0.316, 0.684]
DoReMi
[0.286, 0.714]
DoGE
[0,325, 0,675]
A10L1
[0.456, 0.544]
arXiv/BokssSEE
Grid search
[0.291, 0.306, 0.403]
DML
[0.245, 0.277, 0.477]
Skill-it
[0.292, 0.238, 0.469]
DoReMi
[0.318, 0.180, 0.502] ]
DoGE
[0.592, 0. 132, 0.276]
A10L1
[0.342, 0.275, 0.383]
CCIGitHibbWWii
Grid search
[0.291, 0.306, 0.403]
DML
[0.157, 0.472, 0.371]
Skill-it
[0.275, 0.3, 0.425]
DoReMi
[0.101, 0.714, 0.185] ]
DoGE
[0.536, 0.220, 0.244]
AIOLI
[0.342, 0.325, 0.333]
SlimPajama, full (WB/CAICCIG/SEAN)
Grid search
[0.202, 0.022, 0.28, 0.038, 0.018, 0.376, 0.064)
DML
[0.042, 0, O, 0.579, O, 0.249, 0.013]
Skill-it
[0.098, 0.111, 0.204, O. 103, 0.138, 0.266, 0.076)
DoReMi
[0.08, 0.047, 0.057, 0. 0.467, 0.078, 0.157]
DoGE
[0.056, 0.162, 0.343, 0.28, 0.038, 0.067, 0,051]
AIOLI
[0.142, 0.143, 0.143, 0.144, 0.140, 0.144, 0.143]
Language Model Evaluation Harness (Gao al., 2024) (Table 17). The correlation between perplexity and worse downstream performance. fact, DML obtains the best overall performance, even though it omits three out of seven datasets in SlimPajama (see the average proportions in Table 16).
the macroaverage of our downstream tasks is 0.529, indicating that lower perplexity is predictive of
One potential reason for this disparity the distribution shift between pre-training data and downstream evaluation data; for example, the DML results suggest that training on Books, C4, and Github is not needed to do well on the above selection of downstream tasks. Many recent works have also noted that perplexity and downstream performance are uncorrelated (Liu et al., 2023; Xia et al., 2023; Tay et al., 2023). Furthermore, Levy et al. (2024) proposes a question answering dataset where the perplexity the pretrained model is positively correlated with performance, similar to results. This mismatch between training objective and downstream evaluations also extends to post- post-training, where better learning of humaa preferences does translate better win-rate against other post-trained models (Chen et al., 2024).
38

---

Published as a conference paper at ICLR 2025
ArstviStackeechhng
Gutbahot44
0.32
0.50
0.50
0.65
20gpetho (book)
0.45
0.45
Ainli Dinostricted
0.38
AinlleES
Alcli-ENL EOLL
Abelit Sadi
0.35
Abelit Duliddi
Abili
1000
2000
3000
4000
5000
1000
2000
2000
4000
51000
1000
2000
2003
4000
5000
Atrp
steps
etepe
Figure 8: AIOLI'S proportions throughout training for both unrestricted and restricted settings Arxiv/StackExchang Github/C4, and Book/StackExchang.. These trajectories show that AIOLI meaning- fully alters the mixture proportions time.
Aãoli Unrestricted
Aioli+GS
Aioli+DML.
0.30
0.40
progortin
0.375
0.36
0.350
0.325
0.30
Amttv
0.32
0.300
StackExchange
0.275
1000
2000
2000
4000
5000
1000
2000
3000
4000
5000
1000
2000
2000
4000
5000
Aioli+ Skill-it
Ainli+ DoRMMi
Ainli-DoGE
0.400
0.50
0.50
0.45
argiattion
0.45
0.350
0.40
0.40
0.35
0.325
0.35
0.30
0.30
0.300
0.25
0.275
1000
2000
4000
5000
1000
2000
3000
4000
5000
1000
3000
4000
5000
steps
steps
steps
Figure 9: AIOLI'S proportions throughout training for both unrestricted and restricted settings on Hhtt/Book/StackExchaag These trajectories show that AIOLI meaningfully alters the mixture proportions over time.
Table 17: Downstream evaluation metrics for various data mixing methods after training on Slim Pajaaaama across three random seeds in the unrestricted setting.
Method
Avenge
ARC-C
ARC-E
BoolQ
HellaSwag
LAMBADA
OpenBookQA
PIQA
WinoGrande
Stratfiedd
0.305
0.176
0.314
0.394
0.261
0.116
0.117
0.563
0.499
AIoLI
0.311
0.172
0.315
0.447
0.264
0.114
0.111
0.559
0.504
0.124
DML
0.333
0.181
0.330
0.608
0.261
0.109
0.128
0.554
0,490
Skill-it
0.316
0.182
0.322
0.462
0.261
0.124
0.122
0.559
0.492
DoRMMi
0.324
0.177
0.323
0.507
0.264
0.127
0.122
0.574
0.499
DoGE
0.314
0.173
0.313
0.471
0.262
0.116
0.115
0.557
0.504
Resolving the disconnect between training objective and downstream evaluations is an area of active research. In the case of data mixing, AIOLI remains the only algorithm in tests that robustly minimizes average test perplexity-essenniallyy AIOLI achieves what it sets out to achieve in the LMO framework in (I). Conversely, other data mixing algorithms might be implicitly doing something else with respect to minimizing downstream evaluations. Considering how to incorporate downstream evaluations into data mixing is fruitful area for future work.
39

---

Published as a conference paper at ICLR 2025
Alobi Unrestricted
AIoli+GSS
Aloll+ DML
0.35
0.5
prggetoo
0.33
0.3
CommanCCawel
0.2
0.32
Gitbab Wikipedia
1000
2000
3000
4000
5000
1000
2000
3000
9000
5000
1000
2000
3000
4000
5000
Azolt +Sttill-st
Aioh+ DoReMi
Aboli+
0.40
0.45
0.40
0.35
gropertion
0.40
0.50
0.35
0.25
0.25
0.20
1000
2000
3000
0000
5000
2000
3000
4000
5000
1000
2000
2000
4000
5000
stege
steps
staps
Figure 10: AtoLI's proportions throughout training for both unrestricted and restricted settings Common- www.//GithubWWiiiipeeo These These trajectories show that AIOLI meaningfully alters the mixture proportions over time.
F.2 ABLATIONS
We ablate AIOLI by studying performance when two key properties of A' (Appendix C.2.1) are changed: when 7 (i.e., A only learned once at the beginning of training and used throughout), and when A is assumed be diagonal. We evaluate these ablations in the unrestricted setting presented in Section 6.1 and Table 2:
AIOLI-STATIC: We 7 in Algorithm That is, learn A the beginning of training. We use this A to set p', and use this p for the remainder of the training run. This approach tests if A' needs be adjusted throughout training.
AIOLI-DIAGONAL: We assume that each A' diagonal in this ablation. particular, in LEARN- PARAMS we do A, =Ba/p"* rather than - Br for each i E [m]] in line 11. This approach tests if is sufficient to not model cross-group interactions instead only capture how much group i's performance improves when trained on group i itself.
For both ATOLI-STATIC and AIOLL-DIAGOOAT, we use the same set hyperparameters as AIOLJ as described in Appendix E. For AIOLI-STATIC, additionally sweep over EGD learning rates (1,22,,3,,44) where 17 is the EGD learning rate used by AIOLI.
Our results are in Table 18. We find that AIOLI outperforms both ablations in 3 out of 6 settings, and obtains the lowest test perplexity on average over these settings. This suggests that both T>> and modeling off- off-diagonal entries are important to AIOLI'S consistent performance across datasets.
Table 18: Ablations on ATOLI. The table reports the difference in average test perplexity compared to stratified sampling. Negative values (green) improvement, and bolded = best performing method for given data setting. A=Arxiv, B=Books, GH-GitHub, SE=StackExchange W-Wikipedia. AIOLI outperforms ablations in 3 out settings and attains the lowest test perplexity on average.
Method
A/EE
GH/C4
B/SE
A/B/SE
CC/GH/W
SlimPajama
Average
Stratified
16.532
35.991
47.192
35.114
41.583
26.426
33.806
AIOLI
0.205
0.340
0.439
-0,226
-0.196
-0.240
-0.274
ATOLI-STATIC
-0.065
-0.333
-0.226
-0.117
0.092
-0.330
-0.140
ATOLI-DIAGONAL
-0.182
-0.178
-0.354
-0.246
0.215
-0.202
-0.230
40

---

Published as a conference paper at ICLR 2025
F.3 HYPERPARAMETER SENSITIVITY
We study how robust AIOLI is to changes in its hyperparameters. From the experimental details in Appendix E, the main hyperparameters that we modify are 77 (EGD step size), 6 (proportion of round allocated for learning A"), and (the EMA parameter). In Tables 19, 20, and 21. we report results on AIOLI in the unrestricted setting for ArxivSttacEExchanee and AriiwBooks/SuccXcchhngg 8/m= {0.064,0.8m 0.096,0000000 0....00111 1.1))). and y E (None, O. 0.1,..5.. We find that AIOLI still yields lower test perplexity than stratified sampling across all 17,6, and y we evaluated.
We sweep
E
Table 19:
The difference in average test perplexity of AIOLI with varying step size hyperparameter compared stratified sampling. Bolded result is the original number reported in Table 2.
Method
A/B
A/B/SE
Stratified
16.532
35.114
A1OLI (7==0.1)
-0.110
-02212
A10L. (7/=0.2) 0.2)
0..05
00.21
AIOLI (1)= 0.3)
-0.155
-0.186
AIOLI 0.5)
-0.166
-0.26
Table 20: The difference in average test perplexity of AIOLI with varying 8/m, the fraccton of each round for learning A', compared to stratified sampling. Bolded result is the original number reported Table 2.
Method
A/B
A/B/SIE
Stratified
16.532
35. 35.114
AIOLI (6 0.064)
0.205
-0.152
AIOLI (8 = 0.096) 0.2883
00.266
ATOLI (5/m= 0.1288)
-0.003
-02296
Table 21: The difference in average test perplexity of ATOLI with varying 7 the hyperparameter for computing p' with an exponential moving average, compared to stratified sampling. Bolded result is the original number reported in Table 2.
Method
A/B
A/B/SE
Stratified
16.532
35.114
AIOLI None))
-0.11
0.226
AIOLI (7=0.1)
-0.205
-0.185
AIOLI 0.5)
-0.141
-0.213
F.4 RESULTS ON LARGER MODELS
We examine if our findings-bott both in terms of the mixing law and in terms of AIOLI'S performance- hold larger models. We train -.raeearrmettr models. We learning of 3e-4 and keep all other training details the same. We use a subsample of our data settings, focusing on when we mix Arxiv/StacKExchange (m=2) and Prriv/Book/tuckExxhaap (m=3).
First, we measure if the log-linear static and linear-dynamic mixing laws are well-specified for 1.4B models. We use the same fitting procedure as described Section 4.1 and Appendix C.I. Figure 11 describes the fit of the static and dynamic mixing laws on exxiv/StckExxhhang The full results are in Table 22, which show that the average R² for the static and dynamic mixing laws for the 1.4B model are 0.989 and 0.929, respectively. This accuracy of the mixing law parameterization on the 1.4B model is prerequisite for ATOLI'S performance, which we evaluate next.
Second, we evaluate AIOLI in the unrestricted setting on the 1. 1.4B models. We compare AIOLI to stratified sampling and DoGE, Our results on three random seeds are Table 23. Similar to our results on the 160M models, we find that AIOLI outperforms stratified sampling in both data settings. Moreover, from Table 2, we see that DoGE originally performed worse than stratified sampling at the 160M scale. Our results here confirm that even at the 1.4B model scale, DoGE continues to underperform stratified sampling, Altogether, we see that AIOLI consistently outperforms stratified sampling while existing methods do not-at at both the 160M and 1.4B scale.
41

---

Published as a conference paper at ICLR 2025
Table 22: Comparison log-linear static and linear dynamic mixing law parameterizations when training a 1.4B 1.4B model.
Arxiv/SE MSE
Arxiv/Books/SE
Parameterization
R²
MSE
R²
Log-linear static
2c-4
0.995
le-3
0.984
Linear dynamic
7
0.916
2e-4
0,943
Log-linear static mixing law ArdtrSttackileebbngg 4B model)
Linear dynamic mixing law en Araw/StackExchanp (1.43 model)
10
2.
E
Propaction OE ATMY
Proportion ofsccleeccaeegee
Proportion of ar3oy
Proportion of dcstcckechaaggp
Figure 11: Left: log-lineat static mixing law fit on Arxiv/Stackexchange on 1.4B parameter model, in which each color represents different random seed. Right: linear dynamic mixing law fit on Arxiv\Stuckexchange on 1.4B parameter model on random seed. Each color a different initial mixture pd E P trained for 2000 steps, and the fitting sweeps are done over 100 additional steps.
Table 23:
Difference in average test perplexity compared to stratified sampling the unrestricted setting for 1.4B models. For ATOLI, we use 7 m0.55/mmmmm0066 0.1 for A/SE and 7/= 0.1 0.115/mm mm 0.096,7 0.5 for A/B/SE.
Method
A/SE
A/B/SE
Stratified
15.799
34.733
DoGE
0.551
0.922
AIOLI
-02276
-0.403
F.5 OUT-OF-DOMAIN SETTING
We consider the out out-off do-mann setting, in which the training data groups are disjoint from the groups that the model will be evaluated on. This practical scenario where we have access separate validation dataset that we wish our model to perform well on (Fan et al., 2024; Chen et al., 2023; Xia et al., 2024; Xie et al., 2023b; Engstrom et al., 2024). We will demonstrate how 1) the LMO framework can be adjusted to capture this setting, recovering the out-of-domain versions of Skill-It and DoGE proposed in their respective papers; 2) the linear mixing laws are still well-specified in this setting; and 3) AIOLI adjusted for this setting still more consistently outperform out-of-domant baselines.
LMO framework for OOD setting. Concretely, we suppose we have m training data groups such that Durin is still (D) (Dimaa mam and we have one separate out-of-domain data group that we do not train on; we have IID validation and test datasets for this ont-of-dmminn data group. Let L OOD the validation and test loss on the out-of-domain data group, respectively. Then, the LMO framework can be slightly modified:
minimizepe.a PEAT Tx 1 TT+ val, (p)
(52)
s.t. Lt+1 (10 o
Σ
AboD,PP)
Vt [T].
(53)
where Abon, E E is now vector representing how much each training group influences the validation group. There are two changes to the optimization problem: first, the objective is now to minimize the out-of-domain validation loss; second, the mixing law captures the relationship between the validation loss and the mixture proportions over the training data groups. Note that the DML method can still be applied in the OOD setting by directly minimizing c+ exp( Ex AooD.)))). More importantly, applying
42

---

Published as a conference paper at ICLR 2025
Lemma I to this optimization problem, get the update rule pj ~ p expppnoooo. expression recovers the Skill-It and DoGE OOD update rules, and can be incorporated into AIOLI as demonstrated in Algorithms and 4. These algorithms identical to AIOLI (Alg 1) and LEARNPARAMS (Alg 2), with the exception of lines and lines 3, 8, and 11 respectively, which reflect that 4boD is now matrix.
[m]. This
vector rather than
Algorithm AIOLI-OOD
1: Input: data Duain D model ft. Initial steps Sini. initial proportions p E Am. T rounds over
S Smit remaining steps, fraction per round for learning parameters, learning rate 7,, one-hot smoothing factor
2: If Smn 00, train f on p for Smi steps.
3: Set pm UUni((m).
4: for tert:..,TT 1-...,,Tdo
5: Set Abop. f(6 LEARNPARAMS- Draiin ()) (Alg, 4), and normalize A' to get A
6: plJ xp exp(nAoon,)) for all [m].
7: Train model po+ō with (1 6) steps from mixture y' over Duain- Obtain updated ft+1.
8: end for
Algo Algorithm (.EARRPARAMS-OOO
1: Input: D 8, model f², number of sweeps k, one-hot smoothing factor
2: Split the fraction of a training round 5 into K time segments, where K =nk.
3: Set -= R"
4: Define =1-==1-e)) (niif(m) for e [m], and define P=(p² .....p".. E Amxm
...-.. 5: Randomly shuffle instances of each ie [m]] to create an order I [m]]
6: for L= do
7: Le Train model on mixture p of Duran for one time segment, obtain fo+75/K
8: Update 3j + 33 + wwa.ooo(((t(r->))
( with loss difference on OOD
validation dataset.
9: end for Update 3+
x. Set 100D P-133 E. fe46
Return AboD =RT
Mixing law parameterization results. We study setting where our training data groups are Arxiv, Book, and Github from SlimPajama and our validation data group is StackExchange. Using the same setup as other m=3 3 settings in Section 4.2 (160M model, 5K steps, sweep over 9 runs), we measure the MSE and R2 the log-line linearnear static mixing law, Lval, ooD = =c bexp E 1 AooD.PP) and of the linear L (p) = c' + - AbODD, The MSE and R² for the log-linear static dynamic mixng law, L oOD
mixing law are 1.5x 10-3 and 0.964, respectively. The MSE and R for the linear dynamic mixing law are 1.1 x 10-1 and 0.796. The linear dynamic mixing law fits the true loss-proportion relationship less accurately than the log-linear static law. Nevertheless, both MSEs are low, and the R² still suggests that least 79% of the variability in validation loss can be explained by the mixing law.
AIOLI results. We evaluate stratified sampling, and OOD versions of AIOLI, Skill-It, DoGE, and DML in the unrestricted setting on 3 random seeds. train on Arxiv, Books, and Github and evaluate on StackExchange. Our results are in Table 24.
We find that all methods, including AIOLI, attain lower test perplexity than the stratified sampling baseline, which both the Skill-It and DoGE papers method that achieves this improvement without requiring additional training runs. This improvement over stratified sampling across OOD methods is expected, since stratified sampling include imelevant data
comparison point for the OOD setting. AIOLI is the only
43

---

Published as a conference paper at ICLR 2025
Table 24: Out-of-domain data evaluation, in which mix training data from Arxiv, Books, and Github and evalute on StackExchange data. The table reports the difference in average test perplexity compared to stratified sampling on the training data groups. For AIOLI, we use 008,6 m.006m None,
Method
Arxiv/Book/Gtthuh + StackFxchange
extra runs
Stratified
39.644
GS
-7.244
DML
66.16
Skill-It (OOD)
55.86
3
DoGE (OOD)
-7.626
ATOLI (OOD)
-4.028
due to the distribution shift between training and evaluation. On the other hand, stratified sampling is strong baseline in the in-distribution scenarios studied in the rest of this work.
G ADDITIONAL DISCUSSION
Hyperparameter Optimization and Truncation Bias. Many data mixing methods utilize extra training runs to learn static mixture proportions before the final training run. This allows us to view data mixing hyperparameter optimization problem in Ye et al. (2024) and Liu al. (2024) mitigate the inefficiency grid search in higher dimension by combining it with data mixing laws to impose additional structure. By fitting the loss surface, these methods obtain better sample complexity and performance than grid search alone. However, both grid search and these offline methods can have poor performance when searched for fitted on shorter runs, the restricted setting.
To understand these results, we look to the hyperparameter optimization literature. Many popular hy- perparameter optimization methods carefully control truncation, and some runs allowed continue longer than others (Li et al., 2018; Swersky et al., 2014; Domhan et al., 2015). Thus, generic hyperparam- eter optimization methods may eventually prove effective for tuning data mixes, but for future work we recommend going beyond fixed compute budget for all runs in a hyperparameter sweep, especially in compute-constrainet settings.
H
WHY THE METHOD IS CALLED AIOLI
An aioli is an emulsion, where individual components remain chemically separate from each other, despite being combined into one mixture. Similarly, our A' matrix is formed from separate test runs (the pl1 .....h.m Am Section 5), despite being combined into one update for p².
44


Reviews:
[{"review_id": "lll3pUO2rG", "reviewer": "Reviewer_zQiM", "summary": "This paper targets the challenge of optimizing data mixture groups in training large language models. The authors claims that no existing method consistently outperforms a simple baseline regarding average test perplexity. They propose a unified optimization framework that encompasses several recent methods, demonstrating that while these methods can represent the true loss-proportion relationship, they often fail to set accurate parameters, leading to sub-optimal performance. Leveraging this framework, the authors introduce a new method, AIOLI, which consistently surpasses baselines in both general and practical settings.", "questions": "1 How do you guarantee the \"LearnParams\" will yield accurate estimate given its parameters and potential for error?\n2 The paper employs \"gains in every data group\" as a metric to demonstrate Aioli's effectiveness. However, in practical scenarios, a model need not outperform its counterparts in every data group. How would you address situations where different weights are assigned to various data groups?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1 The unified framework integrates several recent methods about data mixture on large language model. This part is well-written and easy to follow.\n2 The problem it tries to tackle is of great importance. Related work propose different methods, and this paper try to unify them as a framework, which could advance future research in this topic.\n3 The experiments provided are all necessary to support the claims.", "weaknesses": "1 Formulas and letters and extensive details make some contents not easy to follow. For instance, in section 5 (Estimating $A^{t\\star}$), additional clarification on the core concepts prior to presenting numerous notations would enhance understanding. \n2 Experimental settings are insufficient. All experiments utilize a 160M model with a maximum of 50K steps. In Section 4.1, when m=7, 10 different proportions are not enough to find the optimal proportion.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "uNkA4Is5sl", "reviewer": "Reviewer_aN2q", "summary": "The research content of this paper is about how to optimize the data mixing of pre-training corpus. This paper proposes a framework for mixture proportions--LMO, which can incorporate existing methods (this paper lists one offline and three online methods) into this framework. The basic assumption (loss-proportion relationship) of this framework is verified through experiments, and the shortcomings of parameter estimation in existing methods (mainly three offline methods) are pointed out (disparity between estimated parameters and Optimal). This paper improved the parameter estimation. Experimental results have shown that the improved method proposed in this paper can enhance the performance on its test set across two situations.", "questions": "1.In Figure 2, the author used all the points for fitting. Did the author use the fitted mixing data law to predict a certain proportion setting?\n2. May I ask if the author has conducted experiments on larger models or have studied the argument that the conclusions drawn in small models also hold true in larger models?\n3. Would the author mind explaining in layman's terms the basic principles of the effectiveness of Algorithm 2 or the basic principles of the improvement brought by this method (line379-line387) (or whether this method is a universal approach);\n4. Besides PPL, may I ask if the author has tested the basic performance of the model downstream under different data mixing strategies?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "This article is well written and establishes a mathematical expression framework that unifies the data mixing methods mentioned in this paper. It reveals that the different data mixing methods mentioned in this paper are formally unified, with differences mainly concentrated in parameter estimation methods. At the same time, a targeted method for parameter estimation in the framework was proposed, which demonstrated effectiveness in practical experiments.", "weaknesses": "I am willing to acknowledge the contribution of the framework proposed in this paper, but the proposal of the framework focuses more on \"induction\" based on existing methods, which leads to a lack of persuasiveness in the innovation and practicality of the framework. Specifically, 1) the basic assumption of the framework is that the loss-proportion relationship is \"linear\" (or log linear). Although there has been experimental verification, the experiments are very empirical and the experimental setup is relatively simple, involving up to 7 data domains (real pre-training scenarios often contain domains of over 10000 levels). 2) The validation was only conducted at 160M (section 4.1), which is not convincing at all (many abilities cannot emerge on models of this magnitude), especially considering the extensive research on Scaling laws, small models can be used to predict larger models; 3) This paper only needs to consider the impact between up to 2 data domains, and characterizing the impact between more data is limited. However, this is more of a concern in real pre-training scenarios; 4) The parameter improvement method proposed in this paper is also biased towards empirical methods (line 378-line 389), and lacks persuasiveness in terms of generality and effectiveness on a wider scale。", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "ocbqi5DRX6", "reviewer": "Reviewer_HM3q", "summary": "This paper unifies existing methods for data mixing, offline and online, as different parameter estimations of mixing laws. From this perspective, the authors find existing methods to fall short because of inaccurate parameter estimations. With this insight, the author derives a new online data mixing method AIOLI, which updates data proportion according to dynamically fitted data mixing laws along the training process. Experiment results show AIOLI generally works well on 6 out of 6 datasets while other methods underperform some of them.", "questions": "AIOLI only performs one update step on each round (Line 5 in Algorithm 1). Why not iterate for a few steps to approach the optimal mixture proportion given by the estimated mixing laws?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well-organized and easy to follow, with extensive experimental details.\n2. This paper provides a unified framework for data mixing problems that explain the different of different methods, which is a significant conceptual contribution.\n3. The experiment results confirm the method generally works well, shows improvement on 6 out of 6 datasets,", "weaknesses": "1. Although AIOLI improves stratified sampling on all evaluated data, it can underperform baseline methods in some datasets (e.g., A/B/SE in Tab. 2).\n2. The authors propose the LearnParams Algorithm to estimate A* without sweeping, but its accuracy is unclear. It is also unclear whether the simulating process affects model performance.\n3. AIOLI involves a number of hyperparameters while a description of how to decide these hyperparameters is insufficient. This may hinder the application AIOLI.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "4DGNaBLkxQ", "reviewer": "Reviewer_9HhZ", "summary": "This paper unifies all existing data mixture optimization methods within one Linear Mixing Optimization (LMO) framework, and proposes AIOLI as a new online data mixing method. AIOLI iteratively train models by upweighting each training data domains and leverages the loss difference to learn the optimal data mixing coefficients. \nEmpirically, AIOLI is able to consistently outperform or be comparable to the stratified sampling baseline in the unrestricted setting, without extra proxy runs. Additionally, AIOLI can improve on the existing online data mixing methods by a significant margin.", "questions": "1. When the validation distribution is unknown or different from the training domains (e.g. train on the other 6 domains but eval on Books/Wiki/Github etc.)?\n\n2. Does the proposed method robust about the hyperparameter k, $\\delta$ (in Algorithm2)? To understand their impact, could you provide ablation studies or sensitivity analyses for these specific hyperparameters?\n\n3. The finalized domain weights from AIOLI are very close to uniform (Table16), how much the AIOLI framework could help with LLM pretraining above uniform-sampling baseline in practice? \n\n4. Could you show the data mixture weights trajectory from AIOLI? Does it learn from the curriculum like in DOREMI [1] and DOGE [2]  or it is always close to uniform?\n\n[1] DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining\n\n[2] DoGE: Domain Reweighting with Generalization Estimation", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper introduce a novel online data mixing algorithm AIOLI leveraging the EGD optimization framework, without introducing extra runs as previous methods.\n2. The paper also conclude all existing data mixing methodologies within an unified optimization framework, and conducted a comprehensive evaluation with the target of minimizing the average domain losses.  \n3. The proposed method AIOLI demonstrates remarkable robustness comparing to all the other baseline methods, which not requires extra runs with proxy models.", "weaknesses": "1. **The proposed algorithm is very similar to the Skill-it method [1]**, which weakens the algorithmetic novelty. \n- The `LearnParams` algorithm (alg. 2) in the AIOLI paper is very similar to the `LearnGraph` algorithm (Alg. 3 [1]) from the Skill-it paper, where the adjacency matrix A has been updated according to the validation loss difference.  Each \"training domain\" in AIOLI can be considered as each \"skill\" node in the graph, which corresponds to the *Continual pre-training* setting in [1], Table1.\n- The `AIOLI` algorithm (alg.1) is also similar to the `SKILL-IT Online Data Selection Algorithm` (alg. 1, [1]), which applies EGD on the data mixture weights. The author considers this as an unified online optimization framework for data mixture optimization problem.\n\n2. The paper only targets at minimizing the average losses but not consider the other evaluation cases (e.g. out-of-domain genearlization etc). How would AIOLI perform comparing to the other baselines when the validation distribution is unknown (e.g. the OoD setting in DOGE paper [2] and the Skill-it paper [1])? \n\n3. The finalized domain weights from AIOLI are very close to uniform (Table16), which may explain the consistently comparable-to-baseline results in Fig. 3.  Since the uniform sampling across difference domains would be the first-choice for practitioners for minimizing the average domain loss, how much the AIOLI framework could help with LLM pretraining above uniform-sampling baseline in practice?\n\n[1] Skill-it! A Data-Driven Skills Framework for Understanding and Training Language Models\n\n[2] DoGE: Domain Reweighting with Generalization Estimation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)