Source
Porian, Wortsman, Jitsev, Schmidt and Carmon (2024), LIT-690.
What to do
In a sweep meant to fit a compute-optimal allocation law:
- Warmup: keep it well short of the smallest model’s compute-optimal token
count. The source uses
min(N, 0.2·budget)tokens. - Per-size tuning: find learning rate and batch size for each size, and raise AdamW β₂ above 0.95 at small batches (the source’s cut-off is 128 sequences of 2048 tokens).
- Schedule: a constant learning rate after warmup is enough to recover the exponent. Running a cosine per (N, C) pair costs about twice as much and moved the exponent by 0.03.
- Counting: see SOTA-413.
Why
At small scale, fixed warmup and fixed hyperparameters penalize some sizes
more than others. The penalty tilts the fitted slope toward parameters
(THEORY-108).
The protocol is LIT-690’s, which made this recommendation by attributing the
gap between Kaplan’s and Chinchilla’s exponents one step at a time on an
OpenLM sweep from 5M to 901M. Reproducing Kaplan gives a = 0.835, counting
the head’s FLOPs 0.706, warmup tokens = N 0.602, and tuning learning rate,
batch and β₂ per size with a constant learning rate 0.497, Chinchilla’s
equal-proportion allocation. A cosine decay matched to each run moved the
exponent only 0.031, on a side branch, which is why the schedule item allows a
constant rate.
Conditions
- It matters below about 1B. The source expects each factor to fade with scale.
- “Warmup = N tokens” is confounded with “about 5% of a 20:1 run”. It was ablated at one size.
- The β₂ grid stopped at 0.999. A rule holding β₂’s half-life fixed in tokens (LIT-444) would go higher at small batch.
- A constant learning rate is for fitting, not training. The source’s constant runs reach worse loss than cosine. Train the final model with a decay (SOTA-140, SOTA-141).
Known implementations
- The source’s released code and checkpoints.