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

Dear Reviewer Ad96,

As the rebuttal/author discussion period ends in a few days, we sincerely look forward to your feedback. The authors are deeply appreciative of your valuable time and efforts spent reviewing this paper and helping us improve it.

It would be very much appreciated if you could once again help review our responses and additional results and let us know if these address or partially address your concerns and if our explanations are heading in the right direction :)

Please also let us know if there are any further questions or comments about this paper. We strive to consistently improve the paper and it would be our pleasure to have your precious feedback!

Kind Regards,
Authors of Submission1690

We thank the reviewer for the review and appreciate the feedback. Here we provide explanations for the listed weaknesses and responses to the questions.

1. Experiment setup

a. Learning rate schedule

This work follows the standard setup for training GPT-2-style LMs with a linear learning rate schedule and a 10% warmup ratio. The learning rate schedule and warmup steps are relative to the total steps in each training run, automatically scaled for each experiment. We documented the hyperparameters for pre-training GPT-2 Large models in Table 2.

b. Number of epochs for different domains

We agree with the argument that sampling different domains with different numbers of epochs could lead to misaligned results and confusing comparisons. Thus, in this work, we strictly controlled this factor to ensure in every training run, regardless of its mixture, we train on data from each domain with the same number of epochs. The source dataset used in this work, RedPajama, is large enough so we are always able to obtain the target amount of non-repeating samples for every data mixture.

c. Comparative evaluation

To the best of our knowledge, there are no existing solutions for determining the ground truth optimal domain weights, except by computationally intractable exhaustive grid search with training the full model on all combinations of domain weights. [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining, 2023] is one of the earliest works on optimizing training data mixture for pre-training LLMs and has been widely referred to as the main baseline in subsequent works. Thus, we formulate the data optimization problem in Section 3.1 as a bilevel optimization problem where no direct solution is computationally tractable. Consequently, we developed the Direct Data Optimization (DDO) algorithm in Section 3.2 which leverages scaling law analysis to provide a global approximation to the original optimization problem, achieving high-accuracy optimization results with plausible computation complexity. This already pushed forward the frontier of solving the original data optimization problem. But still, we do not settle for a solution that requires training the model with a full data budget multiple times. This led to the development of the novel AutoScale tool, which automatically predicts optimal training data compositions at larger scales based on compositions optimized at smaller scales.

D. Computation overhead

DDO requires training $2M +1$ models where $M$ is the number of data domains. This is conducted on a data scale magnitudes smaller than the target compute budget. AutoScale requires obtaining optimized domain weights with DDO on two different data scales. This allows for predicting the domain weights for training the target model with a much larger data budget. In comparison, the full data budget for 1B-parameter models is 2~3T tokens** (ref: 1.5T tokens in [OpenELM: An Efficient Language Model Family with Open Training and Inference Framework], 2T tokens in [TinyLlama: An Open-Source Small Language Model], 3T tokens in [olmo: Accelerating the science of language models]) where a 0.5% speed-up would justify the cost. We are adding this discussion to the manuscript and the limitations.

Besides, the computation efficiency for AutoScale is advantageous relative to existing works.* In the default setup, in each iteration, DoReMi requires training two models each with 200k steps with token_length 1024 and batch_size 64*8, which equals >100B tokens and ~10x than AutoScale. This process needs to iterate multiple times, resulting in a multitude of compute. This has not included the computation for conducting the grid search to determine which proxy model to use. [Rephrasing the Web: A Recipe for Compute and Data-Efficient Language Modeling] explores improving LLM pretaining efficiency via training on synthetic texts generated by pre-trained LLMs. The computation overhead for generating synthetic texts with 7B-parameter models is in the same magnitude as training the target 1B-parameter model.

Further, as listed in the discussion with Reviewer mxrW, even though the cost of training full-size models on a small proportion of its target data budget should be within a low percentage of the total training cost, there's room for further improvements. The next step could be applying the scaling law solution developed for different data budgets to different model sizes and directly predicting optimal data composition across model scales. The original paper on scaling laws reports that model performance is empirically predictable with power law functions w.r.t. to either the training data size OR the model size–these two were treated as equal, independent factors. Theoretically, the methodology we developed in this work could be applied to scaling with the model size as well. This extension could be an independent work parallel to this manuscript. Should the results turn out to be positive, then the next step will be combining these two methods into a unified framework to predict optimal data mixtures for different compute budgets and model sizes at the same time. We are listing this in future directions and calling for contributions from the community.

Dear Reviewer JgHF,

As the rebuttal/author discussion period ends in a few days, we sincerely look forward to your feedback. The authors are deeply appreciative of your valuable time and efforts spent reviewing this paper and helping us improve it.

It would be very much appreciated if you could once again help review our responses and additional results and let us know if these address or partially address your concerns and if our explanations are heading in the right direction :)

Please also let us know if there are any further questions or comments about this paper. We strive to consistently improve the paper and it would be our pleasure to have your precious feedback!

Kind Regards,
Authors of Submission1690

Dear Reviewer Qo6P,

As the rebuttal/author discussion period ends in a few days, we sincerely look forward to your feedback. The authors are deeply appreciative of your valuable time and efforts spent reviewing this paper and helping us improve it.

It would be very much appreciated if you could once again help review our responses and additional results and let us know if these address or partially address your concerns and if our explanations are heading in the right direction :)

Please also let us know if there are any further questions or comments about this paper. We strive to consistently improve the paper and it would be our pleasure to have your precious feedback!

Kind Regards,
Authors of Submission1690

The authors sincerely appreciate the reviewer's consideration and value the reviewer's feedback! We strive to consistently improve the manuscript and serve its goal to contribute to the research community.

Concern 1(a): accuracy vs. MCC

Re: We fully understand the reviewer's question. In the current literature, there appears no consensus on the primary evaluation metrics for GLUE tasks. The original paper proposing the GLUE benchmark, [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding], lists Matthews corr. as the evaluation metric for "cola", Pearson/Spearman corr. for "stsb", acc./F1 and "mrpc" and "qqp", and accuracy for other tasks. This is how we chose evaluation metrics in this work.

Nonetheless, we did notice the considerable wide range of choices for evaluation metrics of the GLUE benchmark among published works, including seminal papers. For example, in Google's original paper proposing the BERT model, [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding], the evaluation on GLUE benchmark excludes the "wnli" task and reports F1 scores for "qqp" and "mrpc", Spearman correlations for "stsb", and accuracy scores for the other tasks. [Data Selection for Language Models via Importance Resampling] reports accuracies for all tasks. We are neutral about which metric to use and agree that unifying the experiment setup will be beneficial, which would cause less confusion and allow better comparison between different works.

Concern 1(b): low baseline performance for "stsb"

Re: While developing this work, we noticed the performance on task "stsb" demonstrates the phenomenon often referred to as "emergent ability", which describes the scenarios where the task performance of an LLM improves abruptly in an unpredictable manner while increasing the model size/compute budget (ref: [Emergent Abilities of Large Language Models]). In this work, we observed the performance on task "stsb" increases sharply when the training budget reaches certain thresholds. In the experiment, for the same training data budget, DDO-optimized domain weights led to higher training efficiency and the trained model observed a sharp improvement in "stsb" performance. On the contrary, the model trained with uniform domain weights did not reach the threshold for performance hike, making the performance discrepancy appear wide.

It was not the intention of this work to design experiments in this specific manner. Rather, as we stressed in the review responses, we note that task performance is generally less predictable and better used as a qualitative measure. Task performance does not scale smoothly with computing budget/model size and it is currently an open question whether/how task performance for LLMs can be predicted. Perplexity/loss, on the contrary, scales stably and serves as a better indicator to track the model's training progress/capability. We recommend considering validation perplexity as the primary metric for efficiency research on LLM pre-training, which allows more precise quantitative comparisons.

Concern 2: relative improvements in Takeaway 3

Re: We appreciate the feedback. We will revise the manuscript to clearly state that the improvements are in the training efficiency where AutoScale-predicted weights train the model 28%–38% faster than baselines.