Chapter 14 of 16
Scaling laws
Kaplan and Chinchilla, IsoFLOP curves, compute-optimal training.
Slides for this chapter: Transformer Training 13, Scaling Laws, Running a Real Job, and After Pretraining
Given a fixed amount of compute, should you train a bigger model on less data or a smaller model on more? Below, the Decoder's model at eight widths, each trained for exactly the tokens a compute budget allows; four budgets, three seeds each. At every budget there is a best size, and the best size grows with the budget: that trend is a scaling law.
Loading the animation…
Concept
Compute is parameters times tokens. Training costs about FLOPs for parameters and training tokens (chapter 11's per token: 2 for the forward pass, 4 for the backward). A budget fixes the product, so every choice of model size fixes the number of tokens it can afford: the curves above are IsoFLOP profiles, one per budget, each point a model trained for exactly its share. Too small a model learns everything it can and stops improving; too large a model never sees enough tokens. In between is the best size for the budget.
Kaplan and Chinchilla. Kaplan et al. (2020) fitted power laws to many runs and concluded that, as compute grows, most of it should go into parameters: . Hoffmann et al. (2022), training over 400 models with each run's learning-rate schedule matched to its length, found the two should grow about equally: by three different methods (0.50, 0.49 and 0.46, their Table 2), and trained Chinchilla, 70 billion parameters on 1.4 trillion tokens (20 tokens per parameter), to beat much larger models trained on fewer tokens. Their second method is the one above: IsoFLOP profiles, a parabola through each, and a straight line through the minima on log–log axes. Llama 3 did the same to choose its flagship's size (Grattafiori et al. 2024, §3.2.1), fitting for the optimal number of tokens.
The toy sweep. Widths 6 to 64 (1,056 to 70,656 parameters, counting the embedding, which is most of a model this small), budgets from 3000000000 to 100000000000 FLOPs, AdamW at a peak learning rate of 0.01 with each run's cosine schedule stretched to its own length (Chinchilla's point: a schedule cut short favours the smaller models). The rule, fixed before the sweep ran: a parabola in through each budget's lowest mean and its two neighbours; no fit when the lowest is the smallest or largest size tried. The smallest budget has no minimum inside the sizes: every model wider than 8 spends its few hundred steps on chapter 1's plateau and ends near 2.95, an optimisation transient, not a scaling law. At the other three the best size grows with the budget: 2,493 parameters at 10000000000 FLOPs, 3,550 at 30000000000, 6,590 at 100000000000. A line through those three gives 0.42, so 0.58: more of each extra FLOP goes into tokens than into parameters.
How far that is from the real laws. The toy's exponent lands near Chinchilla's, and that should not be over-read. It rests on three budgets spanning one decade, three seeds, one learning rate for every width and a plateau that decides the smallest budgets. Its "tokens" are characters of one 128,000-character book, seen over and over (the best model at the largest budget is trained on 2,529,229 characters, hundreds of times its own parameter count, where Chinchilla's rule of thumb is 20 tokens per parameter of fresh text). Real scaling studies span many decades of compute with fresh data. The toy shows the method and the shape of the trade-off; the GPU notebook runs the same sweep on GPTs of 1–40M parameters (one of the planned GPU replications under Runs to come; every number here comes from the in-browser runs).
Extrapolating to larger models. The point of a scaling law is to predict a run too expensive to try: fit on small runs, extrapolate to the big one. Llama 3 fitted its law on runs of up to FLOPs and extrapolated to its budget, predicting a 402B-parameter model on 16.55T tokens; it trained 405B, because the IsoFLOP curves flatten near their minima as budgets grow, so a nearby size costs little (§3.2.1). Extrapolation is only as good as the assumption that nothing changes between the runs measured and the run predicted: a toy fit over one decade, like the one above, says nothing about a model a million times larger. When the Compute in Transit results site scales its measurements to bigger models or rarer syncs, it labels those numbers as extrapolations for the same reason.
Overtraining for cheaper inference. Compute-optimal minimises training cost for a given loss; it ignores the cost of using the model. A smaller model trained on far more tokens than compute-optimal reaches a loss a bigger model would reach cheaply, but costs less to run for every token it ever serves. Llama 3 trained its smaller models "for much longer than is compute-optimal" for that reason (§1). How model size trades against inference cost is the subject of the Inference Trade-offs site.
Maths
IsoFLOP fit. For each budget and the three sizes around its lowest loss, fit with by least squares; if the minimum is at
and . Then fit ; since , .
Chinchilla's parametric form (their Approach 3) models the loss itself,
an irreducible loss plus a term that shrinks with model size and one that shrinks with data; minimising it at fixed gives .
Code
The parabola fit behind each budget's minimum,
src/lib/anim/training/isoflop-steps.ts (the same arithmetic, in the same
order, as scripts/training/isoflop.py; frame-tested against it):
export function fitParabola(
xs: number[],
ys: number[],
): { c: number[]; xMin: number | null; yMin: number | null } {
const s = [0, 0, 0, 0, 0];
const t = [0, 0, 0];
xs.forEach((x, i) => {
let p = 1;
for (let k = 0; k < 5; k++) {
s[k] = s[k]! + p;
if (k < 3) t[k] = t[k]! + p * ys[i]!;
p *= x;
}
});
const a = [
[s[0]!, s[1]!, s[2]!],
[s[1]!, s[2]!, s[3]!],
[s[2]!, s[3]!, s[4]!],
];
const det3 = (m: number[][]) =>
m[0]![0]! * (m[1]![1]! * m[2]![2]! - m[1]![2]! * m[2]![1]!) -
m[0]![1]! * (m[1]![0]! * m[2]![2]! - m[1]![2]! * m[2]![0]!) +
m[0]![2]! * (m[1]![0]! * m[2]![1]! - m[1]![1]! * m[2]![0]!);
const d = det3(a);
const c: number[] = [];
for (let j = 0; j < 3; j++) {
const mj = a.map((row) => row.slice());
for (let i = 0; i < 3; i++) mj[i]![j] = t[i]!;
c.push(det3(mj) / d);
}
const xMin = c[2]! > 0 ? -c[1]! / (2 * c[2]!) : null;
return {
c,
xMin,
yMin: xMin !== null ? c[0]! + c[1]! * xMin + c[2]! * xMin * xMin : null,
};
}
Comments
Be the first to leave a comment on this section.
Sign in (top-right) to leave a comment.