PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Compute-Optimal Large Language Models

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, Laurent Sifre

arXiv:2203.15556v1cs.CLcs.LG

TL;DR

Choosing model size and training data under a fixed compute budget is critical because large-language-model training is costly. Using over 400 training runs and predictive scaling approaches, the paper finds that Chinchilla, a smaller model trained on more tokens, outperforms Gopher and larger models on nearly every measured evaluation task.

  • Problem

    The paper asks how to choose model size and training-token count for a fixed compute budget, an important question because large-language-model training is costly.

  • Method

    The authors analyze over 400 training runs and predictive compute-optimal approaches that minimize loss under a fixed FLOP budget.

  • Results

    Chinchilla, a 70B-parameter model trained on 1.4 trillion tokens, outperforms Gopher and larger models on nearly every measured evaluation task.

  • Takeaways & Limitations

    A smaller, more-data-trained model can deliver better performance at the same compute budget while reducing inference costs and facilitating downstream use on smaller hardware.

  • Takeaways & Limitations

    At large scale, the analysis has only two comparable training runs and no additional tests at intermediate scales.

Abstract

from arXiv · show

We investigate the optimal model size and number of tokens for training a transformer language model under a given compute budget. We find that current large language models are significantly undertrained, a consequence of the recent focus on scaling language models whilst keeping the amount of training data constant. By training over 400 language models ranging from 70 million to over 16 billion parameters on 5 to 500 billion tokens, we find that for compute-optimal training, the model size and the number of training tokens should be scaled equally: for every doubling of model size the number of training tokens should also be doubled. We test this hypothesis by training a predicted compute-optimal model, Chinchilla, that uses the same compute budget as Gopher but with 70B parameters and 4$\times$ more more data. Chinchilla uniformly and significantly outperforms Gopher (280B), GPT-3 (175B), Jurassic-1 (178B), and Megatron-Turing NLG (530B) on a large range of downstream evaluation tasks. This also means that Chinchilla uses substantially less compute for fine-tuning and inference, greatly facilitating downstream usage. As a highlight, Chinchilla reaches a state-of-the-art average accuracy of 67.5% on the MMLU benchmark, greater than a 7% improvement over Gopher.

1. Introduction

The paper asks how to optimally trade off model parameters and training tokens under a fixed compute budget. Using over 400 models, it predicts that current large models are too large and undertrained, then validates the result with Chinchilla.

  • Motivation: Training compute and energy costs rise with model size, making accurate hyperparameter selection critical when the available compute budget is fixed.The budget may be determined by the number of accelerators and the intended training duration.
  • Problem formulation: The study minimizes pre-training loss L(N, D) subject to a fixed FLOPs budget C, where N is parameter count and D is training-token count.The optimal allocations are represented by N_opt(C) and D_opt(C).
  • Method: Over 400 models spanning under 70M to over 16B parameters and 5B to over 400B tokens provide the empirical basis for estimating the compute-optimal frontier.Each model configuration was trained for several different training horizons.
  • Main result: 4× smaller and 4× more tokens than Gopher’s configuration is predicted to be compute-optimal for the same training budget.The prediction is verified with Chinchilla, a 70B-parameter model trained on 1.4 trillion tokens.
  • Main result: Chinchilla outperforms Gopher while its reduced model size considerably lowers inference cost and facilitates downstream use on smaller hardware.Most recently trained large models use approximately 300B tokens, whereas Chinchilla is trained substantially longer.

2. Related Work

Related work spans scaling laws, hyperparameter selection, and alternative architectures for large language models. Prior studies examine predictable model-size–loss relationships, broader training choices, and conditional computation that increases the effective amount of training data.

  • Large language models: Large language models include dense transformers and mixture-of-expert models, with the largest dense transformers exceeding 500 billion parameters.The cited examples include models from Brown et al., Lieber et al., Rae et al., Smith et al., and Thoppilan et al., alongside MoE models from Du et al., Fedus et al., and Zoph et al.
  • Modelling the scaling behavior: Kaplan et al. showed a predictable relationship between model size and loss across many orders of magnitude, motivating compute-budget-dependent model-size selection.The paper investigates choosing the optimal model size for a given compute budget and builds on prior work on scaling behavior and transfer properties.
  • Modelling the scaling behavior: Mixture-of-expert scaling with the number of experts diminishes as model size increases, but fixed training-token analyses may underestimate branching improvements.Clark et al. model loss using model size and number of experts, while retaining the fixed-token setup used by Kaplan et al.
  • Estimating hyperparameters for large models: Model size and training tokens are only part of language-model design; other choices include learning rate, schedule, batch size, optimiser, and width-to-depth ratio.This work focuses on model size and training steps, using existing research and experimental heuristics for the remaining hyperparameters.
  • Improved model architectures: Conditional-computation architectures include the 1.7 trillion parameter Switch transformer and 1.2 Trillion parameter GLaM, with one approach increasing training tokens by a factor of ∼10.These results suggest that language-model performance may depend more on training-data size than previously thought.

3. Estimating the optimal parameter/training tokens allocation

Three complementary approaches produce comparable compute-optimal scaling predictions: model parameters and training tokens should increase in approximately equal proportions as compute grows. The analysis also indicates that current large language models are oversized for their compute budgets and should instead be trained on substantially more tokens.

  • Approach 1: a = 0.50 and b = 0.50 summarize the estimated compute allocation exponents for model size and training tokens.These estimates come from the first approach’s training-curve analysis.
  • Approach 2: A clear loss valley across IsoFLOP curves identifies an optimal model size for each fixed FLOP budget.The second approach varies model size at fixed training FLOP counts and fits each curve to estimate its minimum.
  • Approach 3: 40B parameters is the projected optimal model size for the Gopher FLOP budget under the parametric loss fit.The third approach models final loss as a function of parameter count and seen tokens.
  • Convergent scaling predictions: Three approaches yield comparable predictions that model size and training data should increase in approximately equal proportions as compute increases.The approaches use different fitting methodologies and different trained models.
  • Implications for existing models: 4.41 × 10^24 FLOPs and over 4.2 trillion tokens are recommended for a 175 billion parameter model.A 280 billion Gopher-like model is instead optimal at approximately 10^25 FLOPs and 6.8 trillion tokens.
  • Additional datasets: C4 and GitHub code reproduce the conclusion that model size and training tokens should be scaled in equal proportions.The IsoFLOP analysis reaches the same conclusion on both additional datasets.

4. Chinchilla

Chinchilla is a 70B-parameter model trained on 1.4T tokens under the same compute budget as Gopher, and it substantially outperforms Gopher and other large language models across diverse evaluations. Its improvements extend to language modeling, multitask reasoning, question answering, and bias-related evaluations, while toxicity statistics remain broadly similar to Gopher’s.

  • Model and training: 70B parameters and 1.4T tokens define Chinchilla’s compute-matched comparison with Gopher.The model size was selected from the predicted optimal range of 40–70 billion parameters, and both models used the same number of FLOPs.
  • Model and training: 4× smaller than Gopher, Chinchilla has a smaller memory footprint and lower inference cost.These reductions facilitate downstream fine-tuning and inference.
  • Evaluation results: 67.6% average 5-shot accuracy on MMLU improves upon Gopher by 7.6%, despite Chinchilla being much smaller.Chinchilla improves performance on the vast majority of MMLU tasks, underperforming Gopher on four and matching it on two.
  • Evaluation results: 77.4% accuracy on LAMBADA exceeds Gopher’s 74.5% and MT-NLG 530B’s 76.6%.Chinchilla also improves RACE-h and RACE-m accuracy by more than 10% over Gopher in both cases.
  • Evaluation results: 31.5% 5-shot and 35.5% 64-shot accuracy on Natural Questions exceed Gopher’s 21% and 28%, respectively.These results establish new closed-book state-of-the-art accuracies on the reported Natural Questions settings.
  • Bias and toxicity: 8.3% and 9.2% gains for female and neutral pronouns exceed the 3.2% gain for male pronouns on Winogender, while toxicity summary statistics show no major differences from Gopher.The largest gender-bias improvement is 10% on female gotcha examples; mean toxicity scores were 0.081 for Gopher and 0.087 for Chinchilla.

5. Discussion & Conclusion · Appendix

The paper argues that compute-optimal language-model training requires balancing model size with training data, and validates this with Chinchilla’s stronger performance at the same compute budget. It also identifies limitations, emphasizes high-quality dataset scaling, and suggests the methodology may extend beyond autoregressive language models.

  • 5. Discussion & Conclusion: Existing large models increasingly expand model size without increasing training tokens, with MT-NLG 530B exceeding GPT-3’s 170 billion parameters while both use around 300 billion tokens.The passage frames this pattern as the motivation for reassessing how compute is allocated during training.
  • 5. Discussion & Conclusion: Three predictive approaches based on over 400 training runs identify Gopher as substantially over-sized and favor smaller models trained on more data at equal compute.The hypothesis was directly tested with Chinchilla, a 70B-parameter model that outperformed Gopher and even larger models on nearly every measured evaluation task.
  • 5. Discussion & Conclusion: The analysis is limited by having only two comparable large-scale training runs and no additional tests at intermediate scales.The methodology also assumes that the efficient computational frontier follows a power-law relationship involving compute budget, model size, and training tokens.
  • 5. Discussion & Conclusion: The authors caution that their methodology may still overestimate the optimal size of large models at high compute budgets.This observation is presented as an additional limitation of the scaling analysis.
  • 5. Discussion & Conclusion: Chinchilla validates the performance predictions by providing a better and more lightweight model at the same compute budget.The comparison with Gopher supports the predicted advantage of reallocating compute toward training data rather than model size.
  • 5. Discussion & Conclusion: The study’s training runs used less than one epoch of data, leaving the multiple-epoch regime for future work.The authors state that these limitations do not negate the Chinchilla–Gopher validation.
  • 5. Discussion & Conclusion: The analysis calls for greater focus on scaling high-quality datasets, with careful accounting for train-test overlap as datasets grow.The authors present this as a responsible approach to collecting larger datasets.
  • 5. Discussion & Conclusion: The authors expect a similar trade-off between model size and data amount in other modalities, making advance choices of model size and training steps essential.They also describe the proposed methods as easy to reproduce in new settings.

A. Training dataset · B. Optimal cosine cycle length

The training dataset uses a specified MassiveText composition, with some subsets repeated across epochs. The training schedule assumes a 10× learning-rate decay and calibrates cosine-cycle length to the maximum steps allowed by the FLOP budget.

  • A. Training dataset: MassiveWeb and Wikipedia are each used for more than one epoch in the training dataset.This reuse applies both to Chinchilla and to the scaling runs.
  • A. Training dataset: The Chinchilla and scaling runs use the MassiveText training dataset described in Table A1.Table A1 reports each subset’s disk size, document count, sampling proportion, and epochs used in 1.4 trillion tokens.
  • A. Training dataset: MassiveText subsets are assigned distinct sampling proportions, with alternative proportions from Rae et al. (2021) shown for comparison.The table lists total disk size and number of documents for each subset alongside these proportions.
  • A. Training dataset: 1.4 trillion tokens is the reference training volume for reporting the number of epochs used by each MassiveText subset.The epoch counts appear in the table’s rightmost column.
  • B. Optimal cosine cycle length: The schedule uses a 10× learning-rate decay, following Rae et al. (2021), together with an assumption about cosine-cycle length.The passage identifies the cosine-cycle length and corresponding learning-rate drop as a key assumption.
  • B. Optimal cosine cycle length: Cosine cycles much longer than the target training steps produce sub-optimally trained models.The study therefore assumes the cosine cycle is calibrated to the maximum number of steps permitted by the FLOP budget.
  • B. Optimal cosine cycle length: The main analysis calibrates cosine-cycle length to the maximum training steps determined by the FLOP budget.This calibration is treated as the condition for optimal training under the assumed schedule.

C. Consistency of scaling results across datasets · D. Details on the scaling analyses

IsoFLOP analyses on C4 and GitHub closely reproduce MassiveText scaling behavior, indicating dataset-independent results when training stays within one epoch. The analyses also show that a 10× learning-rate decay is slightly preferable to decay to zero, while 5× decay performs clearly worse.

  • C. Consistency of scaling results across datasets: C4 and GitHub IsoFLOP experiments use subsets of the alternate datasets to test scaling beyond MassiveText.The analyses use C4 and GitHub data, with the same tokenizer as the MassiveText experiments.
  • C. Consistency of scaling results across datasets: The scaling behavior on C4 and GitHub is very similar to MassiveText’s.The comparison is shown in Figure A2 and Table A2.
  • C. Consistency of scaling results across datasets: One epoch is the stated condition for dataset-independent scaling results.The paper suggests the results remain independent of dataset as long as training does not exceed one epoch.
  • D. Details on the scaling analyses: 10× decay is slightly more performant than decaying to 0.0 over the same number of steps.The difference between the two schedules is small.
  • D. Details on the scaling analyses: 5× decay is clearly worse than the alternatives tested.The passage contrasts decaying by 5× with both 10× decay and decay to zero.
  • D. Details on the scaling analyses: 25% overestimation of training steps leads to clear performance drops when the cosine cycle is too long.The analysis varies cosine cycle length from 1× to 5× the target training steps and attributes the impairment to insufficient learning-rate reduction.
  • D. Details on the scaling analyses: The alternate-dataset scaling analysis estimates exponents a and b for optimal parameters and data as functions of compute.The relationships are N_opt ∝ C^a and D_opt ∝ C^b, estimated from IsoFLOP profiles.

D.1. Approach 1: Fixing model sizes and varying training sequences … D.4. Small-scale comparison to Kaplan et al. (2020)

The paper specifies training-schedule choices, decomposes loss into approximation and stochastic terms, projects compute-optimal frontiers, and validates its approach against Kaplan et al. (2020). At 10^21 FLOPs, the authors’ predicted model outperforms Kaplan et al.’s prediction.

  • D.1. Approach 1: Fixing model sizes and varying training sequences: Learning rates range from 2 × 10^-4 for the smallest models to 1.25 × 10^-4 for the largest, then decay 10× with a cosine schedule.The cosine cycle is approximately matched to the number of training steps.
  • D.1. Approach 1: Fixing model sizes and varying training sequences: A cosine cycle overshooting training by more than 25% noticeably degrades performance.The training curves are smoothed with a 10-step Gaussian window.
  • D.2. Approach 3: Parametric fitting of the loss: The loss decomposition contains Bayes risk, a functional approximation term dependent on parameter count, and a stochastic approximation term from finite-data optimization.The stochastic term also reflects making a single epoch over the provided dataset.
  • D.2. Approach 3: Parametric fitting of the loss: The approximation term is expected to scale as 1/N^1/2, while the stochastic term is lower-bounded by 1/D^1/2.These forms motivate how model size and dataset size enter the fitted loss.
  • D.2. Approach 3: Parametric fitting of the loss: The fitted decomposition uses E = 1.69, A = 406.4, and B = 410.7.The paper notes that both parameter/data coefficients are below 2.
  • D.3. Predicted compute optimal frontier for all three methods: Approaches 2 and 3 estimate optimal model sizes and training-token counts across compute budgets, alongside Approach 1’s predicted frontier.These projections are reported in Table A3 and Figure A3.
  • D.4. Small-scale comparison to Kaplan et al. (2020): At 10^21 FLOPs, Approach 1 predicts 2.86 billion parameters versus Kaplan et al.’s 4.68 billion, and the trained 2.80-billion-parameter model outperforms the 4.74-billion-parameter model.Both models use a batch size of 0.5M tokens and a maximum learning rate of 1.5 × 10^-4 decaying by 10×.

E. Curvature of the FLOP-loss frontier

The FLOP-minimal loss frontier curves as model size increases, so projections from small models differ from those from larger models. The analysis does not account for this curvature, leaving open the possibility that smaller models may be optimal at large FLOP budgets.

  • The FLOP-minimal loss frontier exhibits curvature as models increase.
  • Projections from very small models produce different predictions than projections from larger models.Figure A5 compares linear fits from the first, middle, and final thirds of frontier points.
  • The analysis omits frontier curvature, leaving as future work whether smaller models may be optimal for large FLOP budgets.

F. FLOPs computation

The FLOP accounting includes embeddings and uses detailed attention, dense-block, and logits costs, with a twofold backward-pass assumption. Its results closely match the common 6ND approximation, while yielding a slightly more accurate calculation.

  • F. FLOPs computation: Training FLOPs include embedding matrices, which are also included in the total parameter count; their contribution is small for large models.The calculation uses a factor of 2 for multiply-accumulate cost.
  • F. FLOPs computation: Total forward-pass FLOPs equal embeddings + num_layers × (total_attention + dense_block) + logits.Attention costs include projections, attention logits, softmax, query reductions, and the final linear operation.
  • F. FLOPs computation: 2× forward-pass FLOPs are assumed for the backward pass, following Kaplan et al. (2020).The paper compares this calculation with the common approximation C = 6DN, where C is FLOPs, D is training tokens, and N is parameters.
  • F. FLOPs computation: The detailed and 6ND-based FLOP calculations differ very little and do not impact the analysis.Table A4 reports the ratio of computed FLOPs per sequence to those from the 6ND approximation across model sizes.
  • F. FLOPs computation: 6.3 × 10^23 is the slightly more accurate calculation, compared with 5.76 × 10^23 from the alternative result.This comparison is presented against results from Rae et al. (2021).

G. Other differences between Chinchilla and Gopher … H.4. BIG-bench

The appendix distinguishes Chinchilla from Gopher by optimizer and optimizer-state precision, finding advantages for the Chinchilla setup. It then reports comparative results across The Pile, MMLU, Winogender, and BIG-bench evaluations.

  • G. Other differences between Chinchilla and Gopher: Chinchilla used AdamW and stored a higher-precision weight copy in the sharded optimizer state, whereas Gopher used Adam.These are additional differences beyond model size and training-token count.
  • G. Other differences between Chinchilla and Gopher: AdamW-trained models outperform Adam-trained models independently of the learning-rate schedule.The comparisons are shown in Figures A6 and A7.
  • G. Other differences between Chinchilla and Gopher: The Chinchilla setup clearly outperforms the Gopher setup in an 680 million parameter comparison combining optimizer and weight-copy changes.The experiment compares setups with and without the higher-precision weight copy and with Adam versus AdamW.
  • H.1. The Pile: Chinchilla outperforms Gopher on every subset of The Pile, while Jurassic-1 outperforms Chinchilla on dm_mathematics and ubuntu_irc.The comparison uses bits-per-byte on The Pile.
  • H.2. MMLU: Table A6 reports Chinchilla and Gopher performance on each MMLU subset.The passage identifies the table and comparison but provides no subset-level values.
  • H.3. Winogender Setup: Winogender testing measures pronoun-reference probabilities across female, male, and gender-neutral pronouns, with unbiased models resolving references regardless of pronoun gender.The setup follows Rae et al. (2021) and compares candidate sentence roles.
  • H.4. BIG-bench: Table A7 reports Chinchilla and Gopher performance on each considered BIG-bench subset.The passage specifies the comparison scope but provides no subset-level values.

I. Model Card · J. List of trained models

The model card defines Chinchilla’s intended use, evaluation scope, preprocessing, and documented limitations, while Table A9 lists the configurations of all models trained in the study. Chinchilla is intended for research rather than public release or unmitigated downstream use, and the study reports toxicity, bias, and variability considerations.

  • I. Model Card: Chinchilla is intended for DeepMind researchers and will not be made publicly available.
  • I. Model Card: The model should not be used for harmful or deceitful generation or downstream applications without further safety, fairness, and application-specific analysis.
  • I. Model Card: The model is trained on English data and can show unequal performance across dialects, including African American English.
  • I. Model Card: Evaluation covers language modeling, completion, reading comprehension, MMLU, BIG-bench, fact checking, question answering, toxicity, and gender and occupation bias.
  • I. Model Card: Because Chinchilla was not trained multiple times, evaluation breadth and comparison with Gopher trained on the same dataset provide the main basis for estimating overall performance.
  • I. Model Card: Input text uses a SentencePiece tokenizer with a 32,000-token vocabulary, without NFKC normalization, and the same dataset as Rae et al. (2021) with sampling differences documented in Table A1.
  • I. Model Card: Chinchilla can output toxic language, especially after toxic prompts, and reproduces gender-associated occupational stereotypes and differing sentiment toward some groups.
  • J. List of trained models: Table A9 lists the hyperparameters and sizes of all models trained in the study, including models trained with multiple learning-rate schedules or training-token counts.
Loading 2203.15556v1…