Vol. XXXI · History and Projects Plate 7 of 14project, pretraining, held-out loss, FLOP budget, batch size, learning rate, model sizeDCL

Project: Pretrain a Small Model.

A loss target and a compute budget: choose the model’s size and how to spend the steps.

The brief: a character-level decoder (Vol. V) trained from scratch on the first 90% of 60 000 characters of course text, reaching a held-out loss of at most 3.30 nats per character on the last 10%, with training compute of at most 3 × 10¹⁰ FLOPs, counted as 6 × parameters × tokens (plate V.6), where tokens = steps × batch × 32.

The workshop trains the model in the page each time a setting changes, which takes several seconds. The starting design — width 32, 2 layers, 200 steps of batch 8 at learning rate 0.003 — reaches 3.569 with 9.20 × 10⁹ FLOPs. A project plate grades the design against a brief: every check below the workshop must pass, and a passing design is saved in this browser.

Left: training loss and held-out loss as training proceeds, against the brief’s line. Right: FLOPs spent against the budget.

What you will learn

Train a small model to a loss target.
Count training FLOPs.
Trade model size against steps and batch.
Choose a learning rate.

Workshop

A character-level Transformer with learned positions and tied embeddings, context 32, AdamW at a constant learning rate, trained in the browser; held-out loss measured on the last 6 000 characters. Each change retrains.

Where the FLOPs go

The starting design spends only 31% of the budget and stops at 3.569: it is not that the model is too small but that it has seen too little text. FLOPs scale with parameters times tokens, so the same budget buys a wider model for fewer tokens or a narrower one for more; with this little data and compute, the loss is limited by tokens seen, the regime in which Chinchilla’s rule (plate X.3) says to keep the model small.

Batch and learning rate decide how well the steps are used. Larger batches average the gradient over more text and tolerate a larger step; too high a rate makes a small model’s loss jump instead of fall (plate XI.1). The project is the whole of Parts One and Two in miniature: an architecture, a token budget, an optimiser, and a held-out number that is the only judge.

A 3.3-nat loss is far from good English — this model makes word-shaped strings — but on 54 000 training characters it is near what 30 billion FLOPs can buy.

Terms

Held-out loss: Cross-entropy on text not used for training.
Training FLOPs: About 6 × parameters × tokens.
Batch size: Sequences averaged in one gradient step.
Learning rate: The size of each optimiser step.

A loss target within a compute budget is met by spending the FLOPs on enough tokens, with a batch and learning rate that use them well.

Sources

Kaplan et al., “Scaling Laws for Neural Language Models”, arXiv:2001.08361, 2020 (C ≈ 6ND).
Karpathy, nanoGPT, github.com/karpathy/nanoGPT, 2022.