PaperlayerView exact revision on arXiv

Source-linked AI summary

Resolving Discrepancies in Compute-Optimal Scaling of Language Models

Tomer Porian, Mitchell Wortsman, Jenia Jitsev, Ludwig Schmidt, Yair Carmon

arXiv:2406.19146v4cs.LGcs.CL

TL;DR

Kaplan et al. and Hoffmann et al. give substantially different predictions for compute-optimal language-model scaling. This paper reproduces the discrepancy, isolates its causes, and shows that correcting computational accounting, warmup, and optimizer tuning brings the result into close agreement with Hoffmann et al., without requiring carefully matched learning-rate decay.

  • Problem

    Kaplan et al. and Hoffmann et al. proposed substantially different compute-optimal scaling laws, motivating an explanation of why their conclusions diverge.

  • Method

    The authors reproduce Kaplan et al.’s scaling law on OpenWebText2 and RefinedWeb, then vary last-layer accounting, warmup, learning-rate schedules, and optimizer hyperparameters.

  • Results

    Correcting last-layer cost, warmup duration, and scale-dependent optimizer tuning yields scaling that closely agrees with Hoffmann et al.; matched learning-rate decay has little effect on the scaling law.

  • Takeaways & Limitations

    Compute-optimal scaling can be recovered at smaller scales by correcting computational accounting and tuning optimization hyperparameters, especially AdamW β2 at small batch sizes.

  • Takeaways & Limitations

    The experiments are substantially smaller than Hoffmann et al.’s, and limited compute restricted hyperparameter sweeps to smaller models trained at the Hoffmann token-to-parameter ratio.

Abstract

from arXiv · show

Kaplan et al. and Hoffmann et al. developed influential scaling laws for the optimal model size as a function of the compute budget, but these laws yield substantially different predictions. We explain the discrepancy by reproducing the Kaplan scaling law on two datasets (OpenWebText2 and RefinedWeb) and identifying three factors causing the difference: last layer computational cost, warmup duration, and scale-dependent optimizer tuning. With these factors corrected, we obtain excellent agreement with the Hoffmann et al. (i.e., "Chinchilla") scaling law. Counter to a hypothesis of Hoffmann et al., we find that careful learning rate decay is not essential for the validity of their scaling law. As a secondary result, we derive scaling laws for the optimal learning rate and batch size, finding that tuning the AdamW $\beta_2$ parameter is essential at lower batch sizes.

1 Introduction

The paper explains why Kaplan et al. and Hoffmann et al. predict different compute-optimal scaling by identifying experimental and optimization factors behind the discrepancy. Correcting these factors yields scaling closely aligned with Hoffmann et al., while learning-rate decay is not essential and optimizer hyperparameter scaling matters at small batch sizes.

  • Contributions: Matching learning-rate decay to each token budget has little effect on the compute-optimal scaling law, contrary to Hoffmann et al.’s hypothesis.The hypothesis proposed that Kaplan et al.’s untailored decay schedule caused the discrepancy, but the authors find otherwise.
  • Contributions: Three factors explain the discrepancy: last-layer computational cost, warmup duration, and scale-dependent optimizer tuning.The study reproduces Kaplan et al.’s scaling law and progressively corrects these factors.
  • Contributions: Accounting for the decoding layer shifts compute-optimal scaling toward a more constant token-to-parameter ratio.Kaplan et al. did not include this cost, whereas Hoffmann et al. did.
  • Contributions: Scaling warmup with model size further shifts the scaling law toward Hoffmann et al.’s result, because fixed warmup is too long for smaller models.The authors use warmup tokens equal to model size and report an exponent of roughly 0.6.
  • Contributions: Individually tuning learning rate, batch size, and AdamW β2 produces scaling that closely agrees with Hoffmann et al., while β2=0.95 is suboptimal at batch sizes of 128 or below.Increasing β2 improves performance at smaller batch sizes and reveals clearer hyperparameter trends.
  • Results: Cosine learning-rate decay lowers loss substantially, with benefits persisting at larger scales, whereas the other corrections mainly help at smaller scales.The tuned loss remains broadly consistent with a saturating power law.

2 Preliminaries and experiment design

The paper defines compute-optimal training as choosing model and dataset sizes that minimize loss under a fixed FLOP budget. It estimates optimal scaling by evaluating model-size/loss curves across a FLOP grid and fitting power laws to the resulting optima.

  • 2.1 Notation and problem setting: The analysis treats N as parameters in all linear layers, including the untied output head but excluding embeddings.This definition matters because the paper separately studies the computational cost of the final decoding layer.
  • 2.1 Notation and problem setting: Training compute is approximated as FLOPs(N,D) ≈ 6ND, with D denoting training tokens and L(N,D) the resulting expected log loss.The compute-optimal model minimizes loss subject to this approximate budget constraint.
  • 2.1 Notation and problem setting: The study models optimal model size, token count, and token-to-parameter ratio as power laws of compute, fitting exponents and coefficients from data.The fitted quantities are N⋆(C), D⋆(C), and ρ⋆(C).
  • 2.1 Notation and problem setting: The loss notation abstracts over architecture and training-method choices, although each evaluated (N,D) pair uses one configuration selected through search or prior work.Thus, the practical estimates are not over every possible architecture or optimizer.
  • 2.2 Training setup: Experiments use 16 decoder-only Transformer models ranging from 5M to 901M parameters, with architectures chosen to keep width-to-depth ratios between 32 and 64.The models use four attention heads, based on preliminary findings for smaller models.
  • 2.2 Training setup: The datasets are OpenWebText2, containing roughly 30B tokens, and RefinedWeb, containing roughly 600B tokens.These datasets respectively resemble the data used in Kaplan et al. and Hoffmann et al.
  • 2.3 Scaling-law estimation: For each FLOP-grid value, the authors train several model sizes, estimate the minimum-loss model by noise-and-interpolation, and fit weighted log-space power laws with bootstrap confidence intervals.Validation uses 160M held-out tokens, while smoothed training loss is used where validation loss is unavailable.

3 Main results: settling the scaling law discrepancy

Reproducing Kaplan et al. [30] reveals that last-layer FLOPs, excessive warmup, and fixed optimization hyperparameters distort compute-optimal scaling; correcting them closely reproduces Hoffmann et al..

  • 3.1 Reproducing the Kaplan et al. scaling law: The baseline reproduction matches Kaplan et al.’s [30] setup and obtains a compute-optimal scaling law close to 1.6e9(C/8.64e19)^0.88.
  • 3.2 Counting last layer FLOPs: Counting model-head FLOPs reduces the scaling exponent by more than 0.1 and corrects an underestimation that grows from roughly 10% to 90% as models get smaller.
  • 3.3 Correcting warmup duration: Kaplan et al.’s [30] long warmup makes smaller models reach compute-optimality during warmup; setting warmup tokens equal to model size keeps optimal training at least 5 times longer.
  • 3.4 Learning rate decay has limited impact on compute-optimal allocation: For compute-constrained runs, Kaplan et al.’s [30] 131B-token decay produces virtually no decay because optimal training stays below 10B tokens; per-budget cosine schedules instead decay to 1% of the maximum rate.
  • 3.5 Correcting batch size, learning rate and β2: Tuning batch size, learning rate, and AdamW β2 across model sizes yields clear power-law trends, while β2=0.95 is suboptimal at batch sizes of 128 and below.
  • 3.5 Correcting batch size, learning rate and β2: The corrected setup matches Hoffmann et al.’s [25] exponent of 0.5 within 0.6% and predicts a Chinchilla-scale model within 15%, even with a constant learning rate schedule.

4 Additional Analysis

Additional analyses show that optimizer and warmup corrections mainly improve low-compute loss, while learning-rate decay becomes more beneficial at larger scales and scaling predictions improve with experimental compute.

  • 4.1 Trends in compute-optimal loss: Shortening warmup and tuning hyperparameters each improve low-compute loss by up to 0.5 nat per token, whereas their effect is marginal at larger scales.
  • 4.1 Trends in compute-optimal loss: Learning-rate decay becomes increasingly beneficial as compute grows and appears to improve the rate at which compute-optimal loss decreases.
  • 4.1 Trends in compute-optimal loss: The tuned-optimizer loss follows a saturating power law that extrapolates well, unlike the other experiments; its exponent is approximately −0.1.
  • 4.1 Trends in compute-optimal loss: A 901M-parameter model trained at C+ ≈8e19 FLOPs reaches L+ = 2.943 and falls within the predicted compute-optimal loss trend.
  • 4.2 Experimental compute and prediction confidence: Fixed-schedule experiments cost 1.54e20 FLOPs, while varying-length cosine schedules require 2.99e20 FLOPs, roughly twice as much.
  • 4.2 Experimental compute and prediction confidence: Scaling-law predictions become steadily more accurate and confident as the total scaling-experiment compute increases.

5 Discussion

The discussion situates the paper’s discrepancy analysis within broader scaling-law work, emphasizing tuning, computational scale, and the scope of compute-optimal pretraining loss. It also clarifies how the literature has interpreted Hoffmann et al.’s learning-rate-schedule hypothesis.

  • 5 Discussion: Prior studies broadly support the importance of tuning across scales, especially for smaller models, while related work derives alternative compute-optimal laws and extends scaling beyond the compute-optimal regime.Hu et al. report larger optimal token-to-parameter ratios, whereas Sardana and Frankle incorporate inference cost and favor smaller models.
  • 5.2 Limitations: Computational scale limits the results: the experiments are closer to Kaplan et al.’s scale than Hoffmann et al.’s, and larger scale may reduce the identified effects.The authors note that the model head, fixed warmup, and hyperparameter sensitivity may matter less as scale increases.
  • 5.2 Limitations: Limited compute constrained hyperparameter sweeps to smaller models trained at the Hoffmann et al. token-to-parameter ratio, potentially biasing the selected hyperparameters.The authors acknowledge that perfect tuning across every model size and token-to-parameter ratio could have produced different results.
  • 5.2 Limitations: The paper’s compute-optimal criterion concerns pretraining loss rather than direct capability measures, whose small-scale emergence remains unresolved.Fine-tuning can evaluate downstream performance, but may disrupt the clean scaling patterns observed during pretraining.
  • A Interpreting Hoffmann et al.’s hypothesis: The literature commonly interprets Hoffmann et al.’s hypothesis as attributing the discrepancy to learning-rate decay, although their evidence varied decay while holding warmup fixed.This interpretation appears in subsequent discussions by Hu et al., Hägele et al., and.

B Estimating FLOPs and accounting for attention

The paper compares model-size definitions for estimating training FLOPs and finds that attention and last-layer accounting materially affect smaller-model estimates. Using an effective size that includes attention yields qualitatively similar scaling results, with a modestly larger prediction at Chinchilla/Gopher compute.

  • Counting last layer and attention FLOPs: Attention is the only non-negligible error in the approximation FLOPs(N, D) = 6ND once exact non-embedding parameter counts are considered.In OpenLM, attention adds 6ndl FLOPs per token per layer for forward and backward computation, with n = 2048.
  • Counting last layer and attention FLOPs: The effective size N_eff incorporates attention cost so that 6N_effD captures training FLOPs including attention.The paper defines N_eff by adding an attention-dependent term to N.
  • Counting last layer and attention FLOPs: Because N_eff/N varies only modestly across the model grid, replacing N with N_eff produces qualitatively and quantitatively similar scaling laws.The ratio varies roughly from 1.1 to 1.2 and back to 1.1, limiting differences between the resulting power laws.
  • Counting last layer and attention FLOPs: The appropriate model-size definition remains uncertain, but the final tuned experiment predicts an effective model about 16B parameters larger at Chinchilla/Gopher compute.The comparison is directly valid because both predictions are expressed at the same compute scale.
  • Counting last layer FLOPs: Ignoring the model head produces a large, systematic FLOPs-count error for smaller models and is identified as a primary source of the Kaplan–Hoffmann discrepancy [30].Kaplan et al. approximate compute as 6N_KaplanD, whereas the paper’s standard N includes the model head.

C Additional training setup description

The experiments use decoder-only Transformers trained with OpenLM and a modern optimization and hardware stack. The setup includes limited unintended data repetition in some runs, although the authors expect little impact and report none in the main RefinedWeb experiments.

  • Modeling: The models are decoder-only Transformers trained for next-token prediction with OpenLM using SwiGLU, rotary embeddings, qk-LayerNorm, depth-scaled initialization, and mixed precision.The models use the GPT-NeoX-20B tokenizer with vocabulary size 50,432.
  • Optimization: Training uses AdamW with independent weight decay and an auxiliary z-loss term, while reported analysis uses log loss without the z-loss contribution.The z-loss coefficient is 1e-4 and the independent weight-decay parameter is 1e-4.
  • Hardware and computational cost: The full study uses 22.3K GPU hours and 3.03e21 FLOPs across runs on 40GB A100 GPUs.Individual training runs use 4–32 GPUs in parallel.
  • Hyperparameters: The fixed-hyperparameter experiments use the values specified in Table 3, while later experiments tune selected hyperparameters according to the scaling law in Table 4.Batch sizes are rounded to GPU-count multiples and learning rates to two significant digits.
  • Data repetition: Some hyperparameter-sweep runs may repeat data up to 10 times and some larger OpenWebText2 runs up to four times, but the main RefinedWeb experiments had no repetition.The authors believe the repetition had limited to no impact, citing small sweep models and prior evidence that four repetitions have marginal effects.

D Additional data analysis details

The analysis fits compute-optimal scaling curves from smoothed and interpolated loss measurements, while explicitly modeling validation-set and seed noise. IsoFLOP and loss plots are provided for both datasets, with edge cases excluded or adjusted during optimal-size estimation.

  • Estimating N*(C): The paper fits the N*(C) power law using a procedure analogous to the D* power law, with training-loss smoothing and compute-indexed loss extraction preceding the fit.The N*(C) procedure uses bootstrap-like perturbations, interpolation, and edge handling.
  • Training loss smoothing: Variable-window smoothing and a half-logging-interval shift make smoothed training loss track validation loss within approximately its sampling error.The smoothing parameter is p = 0.05 and losses are logged every k = 20 steps.
  • Fetching loss at C_i: Loss at compute C_i is obtained by log-space interpolation near C_i/(6NB), using validation loss when available and training loss for alternative model-size definitions.The nearest sampled step must lie within 10% of the target compute step.
  • Diagnostic plots: The analysis reports IsoFLOP curves and training-loss-versus-compute plots for both RefinedWeb and OpenWebText2, alongside cross-seed noise diagnostics.The figures overlay smoothed and raw training loss where applicable.
  • Estimating loss noise: Loss uncertainty combines finite-validation-set error, measured at 0.001–0.002 standard deviation, with variation across training seeds.The paper also uses a heuristic bootstrap noise estimate for seed variance.
  • Estimating N*(C): For each bootstrap sample, Gaussian loss noise is added, the loss-versus-size curve is interpolated in log-space, and the optimum is estimated; grid-edge optima are omitted or corrected.Values are omitted when more than half of bootstrap samples land at a grid edge.

E Additional plots for main experiments

Additional plots show the fitted behavior of D*, ρ*, and N* across the main experiments, with qualitatively and quantitatively consistent results on OpenWebText2.

  • Additional plots: Figures 13 and 14 plot observations and power-law fits for D*, ρ*, and N* across the main experiments on RefinedWeb and OpenWebText2.The OpenWebText2 reproduction shows consistent qualitative and quantitative results.

F Ablation of warmup duration and final learning rate value

The ablations show that a warmup of roughly N tokens is nearly optimal, while varying the final learning rate has little effect; tuning β2 is important for clean batch-size trends.

  • Warmup duration: A warmup duration of N tokens is nearly optimal, with durations up to 4N performing similarly or slightly better.The sampling standard deviation is 0.002 across the warmup experiments.
  • Final learning rate value: The loss is barely affected when the final learning rate ranges from 0.1% to 10% of the peak learning rate.This result holds with cosine decay and the Section 3.4 hyperparameters; the sampling standard deviation is 0.002.
  • Hyperparameter sweep: The hyperparameter sweep covers 642 runs across six models, varying learning rate, batch size, and three β2 values.Models range from 5M to 221M parameters and are trained for 20N tokens on RefinedWeb.
  • β2 tuning: Restricting the analysis to β2 = 0.95 breaks part of the clean scaling trend and causes the optimal batch size to saturate for small models.The rounded fit uses all β2 values rather than only β2 = 0.95.

G.4 Estimating scaling law with ideal tuning

Ideal hyperparameter tuning changes losses most for small models and low token-to-parameter ratios, but it barely shifts compute-optimal model sizes or the scaling exponent.

  • Limitation: The learning-rate and batch-size scaling laws are fitted from models no larger than 108M trained for 20N tokens, potentially favoring Hoffmann et al. scaling.Ideal tuning would vary hyperparameters by model size and token-to-parameter ratio, and estimates are unavailable for k > 7.
  • Method: The analysis estimates ideal tuning by interpolating loss across model size and token-to-parameter ratio, updating IsoFLOP curves, and refitting the scaling law.Observed loss is adjusted by subtracting smoothed hyperparameter suboptimality estimates.
  • Conclusions: Ideal tuning changes the estimated compute-optimal exponent by less than 0.032, with the fitted exponent reaching 0.5.The compute-optimal model-size differences are generally within estimated standard deviations, and scaling-law-based choices produce similar results to ideal tuning.
  • Conclusions: For models above 15M parameters and token-to-parameter ratios above 10, the chosen hyperparameters are within 1e−2 loss of ideal tuning.Potential loss reductions are larger for smaller models or lower token-to-parameter ratios, but the optimal model sizes move little.

H Reproducing the adjusted Kaplan et al. scaling law

Reintroducing Kaplan et al.’s FLOP-count and long-warmup issues while retaining optimized hyperparameters recovers the adjusted Kaplan scaling law.

  • Adjusted Kaplan scaling law: 1.3e9 · (C/8.64e19)^0.73 is recovered as the adjusted Kaplan et al. compute-optimal scaling law.The recovery occurs after restoring the FLOP-count and long-warmup issues while decreasing batch size through parameter tuning.
  • Adjusted Kaplan scaling law: The result shows that optimized hyperparameters can recover adjusted Kaplan behavior even when the other identified issues are reintroduced.Figure 18 presents the corresponding optimal model-size relationship.
  • Supporting analyses: The appendix includes expanded saturating power-law fits for the experiments and reproductions on OpenWebText2.These figures extend the scaling-law visualizations across the reported experimental settings.

I The compute-optimal loss

The compute-optimal loss is well described by a saturating power law when hyperparameters are tuned, and the loss-scaling result also holds on OpenWebText2.

  • The compute-optimal loss: The saturating power-law fit is predictive only for the experiment with tuned hyperparameters.The fit has the form L(C) = E + L0C^-ℓ and is obtained by minimizing Huber prediction loss for log L(C).
  • OpenWebText2 results: The loss scaling law is predictive on the OpenWebText2 dataset as well.Figures 20 and 21 reproduce the corresponding fits and analyses for OpenWebText2.
  • Computational cost: Each compute-optimal experiment estimates N⋆(C) over budgets Ck = 2^k · 1.25e16 FLOPs for k from 0 to 11.Cosine schedules require distinct training runs for each budget, whereas constant schedules can reuse runs across smaller FLOP values.
  • Computational cost: For a doubling compute grid, a constant learning-rate schedule costs roughly half as much as a cosine-schedule experiment.The comparison excludes runs with more than 100 times as many tokens as parameters or losses more than 1 nat above optimal.
Loading 2406.19146v4…