AutoScale: Automatic Prediction of Compute-optimal Data Compositions for Training LLMs

Authors: Feiyang Kang, Yifan Sun, Bingbing Wen, Si Chen, Dawn Song, Rafid Mahmood, Ruoxi Jia
Primary area: foundation or frontier models, including LLMs
Keywords: Data Curation, Data Composition, Scaling Laws, Data-centric AI, Large Language Models (LLM)

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 language model pre-training. This paper demonstrates that the optimal composition of training data from different domains is scale-dependent, challenging the existing practice of determining optimal mixtures through small-scale experiments and directly applying them at larger scales. We derive an analytical model for the dependence of optimal weights on data scale and introduce *AutoScale*, a novel, practical approach for optimizing data compositions at potentially large training data scales. *AutoScale* first uses a principled optimization framework to find optimal compositions at smaller, feasible scales, then predicts optimal compositions at larger scales using our derived model. Our evaluation on GPT-2 Large and BERT pre-training demonstrates *AutoScale*'s effectiveness in improving training convergence and downstream performance. Particularly, for GPT-2 Large on RedPajama, *AutoScale* decreases validation perplexity 28% faster than baselines, with up to 38% speed-up over unweighted training, achieving the best performance across downstream tasks. This work provides insights into the varying benefits of data sources across training scales for language models, contributing to the burgeoning research on scale-dependent data curation. Code is open-sourced

Decision: reject
Reviewer scores: 6, 8, 3, 5

Reviews and discussion:
<thread>
creation date:
1730505021458

thread_type:
official_review

<thread_object>
creator: reviewer_ad96

<content>
summary:
This work proposes a method called “AutoScale” that helps predict the optimal composition of pre-training data for LLMs. It challenges the conventional notion of determining this via small scale experiments and simply applying them to a large scale where two axes change (data scale, parameter count). The experiments show a very promising line of research and it was a pleasure to read. 

I couldn’t check the math as well as I would have liked to.

strengths:
* Very strong work in terms of the hypothesis and experimental setup albeit at smaller scales. The promise of finding optimal weights for training large networks without having to guesstimate it is a very attractive proposition. 
* The plots are really well done. They drive the main idea of the paper very well(especially Fig 1 (a, e) )

weaknesses:
I would like to list the following weakness fully ensuring the authors that I am not unreasonable and am completely open to increasing my score if these are addressed/answered satisfactorily.
* The work proposes using a different approach to finding optimal data weights for a given pre-training compute budget. This is well explained via results but does in fact require training the original size model. Given that we obtain suboptimal performance via the conventional way( smaller model, fewer data), an analysis showing how much performance could be gained by spending the compute and training these (equal parameter) networks would be useful. 
* For Takeaway 1, Fig 1(b) only has 2 data points. Additional points would help make the case stronger. It’s a tough sell to make such a bold statement with two data points. But I’m hoping I am wrong :) 
* Maybe I missed this, but the repeated claims that Wikipedia should be sampled less at a higher scale is a result of the OLS fit. But no experiment actually confirmed this fact in the paper, right ? Since the max scale was 1.2B ? Please correct me if I’m wrong.

General Comments/Typos:
* [Section2] :  “this work contribute”  -> “this work contributes”
* [Section 3.1] : wi = Ni/N => wi = Si/N ?
* [Algorithm 1] : Train the model on data S = ({S1 . . . Sm} \ Si) => S = ({S1 . . . Sm} \ Sj) ? 
* Some of the font sizes are very distracting to read.

questions:
* Even at a smaller scale, I see opportunities of clear promise where we could have had more points between 0.3B and 1.2B and show some trend. Any specific reason this was not done/ increased to more than 1.2B ? With scale, a lot of problems disappear that are apparent at lower scales.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730757262739

thread_type:
official_review

<thread_object>
creator: reviewer_jghf

<content>
summary:
This work presents a method to estimate optimal data domain weights at large training-data scales by extrapolating via exponential functions fit to smaller-scale training runs. The proposed method is evaluated on GRT-2 Large + RedPajama and BERT pretraining, and compared to extant method baselines.

strengths:
The general problem the authors attempt to address is important, and the assessment that present methods are limited and that performance headroom is available is well-framed. The code is open-sourced. The evaluations presented are limited, but positive.

weaknesses:
In general, the writing is difficult to parse. It is frequently frustratingly vague, including in the Abstract, where the actual method is alluded to but not elucidated. In the actual methods section, important questions about the method are unanswered, leaving the method underspecified (is the learning rate schedule (linear? presumably linear decay?) the same for the tuning runs as for the final run? Is the decay timing adjusted to the compute budget? The value of the final validation loss hinges critically on this - yet it goes unmentioned). There is no addressing of the profound difficulties this method (and others like it) can be expected to have around epochs for individual datasets. A more thorough analysis would identify and investigate this issue with experiments demonstrating specific datasets being sampled for > 1 epoch, and the subsequent breakdown of the "scaling law" prediction. Evalutation is purely comparative to other methods, and does not assess to what extent the predicted 'optimal' values might differ from more expensively traditionally-derived 'optimal' values. No discussion of the relative cost of the method (with its 'linearly scaling' cost in the number of datasets) is mentioned, though it is clear it would become prohibitively expensive for dataset mixtures with more than a handful of individual datasets. The method proposed is prohibitively expensive at large model sizes, and seems unlikely to scale to larger compute budgets even at small dataset sizes due to the issue of datasets passing through multiple epochs, which is unaddressed in this work. This limitation goes unmentioned.

questions:
Fig 1: [a] suggests that you've tuned 6 models between 30M and 1.2 B tokens, yet [c] shows only three models being used to fit the predictor model. Why is that? where are the other data points? And are *all* of the linear fits R2=0.998? Is that the average R2? Also, [d] shows the predictions of the model extrapolated past 1.2 B to 307 B? Why are you not showing the training data points (between 30M and 1.2 B) as well? And isn't the largest model you look at trained to 10B? why show this extrapolation to so far beyond where you explore? This seems misleading. The x-axis should say (log scale) as well. In [b] the color used for the 1.2B model is the color used for the 0.6B model in [a]. And there is a typo in the title ('scale - depedent' -> "dependent"). In [e] the 38% improvement looks to be overstated due to the noise of those evaluation curves, you could just as easily pick out the peak in Autoscale curve at step 86k and the point in the Uniform curve at step 100k to get a smaller improvment result with the same underlying data. 

Table 1: boolq has the Autoscale value bolded as 'best' but the Data Mixing Laws value is greater. Also, consider place your method on the bottom row separated by a thin line.

Fig2: What is being depicted here? Is this showing power laws being fit to 3 empirical datapoints? Is the first column of points supposed to be at 0? It looks like the points are at [0.2, 1, 3] on the x-axis? 


Nits:

Throughout: "AutoScale" is consistently the wrong font size. Please fix. Similarly, in section 5.2 the font size of the methods needs to be fixed. And in line 418 'from' is included in the method font instead of the default font.
181: work contribute -> contributes
379: N^(1)* is missing the N in summation
465: much lowered -> much lower
155 'a consistent shift can be observed', please be more specific, what is shifting, how is it consistent?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730793713621

thread_type:
official_review

<thread_object>
creator: reviewer_mxrw

<content>
summary:
This paper studies the problem of predicting optimal data mix for a given compute budget (i.e., fixed total token count and model size). A key challenge here is that the optimal domain weighting may change at different scale, hence it is inaccurate to use smaller models to predict large model performance, while solving the optimization problem at the large model scale directly is computationally infeasible (requires multiple retraining).

The paper proposes a method that work on one domain at a time by fixing the rest of the data constant (hence the loss is constant for other domains too), and estimated a scaling law per domain. The power law parameters $\gamma_i$ and $l_i$ can be easily estimated, which approximate a regular data scaling function where $l_i$ is the irreducible loss of that domain. 

After the power law of each domain is found, the final objective is to mix the data so that the loss is minimize while keeping sum of the tokens reaches the budget, which becomes a convex function that can be solved efficiently. This gives the DDO method. The different $\gamma_i$ explains why there is a differet mix at different stage.

A method "AutoScale" is further proposed to obtain the data weight of a larger token budget, by iteratively mxing two data weights at different scale to create the weights of the next one. 

The proposed approach is tested on models like GPT-2 (autoregressive) and BERT (bidirectional), showing improved convergence rates and downstream task performance. Empricially, the results show AutoScale’s ability to shift data weights, favoring diverse data sources like CommonCrawl at larger scales while reducing reliance on traditionally high-quality, standard-format data such as Wikipedia. These findings match the empricial findings of the data weights used for prior succesful models such as Llama.

strengths:
- This paper analyzes an important problem, data weighting of LLM training, which can improve the training efficiency with reasonable cost. It also presents an actionable algorithm for LLM training.
- The proposed method assumes a power law formulation which makes the data weighting problem practically solvable. It is important to point out that data weights is scale dependent.
- The empirical results and findings on the corpus weighting align with common belief of the community, such that further up-weighting high quality source is less effective, and books and web documents continue to be important at larger scale. This shows that the proposed method has strong explanatory ability.
- The experiment is quite thorough, considering the cost for training models is quite high even at small scales.

weaknesses:
- I wonder if carbon footprint of the experiments here should also be reported
- The presentation is good but can still be improved. The core method part can be improved with additional intuitive explanations and better use of notations. Further, I have noted down some minor typo/notation errors:

Typo or notation:
  - L258: $N_i^+2$ should probably be $N_i^+$
  - L526: Ecoder -> Encoder 
  - L200: probably should better use $N_i$ instead of $|S_i|$
  - L281: $w_i'*N'$ appears twice.
  - Figure 2 caption: meaning of (original = 1) is a bit unclear
  - L380: $N$ is missing in the first equation.

questions:
- I am a bit unclear about your definition of "equivalent data size" at L243, what's the equivalence about (i.e., which size and which size)? Note that I understand the meaning of $N_I^0$, just wondering the terminology here.
- Maybe I missed something, but how do one control the budget for the next $N^(3)$? It seems the amount of tokens is defined by the initial weights of  $N^(1)$ and $N^(2)$. Or in other words, say I need to find a optimal weight for a total token of 300B, how should I start with $N^(1)$ and $N^(2)$?
- Adding to the prior question, if the optimal ratio of each domain follows a exponential function, after taking a few data points using AutoScale, can we simply fit the exponential function instead of using the AutoScale iterative method? You seem to be using that in Figure 1 (d). If y es, this simply answer my question above.
- While the problem of different data scale is resolved with a scaling law solution, can we also use a similar approach on model scale? Even though the cost of using a small amount of data for a larger model should be within a low percentage of the total training cost, setting up the experiment for the larger scale is non-trivial. It'd be nice to have a function that can predict the loss across model scales.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731107202067

thread_type:
official_review

<thread_object>
creator: reviewer_qo6p

<content>
summary:
The authors address an interesting topic in this paper: a method for automatically optimizing the mixture proportions of pretraining data domains when training language models. 
They begin by formulating the optimal mixing problem as a bi-level optimization and then propose the Direct Data Optimization (DDO) algorithm to formalize the relationship between optimal data compositions and training data scales. Using DDO, they conduct empirical studies to optimize domain weights at various training data scales, demonstrating that the optimal data composition varies with the scale of the training data. Finally, they introduce AutoScale, which automatically predicts optimal training data compositions at larger scales based on compositions optimized at smaller scales. 
Additionally, their evaluation of AutoScale on both decoder-only and encoder-only models demonstrates its ability to achieve computational savings.

strengths:
1. AutoScale presents an interesting idea that distinguishes it from previous work, demonstrating that the optimal weights are only effective at the scale they were optimized for and become suboptimal when applied to other scales. It offers a practical method for automatically and efficiently determining domain weights when train large language models. 
2. The experiments are conducted on both encoder-only and decoder-only models and shows good results on decoder-only model. 
3. The work is supported by both empirical experiments and mathematical formulations. Additionally, the diagram in the paper is well-designed and effectively conveys the underlying concepts.

weaknesses:
1. The experimental setup is not entirely convincing:

* The models used (a 774M decoder-only model and a 110M encoder-only model) are relatively small compared to today’s large language models, making it difficult to gauge performance at a larger scale.
* The data size is limited to 3B, 5B, and 10B tokens, with results in Table 1 only reflecting the 3B set.
* Figure 3(b) lacks explanation, and the cola baseline and DDO performance seems unusually low, falling below random guessing (0.5). Also, stsb baseline seems low too.

2. The evaluation of downstream tasks could be expanded. It would be helpful to see the models' performance on more complex tasks, such as mathematical problem-solving.

questions:
1. If I understand correctly, for the downstream tasks, the evaluation metric used is perplexity. Why is perplexity chosen as the metric instead of one that is specific to the dataset or task itself?
2. Is there any potential explanation for why AutoScale doesn't perform as well on encoder-only models compared to decoder-only models?
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper attempts to tackle the challenging problem of optimizing data mixing ratios for language model training across different scales. The authors propose a computationally efficient approach by first analyzing domains independently to derive power law scaling parameters, then solving a convex optimization problem to determine optimal mixing ratios. They further introduce AutoScale, an iterative method for predicting mixing ratios at larger scales by interpolating between known solutions, addressing the practical challenge that optimal data mixing strategies often evolve with model scale.

The paper's key strengths lie in its decomposition of an intractable optimization problem into manageable components and its principled mathematical framework. The empirical validation across both autoregressive (GPT-2) and bidirectional (BERT) architectures demonstrates practical utility. The work effectively transforms what has historically been a largely empirical process into a more systematic approach, potentially reducing the computational overhead in developing large language models.

This paper exhibits several notable limitations that could be improved. First, as mentioned in most reviewers, the experimental scope is constrained by using only 3B tokens (in Table 1, the most important table to show the improvements), which is insufficient for a paper studying language model pre-training and should be extended to larger token counts. In practice, we usually use Training-Compute Optimal data allocation for a model pre-training, and it should be nearly 10B tokens for a 774M GPT model. Second, the discussion with previous work is inadequate, particularly in its omission of the Data Mixing Laws (DML) work, which pioneered the application of scaling laws to data mixing - this relationship should be explicitly addressed to properly contextualize the current contributions. Additionally, the methodology used to establish scaling laws raises concerns, as traditional scaling law studies typically analyze thousands of data points to demonstrate consistent patterns, while this study's limited data points (at most 1.2B token budget) make it difficult to confidently characterize the observed relationships as scaling laws. Finally, I suggest reconsidering the term "optimal data composition" used throughout the paper, as it may be over claiming the findings. True optimality would require exhaustive testing of all possible data compositions, which is nearly impossible. A more precise term, such as "effective data composition" or "improved data composition" would better reflect the actual extent of the optimization conducted. Additionally, the absence of comparisons against DoReMi on the original Pile dataset leaves some questions on the experiments, also.

Overall, reviewers have a mixed feeling of this paper. I have carefully read all the reviews and the author rebuttal, and have read the paper again. While the mathematical framework and formulation are commendable, the experimental validation requires substantial strengthening.