Source-linked AI summary

Language models scale reliably with over-training and on downstream tasks

Samir Yitzhak Gadre, Georgios Smyrnis, Vaishaal Shankar, Suchin Gururangan, Mitchell Wortsman, Rulin Shao, Jean Mercat, Alex Fang, Jeffrey Li, Sedrick Keh, Rui Xin, Marianna Nezhurina, Igor Vasiljevic, Jenia Jitsev, Luca Soldaini, Alexandros G. Dimakis, Gabriel Ilharco, Pang Wei Koh, Shuran Song, Thomas Kollar, Yair Carmon, Achal Dave, Reinhard Heckel, Niklas Muennighoff, Ludwig Schmidt

arXiv:2403.08540v2cs.CLcs.LG

TL;DR

Current scaling studies often emphasize compute-optimal training and next-token loss, although models are also over-trained to reduce inference costs and compared on downstream tasks. The paper tests scaling with over-training using a 104-model, three-distribution testbed and proposes a loss-to-downstream-error power law. It predicts over-trained validation loss and aggregate downstream error from substantially cheaper experiments.

  • Problem

    Published scaling laws mainly focus on compute-optimal training and next-token loss, despite practical use of over-trained models and downstream benchmark performance.

  • Method

    The paper trains 104 models across three data distributions, fits scaling laws over parameter count and token multiplier, and chains loss prediction with a power law relating loss to downstream error.

  • Results

    The method predicts a 1.4B, 900B-token model within 0.7% relative error using 300× less compute and predicts aggregate downstream error within 1 percentage point using 20× less compute.

  • Takeaways & Limitations

    The results support using scaling laws to derisk over-training decisions and downstream performance when scaling training recipes.

  • Takeaways & Limitations

    Downstream task performance remains a proxy for in-the-wild user experience, which the paper identifies as an area for further analysis.

Abstract

from arXiv · show

Scaling laws are useful guides for derisking expensive training runs, as they predict performance of large models using cheaper, small-scale experiments. However, there remain gaps between current scaling studies and how language models are ultimately trained and evaluated. For instance, scaling is usually studied in the compute-optimal training regime (i.e., "Chinchilla optimal" regime). In contrast, models are often over-trained to reduce inference costs. Moreover, scaling laws mostly predict loss on next-token prediction, but models are usually compared on downstream task performance. To address both shortcomings, we create a testbed of 104 models with 0.011B to 6.9B parameters trained with various numbers of tokens on three data distributions. First, we fit scaling laws that extrapolate in both the amount of over-training and the number of model parameters. This enables us to predict the validation loss of a 1.4B parameter, 900B token run (i.e., 32$\times$ over-trained) and a 6.9B parameter, 138B token run (i.e., a compute-optimal run)$\unicode{x2014}$each from experiments that take 300$\times$ less compute. Second, we relate the perplexity of a language model to its downstream task performance by proposing a power law. We use this law to predict top-1 error averaged over downstream tasks for the two aforementioned models, using experiments that take 20$\times$ less compute. Our experiments are available at https://github.com/mlfoundations/scaling.

1 Introduction

Existing scaling laws often study compute-optimal training and next-token loss, while practitioners increasingly over-train smaller models and compare them using downstream benchmarks. This paper addresses both gaps by testing extrapolation for over-training and aggregate downstream performance.

  • Scaling studies often target compute-optimal training, whereas over-training smaller models can reduce inference cost and downstream benchmarks are commonly used instead of loss for comparison.The paper identifies both regimes as mismatches between published scaling studies and practical model development.
  • A loss-to-downstream-error scaling law predicts average top-1 error across downstream tasks within 1 percentage point for a 6.9B compute-optimal model using 20× less compute.The prediction concerns aggregate performance rather than reliably predicting each individual task.
  • The results suggest these scaling laws can help derisk over-training choices and downstream performance when scaling training recipes.The experimental results and testbed are released to support further reliable-scaling research.

2 Developing scaling laws for over-training and downstream tasks

The paper develops scaling laws for over-trained models and connects validation loss to aggregate downstream error. It predicts over-trained loss by varying compute and token multiplier, then chains loss prediction with an exponential loss-to-error relation.

  • 2.1 Preliminaries: Reducible loss is modeled as L′(C) = λ · C^-η, where η controls log-log slope and λ controls the intercept.The total loss decomposes into irreducible loss E and reducible loss L′, which should decline with more data and capacity.
  • 2.1 Preliminaries: Over-training allocates compute sub-optimally by training smaller models on disproportionately more tokens, reducing inference cost despite higher loss for a fixed budget.The token multiplier M = D/N measures the ratio of training tokens to parameters, with larger M denoting more over-training.
  • 2.2 Scaling laws for over-training: Increasing the token multiplier changes λ but leaves η approximately constant, producing parallel scaling curves across over-training levels.This empirical pattern motivates a parameterization in compute C and token multiplier M whose exponent remains independent of M.
  • 2.3 Scaling laws for downstream error: Average downstream top-1 error decreases exponentially as validation loss decreases, motivating a power-law relationship between error and perplexity.The relationship is fit over aggregate performance on 17 downstream evaluations, rather than requiring accurate prediction for every individual task.
  • 2.3 Scaling laws for downstream error: The prediction pipeline fits loss from compute and token multiplier, fits downstream error from loss, then chains both models to predict error from compute and over-training.This produces the mapping (C, M) → L → Err for chosen training and validation distributions.

3 Constructing a scaling testbed

The authors build a 104-model testbed across model sizes, token multipliers, and three training distributions, using filtered compute-optimal configurations to fit scaling laws and evaluate extrapolation.

  • Model configurations: A 435-model grid search tunes width, hidden dimension, attention heads, batch size, and warmup steps to identify compute-efficient base configurations.The search uses M=20, which early experiments place near the compute-optimal frontier.
  • Model configurations: Filtering to Pareto-frontier models removes configurations between 5.2 × 10^16 and 5.2 × 10^17 FLOPs that fall below the observed trend.The authors hypothesize that these models over-perform because power-of-two batch sizes give them more optimization steps than neighboring configurations.
  • Model configurations: Four selected configurations closely track the full frontier and extrapolate the performance of 1.4B and 6.9B models in the compute-optimal setting.The trends are fit with a power law including irreducible error, with bootstrap 95% confidence intervals shown in Figure 4.
  • Scaling testbed: A 104-model testbed spans 0.011B–6.9B parameters, token multipliers M=5–640, and C4, RedPajama, and RefinedWeb training data.Runs are omitted when they require more tokens than the dataset contains.
  • Scaling-law fitting: Equation (4) is fit to predict loss across model size and token multiplier, while Equation (5) predicts downstream error; the default fits require about 100 and 1,000 A100 hours, respectively.The target extrapolations are N=1.4B, M=640 and N=6.9B, M=20.
  • Evaluation setup: Evaluation uses C4 validation loss and average top-1 error over a 17-task subset selected to provide measurable signal above chance.Relative prediction error is measured against ground-truth loss or average top-1 error.

4 Results: Reliable extrapolation

The fitted scaling laws reliably extrapolate validation loss and aggregate downstream error across model sizes and token budgets, though individual-task predictions and some out-of-distribution settings remain difficult.

  • Over-trained performance: 0.7% relative error predicts both the 1.4B, 900B-token over-trained run and the 6.9B, 138B-token near-compute-optimal run, using 300× less compute.The loss prediction extrapolates simultaneously in parameter count and token multiplier.
  • Prediction error profile: Relative loss-prediction error remains uniformly low across training distributions, token multipliers, and parameter counts, including interpolation and extrapolation ranges.Figure 5 reports this pattern for C4 evaluation across the testbed.
  • Downstream performance: Within 20× less compute, average top-1 error is predicted within 0.05% for the 6.9B, 138B-token model and 3.6% for the 1.4B, 900B-token model.The prediction chains scaling fits for loss and downstream error as functions of compute and token multiplier.
  • Downstream performance: Average downstream performance is predictable across C4, RedPajama, and RefinedWeb, whereas individual task predictions are substantially noisier.Removing the 1.4B model from the downstream fit raises the RedPajama 6.9B prediction error from 0.05% to 10.64%.
  • Limits and trade-offs: Scaling becomes unreliable at M=5 and for C4-trained models evaluated on code, while German next-token prediction remains reliably predictable.Loss scaling laws can be constructed more cheaply than downstream-error laws.

5 Related work

Prior work established power-law scaling and compute-optimal training, while later studies addressed extrapolation, architecture, inference costs, and downstream metrics; this paper extends those directions to over-training and aggregate task performance.

  • Scaling laws: Early studies observed predictable power-law scaling with training-set size and model parameters, while later work emphasized extrapolation and scaling architectural or hyperparameter choices.The cited studies include Alabdulmohsin et al. [2], Yang et al. [124], and Bi et al. [13].
  • Compute-optimality: Hoffmann et al. [45] formalized compute-optimal scaling with N ∝ C^0.5 and D ∝ C^0.5, but did not verify extrapolation to larger models or token budgets.Their formulation minimizes loss for a fixed compute budget.
  • Inference costs: Sardana & Frankle [98] incorporate inference budgets and find that sufficiently large budgets favor training smaller models for longer than Chinchilla recommendations.This work motivates the paper’s focus on over-training rather than solving a new inference-cost optimization problem.
  • Downstream evaluation: Prior downstream-scaling work relates pre-training tokens to downstream cross-entropy or BLEU, whereas this paper evaluates average top-1 error across many natural-language tasks.Schaeffer et al. [100] argue that emergent abilities can reflect nonlinear metrics and propose smoother alternatives.
  • Over-training: Popular models already span substantial over-training, including Chinchilla 70B at M=20 and LLaMA-2 7B at M=290.The paper examines M=5–640 to cover these regimes and possible future models.

6 Limitations, future work, and conclusion

The paper concludes that over-trained-model loss and average downstream performance are predictable, while identifying boundaries that motivate further validation and broader evaluation. Future work includes cheaper scaling experiments, larger and harder models, post-training effects, per-task prediction, real-world performance, dataset curation, and failure conditions.

  • Scaling up: The reported trends still require validation at larger scales and on harder evaluations such as MMLU.
  • Scaling down: Cheaper experiments remain important, especially for making downstream error prediction more accessible.
  • Post-training: The study does not examine post-training interventions, leaving the benefits of over-training base models after fine-tuning unresolved.
  • Individual downstream task prediction: Average error is predictable across many tasks, but individual downstream-task prediction remains future work.
  • Scope and failure cases: The scope excludes dataset curation and in-the-wild user experience, while conditions causing scaling to break down remain insufficiently understood.
  • Conclusion: Over-trained-model loss is predictable, and a scaling law relates loss to average downstream task performance.The authors provide a public testbed to support further research on reliable scaling.

B Scaling-law derivations

The derivations reparameterize scaling laws using compute and token multipliers to explain over-training. They show that the power-law exponent can remain constant while the scalar changes as token-to-parameter ratios increase.

  • Compute reparameterization: The compute constraint C = 6ND and token multiplier M = D/N reparameterize model size and training tokens.
  • Relation to compute-optimal training: The derivation preserves the power-law exponent under over-training while allowing its scalar to change.
  • Over-training: Over-training modifies model size and token count while keeping compute constant relative to the compute-optimal setting.
  • Token multiplier: The multiplier m is a multiple of the Chinchilla-optimal token multiplier M∗ = D∗/N∗, which varies with compute when α ≠ β.
  • Experimental configuration: The model configurations use parameter counts, architecture widths, depth, attention heads, and batch sizes, with M = 20 corresponding to near-compute-optimal runs.

C Additional training details

Additional training details specify the transformer architecture, data-sequence handling, causal language-modeling objective, optimizer, regularization, and learning-rate schedule.

  • Architecture: The models use decoder-only pre-normalization transformers with LayerNorm, rotary positional embeddings, a 2,048-token sequence length, and a 50k-tokenizer vocabulary.
  • Data handling: Documents are separated with end-of-text tokens, and sequences are packed without attention masking while sampling without replacement.
  • Objectives and optimization: Training uses causal next-token prediction with additive z-loss, AdamW, independent weight decay, linear warmup, cosine decay, and a 3e-5 cooldown learning rate.

D Additional grid search details

The appendix records the final hyperparameters and the architecture configurations used in grid-search selection.

  • Final model configurations: The final hyperparameters are reported in Table 3.
  • Grid search configuration selection: The architecture configurations swept during grid search are reported in Table 4.

E Evaluation dataset details

The evaluation uses a signal-focused subset of downstream tasks and investigates anomalies in small-model grid-search scaling behavior.

  • Evaluation datasets: The evaluation suite comprises 46 downstream tasks, with the 17-task signal-focused subset used for the main analysis.The complete task list is reported in the paper’s evaluation materials.
  • Grid-search behavior: Small grid-search models between 0.011B and 0.079B parameters over-perform relative to neighboring configurations.These models receive more optimization steps than their x-axis neighbors, and the authors exclude this range from the scaling study.

F Additional results

Additional experiments show where scaling predictions remain reliable, how evaluation choices affect them, and where extrapolation or under-training can cause failures.

  • In-distribution scaling: Scaling is largely predictable in-distribution, with the largest reported RedPajama error reaching 15.4% for the N = 1.4B, M = 640 run.The authors hypothesize that near-duplicate sequences may make this model outperform its prediction.
  • Evaluation-set sensitivity: Prediction error remains fairly low and consistent across many downstream evaluation-suite choices, including the full 46-evaluation suite.The 17-task subset is not uniquely responsible for the observed predictability.
  • Failure modes: Scaling can become unreliable under-training, where M = 5 follows a different trend and may require multiplier-specific hyperparameter tuning.The authors also caution that out-of-distribution predictability may depend on the relationship between training and evaluation distributions.
  • Downstream prediction: Average downstream top-1 error is predictable across interpolation and extrapolation, but individual tasks can show exponential, step-function, or chance-level behavior.The aggregate relationship is therefore more reliable than prediction for every individual evaluation.
  • Fit reliability and token multipliers: Increasing the compute or number of data points used to fit a scaling law lowers relative prediction error, while token multipliers from 10 to 80 lie near the compute-optimal frontier.A fitted RefinedWeb optimum of M* = 5.85 conflicts with the empirical frontier, possibly because fits and tuning used M ≥20.

H Broader impact

The paper releases scaling research artifacts while acknowledging model risks and documenting the experimental infrastructure used to fit and evaluate its laws.

  • Broader impact: The released base models are research artifacts, not post-trained or aligned systems, and may still be misused or cause harm.The authors specifically identify harmful language, toxicity, and human automation as risks.
  • Broader impact: The authors position the released testbed as useful for researchers who lack the resources to train comparable model artifacts themselves.They also note that publicly available models outperform the testbed models and may be more widely adopted.
  • Experimental resources: The experimental infrastructure includes 130 grid-search architectures, 46 downstream tasks, and scaling-law coefficients fitted from the reported configurations.Additional tables document token multipliers for GPT-2- and LLaMA-style decoder-only architectures.
Loading 2403.08540v2…