How to Allocate Your Tokens? Scaling Laws with Training Steps and Batch Size

TL;DR

Schaipp proposes a three-term scaling law that splits the data axis of Chinchilla-style fits into training steps T and batch size B as separate variables, rather than lumping them into total tokens D = T·B. Fit on a large set of training runs, the law recovers the empirical scaling of the optimal batch size and — crucially — can be robustly fit from runs at suboptimal batch sizes, which slashes the number of runs needed compared to standard Chinchilla-style sweeps that must sit at the loss-optimal frontier. The paper further shows the three-term form recovers previous empirical critical-batch-size findings as a downstream consequence, unifying two threads (scaling laws for optimal token count and scaling laws for critical batch size) that had been fit separately.

Key claims

The joint loss surface L(N, T, B) is fit as a three-term law that separates training-steps and batch-size dependencies, in contrast to Chinchilla-style laws that use only L(N, D=T·B) [Abstract, §Intro].
The fitted three-term law correctly recovers the scaling of the optimal batch size B*(N, D) as an emergent consequence rather than as a separately-fit relation [Abstract].
Because it consumes runs across a range of batch sizes (including suboptimal ones), the three-term law is fittable from significantly fewer training runs than a standard L(N, D) fit that must lie on the compute-optimal frontier [Abstract].
The three-term law can be used to derive scaling laws for training at suboptimal batch sizes, giving predicted loss for the off-optimal (N, T, B) cells that Chinchilla-style laws cannot address [Abstract].
The three-term law’s predictions match previously reported empirical findings on the critical batch size, unifying two previously separate empirical threads [Abstract].

Method

The core methodological move is to treat total tokens D as a composed quantity D = T·B and fit a scaling law that keeps both factors as explicit variables of the loss surface, rather than absorbing them into D. Standard Chinchilla-style fits collapse the (T, B) plane down to a single axis; Schaipp’s three-term form retains both, which allows: (a) runs at any B to contribute to the fit rather than only near-optimal-B runs, (b) the fitted surface to predict loss for arbitrary (T, B) splits at fixed D, and (c) the batch-size optimum B*(N, D) and the critical-batch-size scaling to be derived from the same underlying fit rather than measured separately.

The precise functional form and empirical protocol (sweep grid, model sizes, tokens per run, optimizer) live in the full PDF; the abstract-level structure is what makes the law interesting: a single fit that closes the gap between the Chinchilla L(N, D) literature and the McCandlish-Kaplan critical-batch-size literature.

Results

Headline empirical results from the abstract:

Optimal batch-size scaling is recovered by the fit, without being an input to the law.
Robust fits from significantly fewer training runs than standard Chinchilla-style methodology, because suboptimal-B runs are informative rather than discarded.
Predicted critical batch size from the three-term law matches previous empirical findings (McCandlish et al. and successors).

Quantitative numbers (R², exact exponents, model-size range, token budgets) are in the full PDF — the abstract is the load-bearing summary at filing time.

Why it’s interesting

This paper is a direct methodological extension of the concept the wiki tracks in Hyperparameter scaling laws: the axes of a scaling-law fit determine what questions it can answer. It sits between two entries already on that concept page. How to Set the Learning Rate for Large-Scale Pre-training? fits η*(N, D) at fixed batch under a WSD schedule; Schaipp adds batch back in as a first-class axis, which is precisely the “batch size as a free hyperparameter” question the concept page’s Open Questions section flags. Seesaw: Accelerating Training by Balancing Learning Rate and Batch Size Scheduling proves an LR/batch schedule-equivalence (B·η² = const for NSGD), but under a fixed total token budget — the three-term law gives a fitting-side counterpart that lets you predict what (T, B) split is optimal for a given (N, D) rather than proving equivalences along a single family. Also complements The Design Space of Tri-Modal Masked Diffusion Models, whose SDE-reparameterization eliminates the batch search below a critical threshold under AdamW+CompleteP — Schaipp’s law would in principle predict where that threshold falls without needing to run the SDE analysis. The fewer-runs-required claim is the immediately actionable part for pretraining planning.