Adaptive Data Optimization: Dynamic Sample Selection with Scaling Laws
Authors: Yiding Jiang
Allan Zhou
Zhili Feng
Sadhika Malladi
J Zico Kolter
Venue: ICLR 2025 Poster (2025)
Keywords: Data selection
Pre-training
Curriculum learning
Language Models
Scaling laws
Primary area: foundation or frontier models, including LLMs
TL;DR: The paper presents a scalable and proxy-free algorithm that dynamically optimizes data distributions during model training.

Abstract: The composition of pretraining data is a key determinant of foundation models' performance, but there is no standard guideline for allocating a limited computational budget across different data sources. Most current approaches either rely on extensive experiments with smaller models or dynamic data adjustments that also require proxy models, both of which significantly increase the workflow complexity and computational overhead. In this paper, we introduce Adaptive Data Optimization (ADO), an algorithm that optimizes data distributions in an online fashion, concurrent with model training. Unlike existing techniques, ADO does not require external knowledge, proxy models, or modifications to the model update. Instead, ADO uses per-domain scaling laws to estimate the learning potential of each domain during training and adjusts the data mixture accordingly, making it more scalable and easier to integrate. Experiments demonstrate that ADO can achieve comparable or better performance than prior methods while maintaining computational efficiency across different computation scales, offering a practical solution for dynamically adjusting data distribution without sacrificing flexibility or increasing costs. Beyond its practical benefits, ADO also provides a new perspective on data collection strategies via scaling laws.

Published as a conference paper at ICLR 2025
ADAPTIVE DATA OPTIMIZATION:
DYNAMIC SAMPLE SELECTION WITH SCALING LAWS
Yiding Jiang Allan Zhou Zhili Feng
Sadhika Malladi
J. Zico Kolter
Carnegie Mellon University! Stanford Universityf Princeton University yidingj jics. cmu. edu,
a sta.ssrd..ffood.o....
ABSTRACT
The composition of pretraining data is a key determinant of foundation models' performance, but there no standard guideline for allocating a limited computa- tional budget across different data sources. Most current approaches either rely extensive experiments with smaller models or dynamic data adjustments that also require proxy models, both of which significantly increase the workflow plexity and computational overhead. In this paper, we introduce Adaptive Data Optimization (ADO), algorithm that optimizes data distributions in online fashion, concurrent with model training. Unlike existing techniques, ADO does not require external knowledge, proxy models, modifications the model up- date. Instead, ADO uses per-domain scaling laws to estimate the learning potential of each domain during training and adjusts the data mixture accordingly, making it more scalable and easier to integrate. Experiments demonstrate that ADO can achieve comparable or better performance than prior methods while maintaining computational efficiency across different computation scales, offering a practical solution for dynamically adjusting data distribution without sacrificing ficxibil- ity or increasing costs. Beyond its practical benefits, ADO also provides new perspective on data collection strategies via scaling laws.
INTRODUCTION
Foundation models, large neural networks pre-trained vast amounts of data, are the backbone for a wide range of today's machine learning workload (Bommasani et al., 2021). It is well es- tablished that the composition of the pretraining data plays crucial role in these models' final performance (Gadre et al., 2024); however, there are currently no standard guidelines for selecting good pretraining data. Most existing datasets are filtered and curated snapshots of the web data (Gao et al., 2020; Penedo et al., 2024), often categorized into distinct heterogeneous domains (e.g.. Com- monCrawl and Github). Even after filtering, a challenge remains: deciding how allocate compu- tational resources across these different domains. Given the ever increasing cost of pretraining large foundation models (Sharir et al., 2020), optimizing the data composition potential avenue for improving performance without increasing computational cost.
There are two popular approaches to adjusting the data distribution. straightforward approach to this problem is to experiment with smaller models, test different data policies, and then apply the best-performing policy to the larger model (Ye et al., 2024; Ge et al., 2024). However, this strategy has several drawbacks. First, the number of domains increases, the cost of training all possible policies can grow linearly or even exponentially (if the level of discretization remains the same), making even small-scale experiments costly. Second, the optimal data policy for smaller model does not necessarily generalize to larger models (Ye et al., 2024; Albalak et al., 2023). Another ap- proach involves dynamically adjusting the data policy during training based on various statistics (Xie et al., 2024; Chen et al., 2024; Qin et al., 2024; Fan et al., 2023). However, these methods often re- quire training smaller proxy model to guide the data policy, whose cost is still nog-neglig
A significant limitation of almost all existing methods (with the exception of Albalak et al. (2023)) is their requirement for additional computational steps beyond standard training. This extra complex- ity entails substantial modifications to the training pipeline, making it challenging to integrate these
Equul contribution.

---

Published as a conference paper at ICLR 2025
Multi-staged data selection proxy models
Adaptive Optimization (ADO)
Data selection
Itl Optimized data distribution
ADO update
Train one more small proxy models
Train the full model
Data distribution at time t
Train the full model for one step
Data distribution at time t+1
Figure 1: ADO is cheap online technique for adjusting the data distribution while training large models. In contrast to prior methods, ADO tailors its data distribution to the model as it is being trained, and does not require training smaller proxy models in advance.
methods when other components are varied (e.g., the architecture, tokenizer, or optimizer). The con- sequence of this usage barrier is that these methods have not been tested widely, making them less popular for those training new model. To enable broader adoption of data mixture optimization, we argue that must be implemented in an online fashion, concurrent with model training, without disrupting the standard training process, which means that the data mixture must adapt itself based on the feedback from the model. This closely related to curriculum learning (Bengio al., 2009), a training strategy where models are progressively exposed to domains in curated order.
At a fundamental level, this work develops and empirically investigates questions about the existence of good pretraining curricula and the feasibility of cheaply identifying them. We first demonstrate in controlled experiments that good curricula can generally be found with more computation and that it is hard to accurately predict larger models' performance with only few small models. This leads to our main contribution, Adaptive Data Optimization (ADO): an algorithm for adaptively adjusting the data mixture online during training (Figure 1). In experiments on language models up to 1.3B parameters trained on the Pile (Gao et al., 2020). we find that ADO improves performance across a variety of common benchmarks and improves validation loss SlimPajama (Soboleva al., 2023) and Fine Web (Penedo et al., 2023), both of which are considered to be higher-quality datasets. Most importantly, ADO achieves these without requiring significant additional computation (less than 0.4% wallclock time for 1.3B), proxy models, or extensive modification to the training pipelines.
MOTIVATIONS
The drawbacks of small proxy models. At first glance, transferring data selection strategies from smaller larger models seems plausible, with positive results in specific cases, such as in Mindermann et al. (2022) and Xia et (2024b). However, as shown Figure 2, using small number of models to predict the behavior of larger models can exhibit high variance on the Pythia model family (Biderman et al., 2023), indicating that they cannot reliably forecast larger models' performance. At a per- domain level (Figure 7), the variation becomes pronounced. While this does not fully rule out the pos- sibility of transfer, it raises doubts about whether small models can reliably select data mixtures for larger mod- els. Another issue that the efficacy of distributions gen- crated by proxy models is highly sensitive to factors like the tokenizer and other hyperparameters (Albalak et al., 2023). Our findings also confirm this brittleness (Sec- tion 4). This means that the optimized distributions from one experiment might not transfer, and they need to be retrained for a new model. This is undesirable because proxy models add aoiviaivial overhead - Ge et al. (2024)
Extrapolating Pile validation loss
Validation loss
2.20.
All
more
4100
1.48
10=
10
Params
Figure 2: Extrapolating larger models' loss using scaling laws up to a threshold model size, indicated by color. Blue dots corre- spond to different models. Extrapolations become more accurate with more models.
estimated that a single round of DoReMi proxy training requires 760 GPU hours on 8 Nvidia A100 GPUs, and several proxy training rounds needed before training the full model.
2

---

Published as a conference paper at ICLR 2025
1.0
Perceptron
Multi-layer perceptron
460
Meta loss: MLP
oetta-optiiized order
meta optimiedd ander ran domn order
0.9
random order
460
0.8
440
0.7
Tost lod5
420
400
0.4
380
260
0.3
340
0.2
0
100
200
300
400
500
0.2
0
200
4.0
600
800
1000
200
400
600
800
1000
Steps
Steps
Steps
Figure 3: The average learning curves of SGD and the meta-optimized curricula for logistic regression (Left) and an MLP (Middle). The meta-optimized curricula consistently outperform SGD. Further, the meta loss for curricula is still decreasing after 1000 steps (Right). The first two figures are averaged over 50 runs and the shaded region corresponds to I standard deviation.
The existence of good curricula. Curriculum learning (Bengio et al., 2009) is a conceptually appealing idea that has been regularly studied over the years, but these methods have not been widely adopted in deep learning outside of few relatively niehe areas (Graves al., 2017; Florensa al., 2017: Jesson et al., 2018; Tang et al., 2018). To the best of our knowledge, there is no consensus on why curriculum learning has not had wide success in deep learning. Some hypothesize that deep learning is different from human learning or that overparameterization makes curriculum less effective (Wu et al., 2021; Mannelli et al., 2024). On the other hand, some experimental evidence suggests that deep neural networks implicitly learn functions of increasing complexity from easy to hard (Kalimeris et al., 2019; Baldock et al., 2021), 1). loosely resembling some aspects of human learning (Lawrence, 1952; Baker & Osgood, 1954; Skinner, 1958). For overparameterization,, it is not clear if modern foundation model pretraining is in the overparameterized regime since dataset size generally scales with the number of model parameters (Hoffmann et al., 2022).
To motivate our later work, we propose an alternative hypothesis: good curricula exist but they are computationally difficult to find. Let's start with an instructive experiment of finding good curricula with more computation. Suppose we have a fixed dataset of size N and want to train SGD at a mini- batch size of for epoch (i.c., each data point is used once, which is common for language model pretraining). We are interested in finding data ordering that is better than a random shuffle when training a model initialized at random. Brute force search intractable even for a small dataset. Unlike most existing curricula that rely on some prior notions of difficulty, we will find good cur- riculum with only training data by optimizing over the space of permutations with meta learning. Due to space constraints, we defer the details of this meta learning procedure to Appendix B.
We apply this procedure to a logistic regression problem as well as a multi-layer perceptron where the goal is to imitate a teacher. We randomly sample the initialization lo. ground truth 0* and the data 2 = {(<,, f(z;; )) For fixed set of data points, run the meta optimization for some steps (200 for logistic regression and 1000 for MLP) and evaluate the performance of learned order- ing on new initialization. We repeat this experiment 50 times and report the results in Figure 3. As can be seen from the results, the performance of meta-optimized data ordering is significantly better than that of random order even though during evaluation the models are initialized at random. In concurrent work, Gu et al. (2024) conducted a similar experiment but their meta objective uses additional validation data and optimizes a distribution over all training data points rather than SGD.
This shows that given a dataset, it is possible to find an ordering that greatly outperforms random ordering without committing to any prior notion of difficulty. This procedure does not leverage ad- ditional data, but it requires significantly more computational resources than random ordering, so is neither practical nor scalable. Moreover, given our limited understanding of the landscape of curriculum optimization, it is highly unlikely that we are identifying the globally optimal curricu- lum. As shown in Figure 3 (right), even this meta-oppiimizatin has not fully converged after 1000 steps. Nonetheless, it demonstrates definitively the existence of good curriculum, which raises the question of how we can find a better curriculum in a computationally efficient m
Since static mixture is nothing but fixed stationary curriculum throughout the training, this computa- tional difficulty applies to finding static mixture too.
3

---

Published as a conference paper at ICLR 2025
3 ADAPTIVE DATA OPTIMIZATION
We demonstrated above a general, albeit costly, strategy for finding good curricula but is too ex- pensive to be practically useful. We believe that an ideal online data selection strategy for pretraining foundation models should incur a minimal computational cost.
Furthermore, computational efficiency alone is not sufficient. Perhaps an equally important con- sideration is that such methods should avoid explicit dependence on any particular specification of downstream tasks. The fundamental premise of foundation model is serve as the basis for all reasonable downstream tasks; explicitly relying on fixed set of downstream tasks to pick pre- training data distribution risks unintentionally overfitting to the downstream tasks (Goodhart, 1985). Indeed, the evaluation of current language models has become increasingly difficult as performance on existing benchmarks becomes saturated. While evaluating downstream tasks is undoubtedly im- portant because they are representative of the models' main use cases, we believe that it is valuable to define an objective for online data selection that is agnostic to the downstream tasks.
Previous works like DoReMi (Xie et al., 2024), and DoGE (Fan et al., 2023) share this task tank-agnostic philosophy of data selection. However, these methods require training proxy models which makes them inconvenient and expensive to use. Lifting the requirement for proxy model would make data selection methods more accessible and practical.
The goal of this work is advocate for data selection methods that satisfy the following desiderata?:
Adaptive Data Optimization
Does not leverage extra information such additional data or existing models.
(ii) Does not depend on multi-staged training with different proxy models.
(iii) Does not require significant computational resources and can be run online.
Based on these criteria, we introduce our method of online data selection, Adaptive Data Optimiza- tion (ADO). In the following sections, we present the algorithmic components of ADO. First, we discuss how fit scaling law for each domain, allowing the data mixture account for variations in the natural diversity across different domains. Next, we describe how the data mixture is dynam- ically adjusted based on these scaling laws during model training. Each domain is characterized by two key quantities: 1. the domain's learning potential, indicating the potential improvement from further optimization in that domain (Section 3.1), and 2. credit assignment score that quantifies the domain's contribution to the reduction of training loss (Section 3.2). To address potential noise from simultaneously updating the data mixture and the model, we employ several time-avee averaging techniques to smooth these quantities (Section 3.3).
3.1 CONSTRUCTION NEURAL SCALING LAWS FOR FACH DOMAIN
To achieve our desiderata, it is crucial to be able to predict the evolution of our model's training online with cheap computational routines that do not scale with the size of the model. To accomplish this, we use neural scaling laws (Kaplan et al., 2020) to extrapolate the loss as a function of the amount of training data observed far, which we denote n. It is important to highlight that this scaling law is for predicting the future training loss within a single training run. In contrast, typical neural scaling laws (Hoffmann et al., 2022) use the results of completed training runs to extrapolate the loss of future training runs, typically with larger models or datasets.
For simplicity and interpretability, we use a standard power law: a(,; (,n) = Bn-", where 0, é are the parameters to be fitted, and L is surrogate loss used model the behavior of the true loss. The parameter è can be seen as the irreducible loss or the "entropy" of the data. Data with high diversity (c.g., CommonCrawl) would naturally have higher irreducible loss and data that are more predictable (c.g., code and math) would tend to have a lower irreducible loss. The parameter a measures how quickly the loss decreases with more training data. Although such scaling laws may be inaccurate at extrapolating training loss far into the future (e.g., due to the learning rate schedule),
Thhese criteria not mutually exclusive with proxy models or other offline data selection methods and be used in tandem if needed (e.g., training model that specializes in coding).
4

---

Published as a conference paper at ICLR 2025
can cheaply refit the parameters on the fly (Appendix A.1). Since the parameters from each fit are only used for a short period before being updated and the learning curves usually do not change abruptly, our scaling law need only be locally accurate.
Instead of fitting a single scaling law for the overall loss (Kaplan et al., 2020; Hoffmann et al., 2022), will fit a separate scaling law for the loss each domain, each of which will guide ADO to adjust the weights of each domain.
Definition 1. A domain scaling law is a training sample scaling law, [x(m) - Z(n; on. 3k, = EN + 1.B that predicts the kthh domain's 's training loss after training on n samples. Eky Bu, Ork are functions of the architecture, training algorithms, and the overall data distribution.
The domain scaling law serves tractable middle ground between modeling the overall loss and modeling how every domain interacts with each other. Modeling the full dependencies faithfully may be challenging while training online because we cannot modify the data distribution drastically. As the first step towards this goal, our method will only model how each domain interacts with itself and leave more complex modeling to future works. Nonetheless, this scaling law reveals important information about each domain. For example, EY corresponds to the estimated irreducible loss of a domain, which could be useful for various applications (Xia al., 2024a). More importantly, we can interpret the derivative of the loss with respect to the number of samples as the following:
dCx(n) -0k32n dn
Ork
(Éx(n) Ex)
(1)
n
Learning speed
Reducible loss
The quantity (n) - EK informs about how far away a particular domain is from the estimated minimum loss, which coincides with the notion of population-level reducible loss (Mindermann al., 2022), and aę indicates how fast the loss of a particular domain is changing with new data. This derivative informs us about how much loss decrease we can expect per data point locally. At high level, the cross entropy loss has an information theoretic interpretation, so this quantity can be understood as a form of information density, or, the amount of information gain per unit of computation. Intuitively, prioritizing data with high information density is a likely better use of computation and could lead to richer representation as it implies more information is "stored in" the model weights. We leave the formalization of this statement to future works. may be useful to contrast this quantity with Albalak et al. (2023, ODM) which only prioritizes domains with a high loss (n) without considering the irreducible loss. This objective could lead to the underrepreseenatioon of domains with inherently lower entropy such as code math?.
3.2 THE CONTRIBUTION OF THE DATA FROM FACH DOMAIN
In contrast to prior works that train multiple models to find optimal data policies offline (Ye et al., 2024: Ge et al., 2024; Liu et al., 2024), our online approach adaptively prioritizes selecting data so that the model learns quickly, as predicted by scaling laws that we fit on the fly. It consists of a data policy (t) E A that specifies sampling distribution over the K domains. An intuitive approach would be to prioritize sampling domains where the model will learn quickly: e.g., we could sample
from domain & in proportion dLz (Equation 1).
dr
However, independent domain scaling laws do not account for how samples from one domain help learning on different domain. Consider thought experiment where the loss La is decreasing rapidly, but we have sampled very little data from domain k. Intuitively, data from domain & should not get much eredit for this loss decrease, which can instead be attributed to transfer learning from other domains. This calls for some form of credit assignment based on whether a domain has actually been sampled recently.
Definition 2. A credit assignment score, (t), is a real positive-valued function that indicates how much data from the 1tth domain contributed to recent changes in the loss Lx-
Ideally, we would like to understand how each domain affects every other domain. Unfortunately, this interaction might be highly complex and not local (i.e., the best solution locally given time
Alhalak et al. (2023) showed that ODM achieves higher loss Github than other methods. In our 124M experiments, ODM achieves 1.54 validation loss on Github while ADO achieves 1.40 validation loss.
5

---

Published as a conference paper at ICLR 2025
Algorithm 1 Adaptive Data Optimization (ADO)
1: Input: prior ja E A update interval /update, warmup duration /warmups 71, 72, 8, 8min
2: Initialize h < M4, f- μs, {1 oss_k < [1]
3: Train with μl for twamnup steps to initialize each losa.k
4: Initialize domain k's scaling law with loss_k, for k E [K]
5: for = 0 + ''dd
> The training loop
6:
p(t) compute the preference distribution according to Equation
7:
T(t) < clip (12 p(t) + (1 (1-7)) 7t t 1), 5min)
8:
{Ls ]E[K]]]] train the model according to ा and add domain & loss to loss_k
9:
h(t) (t) + 1- (t- -1
10:
#(L) < (() 1- ((
11:
mod Lupdate = 0 then
12:
Update domain & scaling law with loss_k, for & [K]
step may not be optimal in the long run). As mentioned earlier, we focus on modeling the domain's 's contribution itself for now. A reasonable assumption that the contribution domain makes to its own learning progress is positively correlated with what proportion of the recent data came from this domain. other words, we assume that if a domain was useful for itself in recent history, it will likely continue to be useful. Based on this assumption, we keep history, h(t), of the exponential moving average (EMA) of the recent data policy (with a coefficient of 71 < 1) and apply a power transformation with smoothing parameter < 1 to smooth the distribution since the assignment can be inaccurate:
(t) (t) + (1 (t 1), Ax(t) hx(t)).
(2)
Here a smaller ?1 puts more weighting on policy history in the EMA, and decreasing s leads to stronger smoothing of the credit assignment across tasks. We find that 7 0.1 and 0.5 work well consistently for all scales we tested. We emphasize that this design choice is simple heuristic that can be easily computed and other choices are possible.
3.3 CONSTRUCTING THE DATA POLICY UPDATE
To define our data policy, we combine the two principles we have seen thus far: prioritize sampling from domain & if the model decreasing its loss La quickly, but only that decrease can be attributed to data from domain k. Additionally, it is common to have natural prior distribution ju E K over the domains (e.g., number of tokens each domain), which we may also incorporate:
Definition 3. Given the learning speed forecast by scaling law, 8n Lx(n), credit assignment score (1), and prior domain weight pis, the preference distribution
① (a(1) -ik On
1
(n)
(n ul Au(t) (Ze(n) Ex).
(3)
Note that the preference distribution is not a perfect indication of the best possible distribution to learn from because we have access the gradient of the data policy. As shown in the definition of the credit assignment score, we have to resort to various approximations.
Temporal average. This local estimate of data policy may not be optimal globally. Both DoReMi (Xie et al., 2024) and DoGE (Fan et al., 2023) found superior performance by averag- ing the proxy model's training distribution over all time steps, to obtain a stationary distribution the actual training run. Since we need to select data policy online, we instead take inspiration from Defazio al. (2024) and use an online analogy to post-hoc averaging:
(1) 41 (1- (
(4)
(t) 122 Pk (t) (1- 72)
(5)
Here R (t) is a temporal moving average the preference policy at every step and Ty (t) is linear combination of the moving average and the chosen policy according to which the data are sampled

---

Published as a conference paper at ICLR 2025
Table Zero-shot performance of different methods across different downstream evaluations. High- lighted cells indicate top-performing methods (deeper color indicates better performance).
HellaSwag
WinoGrandde
PIQA
ARC-E
SeiQ
LegiQA2
LAMBADA
Average
124M-Pile
0.279
0.515
0.615
0.435
0.747
0.2228
0.357
0.454
124M-DoReM1
0.279
0.520
0.609
0.429
0.741
0.237
0.368
0.455
1248--DD
0.285
0,514
0.603
0,453
0.764
0.230
0.374
0,461
0.280
0,511
0.610
0.447
0,762
0.227
0.362
0.457
:249-Natural
0.290
0.503
0.624
0.435
0.755
0.234
0.401
0.463
1248--DO
0.290
0.520
0.635
0.456
0.771
0.244
0.371
0,470
3.33-DDRMMi
0.416
0.582
0.707
0.609
0.870
0.228
0.615
0.575
1.38-Ballanee
0.382
0.546
0.689
0.580
0.862
0.225
0.613
0.557
2. ..38-Natural
0.424
0.584
0.718
0.627
0.886
0.232
0.624
0.585
1.38-ADO
0.442
0.590
0.730
0.625
0.875
0.228
0.638
0.590
from, which combines both the current preference distribution and the time-averaged distribution. Using the coefficient of ensures that all the preference distributions from all time steps contribute equally to the final average, similar to the effect of post-hoc averaging. Intuitively, larger value of 72 makes the algorithm more responsive to the current loss dynamics, while smaller value makes smoother and less responsive. We find that 722 0.1 works well consistently. Additionally, we enforce the smallest domain probability to be no smaller than ómi = 0.01 to avoid domains being unsampled (Appendix C). Algorithm shows the pseudocode for the entire algorithm.
EXPERIMENTS
Evaluation. We conduct all our experiments on the Pile dataset (Gao et al., 2020) with decoder- only transformer language models (Vaswani et al., 2017) of varying sizes. We consider two types of metries: 1. validation loss on the Pile, an unweighted version of the pile validation set where each domain receives equal probability, SlimPajama (Soboleva al., 2023), and billion token subset of Win (Penedo et al., 2024), 2. zero-shot downstream performance on common-sense reasoning domains from the language model evaluation harness (Gao et al., 2024): HellaSwag (Zellers al., 2019), WinoGrande (Sakaguchi et al., 2019), PIQA (Bisk et al., 2019), ARC-E (Clark et al., 2018), SciQ (Welbl al., 2017), LogiQA2 (Liu et al., 2023) and LAMBADA (Paperno al., 2016).
Methods. We conduct experiments 2 computation scales: 124 million parameters (124M) and 1.3 billion parameters (1 3B). For 1 4M models, we use a batch size of 256 with context length of batch size
1024 trained for 60,000 steps (approximately 15 billion tokens), and for 3B, of 2048 with a context length 1024 trained for 60000 steps (approximately 125 billion tokens). Both scales use the natural distribution for f and the same hyperparameters for ADO. Training details can be found in Appendix A.
Comparisons. For comparison, we use the original Pile weights (Gao et al., 2020, Pile) and DoReMi weights (Xie et al., 2024, DoReMi)* as the baselines. We also consider ODM (Albalak et al., 2023), another online method. addition, we introduce a simple baseline that is surprisingly missing from the literature: weighing each domain by the number of tokens in it. This mixture natu- rally takes into account the tokenizer being used and training on this policy essentially corresponds training the natural distribution of the dataset (i.e., corresponds empirical risk minimiza- tion where train, validation, and test sets all follow the same distribution defined by the tokenizer). We estimate the tokens per domain by randomly sampling 1000 documents from each domain to estimate each domain's 's average tokens per document, and then multiply it by the total number of documents (Appendix A). We will refer to this policy ural We also use a second bascline, Balanced, which considers the unweighted mixture of all the domains.
The results for downstream performance are shown in Table 3. We will highlight some key observations here. Observation 1: 1.24 scale, ADO achieves the highest average accuracy; also outperforms all the baselines on 6 out of 7 downstream tasks and competitive on
Observations.
"Since DoReMi is sensitive to the choice of tokenizer, used the DoReMi weights from Albalak al. (2023) who the same tokenizer as do.

---

Published as a conference paper at ICLR 2025
Perplexity on different domains
16
Models
1.38-Natrral
14
1.39 DoRaM
12-
Validation perplexity
Unweightad-Pii Pite
Pite
SimPajanaa
Finewies
BookCorpuc
Arriv
DM Methamatics
Book&
Enron Emaill
EuroPari
Freelaw
Gutenbery (PG-
Github
HackerNoors
NiW ExPorter
OpenSubittes
Open f2
PhiPapees
RubMed Abstracis
Pit--CC
PubMed Contral
StackE
USPTO Back Background grounnss
Uhuntu IRC
Weiperil (on)
SuutttaSuudtter Sudtittess
Figure 4: Perplexity of 1.3B models on different domains, trained on the Pile using either DoReMi, Natural, or ADO. On the left side of the red line, we have the validation perplexity the Pile valida- tion set, the unweighted Pile validation set, the SlimPajama validation sct, and random subset of Fine On the right side of the red line, we show the validation perplexity of cach domain of the Pile.
LAMBADA. At 11.3B 3B scale, ADO also achieves the highest average accuracy and outperforms the baselines on 4 out of downstream tasks. Observation 2: interestingly, Natura turss out to be very competitive baseline. It achieves the second-best average downstream performance both scales and performs well on individual domains, too. To the best of our knowledge, very few works on data selection have benchmarked against using the empirical distribution of tokens given by the dataset and tokenizer. As such, we recommend future works on data selection to compare against this simple baseline. Balanced also turns out be a quite competitive baseline at 1 2 2M scale likely because smaller domains are not extensively repeated given the low total number of tokens processed at this scale. This observation is line with the findings of Goyal al. (2024) which suggest that good data policies different for different compute scales.
Conventionally, validation loss is considered a good indicator of the models' performance if they trained on the same data, but it is less clear whether it remains good indicator when the data policy is dynamically changing throughout the training. Empirically (Figure 4), we found that is difficult to outperform Natural setup and without spending much more computation), even though we do outperform this mixture on the downstream tasks as shown above. Observation 3: for validation loss, ADO slightly under- performs Natural 1 on the Pile validation set, but on the validation set of SlimPajama and a subset of Fine FebeWeb, ADO outperforms Natural. Since SlimPajama and Fine are heavily filtered be "higher quality" data compared to the Pile, we speculate that ADO may be able to implicitly select data aligned with some notion of high quality throughout the training process (e.g.. by prioritizing learnability). Nonetheless, since it can be challenging to quantify or precisely define "high quality" data, we will leave this investigation to future work. We include more analysis in Appendix D.
the Pile validation loss (at least within the same training
4.1 ANALYSIS
Mixture over the training. A natural question about data selection methods is what does the distribution look like. We now show that the curricula learned by ADO differ for models of different scales in Figure 5. We observe that CommonCrawl (P (Pile-m CC) receives the most probability mass in both 124M and 1.3B. This observation is consistent with some previous works that put most of the probability mass on CommonCrawl (Xie et al., 2024). OpenWebText2 also receives higher proportion compared to the empirical mixture. Github receives a higher probability in 124M initially but eventually decays whereas in 1 3B it first decreases rapidly and gradually increases
8

---

Published as a conference paper at ICLR 2025
Data mixture: 124M ADO
Data mixture: 1.3B ADO
0.35
0.25
Probabiiity
Probabilitt
0.20
0.15
0.05
0.00
Training steps
x10*
Training steps
K10
Domains
ArXiv
Books3
Github
OpenWebTex12
Plie-CC
PubMed Central
Figure 5: The sampling distribution produced by our data policy during training on The Pile. For case of visualization, we only highlight the top 6 largest domains. ADO produces qualitatively different strategies at different model scales and adaptively changes its weightings over time.
Loss curve forecasts from scaling laws fit during training
Pil--CC
Wikipedia
Enron Emails
Training loss (log scale)
Trining loss
Fit 10k
22
30k
Fit @ 50k
12027
Training steps (log scale)
Figure 6: The forecasts from scaling laws (fit online throughout training) versus the actual training loss different domains in the Pile. The forecasts shown were fit at 10,000, 30,000, and 50,000 steps (time of shown in dashed lines). observe that the scaling laws: (1) become more accurate over the course of training, and (2) can be surprisingly accurate at forecasting the final loss even very early into training on some domains, such as Pile-CC. Forecasts from all domains are shown in the Appendix, Figure 8.
towards the end of the training. This is consistent with the intuition that code data naturally have much lower entropy so it is "easier" to learn in some sense. ADO does not account for the fact that code could be desirable for general capabilities such as reasoning (Ma al., 2024) since it is agnostic to the downstream tasks. This issue can be partially resolved by assigning higher prior to code but further research on data selection targeting reasoning is likely needed.
Accuracy of the scaling law throughout the training. Unlike most applications of scaling laws, refit the scaling laws the fly the models train and the data mixtures dynamically chang- ing. This introduces new challenges because the forecasts from scaling laws fit early on can be inaccurate for steps far in the future. In Figure 6, we show the predictions of the scaling laws for 3 domains at various points during training (all domains are shown in the Appendix, Figure 8). It can be seen that while the predictions of our scaling laws eventually deviate from the true learning curves, they are accurate locally for much of the training and thus can act learning signal for the data policy. More specifically, we can see that the scaling laws consistently overestimate the final loss for all domains, likely due to the presence of learning rate schedules. Empirically, when the learning rate decays towards the end of the training, the training loss tends to decrease more rapidly, until certain point when the learning rate is too small (Hägele et al., 2024). We believe that incorporating recent techniques for incorporating the effect learning rate schedules into the scaling laws (Tissue al., 2024) could further improve ADO.

---

Published as a conference paper at ICLR 2025
5 RELATED WORKS
Data curation and selection. For current large language models, compute often poses a greater constraint than data availability, making data selection crucial (Albalak et al., 2024). A widely used approach is data filtering (Soboleva et al., 2023; Penedo et al., 2023; 2024; Brandfonbrener et al., 2024), where undesirable data points are removed based on heuristics like perplexity, repetition (Lee et al., 2021), or semantic similarity (Abbas et al., 2023). This filtering process is foundational for constructing most large-scale scale datasets today. After filtering, data is often categorized into subsets or domains (e.g., code, books), and one must decide how much data to use from each domain.
Two primary strategies for data selection are prevalent: one focuses on deciding whether individual data points should be included based on various criteria (Mindermann et al., 2022; Engstrom et al., 2024), and the other uses all available data but samples from different domains with varying proba- bilities (Xie et al., 2024; Fan et al., 2023; Albalak et al., 2023). While data selection aims to enhance training efficiency, these methods may introduce considerable computational overhead (Xie et al., 2024: Chen et al., 2024; Fan et al., 2023). Moreover, Kaddour et al. (2023) show that under the same computational budget, these methods often fail to surpass standard training, and Wang et al. (2024) proves that the data selection's 's efficacy depends on the user's utility function.
Another line of work focuses on selecting pre-training data that aligns more closely with downstream tasks (Kang et al., 2024). Data selection has also been explored for computer vision. For example, Evans et al. (2023) use a small reference model to select data for CLIP, while others propose pruning batches based on diversity criteria (Qin et al., 2024; Hong et al., 2024) to improve training efficiency.
Neural scaling laws. Studies have found that various quantities of interest for large pretrained models (e.g., validation loss) can be reliably predicted with simple statistics such as the model size, dataset size, or the amount of computation (Kaplan et al., 2020; Hoffmann et al., 2022). These findings have been central to the design of training protocols for large language models where trial and error are expensive. More recent works have studied the relationship between data repetition and the performance of the models (Hernandez al., 2022; Muennighoff al., 2024; Goyal et al., 2024). Data curation, in particular pruning, has also been shown to have significant effects to achieve better scaling laws (Sorscher al., 2022). Studying the loss curve via scaling law (Hutter, 2021) is relatively less well-explored because the loss curves do not follow a power law exactly due to the learning rate schedule; however, we found that a power law is still decent model for learning curves with cosine decay for language models. Recently, Tissue et al. (2024) showed that learning rate schedules can be incorporated into scaling laws make even more accurate predictions, though we do not explore this direction in this work.
6 CONCLUSION
Dynamic data selection has the potential to improve the pretraining efficiency of foundation models, but most existing approaches incur additional computation costs. We introduce ADO, a cheap online data selection method that dynamically adjusts data distribution over different domains based on domain scaling laws. The scaling laws forecast the model's loss on different data domains and automatically adjust the training data distribution based on each domain's 's learning potential. ADO does not require a proxy model so it naturally takes into account the architecture, the optimizer, and other hyperparameters of the training. In our experiments, with a single set of hyperparameters, ADO performs comparably to or better than existing methods on scales from 124M 1.3B, with much less additional computation. Our implementation only added 20 minutes of additional wall- clock time to 3.5- training run (~ 0.4%). and can be optimized further. The relative cost would decrease more at larger scales since ADO's cost does not scale with the model size. Overall, we believe ADO represents an important step towards accessible and automated online data selection.
Limitations and future directions. Given further computational resources, it would be interesting to scale ADO up to larger models (such as 8B models in DoReMi (Xie et al., 2024)) and datasets and study its behaviors more thoroughly. On the technical side, we believe there are several promising directions: 1. how to create better and more fine- -raiigd domains, 2. how to design more realistic scaling laws that can model inter-domain interactions and learning rate schedules, 3. whether in- sights of ADO can be applied to other training settings such as continued pretraining or fine-tuning, and 4. why do these domain scaling laws arise from naturally occurring data distributions.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We are grateful to Jack Rae, Stephen McAleer, Zhangir Azerbayev, Minqi Jiang, Roberta Raileanu, Michael Xie, Andrew Jesson, Alex Robey, Sam Sokota, Swaminathan Gurumurthy, Jeremy Cohen, and Marc Finzi for helpful discussions during the early phase of this project. We would also like to thank Mengzhou Xia and Tianyu Gao for discussion Figure are courtesy of po jok d, Muhammad Ali, f aisalovers from flaticon com. YJ supported by the Google PhD Fellowship. AZ would like acknowledge support from Chelsea Finn and the National Science Foundation (GRFP). SM acknowledges funding from NSF, ONR, Simons Foundation, and DARPA. We also thank the Google TPU Research Cloud for generously providing computing for the experiments in this paper.
model evaluation. The assets used in
REPRODUCIBILITY
We have included the main code of the ADO algorithm in the supplementary material along- side this submission, which includes the code for tracking per-domain losses, fitting scaling laws, and calculating the current data policy. This the only non-standard component of our train- ing pipeline. Our experiments utilize and modify existing open-source training pipelines with details such as hyperparameters provided Appendix A. The code for this work is available at https: thub.
REFERENCES
Amro Abbas, Kushal Tirumala, Dániel Simig, Surya Ganguli, and Ari S Morcos. Semdedup: Data- efficient learning at web-scale through semantic deduplication. arXiv preprint 00X1v:2303.05540, 2023.
Ryan Prescott Adams and Richard S Zemel. Ranking via sinkhorn propagation. arXiv preprint arXiv: 10c. 1925, 2011.
Alon Albalak, Liangming Pan, Colin Raffel, and William Yang Wang. Efficient online data mixing for language model pre-training. In RO-FoMo: Robustness of Few-shot and shoo Learning in Large Foundation Models, 2023.
Alon Albalak, Yanai Elazar, Sang Michael Xie, Shayne Longpre, Nathan Lambert, Xinyi Wang, Niklas Muennighoff, Bairu Hou, Liangming Pan, Haewon Jeong, et al. A survey data selection for language models. arXiv preprint arXiv:2402.16827, 2024.
Robert A Baker and Stanley W Osgood. Discrimination transfer along a pitch continuum. Journal of Experimental Psychology, 48(4):241, 1954.
Robert Baldock, Hartmut Maennel, and Behnam Neyshabur. Deep learning through the lens of example difficulty. Advances in Neural Information Processing Systems, 14:10876-10889, 2021.
Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In Proceedings of the 26th annual international conference on machine learning. pp. 41-48, 2009.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle 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 Machine Learning, pp. 2397-2430. PMLR, 2023.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense natural language. AAAI Conference on Artificial Intelligence, 2019.
URL httpt: /pp..pp..///
org/ http:/oopusD:::0220353333
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao. Laurence Golding, Horace He, Connor Leahy. Kyle McDonell, Jason Phang, et al. Gpt-neox-20b: An open-souree autore- gressive language model. arXiv preprint arXiv:2004.06745, 2022.

---

Published as a conference paper at ICLR 2025
Rishi Bommasani, Drew A. Hudson, Ehsan Adeli, Russ B. Altman, Simran Arora, Sydney Arx, Michael S. Bernstein, Jeannette Bohg. Antoine Bosselut, Emma Brunskill, Erik Brynjolfsson, Shyamal Buch, Dallas Card, Rodrigo Castellon, Niladri S. Chatterji, Annie S. Chen, Kathleen Creel, Jared Quiney Davis, Dorottya Demszky, Chris Donahue, Moussa Doumbouya, Esin Dur- mus, Stefano Ermon, John Etchemendy, Kawin Ethayarajh, Li Fei-Fei, Chelsea Finn, Trevor Gale, Lauren Gillespie, Karan Goel, Noah D. Goodman, Shelby Grossman, Neel Guha, Tatsunori Hashimoto, Peter Henderson, John Hewitt, Daniel E. Ho, Jenny Hong, Kyle Hsu, Jing Huang, Thomas Icard, Saahil Jain, Dan Jurafsky, Pratyusha Kalluri, Siddharth Karamcheti, Geoff Keel- ing, Fereshte Khani, Omar Khattab, Pang Wei Koh, Mark Krass, Ranjay Krishna, Rohith Ku- ditipudi, and et al. On the opportunities and risks of foundation models. CoRR, abs/2108.07258, 2021. URL https: //tp///aavvvv/aa////a//eaeeawav arg/abs/////..o..o......a....... 210. 07258.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. com/ pepogle/jaa
David Brandfonbrener, Hanlin Zhang, Andreas Kirsch, Jonathan Richard Schwarz, and Sham Kakade. Color-filte:: Conditional loss reduction filtering for targeted language model pre- training. arXiv preprint arXiv:2406. 10670, 2024.
Mayee Chen, Nicholas Roberts, Kush Bhatia, Jue Wang, Ce Zhang, Frederic Sala, and Christopher Ré. Skill-it! data-driven skills framework for understanding and training language models. Advances in Neural Information Processing Systems, 36, 2024.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. ArXiv, has/1803.05457, 2018. URL pettps:/aapte http://://ap.sssmmtttsssssssss www.sttemstttttaeeeeee CorpusID: 3922816.
.rr/
DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Clau- dio Fantacci, Jonathan Godwin. Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojević, Wojeiech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http: //github. cwm/ http://ooooeedeeeimmmm
Aaron Defazio, Harsh Mchta, Konstantin Mishchenko, Ahmed Khaled, Ashok Cutkosky, et The road less scheduled. arXiv preprint arXiv:2405. 2024.
Logan Engstrom, Axel Feldmann, and Aleksander Madry. Dsdm: Model-aware dataset selection with datamodels. arXiv preprint arXiv:2401.12926, 2024.
Talfan Evans, Shreya Pathak, Hamza Merzic, Jonathan Schwarz, Ryutaro Tanno, and Olivier Henaff. Bad students make great teachers: Active learning accelerates large-scale visual un- derstanding, arXiv preprint arXiv:2312.05322 2023.
Simin Fan, Matteo Pagliardini, and Martin Jaggi. Doge: Domain reweighting with generalization estimation. ArXiv, abs/2310.15392. 2023. URL http://appp..emnnniccchhllllaa .rrg
Carlos Florensa, David Held, Markus Wulfmeier, Michael Zhang, and P. Abbeel. Reverse cur- riculum generation for reinforeement learning. ArXiv, abs/ 1011777.0550 (os/1707.05300, 2017. URL https:
// //api. www.nticschhhol...cc.mm. www.olaa...rrooouuss.......
Samir Yitzhak Gadre, Gabriel Ilharco, Alex Fang, Jonathan Hayase, Georgios Smyrnis, Thao Nguyen, Ryan Marten, Mitchell Wortsman, Dhruba Ghosh, Jieyu Zhang, al. Datacomp: In search of the next generation of multimodal datasets. Advances in Neural Information Processing Systems, 36, 2024.
12

---

Published as a conference paper at ICLR 2025
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2 00027, 2020.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lin- tang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024. URL htRpst:s: /ttp://////oooooo/ org/ records/ 12608602.
Ce Ge, Zhijian Ma, Daoyuan Chen, Yaliang Li, and Bolin Ding. Data mixing made efficient: A bivariate scaling law for language model pretraining. arXiv preprint urXiv:2405.44908 2024.
Charles Albert Eric Goodhart. Monetary control- the british experience. Monetary Conditions for Economic Recovery, pp. 59-84. Springer, 1985.
Sachin Goyal, Pratyush Maini, Zachary C Lipton, Aditi Raghunathan, and 7ico Kolter. Scaling laws for data filtering-date curation cannot be compute agmostic. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22702-22711, 2024.
Alex Graves, Marc G Bellemare, Jacob Menick, Remi Munos, and Koray Kavukcuoglu. Automated curriculum learning for neural networks. In international conference on machine learning. pp- 1311-1320. Pmlr, 2017.
Yuxian Gu, Li Dong, Yaru Hao, Qingxiu Dong, Minlie Huang, and Furu Wei. Towards optimal learning language models. arXiv preprint arXiv:2402.11755, 2024.
Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Von Werra, and Martin Jaggi. Scaling laws and compute-optimal training beyond fixed training durations. arXiv preprint arXiv:2405. 18392, 2024.
Danny Hernandez, Tom Brown, Tom Conerly, Nova DasSarma, Dawn Drain, Sheer El-Showk, Nel- son Elhage, Zac Hatfield-Dodds Tom Henighan, Tristan Hume, et al. Scaling laws and inter- pretability of learning from repeated data. arXiv preprint arXiv:2205. /0487, 2022.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aure- lia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and L. Sifre. Training compute-optimal large language models. ArXiv, abs/2203. 2022. URL https: apii
Org/CorpusID 2
Feng Hong, Yueming Lyu, Jiangchao Yao. Ya Zhang, Ivor W Tsang, and Yanfeng Wang. Diversified batch selection for training accelcration. arXiv preprint arXiv:2406.0453 2024.
Mareus Hutter. Learning curve theory. arXiv preprint arXiv:2102.044 00.000000444..... 2021.
Andrew Jesson, Nicolas Guizard, Sina Hamidi Ghalehjegh, Damien Goblot, Florian Soudan, and Nicolas Chapados. CASED: curriculum adaptive sampling for extreme data imbalance. CoRR, 10b//1807.1889, 2018 URL http: //trxvv//axvv... org. org/ halp//////aa /107/ 10819.
Jean Kaddour, Osear Key, Piotr Nawrot, Pasquale Minervini, and Matt J. Kusner. No Train No Gain: Revisiting Efficient Training Algorithms For Transformer-based Language Models, July 2023. URL http: /arxiv..vvv org//bgg//.s.. 2207 06440. arXiv:2307.0644 7.6440 [cs].
Dimitris Kalimeris, Gal Kaplun, Preetum Nakkiran, Benjamin Edelman, Tristan Yang, Boaz Barak, and Haofeng Zhang. Sgd on neural networks learns functions of increasing complexity. Advances in neural information processing systems, 32, 2019.
Feiyang Kang, Hoang Anh Just, Yifan Sun, Himanshu Jahagirdar, Yuanzhi Zhang, Rongxing Du, Anit Kumar Sahu, and Ruoxi Jia. Get more for less: Principled data selection for warming up fine-tuning in llms. In ICLR, 2024. URL https:/ /wwpeeeiieeeeeeeeeee///ee net wtt./ooormi///// id= YNYNBVkee..
13

---

Published as a conference paper at ICLR 2025
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
Patrick Kidger and Cristian Garcia. Equinox: neural networks in JAX via callable PyTrees and filtered transformations. Differentiable Programming workshop at Neural Information Processing Systems 2021, 2021.
Douglas H Lawrence. The transfer of discrimination along continuum. Journal of Comparative and Physiological Psychology, 45(6):511, 1952.
Katherine Lee. Daphne Ippolito, Andrew Nystrom, Chiyuan Zhang, Douglas Eck, Chris Callison- Burch, and Nicholas Carlini. Deduplicating training data makes language models better. arXiv preprint arXiv:2 10022107.0000 107.06499, 2021.
Hanmeng Liu, Jian Liu, Leyang Cui, Zhiyang Teng. Nan Duan, Ming Zhou, and Yue Zhang. Logiqa 2.0-an an improved dataset for logical reasoning in natural language understanding. IEEE/ACM Transactions on Audio, Speech, and Language Processing, 31:2947-2962, 2023. URL https:
/ //api. www.sseeaannnnoltie.e.... org org CorpusID: ott.eechh..a.....ooooooooaa::::. 2595151154
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. arXiv preprint urXiv:2407.01499 2024.
Loshchilov. Decoupled weight decay regularization, arXiv preprint arXiv: 10t.::111..,111,, 2017.
Yingwei Ma, Yue Liu, Yue Yu, Yuanliang Zhang, Yu Jiang, Changjian Wang, and Shanshan Li. At which training stage does code data help LLMs reasoning? In The Twelfth International Conference on Learning Representatinss 2024. URL https::/ http///pppnneeeeeeeeeeeeeeeeeeea net forum? 1d=KIPJKST4gW.
Stefano Sarao Mannelli, Yaraslau Ivashinka, Andrew Saxe, and Luca Saglietti. Tilting the odds the lottery: the interplay of overparamettrisation and curricula in neural networks. arXiv preprint 58Xiv:2406.01589, 2024.
Gonzalo Mena, David Belanger, Scott Linderman, and Jasper Snoek. Learning latent permutations with gumbel-sinkhorn networks. arXiv preprint
2018.
Sören Mindermann, Jan M Brauner, Muhammed Razzak, Mrinank Sharma, Andreas Kirsch, Win- nie Xu, Benedikt Höltgen, Aidan N Gomez, Adrien Morisot, Sehastian Farquhar, et al. Prioritized training on points that are learnable, worth learning, and not yet learnt. In International Confer- ence on Machine Learning, pp. 15630-15649, PMLR, 2022.
Niklas Muennighoff, Alexander Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin A Raffel. Scaling data-constrained language models. Advances in Neural Information Processing Systems, 36, 2024.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and R. Fernández. The lambada dataset: Word prediction requiring broad discourse context. ArXiv, abs/1606.06033 2016. URL https: / //api. semant wwwchhaacarr.ooooo www.oolrrrrrroooooooooooo Corp/orruuuusc : 2351
Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay. The refinedweb dataset for falcon llm: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.0116,, 2023.
Guilherme Penedo, Hynek Kydlfček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data scale, 2024. URL https: /taaaiiivoorr /240 7557.
14

---

Published as a conference paper at ICLR 2025
Ziheng Qin, Kai Wang, Zangwei Zheng, Jianyang Gu, Xiangyu Peng, Zhao Pan, Daquan Zhou, Lei Shang, Baigui Sun, Xuansong Xie, and Yang You. Infobatch: Lossless training speed up by unbiased dynamic data pruning. In The Twelfth International Conference Learning Represen- tations, 2024, URL https::////... paappp////// //wprrveeeeeeeeeeeeeeeeeeeeeeeee ew net/ forum?id=C6i 5k51.K6.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. An adversarial wino- grad schema challenge seale. 2019. URL https: /app//ppp.... Corpus 199370376.
Or Sharir, Barak Peleg, and Yoav Shoham. The cost of training NLP models: A concise overview. CoRR, abs/2004.08900, 2020. URL https: / //wxiviivv/v./////////////aaaav/ wttgabs////eoooooogogogoo/.ee 2004 08900
Noam Shazeer. Glu variants improve transformer. arXiv preprini arXiv: arXiv:2002.05202. 2020.
Richard Sinkhorn. A relationship between arbitrary positive matrices and doubly stochastic matri- ces. The annals of mathematical statistics, 35(2):876-879 1964.
Burrhus F Skinner. Reinforcement today. American Psychologist, 13(3):94, 1958.
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel A 627B token cleaned and
Hestness, and Nolan Dey. SlimPajama: deduplicated version of RedPajama. s1 slimpa jampaaama--ae jama 2023. URL https: h/thgginggaacccee..ee//// wttpaaaaccccddaaaaaaaeeeeeee /tt.///tteeetstrrrrrraaa/eeeeeee
ptttss//cerebras.aa/// otw.ooee---eeeeeedddddeeeeeeeeee
627B.
Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos. Beyond neu- ral scaling laws: beating power law scaling via data pruning. Advances in Neural Information Processing Systems, 35:19523-19536 2022.
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: enhanced transformer with rotary position embedding. corr abs/2104.09866 (2021). arXiv preprint rrXiv2210.009644 2021.
Yuxing Tang, Xiaosong Wang. Adam P. Harrison, Le Lu, Jing Xiao, and Ronald M. Summers. Attention-guided curriculum learning for weakly supervised classification and localization of thoracic diseases on chest radiographs. ArXiv, abs/1807.07555 2018. URL https: : /api. semant cscholar. wwwpoolarrrrg//ooopuuuoe C/rpporrrppp/cc 49882848.
Howe Tissue, Venus Wang, and Lu Wang. Scaling law with learning rate annealing, 2024. URL /200 11029.
httts::/erxiv..rrtt
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov. Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.002886 2023.
Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Neural Information Processing Systems, 2017. URL https: /api semant: i cs cscholar org/ 1:w.
Jiachen T Wang, Tianji Yang, James Zou, Yongchan Kwon, and Ruoxi Jia. Rethinking data shapley for data selection tasks: Misleads and merits. arXiv preprint arXiv:2405.03875 2024.
Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science ques- tions. ArXiv, 0bs/1707.06209 2017. URL https:: https://ppp..eematticchhho.. semanticscholl.rr.......c..cc... TD:
Xiaoxia Wu, Ethan Dyer, and Behnam Neyshabur. When do curricula work? In International Conference on Learning Representations, 2021. URL httLttts/// wwwp/ewweeeew..co///o. new./eee/eeeee /forum? id=t#4QEInpni.
/
Mengzhou Xia, Tianyu Gao. Zhiyuan Zeng, and Danqi Chen. Sheared LLaMA: Accelerat- ing language model pre-training via structured pruning. In The Twelfth International Confer- ence on Learning Representations, 2024a. URL https::////// /pwp///iep/eeeeeee/ee///e. 091 OdaeOzp.
/ net fftfroorrddd
15

---

Published as a conference paper at ICLR 2025
Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. LESS: Selecting influential data for targeted instruction tuning. In Forty-fira International Conference on Machine Learning. 2024b. URL
net wttt://oooommdod/eeSeeeeee
Sang Michael Xie and Stefano Ermon. Reparameteriizble subset sampling via continuous relax- ations, 2021. URL https::////.
10517.......
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy S Liang, Quoc V Le, Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining- Advances in Neural Information Processing Systems, 36, 2024.
Jiasheng Ye, Peiju Liu, Tianxiang Sun, Yunhua Zhou, Jun Zhan, and Xipeng Qiu. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. arXiv preprint arXiv:2403.16952 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? In Annual Meeting of the Association for Computational Lin- guistics, 2019. URL https::/ /app:///pp.p/ semant eeschllar org/ http///oo:::t1004//1111
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Infor- mation Processing Systems, 32, 2019.
Allan Zhou, Nicholas C. Landolfi, and Yiding Jiang. midGPT: simple and hackable repository for IIm pretraining. 2023. URL https: / /g1 thub. com /al.llllaannngnnn wnw.annhooommiiiioo dGPT.
16

---

Published as a conference paper at ICLR 2025
A TRAINING DETAILS
Architecture. All experiments use the same architecture, which follows the Llama 2 family (Tou- al., 2023): Transformer-based language model with key components including SwiGLU MLP layers (Shazeer, 2020), RMS normalization (Zhang & Sennrich, 2019), and rotary embed- dings (Su al., 2021). We use the GPT-NeoX-20I tokenizer (Black al., 2022), and tie the weights in the embedding and final layers. Our experiments contain models trained at two scales: 12 parameters and 1. 3B parameters.
Training, All models were trained for 60,000 steps at batch size 2048 (1B) or 256 (124M), using AdamW (Loshchilov, 2017). We use decoupled weight decay with 10 set B2 = 0.05, and otherwise use default hyperparameters as specified by Optax (DeepMind et al., 2020). For ADO, we fit scaling laws for each domain every 1, 1,000 training steps starting at step 5,000 (we run an empirical sampling strategy for the first 5,000 steps).
Compute and time expenditure. We ran our experiments on TPUs using the open-source midGPT library (Zhou et al., 2023), which is based on JAX (Bradbury et al., 2018) and Equinox (Kidger & Garcia, 2021). All experiments were run on Google Cloud TPUs. On a TPU v3-128, a IB model can be trained for 60,000 steps in 3.5 days. Fitting scaling laws for all domains takes less than 20 seconds-over the course of a training run, this amounts to under 19 minutes additional time spent fitting scaling laws. This number can likely be further improved by using smaller smarter non-uniform grid (we found the results to be not sensitive to the size of the grid search) or smarter optimization custom-made for power law. Although not present in our particular implementation, the fitting of the scaling laws could, in principle, take place asynchronously in parallel with model training since it does not require any information from the model other than the loss, which can be staggered as the scaling laws remain accurate for some duration into the future.
A.1 FITTING SCALING LAWS
We fit an individual loss curve scaling law to each domain where the input is the total number of samples processed so far, and the prediction is the training loss of that particular domain. Namely, suppose the dataset is broken down into K distinet heterogeneous domains, Ex + Ban n -kk To estimate these parameters of the scaling law, we use the standard procedure for fitting power law (Hoffmann et al., 2022) which minimizes the Huber loss with ó 0.001 between the predicted log loss and observed log loss with L-BFGS:
would learn (x(n)
min Ek.P
Σ
Hubers (m)) log La (n))
Since learning curves have a large number of points, we subset the learning curve at regular intervals obtain the data for fitting. Similar prior works, we use grid search with different initialization, although we found that usually a smaller grid search suffices since the individual scaling laws do need to be very accurate. The losses of the first 500 steps are not included as the weights have not stabilized yet and we subsample the trajectory at 10-step intervals to speed up the training. This does not significantly impact the algorithm since scaling laws do not change drastically. We also apply a threshold to the parameters' values to mitigate potential instability due to online fitting.
and e are both parameterized their log forms. We use grid search over the following parameters initialization:
ann E
0.6, 0.7)
log 3o E 6--2,-1,,0,1,,33344555
. log co E --,,,-15,-1,-05,,
Further, we enforce 0 a < 0.8, 3 6.5 and log > 0.5. These bounds are purely preventative in cases of numerical instability and are almost never saturated.
The computation complexity of fitting scaling law depends on the rate of convergence of L-BFGS, which can be problem-dependent but is generally efficient and stable. The complexity of fitting dom-dn scaling laws scales with the number of domains and the granularity of the initialization search grid. Suppose G is the granularity of discretization for parameter (e.g., 7 for og), and K
17

---

Published as a conference paper at ICLR 2025
Table 2: Different mixtures used in this paper.
Dataset
Natural
DeReMi
Pile
Freel.aw
0.0449
0.0380
0.0612
Enron Emails
0.0010
0.0040
0.0014
Github
0.1227
0.0325
0.0759
OpenSubtitles
0,0158
0,0032
0.0155
PubMed Central
0,1215
0.0608
0.1440
Open OpenWeTTex12
0.1096
0.1905
0.1001
StackExchange
0.0491
0.0746
0.0513
Pile-CC
0.1825
0.1379
0.1811
ArXiv
0.0886
0.0535
0.0896
USPTO Backgrounds
0.0262
0.0327
0.0365
Books3
0.1046
0.0757
0.1207
Wikipedia (en)
0.0402
0.1068
0.0153
PubMed Abstracts
0.0221
0.0970
0.0307
NIH ExPorter
0.0019
0.0084
0.0030
BookCorpus2
0.0063
0.0037
0.0075
EuroParl
0.0081
0.0120
0.0073
HackerNews
0.0047
0.0084
0.0062
DM Mathematics
0.0191
0,0019
0.0124
YoutubeSubtitles
0.0040
0.0117
0.0060
PhilPapers
0.0027
0.0093
0.0038
Ubuntu IRC
0.0049
0.0083
0.0088
Gutenberg (PG-19)
0.0195
0.0292
0.0217
the number tasks, the complexity of fitting would O(G² K), which grows linearly with the number of domains. However, can easily parallelize multiple L-BFGS runs TPUs (or GPUs), which is why ADO does not add much time in practice. As we mentioned above, 0.04% training time in 3-day which less than 20 minutes, number that can likely be optimized further.
A.2 ESTIMATED EMPIRICAL MIXTURE
In Figure 5, we show various data mixtures used this paper. For the natural distribution, we sampled 1000 documents from each domain, tokenized them with chosen tokenizer, and then computed the average number of tokens per document. Using the estimated average, we can estimate the total number of tokens within domain by multiplying by the number of documents.
Note that the DoReMi differs from the number of Xie et al. (2024) due to the use of different tokenizers. The numbers tokenizer (Black et al., 2022) like us. Also, note that the estimated empirical distribution differs from the default Pile weights. Notably, Github and Wikipedia weighted significantly higher.
here are from Albalak et al. (2023) which uses GPT-NeoX- X-B
B META-LEARNING PERMUTATIONS
Suppose we have dataset 2 = {2), 1,,,,0 22,... n drawn from the distribution Pz and we wish to train a model in the one-epoch setting, that is, the model gets to see each example only once". In this setting, the space of curricula is SN. all permutations of [N]. Note that the size of this search space [Sw] N! is extremely large for even a small N, and would
allowed data repetition. Given curriculum a [N] + [N] and the stochastic gradient descent update rule UO0, 2) f- (2; 0), we initialize the model parameters 9g and iteratively update them: 0,41 = U0,, ). We can denote the entire training process from 0g to 0N by ON SGD SGD(Oo, N).
even larger
In standard SGD, the curriculum is randomly sampled permutation, Frandom Unif (Sx). Given the size of the search space, we can instead formulate the problem of finding curricula
Thii common for the training regime of the modern language model.
18

---

Published as a conference paper at ICLR 2025
optimization problem over Sy. An optimization problem needs objective, but what should the objective be for finding a good curriculum? The ultimate goal of this problem to find curriculum such that the final parameter 8, achieves low population loss, ie., o* arg min, Ez-r, [e(=; SCD a, N) However, we cannot general assume access to the data distribution need surrogate objective. In this experiment, meta objective that depends only on the training data:
the following surrogate
N-1 N
1
o² = arg min
E 20 2m1
l(zz; 0z) arg min ((a); a
N
where the dependence of 0, implicit since 0.+1 ((),,
This objective encourages the curriculum to minimize the total training loss after each update (i.e., training fast possible). Since we are in the one-epoch setting, the parameters are updated with the gradient of each example only once so there is less, if any, risk of overfitting.
One obstacle directly learning a good curriculum here is that gradient-basee baseed optimizers not suitable for searching over permutations. We follow Mena et al. (2018) and parameterize the cur- riculum with matrix Z N We then the Sinkhorn operator S(.) to project Z to doubly stochastic matrix X = S(Z) (Sinkhorn, 1964; Adams Zemel, 201 1), which we can think of as "soft" approximation
permutation. Since X doubly stochastic but not an actual permutation, we use a weighted batch generalization of our update:
N
0++1 = 02
E 1=1
X1
Vol(zz;;)
(8)
N
This is generalization in the sense that, if X is only doubly stochastic but also a permutation a, then it reduces to SGD with curriculum a.
After optimizing Z against the meta objective, we can project Z to true permutation which serves as our final learned curriculum:
Zi+1 = Z, - (V22((())
= arg PEPN
P).
(9)
This procedure can further generalized to selecting a minibatch. To do so, we relax the con- straints X from doubly stochastic matrix matrix with differentiable top-K relaxation each row (Xie & Ermon, 2021). We apply this to the MLP experiments in Figure 3 with minibatch size 10.
C CLIPPING FUNCTION
Below is the pseudocode for clipping the probability. It evenly distributes the excess proba- bility amongst the other categories.
Algorithm Clip Minimum Probability
1: Input: probs E A min prob E [0.
2: Output: clipped_probs
3: total.deficit + mwx
Σ probs,0)
Compute total deficit
4: scale.factor + (1-tooaa..ffiicc
Compute scale factor Scale the probabilities b Clip probabilities Normalize
5: scaled_probs t- probs scale.factor
6: clipped.probs -
7: clipped + probss) ppoppepdoposs
8: return clipped probs
19

---

Published as a conference paper at ICLR 2025
D ADDITIONAL EMPIRICAL RESULTS
D.1 ADDITIONAL ANALYSIS
In terms of individual tasks, we observe that ADO outperforms ARC-E, SciQ, and LogiQA at 124M scale but not 1.3B scale. The deficit for ARC-E is 0.2% which is relatively small and the performance far above the other points of comparison. For LogiQA, the performance of DoReMi, Natural, and ADO are all within 0.4% of each other, which is relatively marginal. For SciQ, the difference in performance is more pronounced at 1.1 %, but still outperforms other baselines. This may suggest that ADO has put less focus scientific knowledge during pretraining. We believe this could be alleviated by either putting higher prior weights on scientific data such as wikipedia or mixing in high-quality scientific corpus at the later stage of training.
On the other hand, for LAMBADA, ADO outperforms Natural at 1.3B scale by a large margin (1.4%) but not the 124M scale. This highlights the fact that, for comparing different data mixtures, is very hard to obtain uniform or monotonic improvement on all downstream tasks the models scale up. Note that downstream tasks are the proxy the model's capabilities. If there is particular downstream task of interest, it would be advisable to collect data for the particular task and finetune further. As an evaluation for generic pretraining, aggregated metries such as average would be more reliable measurement of the model's 's performance.
The performance metrics would naturally become more difficult to improve with larger scales because the room to improve becomes smaller. The average error for "natural" at 124M is 1 0.463 0.537 and ADO improved over "natural" by 0.07 so the relative improvement is 1.3%. The average error for "natural" at 1.3B is 1 0.585 = 0.415 and ADO improved over "natural" by 0.05 so the relative improvement 1. 1.2%. This shows that the relative improvement of ADO in average performance at different scales is consistent.
D.2 ADDITIONAL PLOTS
Table 3: Ablation for the zero-shot performance of different smoothing parameter values, 8. We observe that 3 out 4 values between 0 and outperform the natural baselines, but larger values of tend to perform better. This is expected since smaller s is less effective in assigning credit to different domains (s would make X(t) uniform). On the other hand, since the credit assignment imperfect, using an overly large could also be undesirable. In our experiments, a value smaller than but not too small achieves the best result.
HellaSwag
WinoGrande
PIQA
ARC-E
SciQ
LogiQA2
LAMBADA
Average
Natural Balanced
0.2897
0.5028
0.6235
0.4348
0.7550
0,2341
0.4007
0.4629
0.2800
0,5114
0.6099
0.4470
0.7620
0.2271
0,3623
0.4571
1.0
0.2951
0.5067
0.6295
0.4747
0.7460
0.2360
0.3635
0.4645
0.5
0.2901
0.5201
0.6349
0.4562
0.7710
0.2436
0.3707
0.4695
0.3
0.2854
0.5130
0.6311
0.4369
0.7630
0.2277
0.3742
0.4616
0.1
0.2894
0.5328
0.6170
0.4457
0.7650
0.2290
0.3625
0.4631
20

---

Published as a conference paper at ICLR 2025
GacrrWbbioo22
3tackechargg:
cxbruprate etc
4x10
cebupoiate
2x
A1
A11
CAU apoluu
3x
Log La
30
20
2x
2:
Emals
00000
4x10'
[d
(xl
(11 addraphbos
estrapo ate
octropoiate
ectrapeirma
3x10'
2x100
2xh
2x30
extrapo ato
4x1
ortropoiate
scrapelezz
20
u00 Lone
3x110
5x 10
2x10
2X"00
Gulanbory (70121
Muhomaakks
NHH EaPtdeer
2x"10
3x 3x:10
oetrupiiate
1x100
eecopolez
22x30'
L4xIp
13x10
Corpost
4x11P
Oper Sustilos
4x36
20
Ciboor de
cotopgiate
15300
cecopolas
$K.00
apoluu
2x110
2x100
Pr Papers
5x010
estrapo ate
ectropoiate
4'100
10g 010
2x
$
Figure 7: Exrapolating the loss of larger models on each domain of the Pile with the loss of models with fewer than IB parameters.
21

---

Published as a conference paper at ICLR 2025
Loss forecasts from scaling laws during training
ArXiv
BookCarpus?
Books3
DM Mathematics
12x
2407
18x
202
10
10
10
10a
Enron Emails
EuroParl
FreeLaw
Github
12x00
13x 17
10a
10
100
Gutenberg (PG-19)
HockerNews
NIH ExPorter
OpenSubtities
121.07
322
34.07
17
Training loss (log scale)
282-27
10:17
24
24
10'
10
16
30
10
10"
Open/ee Tex12
PhilPapers
Pile-CC
PubMed Abstracts
37
30x00
121.7
1415
20
27
18
10x10
10'
10
10
30a
10'
10
20
10'
PubMed Central
StackExchange
USPTO Backgrounds
Ubuntu
30x02
241
11
10
10
10
10
10
101
107
10'
Wikipedia (en)
YoutubeSubbitles
1412
Training loss
-ii @ 30k
Fit @ 10k
Filt 50k
100
107
10a
16a
10a
Training steps (lng scale]
Figure& 8: ADO loss forecasts produced by our scaling laws on each domain in The Pile while training 1B model. Scaling laws are fit regularly throughout training, with the forecasts shown being from fits at 10,000, 30,000, and 50,000 steps (dashed line shows when each scaling law was fit).
22

---

Published as a conference paper at ICLR 2025
Training loss of different data policies
1.3B-Natural
3.00
1.3B-DoReMi
1.3B-ADO
2.75
..3B-Balanced
2.50
2.25
2.00
1.75
10000 20000 30000 40000 50000 60000
Figure 9: The training losses of different methods differ significantly from each other. ADO has much higher training loss compared to the other methods because ADO selects data with high learn- ing potential, which seem be correlated with higher entropy.
A(t): ADO
A(d) 1.38 ADO
0.175
0.125
Probabtty
0.100
0.075
0.050
0.000
Training sleps
Training
Domains
Booksb
Givub
Opentiees lest2
Puutiedd Centra
Figure 10: The credit assignment X(t) over the course of training for s 0.5. For clarity, only the top six domains colored.
23

---

Published as a conference paper at ICLR 2025
BookCorpus2
Enron Emails
1e-8 ArXiv
1e-6
1e-6
Books3
1e-7
1
2
4
4
2
2
2
25000
50000
25000
50000
25000
50000
25000
50000
25000
50000
Gutenberg (PG-19)
Hacke nNews
EuroParl
FreeLaw
Github
4
2
2
25000 50000
25000 50000
25000 50000
25000
50000
0
25000
50000
NIH ExPorter
OpenSubiites
OpenWebTex12
PhilPapers
1
6
1e
16
6
1e 6
1e 6 Pile-CC
3
4
2
2
2
2
2
25000 50000
25000
25000
25000
25000
PubMed Abstracts
PubMed Central
StackExchange
USPTO Backgrounds 16-6
Ubuntu
fa
3
1e-6 6
3
1e-6
6
1e-6
3
2
2
2
1
1
1
25000 50000
25000
50000
25000
50000
25000
50000
25000
50000
Wikipedia (en)
YoutubeSubiiiles
3
3
2
2
25000
50000
25000
50000
Figure 11: The slope of the scaling law, on (n))), for each domain throughout 1.3B ADO training.
24


Reviews:
[{"review_id": "lQNW8LAjO0", "reviewer": "Reviewer_26JP", "summary": "The authors propose ADO, an online data selection method that adjusts data distribution dynamically across various domains by leveraging domain scaling laws. Without requiring a proxy model or external knowledge, ADO forecasts the model's loss on different data domains and automatically modifies the training data distribution according to each domain's learning potential. The experiments indicate that ADO surpasses baseline methods overall, with only a minimal increase in clock time.", "questions": "1. Figures 6, 7, and 8 show the training loss for ADO only. I'd like to see how quickly ADO's training loss decreases compared to the baselines. I assume that ADO converges faster, but it would be interesting to observe how it performs relative to the baselines.\n\n2. Can ADO be effectively applied to vision data and tasks? It would be of great importance if ADO could be used in training vision models that require multiple domains, such as large vision-language models (LVLMs) and diffusion models for image generation.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The motivations presented in Section 2 highlight significant concerns, and the proposed method properly addresses these issues.\n2. ADO outperforms the baselines overall, offering online data selection without the requirement of external knowledge or proxy models.\n3. The authors present several interesting observations that align with our intuition and the literature (e.g. line 428, 409, 431).", "weaknesses": "1. Although the 1.3B model is not small, ADO is especially practical and relevant for LLMs with at least 8B parameters—the size typical of many popular \"small\" LLMs such as Llama 3.1 8B. The paper's relevance and importance might have been enhanced if the authors had conducted experiments with language models of at least 8B parameters\n2. The performance improvement achieved by ADO does not seem significant enough, especially on the 1.3B model. More problematically, Table 1 shows diminishing returns as the model scale increases from 124M to 1.3B parameters. For example, 1.3B-ADO only outperforms 4 out of 7 downstream tasks, whereas 124M-ADO outperforms 6 out of 7 tasks. Additionally, the gap between the average score of ADO and the second-best baseline is less pronounced in the 1.3B model. This suggests that ADO may not scale well as the number of parameters increases.\n3. Although the authors propose interesting future research directions (e.g. Section 6), incorporating some of these suggestions—such as learning rate scheduling—into the current paper would make it more complete and thorough.", "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": "KVFLO6MUN9", "reviewer": "Reviewer_monZ", "summary": "This paper proposes a novel method ADO to adaptively sample data from different domains based on their contribution to model training, which is estimated from scaling law in an online manner. Empirical results on different data sets demonstrate the effectiveness of the proposed method.", "questions": "- The smoothing parameter $s$ for computing the credit assignment scores seems a bit confusing. By using $s<1$, I suppose $\\lambda_k$ indeed changes drastically when $h_k$ is close to 0? Should using $s<1$ encourages $\\lambda_k$ to be larger than $h_k$ especially when $h_k$ is originally close to 0? Some further discussion on the effect of this hyper-parameter may be necessary. \n- Also, can we consider other types of functions to compute $\\lambda_k$ from $h_k$? Some empirical comparison on different types of functions may help others better understand how the credit assignment score affects model training, and why the authors reject simply using $\\lambda_k=h_k$ here. \n- Experiments are limited to medium scale models (124M and 1.3B), and the performance gain seems to decrease with increasing model size. Such a tendency can naturally make one wonder if the performance gain can be smaller or even diminish on larger models (e.g., 2.3B or 7/8B)? Some more empirical results may be necessary to better support the proposed method. \n- While the proposed method includes many hyper-parameters (e.g., the smoothing parameter $s$ mentioned above), the authors did not provide sufficient analysis on how these hyper-parameters can affect the final performance. These results should be necessary to demonstrate that ADO can be robust to different values of these hyper-parameters. \n- Despite showing the final sampling probabilities in Figure 5, the authors may also consider further showing how $\\lambda_k(t)$ and $\\frac{\\partial}{\\partial n} \\hat{L}_k(n)$ changes with training iteration $t$. These results may help others better understand how these two terms affect the sampling probabilities of different domains and how they can affect model training.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The proposed method is introduced clearly and easy to understand", "weaknesses": "- Several parts of ADO lack sufficient motivation\n- Empirical results seem not supportive enough to fully evaluate the proposed method", "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": "HMEannbaL0", "reviewer": "Reviewer_7Mvd", "summary": "The paper introduces ADO, a novel online algorithm for optimizing data distributions during the training of large-scale foundation models. ADO addresses the challenge of efficiently managing data mixtures in pretraining by continuously adapting the data distribution based on scaling laws, without requiring proxy models or significant computational overhead. The primary contribution lies in the use of per-domain scaling laws to estimate learning potential and dynamically adjust the data mixture during model training. Experimental results demonstrate that ADO can achieve comparable or better performance than prior methods, while maintaining computational efficiency across various model scales.", "questions": "1.\tIn ADO's data selection process, some datasets might receive minimal training due to low selection probabilities, potentially leading to under-representation. Have any additional measures been implemented to balance data distribution and ensure adequate representation across datasets?\n2.\tHow sensitive is ADO to changes in its hyperparameters, such as the smoothing coefficient (γ1) and update intervals? Would different tasks or datasets require significant parameter tuning?\n3.\tGiven the current limitation in modeling solely intra-domain interactions, have the authors considered potential methodologies to incorporate inter-domain interactions within the ADO framework? What approaches might be feasible for capturing these cross-domain dynamics to enhance the model’s adaptability and performance in tasks where inter-domain relationships play a significant role?\n4.\tADO is designed to be task-agnostic, but certain applications may benefit from targeted data selection. Could ADO be extended to incorporate task-specific objectives without compromising its efficiency?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1.\tADO utilizes online scaling laws to dynamically adapt data selection during training, replacing the need for pre-trained proxy models or multi-staged processes. It requires no prior domain knowledge or external models, making it highly practical and versatile.\n2.\tThe paper presents empirical evaluations demonstrating ADO's effectiveness across multiple benchmarks and datasets. ADO outperforms or matches existing baselines on several downstream tasks while incurring a small computational overhead (e.g., less than 0.4% additional wall-clock time).\n3.\tADO offers a practical, scalable solution for data selection in large model training, with the potential to significantly reduce computational waste. The method aligns with the growing emphasis on efficient resource utilization in AI research, especially as models scale in size and cost.", "weaknesses": "1.\tFigure 4 reveals that fewer than half of the datasets achieve minimal perplexity under the ADO algorithm, indicating potential limitations in the algorithm’s generalization capability across diverse scenarios. The paper would benefit from a more detailed analysis of ADO's applicability.\n2.\tThe algorithm employs heuristic choices for key parameters, such as the exponential moving average in credit assignment and smoothing coefficients, which may impact the robustness and consistency of ADO across different datasets and tasks due to the lack of a more systematic basis for parameter selection.\n3.\tSymbols such as $\\gamma_2$ appears in key formulas but might be clarified further with more detailed context, especially for readers unfamiliar with its role in the algorithm.\n4.\tAlthough ADO highlights low computational overhead as a strength, the paper lacks a systematic computational complexity analysis and comparative experiments. The absence of quantitative and visual evidence of its efficiency limits the clarity of its computational advantages.", "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": "wXvEmAUz2q", "reviewer": "Reviewer_Axts", "summary": "This work proposes an online adaptive data optimization (ADO) method for finding a policy (weighting) for different domains when pre-training a language model. The core idea is to use scaling law to model training loss. The derivative of the loss w.r.t. to the number of data points can be understood as loss reduction per data point locally, i.e. outline the potential for a domain. The final scoring also considers the contribution of the domain itself from the last update and makes it stable with a moving average. Their experiments show the ADO brings improved downstream performances on 7 common-sense reasoning tasks.", "questions": "- About Figure 2, I think the point on the variance is clear. But what’s more important is the agreement of the relative order of data strategies: if a data strategy is better than another on a smaller model, is it also better on a larger model? Did the authors study that? Also, what are the blue dots? I guess they are the actual validation loss of different model sizes, but would be nice to make it clear in the caption.\n- The authors also write “they find $\\gamma_1 = 0.1$, $s = 0.5$, $\\gamma_2 = 0.1$” works well. How exactly is the process of finding those values?\n- The authors write in line 480: “they are accurate locally for much of training and thus can act as a learning signal for the data policy”. On what basis do the claims are made?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- Novelty. It’s a novel idea to improve over online data optimization (ODM, Albalak et al. (2023)) on 3 aspects. 1) Using scaling law to predict training loss on each domain and the derivative of the loss w.r.t. to the number of data points has intuitive interpretation as the loss reduction per data point. 2) The weighting also considers the number of data points used for a domain. 3) A combined weighting with the above two and temporal average.\n- Readability. The paper not only shows their method works but also spends a reasonable amount of space discussing perspectives, hence better readability. For example on the curriculum learning (Section 2) and on the requirements (Section 3, online and agnostic to downstream tasks) of data optimization for pretrained models.\n- Besides their proposed method, the authors also discover a strong baseline called “natural” policy, that depends on the number of tokens in each domain.", "weaknesses": "- Many heuristics are used such as fitting of scaling law (Section 3.1 eq1), credit assignment (Section 3.2, eq2 ), preference distribution (Section 3.3, eq3) and temporal average (eq4 and eq5). The authors try to motivate those choices from related works and intuitions, but only eq1 is adequately explained and validated. \n- Following the above, are all the heuristics (eq 1- 5) necessary? How important are they? This is a missing part in the paper. An ablation on them can validate if the invented heuristics are actually all useful.", "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)