AutoScale: Scale-Aware Data Mixing for Pre-Training LLMs
Authors: Feiyang Kang
Yifan Sun
Bingbing Wen
Si Chen
Dawn Song
Rafid Mahmood
Ruoxi Jia
Venue: COLM 2025 (2025)
Keywords: Large Language Models (LLM)
Data Curation
Domain Reweighting
Scaling Laws
Data-centric AI
TL;DR: We propose AutoScale, which automatically predicts compute-optimal data compositions for training LLMs at the target training data scale.

Abstract: Domain reweighting is an emerging research area aimed at adjusting the relative weights of different data sources to improve the effectiveness and efficiency of LLM pre-training. We show that data mixtures that perform well at smaller scales may not retain their advantage at larger scales, challenging the existing practice of determining competitive mixtures in small-scale experiments and *directly* applying them at much larger scales. To address this, we propose AutoScale, a two-stage, scale-aware data composition framework. First, AutoScale fits a parametric model that predicts the model’s loss under different data compositions, then uses it to find an approximate best allocation at smaller, more manageable budgets. Next, leveraging a novel theoretical analysis of how optimal compositions evolve with scale, AutoScale extrapolates that composition to larger budgets without further retraining. Empirically, AutoScale accelerates convergence and improves downstream performance.
For instance, when pre-training GPT-2 Large, it achieves a 28\% faster perplexity reduction than baselines and up to a 38\% speed-up over unweighted training, while yielding best-average results on various downstream tasks. Overall, our findings illustrate how domain importance shifts with training scale, underscoring the need for scale-dependent data curation in LLM training. 
Our code is open-sourced.

Published as a conference paper at COLM 2025
AutoScale: Scale-Aware Data Mixing for Pre-Training LLMs
Feiyang Kang"t Virginia Tech fykēvt.edu edu
Yifan Sun" UICC yi yifan500illinois.com
Bingbing Wen
Si Chen Virginia Tech
University of Washington
bingbweuw. edu
chensi@vt.en
Dawn Song UC Berkeley
Rafid Mahmood
Ruoxi Jia'
University of Ottawa & NVIDIA
Virginia Tech
dawnsong@gmaill.cm
mahmood@telfer. www..oottttt....
ruoxi jiavvt.eeee edu
Abstract
Domain reweighting is an emerging research area aimed at adjusting the relative weights of different data sources to improve the effectiveness and ef- ficiency of LLM pre-training. We show that data mixtures that perform well at smaller scales may not retain their advantage at larger scales, challenging the existing practice of determining competitive mixtures in small-scale ex- periments and directly applying them at much larger scales. To address this, we propose AutoScale, a two-stage, scale-aware data composition frame- work. First, AutoScale fits parametric model that predicts the model's loss under different data compositions, then uses it to find an approximate best allocation smaller, more manageable budgets. Next, leveraging a novel theoretical analysis of how optimal compositions evolve with scale, AutoScale extrapolates that composition to larger budgets without further retraining. Empirically, AutoScale accelerates convergence and improves downstream performance. For instance, when pre-training GPT-2 Large, it achieves 28% faster perplexity reduction than baselines and up to a 38% speed-up over unweighted training, while yielding best-average re- sults various downstream tasks. Overall, our findings illustrate how domain importance shifts with training scale, underscoring the need for scale-dependent data curation in LLM training. Our code is open-sourced!.
1 Introduction
Large language models (LLMs) are pre-trained on vast datasets sourced from diverse domains. However, the immense computational demands of this process, coupled with limited resources, create pressing need to enhance the effectiveness and efficiency of pre-training. A promising approach to address this challenge is through domain reweighting- adjusting the ratio (or weights) of data from different sources.
However, developing a principled and efficient framework for determining an optimal data mix remains challenging. Many industry pipelines still rely on trial-and-error heuris- tics (Rae et al., 2021; Grattafiori et al., 2024) or reuse domain weights designed for previous models (Mehta et al., 2024), without systematic approach for deciding how much of each domain to include. The seminal domain-optimizaaton work by Xie et al. (2024) attempted to upweight "difficult" domains, but later work (Fan et al., 2023) reported instability and only limited validation-loss improvements, partly because the chosen optimization objective does not robustly align with the model's ultimate test-time performance. Therefore, recent methods (Liu et al., 2024; Ye et al., 2024) focus on directly optimizing domain weights for lower validation loss. However, the highly complex relationship between domain weights and model performance makes such optimization expensive. Á common strategy of these works to reduce costs is train multiple times at smaller scales, identify "best" mix,
"Equal contribution. *Correspondence to: Feiyang Kang and Ruoxi Jia <fyk, ruoxijia>@vt.edu. htttps://gi https////giitub...caffiyggggg/// A
1

---

Published as a conference paper at COLM 2025
and then assume it transfers to large-scale pre-training. Yet our experiments show that compositions found at smaller scales may not remain competitive when training is scaled up, whereas directly optimizing at full scale is infeasible. This yields dilemma: either accept ssall-scale solutions that may not transfer or attempt large-scale optimization that is prohibitively costly.
To resolve this dilemma, we develop novel theoretical analysis that shows how the optimal domain composition evolves at different scales. Building on this insight, we propose two- stage framework, AutoScale, for scale-aware domain reweighting. In the first stage, we fit parametric model that predicts the model's loss under different data compositions, then use it to discover an approximate optimal allocation at smaller, more manageable budgets. Next, we apply our theoretical result extrapolate that allocation to larger budgels-without wwihhut re-optimizing at full scale-thus thus bridging the gap between small-scale optimization and full-scale pre-training.
Our experiments, conducted on both dnloder-only and encoder-only architectures, con- sistently that AutoScale speeds up convergence and yields favorable downstream task performance. For instance, in pre-training GPT-2 Large on the RedPajama dataset, our approach achieves a 28% faster perplexity reduction compared with any baseline and up to 38% speed-up over unweighted training, while also delivering the best downstream-tase performance. Moreover, we made surprising empirical observations that data sources tra- ditionally viewed as "high-quality" (e.g., Wikipedia and scientific papers) excel at smaller scales but exhibit sharp diminishing returns as the training grows. Meanwhile, domains containing more diverse examples (e.g., CommonCrawl) continue offering loss reductions at larger scales, underscoring the importance of scale-aware data curation.
Related Work
Principled training data curation for LLMs is an emerging research area, aiming to strate- gically select data that improves model performance. can be performed multiple levels-from token-level (Lin et al., 2024) or point-level (Wang et al., 2024) up to domain- level selection. Domain-level approaches are often more efficient because they operate coarser granularity, typically applying soft selection iee, upweighting or downweighting entire data domains.
Domain reweighting can generally be viewed as two-step process: (i) define an objective that captures the goal of improving test loss or other model performance measures, and (ii) optimize domain weights according to that objective. DoReMi (Xie et al., 2024), seminal paper in the space, adopted GroupDRO (Sagawa et al., 2019) as the objective, which implicitly upweights "difficult" domains. However, subsequent studies (Fan et al., 2023) found that the performance gains are unstable and limited, partly because the chosen objective does not align with the metrics ultimately used to evaluate the model at test time. Recent methods (Liu et al., 2024; Ye et al., 2024; Fan et al., 2023) attempt to directly optimize validation loss, which serves as closer proxy for the test-time model performance metrics we actually care about. We note that there can still be misalignments between validation loss and test-time performance metrics (Barton, 2024), which an active research area in itself, but validation loss remains widely accepted objective for model selection in pre-training.
With validation loss as the objective, the next challenge is how to optimize it. Evaluating the objective for a given set of weights is computationally expensive, as it requires training a
model from scratch on the weighted data set and evaluating the corresponding validation loss. Existing approaches tackle this in two broad ways. One line uses surrogate models to approximate the mapping from domain weights to model performance (Ye al., 2024; Liu et al., 2024); fitting such models can still require large amounts of retrainings. For instance, in Liu et al. (2024), fitting the surrogate requires more than ten times as many training runs as there are domains. Another line performs local approximations, assuming only single gradient step for the underlying model (Fan et al., 2023), which may not hold for practical learning rates.
Our work follows the surrogate-modeling line but differs by proposing a new parametric function to model performance versus domain weights, which can be reliably fit with only
2

---

Published as a conference paper at COLM 2025
about twice as many retraining runs the number of domains. Crucially, existing methods often conduct this optimization at a smaller data scale to keep cost manageable, then directly apply the small-scale "best" mixture at large scale. We show that this scale-invariance assumption can break down in practice. To address this, we contribute a novel theoretical analysis that characterizes how the optimal domain ratio shifts across scales, enabling us to extend small-scale insights effectively to larger budgets.
Finally, the general notion that data curation should be scale-dependent has appeared in prior works (Sorscher et al., 2022; Goyal et al.). However, Sorscher et al. (2022) argues this point using simplified analysis with perception model- that larger scales favor "harder" samples while smaller scales favor "easier" samples-bur but does not propose practical pipeline for scale-aware data selection. Our work addresses this gap by introducing concrete method for scale-dependent curation at foundation-mode! scale. Meanwhile, Goyal et al. focuses on the CLIP model (Radford et al., 2021), finding that different training epochs call for different data-selection thresholds. By contrast, our setting involves (i) LLM pre-training (with typically one epoch), and (ii) different data modality and selection granularity.
Perplexity Reduction (vs. unifoorm
3 Methodology
2.0
weights Wi weights W2
Evidence of scale dependence. A simple experi- ment illustrates that domain weighting is not one- size-fits-all: we derive two data mixes (with a proce- dure introduced later) and compare them at different training budgets. As shown in Figure 1, when tested at 0.3B tokens, Mix beats Mix B as measured by validation perplexity reduction compared against uniform weights, but at 1.2 B tokens Mix B outper- form Mix A. This flip indicates that a domain which
1.8
1.6
1.4
1.2
1.0
helps more at a smaller scale may not remain a bet- ter choice at larger scale, while another domain initially less impactful can become more valuable as training grows. Consequently, scale-aware ap- proach is needed so that domain weights adapt more tokens are introduced.
1.2B Training Data Scale/billion tokens
0.3B
Domain weights that excel one scale may underperform at another. Weights wi and w2 are obtained by run- ning DDO (as introduced in Section 3.2.1) 0.3B and 1.2B, respectively.
Figure 1:
3.1 Problem Formulation
To capture how domain importance shifts with the total training budget, we now formalize the scale-dependent domain reweighting problem. This framework lets us solve for a better domain mixture at any given data scale.
Notations and setup. We consider m domains (Di,..., Dm each with a large pool of training examples. A domain mix is specified by a weight vector w [w]..... Wm] on the probability simplex l := (w E R" Σm = 1, 2 for all i}. Given a total budget of N tokens, let N; [w; NJ be the number of tokens chosen from domain D,. We denote this resulting dataset by S(N, w)
{S,,..,, Sm. where S; S. | |S N: Training model 0 on S(N, w) means solving an empirical risk minimization (ERM) objective: 0* (NNNw)
arg mino ((, (N,,)), where L is next-token prediction loss.
Objective function. We assess a domain mix w by measuring the validation loss _ (N,,)) L( (N, w), D on held-out dataset DP, We then seek the mix w that minimizes this validation metric at scale N:
w
= arg min wsW
(N,((N))))
(1)
Because 0' depends on w via ERM, this becomes a bi-level optimization problem. Since closed-form expression exists for 0t, any gradient-based approach must approximate
3

---

Published as a conference paper at COLM 2025
0. ("" w)). Traditional bi-level methods rely on higher-order derivatives with respect to model parameters for this approximation, but such techniques become infeasible at the scale of modern foundation models (Liu et al., 2021).
3.2 Our Solution
We propose two-stage framework, AutoScale, for finding scale-aware data compositions, which first approximates the optimal data mix at small scales and extrapolates to a larger target scale:
1. Direct Data Optimization (DDO): At smaller, computationally feasible scales, learn a mapping from domain weights to validation loss. This reduces the original bi-level problem
single-level convex optimization-letting us approximate the "best" domain mix for that smaller budget.
2. Optimal Mix Projection: Building on theoretical analysis of how domain im- portance changes with total tokens, we then extrapolate those small-scale DDO solutions to a larger data budget.
3.2.1 Direct data optimization
Direct Data Optimization (DDO) is practical method for approximating the solution to the bi-level domain- problem at relatively small data scales. The key idea of DDO is to approximate the validation loss L (" (N, w)) as parametric function of the domain-weight vector w. This effectively reduces our bi-level objective (choose w while also training 0) to a single-level optimization, which can be solved efficiently via standard gradient-based methods.
We begin by noting that the validation loss can be represented by scaling laws as function of data size for each individual domain. We model the dependence of validation loss on the size of data from each domain, then aggregate these functions to derive the final approximation for the validation loss on w.
Drawing inspiration from neural scaling laws-which indicate a power-law relationship between training data scale and validation loss (Kaplan et al., 2020) assume that validation loss as function of domain i's data size follows
L C" (N, w))
(Ng + w. N) 7: 4.
Here, wi denotes the fraction of the total token budget N allocated to domain i. The term Ný represents an the "equivalent data size" contributed by domains other than i, while Y governs how quickly domain reaches a point of diminishing returns. Lastly, e represents the irreducible term in the loss function.
To learn these parameters [Ní, Yir li for each domain i, we retrain the model after perturb- ing w, upward and downward, measure the change in total validation loss, and then fit (Ng + w0 - + E; via least squares. Because L aggregates the effects of data size for each domain, our final approximation for the validation loss is:
m
w w. N) -7
(N,
Σ
No +
+ P,.
Once we have fitted the parameters, we can directly optimize over w subject to E," w to approximate the optimal domain mix under the total token budget N.
Because DDO only requires retraining at (2m + 1) mixes (one baseline plus up/ up/dows down perturbations for each of the m domains), it is far cheaper than naive zero-order method that retrains the model at every weight update. Nevertheless, DDO is best suited for moderate domain counts (m) and data scales (N). For much larger target scales, introduce a second stage that extrapolates the "best" DDO mix from smaller scales to significantly bigger budgets, all without additional retraining.
4

---

Published as a conference paper at COLM 2025
3.2.2 Optimal mix projection
Our method for extrapolating domain mixes to much larger training budgets hinges on novel theoretical result that we developed to characterize how the optimal mix ratio depends on the total data scale. Note that all inverses (-) products, and exponentiations on vectors below are understood elementwose.
Theorem 1: Scale-Dependent Optimal Composition
Consider the optimization problem
MK min N
MN
N = N
where 8; 2 O and Yi 2 for all i, and N = (N...., Nm. denotes the domain allocations. Let N* (N) be the optimal allocation that minimizes the sum above for total budget N. For two distinct budgets N(1) * N and any larger budget N(3),
suppose there is constant k > 0 such that
N(N(O)) = (N(2) (NO()) (N())) with Σ ((( = N133
Then N ((N(3)) is also the optimal allocation for the budget N133), i..,
H!
mt Σ N; = N(3I N (N)3)).
N(N(3)) = arg min
Σ
B N,
"
Proof overview (high-level). At optimality, the first-order (KKT) conditions impose that each domain's ' partial derivative of the loss matches up to a single Lagrange multiplier. From this, we can derive how each domain's 's optimal allocation N; (N) scales when transitioning from budget (1) to another budget N2(2) These domain-by-domain scaling factors do not depend on the absolute size of N, only on the relative shifts between domains, which in turn yields an exponential-style expression for the optimal allocation at third budget N3)), Thus, once we know the optimal allocations two budgets, we can directly construct the optimal allocation for any larger budget without re-solving the entire optimization.
Interpretation of the theory. The statement above assumess each domain D; contributes
B; N; T! independently to the total validation loss. In Appendix B.2, we generalize this to cases where domains may overlap by treating the evaluation as composed of multiple "latent skills" (Tiong et al.); the same exponential-style scaling behavior still emerges.
We defer the full proof Appendix B.1 (where we employ first-order optimality /KKT conditions), but the key insight is that domains saturate at different rates depending on their exponents Yi Specifically, domain D, with a small 7i saturates more slowly and thus continues to yield benefits at larger budgets, receiving an increasingly bigger fraction of tokens as grows. In contrast, large y indicates that D, quickly saturates, so it is favored primarily at smaller scales.
Concretely, this means the "optimal mix ratio" is not constant across scales. As the total budget N increases, domains with smaller 7 allocated larger share. The theorem's 's exponential-style update precisely captures these changing allocations, enabling us to predict the best mix at higher budget given the solutions two smaller buthets-without re-solving the entire optimization problem.
3.2.3 Overall algorithm
Having established that the optimal domain mix varies predictably with training budget, now summarize AutoScale, our proposed two-stage approach to optimize data mix.

---

Published as a conference paper at COLM 2025
Stage 1 (pseudocode provided in Algorithm 1): N(1 < N(2)), where retraining the model still affordable. Run DDO to obtain optimal allocations N* (N(1) and No (.
Pick two feasible scales N(1) and N(2) (with
Stage 2 (pseudocode provided in Algorithm 2): Leveraging theoretical result, we automatically predict the optimal domain mix for any larger scale. Specifically, starting from the optimal domain allocation No No (N(2)), we repeatedly "scale up" by multiplying by
N
Each such update yields new allocation at larger budget
than before. We continue until reaching or exceeding the target budget Nést, By adjusting the resolution 8, we control the granularity of these updates, thus reaching the target scale with any desired accuracy.
Evaluation
Our evaluation aims to address the following questions:
(RQ1) Does DDO yield better domain weighting at smaller scales? In our approach, DDO is designed to approximate the best data mix at given scale. While we cannot verify its optimality, we want to see if DDO meaningfully improves domain weighting compared to baseline methods (Section 4.1).
(RQ2) Can AutoScale-DDOO at smaller scales plus our theoretical projection to larger training efficiency and performance benefits when direct DDO at large scale is prohibitively "npenside((Section 4.2)
Overview of experimental settings. We provide an overview here and defer the full details to Appendix C. (I) Models and datasets. We experiment with two architectures GPT- and BERT (110M), training on up to 10B tokens. While this budget is comparatively small for the latest LLM regimes, it already exceeds the data scales used in many existing domain-
2 Large (774ppaameest
studies (Fan et al., 2023; Chen et al., 2024a), serving as a feasible testbed for prototype ideas in non-commercial settings. Specifically, we pre-train GPT-2 Large on the RedPajama dataset (Computer, 2023),which spans seven domains (e-g., Common Crawl, C4 (Raffel et al., 2020), GitHub, Wikipedia, ArXiv, StackExchange). For Reviews, Arxiv, Books, Wikipedia, and Open WebText (Gokaslan & Cohen, 2019). (II) Baselines. We compare our methods against several baseline strategies. Uniform samples data from each domain, leading the same token count per domain. Llama weights are curated set of heuristically tuned domain weights from the LLaMA-1/2 models (Touvron al., 2023). DoReMi (Xie et al., 2024) is seminal paper in this domain-reweightinng space, offering an early, principled approach to finding domain weights. Data mixing law (Ye et al., 2024) and RegMix (Liu et al., 2024) represent the latest ofate-of-the-art. (III) Metrics. We measure test perplexity and also evaluate downstream performance to confirm that improvements extend to practical tasks.
BERT, we use data from five nources-Amazee
4.1 Evaluating DDO
Effectiveness of DDOOoptimized weights. We perform DDO on GPT-2 Large at two different data scales (0.3B and 1.2B tokens) to obtain DDO weights specifically optimized for each scale. We then retrain the model under these DDO-derived weights and compare the evaluation loss against two baselines: (1) Uniform (no reweighting), and (2) RegMix, the latest state-of-the-art approach. (We omit DoReMi here, as has been surpassed by RegMix (Liu et al., 2024).) For each set of weights, the model is trained at both 0.3B and 1.2B tokens, with results in Table 1. At both scales, DDO-optimized domain weights significantly outperform the Uniform baseline, achieving a notably lower evaluation loss. DDO-optimized weights also surpass RegMix when models are trained at the same scale as the demain-weight optimization, indicating that DDO finds more effective domain weights than RegMix. Further, RegMix does not consider adaptation for training models at different scales. Applying RegMix optimized weights on larger data scales appears less effective,

---

Published as a conference paper at COLM 2025
evident by the widening gap between its performance from DDO's. Notably, the DDO- derived weights yield the strongest gains at the scale for which they were optimized, while showing less advantage when used at different scale, highlighting the scale-dependent nature of domain weighting.
Weights/Actual training scale
0.3B training tokens
1.2B training tokens
Uniform Weights
48.04 46.56 (-1.48)
28.11 27.86 (-0.25)
RegMix Weights (optimized at 0.3B)
DDO Weights (optimized at 0.3B)
46.13 (-1.91)
27.09 (-1.02)
DDO Weights (optimized at 1.2B)
46.31(-1.73)
26.97 (-1.14)
Table 1: GPT-2 Large trained with DDO optimized domain weights achieve significantly reduced test perplexity compared to with non-optimized, uniform weights, also outperforming RegMix. DDO optimized weights appear most performant the data scale they were optimized.
In addition, we apply DDO to BERT at 0.3B tokens; the resulting model performance from the DDO-optimized weights is shown in Fig. 2. These weights reduce the model's validation loss on all training domains and on held-out non-training domains, demonstrating DDO's effectiveness in improving training efficiency. Furthermore, when evaluated on the GLUE benchmark and the SQuAD dataset, the DDO-optimized weights also yield a notable improvement in downstream task performance.
Velidation Reweighting
Task Performance after Reweighting
Domein expiloss) (lower is better)
70-
Uniform Weights
DOO Optimized
Task Performance (ligher is better)
improved
Uniform Weights
Optimized
0.7
improved I
0.6
40-
30
0.3
20-
0.2
amazon andi
In-distribution Domains
books
wiki
owtc
news pubmeddthhub Out-of-distribution
09
138
6872
mmpc
mmll
sbub
GLUE
avg
SQWAD
Downstream
Tasks
(a) Validation Loss lower is better)
(b) Task Performance (f higher is better)
Figure 2: Optimizing domain weights with DDO algorithm for pre-training Encoder-only LMs (BERT).
DDO substantially reduces validation loss. After reweighting, all training domains' loss has decreased or remained unchanged. Out-of-domailo loss on non-training domains also decreased considerably. Enhanced performance is observed on all GLUE tasks (eval metric: cola: Matt. corr, stsb: Pearson corr., rest: acc.) and SQuAD (acc.).
Analyzing DDO's effectiveness. Recall that the key idea of DDO is use a power-law- based parametric function to predict validation loss from domain weights. A major factor in DDO's effectiveness lies in the accuracy of this function. We quantify its predictive power via the average absolute relative error (AAR) between the predicted and actual losses. In our experiments, the AAR is 1.00%, indicating that DDO's modeling closely reflects actual loss.
4.2 Evaluating AutoScale
Effectiveness of our extrapolated weights. Recall that AutoScale is a two-stage pipeline: first, run DDO at smaller scales to identify domain weights, then extrapolate those weights to a larger scale. We call the resulting allocation the AutoScale weights. For GPT-2 Large, we run DDO up to 0.6B tokens, then extrapolate to 3B and 10B tokens. Figure 3 shows the change of test perplexity during training for models trained with 10B tokens using AutoScale weights versus baseline allocations. AutoScale consistently outperforms every baseline by a 28-38% margin and also demonstrates advantageous downstream perfor- mance. Table 2 demonstrates the results 3B tokens, revealing that AutoScale maintains its superiority in both final loss achieved and faster convergence. Table 3 examines domain-wise test perplexities; AutoScale weights significantly reduce the loss on the Books domain and
7

---

Published as a conference paper at COLM 2025
improve worst-domain perplexity, also yielding better average across domains. Finally, Ta- ble 4 evaluates eight downstream tasks. The model trained with AutoScale weights achieves the best overall performance, further underscoring the effectiveness of our extrapolated domain weights.
For BERT, we train up to 288k steps (approximately 120% of the original BERT-base bud- get (Devlin et al., 2018)). Table 10 shows that, compared to uniform (no reweighting), AutoScale yields 16.7% speed-up most data scales and 10% speed-up at the largest scale, demonstrating consistent effectiveness. However, these gains are smaller than those observed for GPT-2 Large, indicating that different architectures and training objectives may respond differently to domain reweighting. This is also hinted at in Figure 10, where the evaluation loss shows a more uniform response to each domain, suggesting fewer benefits from reweighting in BERT's setup.
AutoScale Weights v. Baselines
20
Few-shot Performance (avg of 9 tasks)
AutoScale
AutoScale
0.45
Test perplexity (PPL)
19
DoReMi
DoReMi
LLaMA
Few Rew-shot Performance
0.44
LLaMA
Uniform
18
Uniform
0.43
31% faster
38% faster
0.41
28% faster
15
40000
60000
80000
100000
120000
0.40
Domain Weights
Training Steps
Figure 3: Training 774M Decoder-only LMs (GPT-2 Large) for 10B tokens (96k steps). AutoScale- predicted domain weights decrease test perplexity at least 28% faster than any baseline with up to 38% speed up, achieving best overall task performance.
Weights
Final Perplexity (PPL)
AutoScale Speed Improvement (% steps saved to final PPL)
(3B training tokens)
AutoScale (ours)
21.123
DoReMi
21.676
25%
Data Mixing Laws
23.333
37%
LLaMA
22.944
31%
RegMix
21.740
28%
Uniform (30% more tokens)
21.736
37%
Table 2: Domain perplexity for 774M Decoder-only LMs (GPT-2 Large) trained for 3B tokens. Au- toScale -predicted weights decreases val loss at least 25% faster than any baseline with up to 37% speed up. Despite LLaMa weights being very different from uniform weights, they yield highly similar training efficiency at these data scales.
Domain/Method
AutoScale (ours)
DoReMi
Mixing Laws
LLaMA
RegMix
Uniform (30% more takens)
Common Crawl
25.598
24.116
30.824
21.464
24.430
28,351
Github
7.482
6.678
5.845
7.376
6.145
5.784
Books
29.162
33.324
34.450
35.533
32,985
31.140
Wikipedia
18.828
17.154
26.795
21.110
21L.777
19.570
34.242
39.429
38.521
37.393
39.654
40.323
Stack Exchange
15.991
15.393
14,519
20.133
15.225
13,890
Anxv
16.558
15.638
12.372
17.598
13.563
13.082
Average
21.123
21.676
23.333
22.944
21.740
21.736
Worst-domain
34.242
37.393
39.654
40.323
Table 3: Domain perplexity for 774M GPT-2 Large trained for 3B tokens. AutoScale notably achieves the lowest average test perplexity while also significantly decreasing worse-domain perplexity.
Examining how domain importance evolves with scale. To illustrate the shift in domain importance, we first run DDO GPT-2 Large across scales ranging from 30M to 1.2B tokens.
8

---

Published as a conference paper at COLM 2025
Method/Task
Avg
pubmedqa
piça
hellaswag (10-shot)
crows, pairs _english
boolg
anc_easy
truthfulqa
hellaswag
_mc2
(zero-shot)
AutoScale (ours)
0.4786
0.536
0.6202
0.3021
0.5850
0.6141
0.3977
0.4385
0.3030
Uniform Weights
0.4514
0.2923
0.5886
0.5636
0.3742
0.4526
LLaMA Weights
0.4585
0.492
0.6055
0.2944
0.5903
0.5612
0.3956
0.434
0,2952
Mixing Laws
0.5778
0.6162
0.3771
DuReMi RegMix
0,4482
0.468
0.5985
0,2886
0.5742
0.5410
0.3750
0.4505
0,2896
0.526
0.3721
Table 4: Task performance for 774M GPT-2 Large trained for 3B tokens. Models trained with
AutoScale predicted weights achieve the best overall performance across the tasks.
Scale/steps
36k
144k
288k
Final Loss (exp)
38.32
16.94
10.97
8.13
6.30
Steps Saved
5k (28%)
5k (14%)
10k
(14%)
20k (14%)
20k (10%)
Table 5: AutoScale notably improving training efficiency for BERT models on all scales- even for considerably large scale, 288k steps, the speedup margin remains visible.
Figure 4(a) shows that the DDO-optimized weights differ visibly at each scale, highlighting clear shifting pattern. Data sources with more standardized formats (Wikipedia scientific papers) often regarded as "high quality"-dominae" diminishing returns as the data budget grows. By contrast, domains with more diverse examples (C4, CommonCrawl) continue to lower training loss even at higher scales.
at smaller scales but exhibit sharp
Consistently, taking DDO-optimized weights from up to 0.6B tokens, we use our theory to project how the composition would shift at scales beyond 1.2B. Figures 4bd) and show that as the training data scale grows, diverse domains (C4, CommonCrawl) command larger share of the mix compared to "standard" domains. We observe a similar pattern with BERT, where extrapolate the DDO-optimizer that domains like WebText and Amazon Reviews gain significance over clean, standardized data (Wikipedia, Arxiv) (see Fig. 11). A plausible explanation is that "diverse" data provides broader topical coverage and linguistic styles, while "standard" data saturates more quickly.
weights at 0.5B tokens to even larger scales, revealing
2014 tnkers
0.30 TDARTS
0,5
0.25
Domain Weights
214
Domenn Noight
stack exchange endv
0.15
0.10
0.00
common glthb books
wiki
C4
stack
andlv
crawl
eschange
Data Domains
Data Scaledillion
(a) DDO optimized domain weights.
(b) AutoScale projected domain weights.
Figure 4: Domain importance evolves with training data scales. (GPT-2 Large)
Note that these trends show how our approach predicts domain importance may evolve, not a proof that each extrapolated mix guarantees the best performance at its target scale. Nonetheless, the consistent shifting patterns across GPT-2 Large and BERT reinforce the idea that domain importance is scale-dependent.
Conclusions
This paper explores how the importance of each training domain shifts across different scales and proposes scale-aware framework (AutoScale) that outperforms existing approaches across various architectures, datasets, and training scales. Still, our experimental settings remain limited in scale and the diversity of evaluations. Extending this work to larger

---

Published as a conference paper at COLM 2025
training budgets, additional data modalities, and broader benchmarks would clarify how well our insights generalize beyond the current scope. Another exciting next step is to adapt AutoScale for directly optimizing downstream metrics, moving beyond perplexity as a rough proxy for language-model quality.
6 Discussions
Online vs. offline data mixing methods. These represent two independent paths to approach data mixing problems in LM pre-training. Current industrial practices are typ- ically offline. Offline methods, such as DoReMi (Xie et al., 2024), Doge (Fan et al., 2023), RegMix (Liu et al., 2024), DML (Data Mixing Laws, Ye et al. (2024)), and our proposed AutoScale, decouple data curation from the final training process. This separation allows for the distribution of model development workloads across different teams (e.g., data, model, and training) working in parallel. This approach is common in industrial settings for developing state-of-the-arr LMs (e.g,, Phi-4 (Abdin et al., 2024), Llama4 (Meta, 2025)). Notably, many offline methods like DoReMi, DML, and BIMIX (Ge et al., 2024) originate from industrial research, whereas current online methods (ODM (Albalak et al., 2023), Skill-it (Chen et al., 2024a), Aioli (Chen et al., 2024b), ADO (Adaptive Data Optimization, Jiang et al. (2024))) are primarily academic contributions.
Offline methods treat the training process integrally, primarily adjusting domain weights and training budgets before training commences. In contrast, online methods (e.g., ODM, Skill-it, Aioli, ADO) adjust domain weights dynamically during the training process. ODM, for instance, models online domain weighting as a bandit problem and updates domain weights based on local gradients. Considering the learning rate often changes throughout training, the loss landscape may be non-convex where a greedy approach cannot guarantee global optimality. Skill-It primarily considers fine-tuning scenarios, while Aioli extends this to small-scale LM pre-training relationship between each pair of training and validation domains (in contrast to DML's offline estimation), while ADO considers scaling laws for individual training domains. The scaling law relationships these methods were originally proposed for adjusting training data size within a fixed training pipeline (Kaplan et al., 2020; Hoffmann et al., 2022). They may not always accurately track training process and its dynamics, especially with varying learning rate, complicating their conceptual grounding. This creates space for future research contributions. For deployment in production pipelines, better understanding on the predictability of the outcome under different training paradigms is desirable.
Aioli performs online estimation of the scaling law
Effect of repeated documents and 'data--onstraiiee setting" The data constraint is relevant for real-world LM pre-training, particularly for low-resource domains. Beyond con- straints on domain weights, practical limitations include the amount of available data and the maximum number of repetitions tolerated before significant performance degradation. Existing research on domain weighting has not explicitly modeled this factor. Despite this research gap, Muennighoff et al. (2023) offers positive evidence that current paradigms can still be effective under such constraints. The paper indicates that repeating data for up to 3 epochs has almost no negative impact, and models can still train reasonably well with up to 7 epochs of repetition. This suggests that in practice, repeating data from some domains a few times can be acceptable without significant detriment to performance, Current data mixing methods may be applied to these scenarios without major modifications.
Data mixing for model-size scaling Applying the data- sataa-scaling methodology to model-size scaling appears plausible. Scaling law papers (Kaplan et al., 2020; Hoffmann et al., 2022) reported similar scaling trends between validation loss and training data budget/ model size. (Shukor et al., 2025) proposes to model the combined effects of data and model-size scaling as a linear addition and validates its effectiveness in optimizing domain weights for LLM, native multimodal model (NMM), and large vision models (LVM) pretraining.
Additional discussions on Intuition for Stage 2 of AutoScale, Optimal mix projection and Number of predictors (scaling law components) can be found in Appendix D.
10

---

Published as a conference paper at COLM 2025
Impact Statement
Reducing the complexity and resource requirements associated with pretraining LLMs, AutoScale contributes to the democratization of AI. Smaller organizations, academic institu- tions, and individual researchers can more easily participate in cutting-edge AI research and development, fostering innovation and collaboration across the Al community, Moreover, learning from massive amounts of data requires large and costly computational resources, which not only consume substantial energy but also generate a significant carbon footprint, contributing to environmental issues. Furthermore, these resources quickly become obsolete due to the rapid pace of technological advancements, leading to e-waste. This research makes contributions to mitigating these issues by improving the efficiency of resource utilization in Al training.
Acknowledgement
This work is supported in part by the National Science Foundation under grants IIS-2312794, IIS2313130, OAC-2239622, Amazon-Virginia Tech Initiative in Efficient and Robust Machine Learning, AWS computational credits, and the Commonwealth Cyber Initiative. The authors are grateful for Ankit Battawar and Alix Delgado from AWS, whose dedicated help and support were crucial for securing computing resources and implementing empirical studies.
11

---

Published as a conference paper at COLM 2025
References
Abdin, M., Jn Behl, H., Bubeck, S., Eldan, R., Gunasekar, S., Harrison, M., Hewett, R. J., Javaheripi, M., Kauffmann, P., al. Phi-4 technical report. ariio preprint arXiv:2412.08905, 2024.
Albalak, A., Pan, L., Raffel, C., and Wang, W. Y. Efficient online data mixing for language model pre-training. arXiv preprint ar arXio:2312.02406 2023.
Barton, I. Calibrating mosaic's evaluation gauntlet. https: /www. databricks com/ hom/llgg/// calibrating mosaic m evaluatiion gauntlet, 2024. [Online; accessed 03-28-2025).
Bisk, Y. Zellers, R., Gao, J., Choi, Y., et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, Pp. 7432-7439, 2020.
Chen, M., Roberts, N., Bhatia, K., Wang, J., Zhang, C., Sala, F,, and Ré, C. Skill-it! data- driven skills framework for understanding and training language models. Advances in Neural Information Processing Systems, 36, 2024a.
Chen, M. F., Hu, M. Y., Lourie, N., Cho, K., and Ré, C. Aioli: A unified optimization framework for language model data mixing. Xiv preprint arx:::2411 Ktt::241..05735, 2024b.
Clark, C., Lee, K., Chang, M.-W., Kwiatkowski, T., Collins, M., and Toutanova, K. Boolq: Exploring the surprising difficulty of natural yes/no no questions. arv preprint Xii:19955.1044 2019.
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the ai2 reasoning challenge. arxiv preprint ar Nlt:1803.05457, 2018.
Computer, Redpajama: An open source recipe to reproduce llama training dataset, 2023. Data.
URL https::/githhb.ccmm wtw.:tttt..cooooottetteeeeeeeeee htwp
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint wxXi::1100004805 2018.
Fan, S., Pagliardini, M., and Jaggi, M. Doge: Domain reweighting with generalization estimation. arv preprint ar arXio:2310.1539, 2023.
Gadre, S. Y., Smyrnis, G., Shankar, V., Gururangan, S., Wortsman, M., Shao, R., Mercat, J., Fang, A., Li, J., Keh, S., et al. Language models scale reliably with over-training and on downstream tasks. Xrviv preprint arX:::2403.08540, 2024.
Gao, L., Tow, J., Biderman, S., Black, S., DiPofi, A., Foster, C., Golding, I.,, Hsu, J., McDonell, K., Muennighoff, N., et al. A framework for few-shot language model evaluation. Version v0. 0.1. Sept, 2021.
Ge, C., Ma, Z., Chen, D., Li, Y., and Ding, B. Bimix: A bivariate data mixing law for language model pretraining. arXiv preprint arXiv:2405.14908, 2024.
Gokaslan, A. and Cohen, OpenWebTextCCrpus, 2019. 2019.
Openwebtext corpus. http://Skylionee) .11007..tthuu -it.gt..hhh.... io/
Goyal, S., Maini, P., Lipton, Z. C., Raghunathan, A., and Kolter, J. Z. The science of data filtering: Data curation cannot compute agnostic. In ICLR 2024 Workshop on Navigating and Addressing Data Problems for Foundation Models.
Grattafiori, A., Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Vaughan, A., et al. The llama herd of models. aiv preprint MxXi::240.211783 2024.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. ar Xiv preprint ar X20: 15556, 2022.
12

---

Published as a conference paper at COLM 2025
Jiang, Y., Zhou, A., Feng, Z., Malladi, S., and Kolter, J. Z. Adaptive data optimization: Dynamic sample selection with scaling laws. arXiv preprint ar orii:24400.11200, 2024.
Jin, Q Dhingra, B., Liu, Z., Cohen, W. W., and Lu, X. Pubmedqa: A dataset for biomedical research question answering. arXiv preprint ar arXiv:1909.0611111 2019.
Kaplan, J., McCandlish,,, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford,
A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 2020.
Lin, S., Hilton, J., and Evans, O. Truthfulqa: Measuring how models mimic human false- hoods. Xrv preprint ar arXio:2109.07958 2021.
Lin, Z., Gou, Z., Gong, Y., Liu, X., Shen, Y., Xu, R., Lin, C., Yang, Y., Jiao, J., Duan, N., et al. Rho-1: Not all tokens are what you need. arXiv preprint ar xio::2404.07765, 2024.
Liu, Q,, Zheng, X., Muennighoff, N., Zeng, G., Dou, L., Pang, T., Jiang, J., and Lin, M. Regmix: Data mixture regression for language model pre-training. arXiv preprint arXiv:2407001492, 2024.
Liu, R., Gao, J., Zhang, J., Meng, D., and Lin, Z. Investigating bi-level optimization for learning and vision from unified perspective: survey and beyond. IEEE Transactions on Pattern Analysis and Machine Intelligence, 4A(12):0045-10007,, 2021.
Mehta, S., Sekhavat, M. H., Cao, Q,, Horton, M., Jin, Y., Sun, C., Mirzadeh, I., Najibi, M., Belenko, D., Zatloukal, P., et al. Openelm: An efficient language model family with open-source training and inference framework. ar preprint ar Xrx::22404 14419, 2024.
Meta, A. The Ilama 4 herd: The beginning of new era of natively multimodal ai innovation. htttps//ai. meta.
checked on, 4(7):2025, 2025.
Muennighoff, N., Rush, A., Barak, B., Le Scao, T., Tazi, N., Piktus, A., Pyysalo, S., Wolf, T., and Raffel, C. A. Scaling conaa-conttained language models. Adnances Neural Information Processing Systems, 36:50358-50333 2023.
Nangia, N., Vania, C., Bhalerao, R., and Bowman, S. R. Crows-pairs: A challenge dataset for measuring social biases in masked language models. arXiv preprint ar xix::2010.00133,, 2020.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8763. PmLR, 2021.
Rae, J. W., Borgeaud, S., Cai, T., Millican, K., Hoffmann, J., Song, F,, Aslanides, J., Henderson, S., Ring, R., Young, S., et al. Scaling language models: Methods, analysis & insights from training gopher. ar Xiv preprint ar axX::2212..11445 2021.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, J. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21(140):1-67, 2020.
Rajpurkar, P., Zhang, J., Lopyrev, K., and Liang, P. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint ar arXiv:1606.05555 2016.
Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang, P. Distributionally robust neural net- works for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint ar 11111991,.,8711, 2019.
Shukor, M., Bethune, I.., Busbridge, D., Grangier, D., Fini, E., El-Nouby, A., and Ablin, P.
Scaling laws for optimal data mixtures. arXiv preprint ar Xiv:2500 .04044, 2025.
Sorscher, B., Geirhos, R. Shekhar, S., Ganguli, S., and Morcos, A. Beyond neural scaling laws: beating power law scaling via data pruning. Advances in Neural Information Processing Systems, 15-19523-19534 2022.
13

---

Published as a conference paper at COLM 2025
Srivastava, A., Rastogi, A., Rao, A., Shoeb, A. A. M., Abid, A., Fisch, A., Brown, A. R., Santoro, A., Gupta, A., Garriga-Alonso, A., et al. Beyond the imitation game: Quantifying and extrapolating the capabilitiis of language models. arXiv preprint ar arXio:2206.04615, 2022.
Tiong, A., Zhao, J., Ji, J, Hoi, S., Xiong, C., and Li, B. Toward data-driven skill identification for general-purpose vision-language models. In ICLR 2024 Workshop on Navigating and Addressing Data Problems for Foundation Models.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation language models. arXiv preprint ar arXio:2302.13977, 2023.
Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. Glue: A multi-task benchmark and analysis platform for natural language understanding, In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, Pp. 353-355, 2018.
Wang, J. T., Wu, T., Song, D., Mittal, P., and Jia, R. Greats: Online selection of high-quality data for 11m training in every iteration. Advances in Neural Information Processing Systems, 37:131197-131223, 2024.
Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., et al. Huggingface's s transformers: State-of-thh-art natural language processing. Xiv preprint ar arX1o:1910.03771, 2019.
Xie, S. M., Pham, H., Dong, X., Du, N., Liu, H., Lu, Y., Liang, S., Le, Q. V., Ma, T., and Yu, A. W. Doremi: Optimizing data mixtures speeds up language model pretraining. Advances in Neural Information Processing Systems, 36, 2024.
Ye, J,, Liu, P., Sun, T., Zhou, Y., Zhan, J., and Qiu, X. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. arXiv preprint arXiv:2403.16952, 2024.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can a machine really finish your sentence? Xiv preprint Xiv::1995.00883, 2019.
14

---

Published as a conference paper at COLM 2025
Appendices
A Algorithms and Operational Pipeline
16
Proofs for Section 3.2.2, Optimal mix projection
18
B.1 Theorem 1: Scale-Dependen Optimal Composition
18
B.2 Scaling Latent Skills
20
C Experimental Details and Additional Results for Section 4, Evaluation
23
Experimental Details on GPT-2 Large (774M)
23
C.2 Experimental Details on BERT (110M)
24
C.3 Implementation Details for Baselines
25
C.4 Evaluation Details
26
C.5 Additional Results on GPT-2 Large (774M)
27
C.6 Additional Results on BERT (110M)
27
Extended Discussions
28
15

---

Published as a conference paper at COLM 2025
A Algorithms and Operational Pipeline
Algorithm 1 Direct Data Optimization (DDO)
Require: m domains (data sources) with data D1 ... Dm, data budget No (< for full-scale
training), training dataset S, model parameters 6, validation loss L,, perturbation ratio
y> 1 (e-g.r
Initialize weights for all domains vi E "I,...m]:: wi + 1/m;
Initialize training data for all domains Vi E {1, m): sample S; D; where |S w; NN
Train the model on data {S ... ) and evaluate its loss C +- (o (5)); for from 1 to m do
w 4 10;;
Perturb domain weights (+)
Resample s/t C Dj where | = w -N; Train the model on data S = (1S) Sm \ Sj) ust and evaluate its loss 2 +
C. (O));); " + 7
> Perturb domain weights (-)
Resample S, C Dj where |s; = w N;
Train the model on data (15) ... Sm) \ 5/) U and evaluate its loss Lj (0))))))
OLS fit for scaling functions Ný Yi, lj arg min Ng. C9 (Ng + N:) e,1
[C++)) - (Ng + N L/] [C(-)) - (Nj N;) E
end for
Output optimized domain weights w arg w'EW" Σ E (Ng wt N
Algorithm 2 AutoScale
Require: Optimal domain weights (obtained from DDO) (I)) at data scale N(1) and (2)+
at data scale N(2), target data scale N(t), where N(1) N(2) < N(1); resolution 5.
Optimal domain data N (N())) w(])+ NI);
Optimal domain data N (2 < (2)+ N(2);
Current data budget N + Σ N(
Optimal domain data under current data budget NO (N < N'(22))); while N < NO) do Compute optimal domain data under the next data budget: N N*(Nnext) (N(()) (N(2)) Compute the next data budget Nmext +- N Update current data budget N +- Nmext,
+
end while
Output predicted optimal domain weights: (()++ <- (N)/NN.
Operational Pipeline (DDO)
1. Train base proxy model with uniform weights (or reference weights, if available);
2. At each time, add/reduce reduce data quantity for one domain and re-train the proxy model;
3. Fit power law scaling functions and solve the optimization problem;
4. Iterate the process if necessary.
16

---

Published as a conference paper at COLM 2025
Operational Pipeline (AutoScale) )
1. For two smaller training data scales N(1) and N where re-training the model is affordable, find their corresponding optimal training data compositions N (N())
and N (N(2) using DDO Algorithm described above;
2. Initialize current data budget at N = N(
3. With the chosen resolution predict the next optimal training data composition as
N (Nmex)) = (*) (NI (N yielding optimal domain weights
w = Nj (Nmex)) / Ninxl at new training data scale Nmext = Σ: Nj (Nnext);
4. Update current data budget to N Nmext Repeat this process until the target training data scale is reached.
17

---

Published as a conference paper at COLM 2025
B Proofs for Section 3.2.2, Optimal mix projection
B.1 Theorem 1: Scale-Dependent Optimal Composition
Theorem B.1 (Scaling Law for Optimal Data Compositions (restated)). Consider the following
optimization problem
min Σ
m BNN
T
Σ N N
m
N
For any two compute budgets N(1) * 2(2), let N (N() and N (N be their respective min- imizers. For any third data composition N(N(3)), if there exists some constant k E R such that
N(N(3)) = (NO))
then, ((N(3)) is the minimizer for data budget N(3) = - M. given as
((N(3)) = arg min N
E
P
T
Σ
m
N; = N(3) = (N(3))
We begin with the following optimization problem, defined at a given total training data scale N(1;
Proof. Setup:
min N
Σ
N,N,
7
Σ N = N(1)
Here, N diag N1, N2,..., Nn is diagonal matrix whose diagonal entries are the amounts of data allocated to each of the mı domains.
For this problem, there exists
unique optimal solution N*(N())
diag( (1(1)= ....,,N(') (1)+).
This (*((()) represents the compute-optimal data
composition at the data scale N(1.
First-Order Conditions (KKT): At optimality, the Kurush-Kuhn-Tuukke (KKT) conditions ensure that the partial derivatives of the objective function with respect to each N; are equal (up to the same Lagrange multiplier for the equality constraint E; N; NT()). For any pair of domains a and b, we must have:
m
m
Σ CAN
Σ
B,N, 7
aN,
No No-
aN,
No
Computing these derivatives, get:
-o-1 -Naro - - - Bo (N((1)) -
From this equality:
BaTa (N (N
Brr
(N (N("+)>>
Rearranging, we obtain fundamental scaling relationship:
(1)
(Prta (N (1)+ ) BrYb
tot1
Scaling to Second Data Scale N(2) Now consider different total data scale (22 * N(1,, with the corresponding compute-optimal solution (N(2)) (
(2) (2)= ...., Nm (2))
18

---

Published as a conference paper at COLM 2025
Suppose know how the optimal amount of data for domain b changes from N to N(2). Specifically, let:
N = m Nà (1)
for some scaling factor
Applying the same relationship used for the first scale, but now at the second scale, we find that for domain
2/2)+
= (BaYaa (N(
Ta77
PaYa (m N₂ N T Yb
BoYb
This simplifies to:
N₂ (2)
FT No
7h+1 (1)-
7
Notice that m Ta o m in general. Thus, when the budget scales by a factor m in domain b, the optimal amount for domain scales by different factor. This shows that the optimal composition is scale-dependent.
Predicting a Third Scale N3);; We now know the optimal compositions at two scales N(1) and N(2), Consider a third scale N(3) and its optimal solution (*((8)) =
diag (N(3), N/! ,...,NM N,,))
we choose (²(3) such that:
(3)a
N(2)*
(3)
N222* (2)
N (1)='
then the change in N₂ from (2)2 to (3) mirrors the change from N(1) to N(22
Since the scaling exponent
Ta TI
remains the same, this symmetrical setup leads to:
(3
N (1)
Matrix Form: Because all domains scale in similar fashion, can write this relationship compactly using diagonal matrices. Define:
N (N(0)= = diag' N (0) ,...,,N, N (0)
The element-wise relationship (N/
for each domain a can be expressed as: (1))
(N (N
Here, (N' (N+ the inverse of the diagonal matrix N N (N(()), obtained by taking the reciprocal of each diagonal element (1)+
We have shown that given two distinct data scales N(1) and N and their correspond- ing optimal solutions N' (N()) and (N((2)),, one can construct a third optimal solution N* (N))) using the formula:
N (N(3)) = N' (N
(N(2).
This relationship holds without needing to explicitly estimate the parameters y; Bü, and it confirms that the optimal data composition is scale-dependent. Thus, the given scaling law for optimal data compositions is established.
19

---

Published as a conference paper at COLM 2025
Generalization to Prediction for Any Data Scale: Finally, we generalize from the case in Eq. (3) to allow prediction of optimal data composition for any data scale. Consider for some constant 7k R*, we choose N such that
N (k)
(2)+ k
m
2
(1)
Same as the procedure in Eq. (2), KKT optimality conditions yield the corresponding optimal data quantity for domain a at the same scale as N(* as
(2)0 k
(k) = (mr") TOTT Na
(2)+
N/22
N:
Rearranging in the matrix form, we have the following formula
No No (N(()) = No (N
which concludes the proof.
Application in AUTOSCALE: Note that this formulate holds for any k E R+ Thus, by scanning through the values of k, one can find optimal data composition N for any target
data scale N = N In practice, for a target data scale N > N(2) Li (2) one only needs
to conduct a line search along k > 1 to find the value of k where Σ N N to determine its corresponding optimal data composition N
Remark B.2 (An example). This example helps visualize the operation pipeline.
If at training data scale ((1) (1) + N 200, we have optimal domain data compo-
sition as (1)= 100, (1) 100 (50%- 50%); and at scale N(2) N N₂?) 500,
we have optimal domain data composition as N2 300, N!2)= 200 (60% 40%).
Then, from the theorem, when the optimal domain data composition has (3(3)
(N! / 900, we can predict N (2/
400, which gives the
optimal ratio at (3) (3 + N/ = 1300 as 69% - 31%.
Similarly,
For N(4)+ 2700, we have N (4)+ 800, which gives the optimal ratio at N(4) = 3500 as 77%
23%
For NJ(5)* 8100, we have Ng5)= [5) 1600, which gives the optimal ratio N(5) 9700 as 84%
16%
For (((6)= 24300, we have ((6)= 3200, which gives the optimal ratio at (66) 27500 88%
12%
For Ni( 72900, we have NJ? 6400, which gives the optimal ratio at N(7) 79300 92% 8%
For ((8)-
218700, we have N (8)+ 12800, which gives the optimal ratio at ( 231500 as
94%
6%
For N/9)" 656100, have N/9;- 25600, which gives the optimal ratio at (9 681700 as
96% 4%
We visualize it in Fig. 5.
B.2 Scaling Latent Skills
We extend this theory to a general case where the evaluation loss is the perplexity averaged over training domains. Consider the evaluation is composed of a number of independent sub-tasks ("latent skills" (Tiong et al.)) which are hidden variables, where each of them observes a power law scaling law relationship with the amount of data contributing to
this task ("equivalent data size"), , L to Ba K₁ + BB K2 + B. K_ +... where
20

---

Published as a conference paper at COLM 2025
100
a ratio
b ratio
60
optimal
40
2.5
3.5
5.0
log_10
Figure 5: Illustration: optimal data composition scales in exponential-style functions with training data quantity.
scalar Kj 2 0 denote equivalent data size for skill,,, and constants (B; 2 0 are coefficients associated with skill,, respectively. Mathematically, these latent skills can be seen orthogonal basis that spans the space of evaluation loss.
Consider training data from each domain D, contributes to these skills to varying degrees, where Equivalent data size for skill,, Kj, is given Kj- C₁,1 N1 + C;22' N2 + where N; wi N denotes the amount of training data from domain D; and constant C₁ is the coefficient measuring the degree of contribution between domain D, and skill Defining diagonal matrices for training data composition N diag { N₂, N2, ... } and skill data composition K= = diagi Ka, Kp, ... }, we have K = AN, where A = 9j, is the matrix for coefficients. For simplicity, we consider training data from each domain will be distributed to the skills such that vi, E N; 1. This gives the amount of total training data from all domains is identical to the amount of total equivalent data for all skills, Σ, K; = Li Nj. For training data scale N = Σ, N; = Σj Kú, define optimal skill data composition K = diag K;, Kū ... as the minimizer of L., given as K = arg min K k, fa- K, 7a + fb. K2 Ti + ..... Theoretically, there can be an infinite number of latent skills. For analysis, we consider a finite number of k independent skills most important for the evaluation. This can considered performing Principal Components Analysis (PCA) with orthogonal transformation and selecting the first k independent components. We consider the standard scenario with an equal number of relevant skills and data domains where & = m and A a square matrix with full rank. This describes the where this optimization problem is well-defined. We discuss in App. B.2 what will happen in other scenarios. In this case, A is invertible and the corresponding optimal training data composition for K* can be given as N = A'K
We provide the following theorem, which states that for the scenario described above, optimal training data composition scales in exponential-style functions with training data quantity and can be directly predictable from that of smaller scales without needing identify the latent skills.
Theorem 2 (Scaling Latent Skills). Consider the evaluation is composed of a number of independent sub-tasks (" ("latent skills") where each of them observes a power law scaling law relationship with the amount of data contributing to this task (" ("equivalent data size") Namely,
L to Ba K1
Be Ke Bb K
21

---

Published as a conference paper at COLM 2025
where scalar K; 2 0 denote equivalent data size for skill,, and constants (B 7;) 2 0 are coefficientss associated with s respectively. Define diagonal matrices for training data composition N diag{ N1, N2. } and skill data composition K diagl Ka, Kg. }. Consider training data from each domain D, contributes to these skills to varying degrees, given as K = AN where A Cj. is the matrix for coefficients. Assume the amount of total training data from all domains is identical to the amount of total equivalent data for all skills, Σ K, E N₁. Assume there is a finite number of latent skills and data domains and A is square matrix with full rank.
For a training data scale N = NN = K, define optimal skill data composition K' = diag K2r Kg, .... as the minimizer of C s.t. K, N with corresponding optimal training data composition. If we have optimal data compositions NO (N()) diag (N(")) No ((()) ....)
where its corresponding skill data composition K(T)* = diag ( K(1), ..... = (NO)) minimizes L s.t. K; = N, NNO), and No = diag N((2** N222x ,...} where its corresponding skill data composition K(2)* diag K!?"", (...} AN (N(2)) minimizes
C s.t. E K D -(2() = N(2) where N(2) + NO), then, other optimal data compositions (N diag (3)-
(N _N(3)),,.. .... where the corresponding skill data composition K(3)*
diag( ((0)=, (33). ,.... AN minimiizes s.l. Σ = Σ N(3)= N(3) where
N(3) + N(2) + N(1) must satisfy
(N = NN (NVP)) (N
for some k E R
Proof. By definition, we have
AN = k(), AN = ((2)- AN(3)) = K(3)
From results of Theorem 1 in Section 3.2.2, we have
K(3)+ (
for some E Rt, which gives
AN (N(3)) = (AN' (N (N(( -1 AN (N
Since A is invertible and N and K are diagonal matrices, naturally,
(AN' (N(")) = -A-
and we have AN" (N(0)) = AN (N 'A 'N = AN (NI)) (N'2)))
This directly gives
(²(3))) A (NN)
(N (N2R)))
which completes the proof.
The above result does not require identifying the latent skills or observing skill data compo- sitions K. Rather, the theorem gives that as long as the coefficient matrix A is invertible, the scaling of N complies to the same scaling law as in Sec. 3.2.2.
Remark 2 (what happens when A is not invertible.). In general, if A is not invertible, scaling for optimal training data composition is not directly predictable. Specifically, if A does not have full rank, there exists redundant domains/ data sources where their contribution to the skills are identical/ exact multipliers of each other. Some data sources may not be needed at any scale; if A has more rows than columns (more domains than skills), this suggests multiple training data compositions can achieve the same skills data composition and the optimal training data compositions are non-unique (infinitely many). If A has more columns than rows (more skills than domains), this means there are too many skills to optimize for. No optimal training data composition exists and one has to make trade-offs. If this is relevant to the practical needs, training data may be processed with additional techniques such as clustering and split into more different domains.
22

---

Published as a conference paper at COLM 2025
C Experimental Details and Additional Results for Section 4, Evaluation
C.1 Experimental Details on GPT-2 Large (774M)
Evaluation We test the perplexity on the held-out dataset, comprising 10K samples each from the 7 domains. For downstream tasks, we include: BoolQ (Clark et al., 2019) (zero- shot), HellaSwag (Zellers et al., 2019) (zero-shot, 10-shot), PIQA (Bisk et al., 2020) (zero-shot), TruthfulQA (Lin et al., 2021) (zero-shot), PubMedQA (Jin et al., 2019) (10-shot), CrowsPairs (Nangia et al., 2020) (25-shot), and ARC-Easy (Clark et al., 2018) (zero-shot). Additionally, BBH Novel Concepts (Srivastava et al., 2022) task is added to the aggregated results for models trained beyond 10B tokens, making total of 9 tasks. We select tasks that ensure the model's performance surpasses random guessing, spanning from question answering and commonsense inference bias identification and scientific problem solving. These tasks provide a comprehensive assessment of model performance (Mehta et al., 2024; Gadre et al., 2024). We adopt the evaluation framework from (Gao et al., 2021).
Baselines We report results for our methods (DDO and AutoScale ) and 6 base- Lnes-UNIFORMA LLAMA WEIGHTS (curated), DOREMI (LLaMA weights initialization), DATA MIXING LAWS FROM (YE ET AL., 2024), DOREMT from Xie et al. (2024) (uniform initialization), and REGMIX from Liu et al. (2024). Uniform weights uniformly sample data from all domains, resulting in the same number of training tokens from each domain. LLaMA weights are a set of curated domain weights heuristically tuned for training LLaMA- 1/2 models. We implemented DOREMI proposed in (Xie et al., 2024). DOREMI trains two smaller-scale auxiliary models (proxy models). First, a reference model trained with the dataset's original domain weights, which are the LLaMA weights for RedPajama dataset. Then, optimized domain weights are obtained by using a proxy model to minimize the worst-case excess loss across different domains. We train both auxiliary models for 50K steps. Implementation details are available in App. C.3. Besides, we compare with 2 do- main weights from existing literature, which are optimized on the same dataset, RedPajama, with similar Decoder-only LMs. DATA MIXING LAWS (Ye et al., 2024) first performs grid search on the space of possible data mixtures and records evaluation loss for proxy models trained on these mixtures. Then, the loss is interpolated with exponential functions to find the optimal domain weights for the proxy model. DOGE (Fan et al., 2023) also implements DOREMI (Xie et al., 2024) with auxiliary models trained for 50K steps but with the reference model trained with uniform weights. REGMIX (Liu et al., 2024) first trains an array of smaller, proxy models different data mix and small data scales, abd fits regression model between domain weights and evaluation loss. Then, the fitted regression model is used to predict the evaluation loss for all feasible domain weights to find the best-performing weights. We use the same pairs of domain weights and evaluation loss DDO used in optimizing domain weights for 774M Decoder-only LMs at 0.3B tokens to fit RECMIX's LightGBM regressor. The fitted LightGBM model is then used to optimize the evaluation loss over domain weights. We evaluate the model trained on these domain weights to present a complete landscape.
Model Training GPT-2 Large is variant of the GPT-2 architecture, featuring embed- ding dimension of 1280, 36 transformer layers, and 20 attention heads. We rely on the Hugging Face Transformers library for implementation (Wolf et al., 2019). Specific training hyperparameters are detailed in Table 6.
Dataset Details The RedPajama dataset is available at: https://huggingface co/ httpee/sstttggeeerroooppe//////p 11. The 7 domains involved are character- ized as follows:
Commoncrawl: A vast repository of web-crawled data, providing heterogeneous mix of internet text.
C4: The Colossal Clean Crawled Corpus, filtered to remove low-quality content, thus ensuring the reliability and cleanliness of the data.
23

---

Published as a conference paper at COLM 2025
Architecture
gpt2
Optimizer
W
Tokenizer Vocabulary Size
50257
Batch Size Per Device
1
Gradient Accumulation Steps
10
Maximum Learning Rate
2e-4
LR Schedule
Linear le-2
Weight Decay
Warm-up Ratio
10%
Epochs
3
GPU Hardware
8x NVIDIA A100/8x NVIDIA H100
Table 6: The list of hyperparameters for GPT-2 Large pretraining.
This domain includes a compilation of publicly available code repositories, offering a rich source of syntactic and semantic patterns inherent in programming languages.
GitHub:
A collection of textual content from published books, providing diverse narrative styles and complex character developments.
Books:
Comprising scientific papers primarily from the fields of physics, mathe- matics, computer science, and quantitative biology, this domain offers high-quality, scholarly content.
ArXiv:
Wikipedia: A well-organized and meticulously curated dataset of encyclopedia articles, delivering broad spectrum of knowledge across multiple disciplines. We only use English samples with 'en' meta-data.
StackExchange: This domain captures a variety of user-generated content from discussions and question-answer sessions across numerous technical topics.
Given copyright restrictions with the Books domain on Hugging Face, we have opted for an alternative source available https: /yknzhu....///// wixsite te. com/mbweb.
For each domain, we ensure only samples with more than 1000 characters are retained. For each sample, the first 1000 characters are truncated, with the exception of the ArXiv and Gi tHub domains where we randomly extract continuous block of 1000 characters. For the Wikipedia domain, we keep only those samples that are in English. Samples are selected without replacement, based on the computed data volume for each domain. Additionally, for each domain, a held-out dataset comprising 10K samples is reserved to evaluate the perplexity of the pretrained model.
C.2 Experimental Details on BERT (110M)
We evaluate the model's MLM loss on held-out validation datasets, comprising 10K samples each from the 5 training domains. Additionally, loss on 3 non-training held-out domains. To be consistent with the perplexity loss used in CLM, we report the exponential cross-entropy loss for MLM. We evaluate the model's task performance on GLUE benchmark (Wang et al., 2018) (with 8 diverse tasks for natural language understanding (NLU)) and SQUAD (Rajpurkar et al., 2016) (a large-scale QA dataset). Uniform weights are used as the baseline.
auxiliary evaluation, we test the MLM
Model Training We employ the BERT-base-uncased model from the Hugging Face Trans- formers library. Originally, BERT's pretraining scheme involved MLM and next sentence prediction (NSP); however, our experiments, we exclusively utilize MLM. Detailed training hyperparameters can be found in Table
Dataset Details
The 5 domains of training data utilized listed follows:
24

---

Published as a conference paper at COLM 2025
Architecture
bert-base-uncased
Max Token Length
300
Mask Token Percentage
15%
Optimizer
AdamW
Batch Size Per Device
12
Devices
4
Maximum Learning Rate
le44
LR Schedule
Linear
Weight Decay
le-2
Warm-up Steps
3000
Epochs
1 1~4 4
GPU Hardware
4x NVIDIA RTX A6000
Table 7: The list of hyperparameters for BERT pretraining.
Amazon Reviews: compilation of customer reviews from Amazon, widely utilized in sentiment analysis studies, available at: https://hugginggace. ootp://uuggggg.....ooootttt//// amazon_us_reviews.
Arxiv: Comprises 1.7 million articles from arXiv, available at: https://www.
tensorflow. http://ttssestttttaaaaa/sseeeete
Books: A corpus of 11,038 novels by unpublished authors across 16 genres, available at: https://yknzhu. wixsite. hbw.://.ww..iiiiiiieooooneobbb
Wikipedia: Offers datasets extracted from Wikipedia in various languages, avail- able at: https://www. tensorflow www.ooooww.oooo/tteeeeeeeeeeeeee English samples with 'en' in meta-data.
We only use
Open WebText Corpus (OWTC): A corpus of English web texts from Reddit posts, available at: https://skylioneee
github. to/OpenweeTeextCoppss/
3 held-out non-training domains used the evaluation include:
Features 19,717 diabetes-related publications from the PubMed database, organized into three classes and linked by network of 44,338 citations, available at: httpp::/www/ tensorflow org ceaatasets/atalloosccieettet wttp://g//iiinniinieiiiiiieeeece
Pubmed:
News: Comprises a significant collection of news articles derived from CommonCrawl, specifically from 5000 news domains indexed by Google News, available at: https:
GitHub: A curated selection from the RedPajama dataset, this segment includes ar- ray of open-source code projects, available at: https://uuggingef www.aacc........e.eeeeeeeeeeee.e co/datastss/ together RRdPajaaa--aal
C.3 Implementation Details for Baselines
Implementation details We followed the official implementation? of DOREMI for our experiments. We evaluated two sets of reference domain weights: (1) the domain weights utilized in the LLaMA-2 paper Touvron et al. (2023) (referred to as LLaMA weights), and (2) uniform weights. Both the reference and proxy models have 120M parameters and are trained from scratch. We use GPT-2 tokenizer with vocabulary size of roughly 50K. For LLaMA weights, we train each model for 20K, 50K and 200K steps for comparison. For uniform weights, train each model for 10K, 20K and 50K steps. Refer to Table 8 for detailed hyperparameters. The effect of reference weights on the output DOREMI is discussed in Fig.9.
25

---

Published as a conference paper at COLM 2025
Architecture
Decoder-only LM
Max Token Length
1024
Optimizer
Adam
Batch Size Per Device
8
Devices
8
Maximum Learning Rate
2e-4
LR Schedule
Linear
Weight Decay
1e-2
Warm-up Steps
3000
Epochs
GPU Hardware
8x NVIDIA RTX A6000
Table The list hyperparameters for DOREMI.
C.4 Evaluation Details
The following tasks are considered for downstream performance evaluation, in line with the setup from (Mehta al., 2024; Gadre et al., 2024). For few-shot tasks, the demonstrations are sampled at random.
GPT/CLM
BoolQ (Clark et al., 2019) consists of a question-answering format that requires binary yes/no no answers.
HellaSwag (Zellers et al., 2019) challenges models on their ability to make common- sense inferences.
PIQA (Bisk et al., 2020) focuses on evaluating model's commonsense reasoning regarding physical interactions.
TruthfulQA (Lin et al., 2021) is designed to assess the ability of models to generate truthful and factual responses.
PubMedOA (Jin et al., 2019) offers a dataset for evaluating question-answering in the biomedical domain.
CrowsPairs-English (Nangia et al., 2020) tests models on their ability to identify and correct stereotypical biases in English text.
ARC-Easy (Clark et al., 2018) presents set of relatively simpler scientific reasoning questions, aimed at evaluating model's basic understanding of scientific principles.
BigBench-Novel Concepts (Srivastava et al., 2022) serves as test of the model's creative abstraction skills, challenging it to make sense of scenarios that could not have memorized during training.
BERT/MLM For each task, we conduct supervised fine-tuning the corresponding training data and test the fine-tuned model on the validation data. The hyperparameters for supervised fine-tuning are given in Table 9.
Architecture
bert-base-uncased
Max Token Length Batch Size Per Device
128
8 or 300
Optimizer
AdamW
Devices
Maximum Learning Rate
2e-5 or 5e-5
Epochs
GPU Hardware
4x NVIDIA RTX A6000
Table The list hyperparameters for supervised fine-tuning of BERT.
26

---

Published as a conference paper at COLM 2025
135.0
common crawi
github
common crawi gthub
135.0
books
Domain Data Size/billion
116.0
Wiki
c4
wkii
97.0
stack exchange
stack exchenge
andv
Domein Weights
andio
0.3
58.0
39.0
0.2
19.0
0.1
0.0
0.3
0.6
1.2
3
5
10
19
38
77
154
307
Data Scale/billion tokens
0.3
0.6
1.2
10
19
30
77
124
307
Scale/billion
(a) AutoScale -predicted optimal data quan- tity for each domain as training data scales up.
(b) AutoScale -predicted optimal domain weights as training data scales up.
Figure 6: AutoScale -predicted domain weights for training 774M Decoder-only LMs. Opti- mal data quantity for each domain grows in exponential-style functions with training data scale (left) where data sources with diverse samples (e-g., C4) are upweighted relative to domains with standard format (e.g,, Wikipedia).
C.5 Additional Results on GPT-2 Large (774M)
Fig. depicts AutoScale-predictel domain weights for training 774M Decoder-only LMs (GPT-2 Large). Optimal data quantity for each domain grows in exponential-style func- tions with training data scale (left) where data sources with diverse samples (e.g., C4) are upweighted relative to domains with standard format (e.g., Wikipedia).
Fig. 7 shows that when training on up to 5B tokens, AutoScale AutoScale-predicted weights decreases val loss at least 25% faster than any baseline with up to 37% speed up.
Fig. 8 visualizes domain weights used for training GPT-2 Large, given by different methods.
Fig. 9 visualizes DOREMI optimized domain weights with different reference weights and training steps. Training proxy /reference models for different steps gives different weights It is unclear which weights are optimal. DOREMI recommends 200k steps, which equals >10BB tokens the default setup. Since optimization was conducted relative the reference weights, reference weights have profound impact on DoREMI's output.
C.6 Additional Results on BERT (110M)
Fig. 10(b) shows the results on fitting validation loss with power-law functions, directly approximating how loss changes with each domain's data quantity. Compared BERT models trained with MLM (right), GPT models trained with CLM (left) demonstrate much stronger response to domain reweighting. In final results, GPT GPT/CLM achieved > 2x speed- up margins relative to uniform weights compared to BERT DER//MMM.
Fig. 11 depicts the AutoScale -predicted domain weights for training BERT. is evident that optimal data quantity for each domain grows in exponential-style functions with training data scale where data sources with diverse samples (e.g., WebText) are upweighted relative to domains with standard format (e.g., ArXiv).
Table 10 shows AutoScale notably improving training efficiency for BERT models on all sceles-even for considerably large scale, 288k steps, the speedup margin remains visible.
27

---

Published as a conference paper at COLM 2025
AutoScale Predicted Weights Baselines
AutoScale Predicted Weights Baselines
30
AutoScale Predicted Weights
AutoScale Predicted Weights
DofteMi Ref Weights
24
LLaMA Weights
28
Data Mixing Law
Uniform Weights
LLaMA Weights
Uniform Weights
22
Validation
26
Validaatio
27% faster
20-
24
37% faster
25% faster
-X
18
31% Taster
22-
33%
25% faster
16
20
15000
20000
25000
30000
35000
40000
25000 30000 35000 40000 45000 50000 55000 60000 65000 Training Steps (x10 gradient accumulation)
Training Steps (x10 gradient accumulation)
(a) Training Decoder-only LMs for 3B tokens.
(b) Training Decoder-only LMs for 5B tokens.
AutoScale prrdicctdd weights decreases val loss least 25% faster than any baseline with up to 37% speed up. Despite LLaMa weights being very different from uniform weights, they yield highly similar training efficiency at these data scales.
Figure 7:
Domain Weights
AutoScale
0.6
DoReMi Ref
Data Mixing Laws
0.5
LLaMA
Domain Weights
Uniform
0.4
0.3
0.2
commongithub
books
wiki
stack
arxiv
crawl
exchange
Data Domains
Figure 8: Domain Weights used for training 774M Decoder-only LMs for 3B tokens. (Domain weights for DATA MIXING LAWS and DOREMI are from references (Ye et al., 2024) and (Fan et al., 2023), respectively, which are implemented on the same datasets/ / data domains with highly similar model architecture/ mrchttecture/moeel model size tize/tokenizers)
Data Scale/steps
18k
36k
72k
144k
288k
Final Loss (exp)
38.32
16.94
10,97
8.13
6.30
Steps Saved.
5k (28%)
5k
(14%)
10k
(14%)
20k
(14%)
20k (10%)
Table 10: AutoScale notably improving training efficiency for BERT models on all scales-even for a considerably large scale, 288k steps, the speedup margin remains visible.
D Extended Discussions
Intuition for Stage 2 of AutoScale, Optimal mix projection. Consider a stylized case where each training domain is independent. In this scenario, the validation performance attributable to each domain scales with the amount of training data from that same domain. Domains exhibit different scaling behaviors:
As the training budget increases, the loss associated with some domains (e.g., 'easier' domains like standard texts from Wikipedia) decreases rapidly at first and then
28

---

Published as a conference paper at COLM 2025
DoReMi Weights Steps for Proxy Models
DoReMi Weights Steps for Proxy Models
1.0
0.30
Uniform Init.
LLaMA Init.
DolteMi 10k steps
DofeMi 20k steps
0.25
DoReMi 50k steps
DoReMi 20k
0.8
DofteMi 200k
DoReMi 50k steps
Domain Weights
Domain Weights
0.20
0.15
0.4
0.10
0.2
0.05
0.00
0.0
commongttub books
wild
C4
stack
andv
commongithub books
wiki
C4
stack andiv
crawl
exchange
crawl
exchange
Data Domains
Data Domains
(a) with Uniform Reference Weights
(b) with LLaMA Reference Weights (Default)
Figure 9: DOREMI with different reference weights and steps. Training proxy / reference models for different steps gives different weights. unclear which weights are optimal. DORFMI recommends 200k steps, which equals >100B tokens in the default setup. Since optimization was conducted relative to the reference weights, reference weights have profound impact on DOREMI's output.
common_craat
gthus
stackeschange
vallddadenn uss teea)
29
20-
26-
23
0.0
2.5
1,0
Domain Change (Orighale1)
Domain Eata Change (Originale1)
(a) 774M Decoder-only LMs (GPT-2 Large)
(b) Encoder-only LMs (BERT-caSe)
Figure 10: Fitting validation loss with plwer-law functions, directly approximating how loss changes with each domain's data quantity. Compared to BERT models trained with MLM (right), GPT models trained with CLM (left) demonstrate much stronger response to domain reweighting. In final results, GPT/CLM achieved 2x speed-up margins relative to uniform weights compared to BERT/MA BERT/MLM.
plateaus. the total training data budget small, sensible allocate more budget to such domains, as they offer a greater initial reduction in validation loss.
Conversely, for other domains (e.g., diverse sources like Common Crawl), the loss decreases more slowly but steadily with an increasing training budget, continuing to provide benefits even when the loss from other domains has plateaued.
With larger compute budgets, it becomes more beneficial to allocate resources to these steadily improving domains rather than adding more data to 'easy' domains that are already saturated. (As Figure illustrates, optimal compositions are scale- dependent.) The key insight, which formalize in theorem, is that under certain independence assumptions, the evolution of optimal domain compositions with increasing training data budgets be analytically derived from their individual scaling laws. This allows us to predict optimal compositions at larger data budgets.
We then relax the assumption of strict domain independence. We theoretically prove that if the validation data requires number of independent 'latent skills,' and data from
29

---

Published as a conference paper at COLM 2025
amazon
Domain Data Size/billion tokens
6.0
books
5.0
wiki
books
web text
5.0
web text
4.0
3.0
0.15
1.0
010
0.0
0.3
0,50.811
2.2
4.3
8.6
17.3
005
Scale/billion tokens
0.3
0.5
0.8
1.1
2.2
0.0
17.3
Data Scale/bilion tokens
(a) AutoScale -predicted optimal data quantity for each domain as training data scales up.
(b) AutoScale -predicted optimal domain weights training data scales up.
Figure 11: AutoScale -predicted domain weights for training Encoder-only LMs (BERT). Optimal data quantity for each domain grows in exponential-style functions with training data scale (left) where data sources with diverse samples (e.g., WebText) are upweighted relative domains with standard format (e.g,, ArXiv).
each training domain contributes to one more of these skills, our previous theory still holds in the same form. Thus, direct domain independence is not a strict requirement. This provides complete theoretical grounding for our AutoScaletool, which captures the evolution of optimal domain weights with the training data budget and offers predictability for larger scales. Finally, our empirical evaluations demonstrate that the weights predicted by AutoScale offer clear advantages over domain weights optimized without considering this scale-dependence. the discussions and future work.
Number of predictors (scaling law components) We conducted ablation studies with DML on optimizing domain weights under a fixed compute dudget, following the same experiment procedure in Section 4.1 to pre-train GPT-2 Large models the RedPajama dataset. According to DML's procedure, we fit a separate exponential function between domain weights and validation perplexity for each of the 7 domains, resulting in 7 separate predictors. When predicting for the average validation loss, we use each predictor to predict the validation perplexity for the respective domain and take the average of these predictions as the final output. When fitting on proxy training on random data mixtures, this prediction approach yields an average absolute relative error (AAR) = 4.46%. In comparison, similar to the treatment in the proposed DDO, we also fitted single function between domain weights and the average validation perplexity using DML's exponential formula. We experimented with using this function to directly predict the average validation perplexity. When fitting on proxy training on random data mixtures, this prediction approach yields an AAR 1.61%.
When using the fitted predictor to derive optimal data mixtures, we realized DML's formula does not allow optimize domain weights over single predictor. For validation domain, the predicted validation loss is given as the exponentiation of weighted average of domain weights plus some constant (Éq. (1) from Ye et al. (2024)). Since the domain weights are combined linearly before the exponentiation, if there is a single validation domain, minimizing the predicted loss will lead to degenerate solution where all training data budgets is assigned to single domain. Thus, it is necessary for DML to incorporate number of separate predictors introduce the crucial nonlinearity needed to model the interactions between different domains. The power-law formula in DDO is able to capture the nonlinearity domain interactions in a single predictor, leading to lighter parametrization and enhanced efficiency
30


Reviews:
[{"review_id": "otZgjqgDI0", "reviewer": "Reviewer_HTn5", "summary": "Summary: many works on data mixing utilize proxy training runs (smaller model, fewer tokens) to determine the optimal data mixture at a larger scale. However, this paper shows that the optimal data mixture can change depending on the token budget. The paper proposes a two stage approach: first, a regression-based mixing algorithm that requires fewer proxy runs than Data Mixing Laws and RegMix; second, a theoretical analysis that leads to a way of scaling up the optimal mix at a small token budget to a larger token budget. Experiments show that AutoScale improves over RegMix and Uniform weights.", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "- Data mixing is a critical step of the pretraining curation pipeline, and the current reliance on smaller and shorter proxy runs should be studied more closely, which this paper does. \n- Generally thorough evaluation, including average test perplexity, downstream tasks, reporting the mixture weights, etc. \n- Paper provides a rigorous theory for how to scale across token budgets and provides a motivating example for why this is important in Figure 1.", "reasons_to_reject": "**Related work:** since this paper mentions how $\\gamma_i$ changes the importance of each domain over time, online data mixing methods should also be discussed in the related work, perhaps as an alternative to determining how to scale mixes across token budgets. An incomplete list of these include: ODM (https://arxiv.org/abs/2312.02406), Skill-It (https://arxiv.org/abs/2307.14430), Aioli (https://arxiv.org/abs/2411.05735), and ADO (https://arxiv.org/abs/2410.11820).\n\n**Effect of repeated documents:** for the paper's setting, scaling up the token budget to 10B should not result in any repeated documents for any mixture over RedPajama. However, in practice it is conventional to train for much longer and utilize all of the unique documents available at least once (i.e., the \"data-constrained setting\" of https://arxiv.org/abs/2501.11747). Repeating documents could significantly change the scaling law and theory of this work, although I do not think this significantly diminishes the paper's contributions. \n\n**Statistical significance of Figure 1:** It is natural that N* @ 1.2B is going to be different from N* @ 0.3B, but it is unclear how much of this is due to noise. While Figure 4 shows that Figure 1's $w_1$ and $w_2$ are different---the weight for Github only slightly changes, while the weight for C4 changes significantly---it would be helpful to ground Figure 1's perplexity reductions by comparing it to other interventions, especially because the changes in perplexity reductions are not significant enough for DDO to perform worse than RegMix or uniform weights. Some suggestions:\n- Compare the changes in downstream performance of $w_1$ vs $w_2$\n- Measure how much changing the random seed (shuffling data/sampling \"fresh\" documents from each domain) impacts perplexity\n- Measure how much a different fitting approach (i.e., using different baseline and perturbation) impacts perplexity\n\n**Justification for scaling law:** in line 175, the approximation of $\\mathcal{L}^v$ using a constant $N_0^i$ is interesting; why not make the $N_0^i$ term dependent on $w_1, w_2, \\dots, w_{i-1}, w_{i+1}, \\dots, w_m$? I.e., something of the form $\\sum_{i=1}^m \\big(\\sum_{j \\neq i}^m w_j \\cdot N_0^{ij} + w_i \\cdot N \\big)^{-\\gamma_i} + \\ell_i$? Note that this reduces to your scaling law if $N_0^{ij} = N_0^i$ for all $j$, i.e., the equivalent data size is constant for each domain (I am skeptical that this is constant, since intuitively, e.g., CC and C4 interact with each other more than some other domain pairs) \n- The result on AAR in line 275 should be explored more. Namely, reporting predicted loss on \"training data\" (baseline mix and perturbations) vs test data, and also reporting AAR for RegMix and Data Mixing Laws. A discrepancy in test data prediction accuracy between DDO and RegMix would justify why RegMix performs more poorly. \n\n**More evaluation of Data Mixing Laws (Ye et. al.):** Data Mixing Laws has a two stage approach where it fits a mixing law and then uses a scaling law to determine the mix at a larger token budget. It would be interesting to try DDO + Data Mixing Law's token budget scaling law, and Data Mixing Law's mixing law + AutoScale's token budget scaling.", "questions_to_authors": "See \"Reasons to reject\" above, in particular:\n\n1) What are the downstream performance changes when using $w_1$ vs $w_2$?\n2) How much does changing the random seed (in data) impact perplexity?\n3) How much does a different fitting approach (using a different baseline mix and perturbation size) impact perplexity?\n4) Can you report AAR on train/test mixes for DDO and RegMix? Can you report AAR when using a scaling law that has $N_0^{ij}$?  \n5) What happens when you use one stage of AutoScale and another stage of Data Mixing Laws? I am curious how Data Mixing Law's token budget scaling performs in conjunction with DDO.", "extra_scores": {}}, {"review_id": "ZGDJ4srBoI", "reviewer": "Reviewer_jtwN", "summary": "The authors develop and test a novel method, \"AutoScale\", for choosing an optimal mixture of domains to use in a training data set for pretraining an LLM at a given data scale, where data scale refers to the total number of tokens that will be present in the training data set.\nThe authors show that mixtures which are optimal at small scale often are not optimal at large scale, but testing multiple mixtures directly at large scale is prohibitively expensive. This is the problem that AutoScale is intended to address.\n\nThe AutoScale method rests on the assumption that the size of each domain, as a subset of a training data set, individually is governed by the empirical power scaling law identified by Kaplan et al. (2020), which defines a functional relationship between the amount of training data and validation loss for neural models. This assumption is then combined with the claimed original \"Scale-Dependent Optimal Composition\" theorem, which the authors prove in the paper (I do not have the expertise to verify the proof), to yield the \"AutoScale\" method for choosing the optimal mixture of domains in the training data set for a given training scale.\n\nThe authors present empirical results in which the application of the AutoScale method for choosing domain mixtures at large scale reduces training time and improves model performance compared to baselines.", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "This is a new method that can accelerate convergence and provide better scalability for training LLMs.", "reasons_to_reject": "The paper is very dense, and much of the content required to understand the approach is in the appendix. I could not understand how the second stage of the method works or oven what the intuition behind it is. I understand that the authors are limited by the 9 page max, so maybe a conference paper is not the best venue for this. However, there are some redundancies between the intro and the lit review, so maybe the authors can tighten these sections and use the freed up space to explain at least the intuition behind stage 2.", "questions_to_authors": "Do you anticipate that the scaling relationships predicted by AutoScale will continue to hold as larger budgets are used (e.g., greater than 100 B tokens)? Is it possible that a domain can appear to be headed toward saturation at small or medium scale but at large or very large scale the saturation does not continue or reverses?\n\nHow does AutoScale compare to recent methods that adjust the domain weight vector dynamically during training, instead of computing a fixed data mixture for a given training budget before starting a large-scale training run?\nE.g., Jiang et al. (2024) [https://arxiv.org/abs/2410.11820], Luo et al. (2024) [https://arxiv.org/pdf/2411.14318]", "extra_scores": {}}, {"review_id": "3SC3Y7JnAp", "reviewer": "Reviewer_TYrQ", "summary": "This paper focuses on the scenario of data mixture for LLM pre-training. Inspired by scaling-law, it verifies that the optimal mixture rates are not uniform across different scales. By leveraging DDO and iteratively scaling based on the theoretical analysis, the paper proposes AutoScale, which seeks larger-scale data mixture at a lower cost using smaller models. Extensive experiments conducted on GPT-2 (774M) and BERT (110M) with ~10B tokens validate the effectiveness of the proposed method.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "1. *Valuable Scenario:* This paper addresses an important, interesting, and practical problem of exploring large-scale data mixtures with minimal cost, which is crucial for LLM pre-training, given that trial-and-error can be an expensive process.\n2. *Well-Grounded and Solid Clarification:* The insights presented in the paper are empirically driven, while the proposed method is theoretically grounded. The dual foundation enriches the content and provides solid support for the findings. Moreover, the conclusions align with existing discourse in the community, offering additional insights.\n3. *Easy to Follow:* The paper features a clear structure and is well-written, making the methods easy to understand. The conclusions drawn are also clear, with comprehensive experiments that exhibit coherent logic.", "reasons_to_reject": "1. *Stability and Robustness of the Method:* The stability of Stage 2 of the proposed method, as discussed in lines 219-225, under multi-round iterations has not been adequately considered. The statement regarding \"… for any larger scale\" and \"… any desired accuracy\" may be overly optimistic. While the initial scaling-up iterations might yield accurate results, the potential for error could gradually accumulate as the scale increases. This aspect has not been addressed in the paper.\n2. *Experimental Scale and Model Selection:* Given that the data mixing scenario in the paper focuses on pre-training, the current volume of data and model-size is still insufficient to support prevailing standards in the field, despite the authors' explanations in lines 237-239. Although I understand that as a validation method, large-scale experiments may not be necessary, conducting experiments at a larger scale would enhance the reliability of the paper's conclusions (related to my previous concerns about stability).\n   - The volume of data cited in the paper is limited to 10B tokens, which is quite small compared to the current mainstream pre-training requirements (for instance, the Llama3 series uses 15T tokens).\n   - The models used in the study, GPT-2 (774M) and BERT (110M), are relatively outdated frameworks, and their sizes do not align with the scale of contemporary models.\n3. *Iterative Scaling Curve:* Still considering Stage 2 of the proposed method, the iterative scaling approach allows for the evaluation of larger-scale performance. This raises the question of whether additional data points could be used to plot an exponential-like curve that describes the properties of mixture scaling. Utilizing this \"ground truth\" curve would facilitate a better assessment of the accuracy and effectiveness of this iterative generation process.", "questions_to_authors": "1. Following the scaling-law, I am curious about whether the proposed method can be applied to model-size scaling as well?\n2. Additional reference: BiMIX [1] also discusses data-mixing scaling in the context of pre-training, it is suggested to be included in the discussion and comparison in this paper.\n3. Minor notation: In Theorem 1, the two variables \\$N^\\*\\$ and N in the optimal domain allocation are somewhat difficult to distinguish. A different notation could enhance clarity.\n4. Minor typo: In line 286, there should be a space between \"AutoScale\" and \"consistently.\"\n5. Minor figure refinement: In Figure 3, the image is not a vector graphic, the axis font is overly large, and the background color at '38%' obscures the subsequent line variation of AutoScale.\n\n[1] A Bivariate Scaling Law for Language Model Pretraining", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces AutoScale, a two-stage, scale-aware data composition framework for pre-training. Stage 1 fits a parametric model to predict the model's loss under different data compositions --- used to find an approximate best allocation with less budget. Stage 2 extrapolates the optimal composition to larger budgets without further retraining. The authors demonstrate that AutoScale accelerates convergence and improves downstream performance, achieving a 28% faster perplexity drop and up to ~38% speed gain when pre-training GPT-2 large.

The reviews were generally sensible and asked some insightful and valuable questions which generated productive discussions and valuable improvements to the paper. All discussions should be incorporated into the revisions of the submission. The paper includes some solid theoretical and empirical contributions which not all reviewers were comfortable to assess. I do not think the paper should be penalised because there is not a directly appropriate journal venue for this work. One reviewer did not respond to the discussion, but I would consider their queries extremely valuable and the author's response to be thorough and highly valuable for paper improvements. Reviewers highlight the significance, volume and importance of the contributions to the paper. My opinion is that this will be a valuable feature in a conference agenda.
Decision: Accept