[Paper Review] Training Compute-Optimal Large Language Models

Jordan Hoffmann, Sebastian Borgeaud+20 more|arXiv (Cornell University)|Mar 29, 2022

Topic ModelingComputer Science651 citations

TL;DR

The paper shows that for compute-optimal training of transformer LLMs, model size and training tokens should scale in roughly equal proportions; they validate this by training a compute-optimal 70B model (Chinchilla) that outperforms larger models like Gopher and GPT-3 on downstream tasks.

ABSTRACT

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$ imes$ 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.

Motivation & Objective

• Determine the optimal allocation of compute between model size and training data for pre-training language models under a fixed FLOPs budget.
• Test whether current large models are undertrained and whether equal scaling of parameters and data yields better performance.
• Develop and evaluate a compute-optimal model (Chinchilla) that follows the predicted scaling and compare to existing large models.

Proposed method

• Train over 400 models ranging from 70M to >16B parameters on 5B to >400B tokens to map the FLOPs–loss frontier.
• Fit three approaches (fixed model size with varying tokens; IsoFLOP profiles; parametric loss model) to estimate optimal N and D as functions of compute C.
• Assume a power-law compute relationship with model size and data; derive N_opt(C) ~ C^a and D_opt(C) ~ C^b with a ≈ b ≈ 0.5.
• Train Chinchilla (70B parameters, 1.4T tokens) under the same compute budget as Gopher to validate the predicted frontier.
• Compare Chinchilla to Gopher, GPT-3, Jurassic-1, and MT-NLG across multiple benchmarks including MMLU, BIG-bench, LAMBADA, and closed-book QA.
• Report that Chinchilla achieves state-of-the-art performance on several metrics, e.g., 67.5% average on MMLU.

Experimental results

Research questions

• RQ1 What is the optimal allocation of training compute between model size and training tokens for pre-training transformers under a fixed FLOPs budget?
• RQ2 Do model size and data scale equally with increased compute to achieve compute-optimal performance?
• RQ3 Does training a compute-optimal smaller model on more data (Chinchilla) outperform larger, more data-starved models on downstream tasks?

Key findings

• Optimal scaling: both model size and training tokens should scale in near-equal proportions with compute (N_opt ∝ C^a, D_opt ∝ C^b with a≈0.5, b≈0.5).
• Current large LLMs are undertrained for their compute budgets; compute-optimal models require more training tokens than commonly used with large parameter counts.
• Chinchilla (70B, 1.4T tokens) outperforms Gopher (280B, 300B tokens), GPT-3 (175B, 300B tokens), Jurassic-1 (178B, 300B tokens), and MT-NLG (530B, 270B tokens) on many downstream tasks.
• Chinchilla achieves a state-of-the-art average 67.5% on MMLU, exceeding Gopher by over 7 percentage points.
• Chinchilla reduces inference and fine-tuning compute due to smaller parameter count while delivering better performance across diverse tasks (The Pile, MMLU, BIG-bench, LAMBADA, RACE).
• Paper also discusses risks such as bias and toxicity associated with large models.