Optimisers, Schedules & Regularisation That Actually Work

The division by $\sqrt{{\hat{v}}_t}$ is the whole idea: parameters whose gradients have been consistently large get smaller steps, and parameters with tiny gradients get relatively larger ones. It is an automatic per-parameter learning rate, which is why Adam works out of the box on architectures where a single global LR would be hopeless.

The AdamW fix, precisely

Classic L2 regularisation adds $\lambda \theta$ to the gradient. Under Adam that term flows into both $m_t$ and $v_t$, so the effective decay becomes $\frac{\eta \lambda \theta }{\sqrt{{\hat{v}}_t}}$ — inversely proportional to the parameter's gradient magnitude. Parameters that need regularising most (large, active weights) receive the least. AdamW removes the term from the gradient and applies it directly to the update:

What warmup actually fixes

At step 1, $v_1=(1-{\beta }_2)g_1^2$. Bias correction divides by $(1-{\beta }_2^1)=1-{\beta }_2$, giving ${\hat{v}}_1=g_1^2$ — an estimate based on a single sample. If that one gradient happens to be small, the update $\eta \hat{m}/\sqrt{\hat{v}}$ is enormous. Early large steps in a randomly initialised network push it into a region it may never recover from, which is why the failure shows up as a run that never converges rather than one that recovers.

There is a second, architecture-level reason. In post-norm transformers, gradients at the early layers are poorly scaled at initialisation, and a large first step produces exactly the divergence the original paper worked around with its inverse-square-root schedule. Pre-norm reduces but does not eliminate this.

• Warmup length: 1–5% of total steps, or a flat 2,000 for large runs. Longer warmup for larger batches.
• Peak LR: run an LR range test and take roughly one order of magnitude below divergence, or start from a known recipe and scale with batch size.
• LR ∝ 1/√width is a reasonable heuristic when scaling model size; μP (maximal update parametrisation) makes this transfer principled so you can tune on a small model and reuse the LR on a large one.
• Layer-wise LR decay (lower LR for earlier layers) is standard for fine-tuning pre-trained encoders and reliably beats a uniform LR.

The diagnostic ladder

1. 1

   1 — Inspect the offending batch

   Log the batch index deterministically. Look for: NaN/Inf in the inputs, a sequence that is entirely padding (mean over zero valid tokens → 0/0), a duplicated document repeated thousands of times, or a token id outside the embedding range. Corrupt data is by far the most common cause and the cheapest to rule out.
2. 2

   2 — Check for fp16 overflow

   fp16's maximum is 65,504. Attention logits before softmax, and the sum inside a normalisation, routinely exceed it in a large model. bf16 has the same exponent range as fp32 (max ≈ 3×10³⁸) and simply does not have this failure mode — which is why every modern large-scale run uses bf16.
3. 3

   3 — Look at the gradient norm history

   Log global grad norm every step. A spike almost never comes from nowhere: you usually see the norm climbing over the preceding hundreds of steps. If so, the problem started long before the NaN and clipping alone is treating the symptom.
4. 4

   4 — Suspect the Adam second moment

   With β₂ = 0.999 the effective window is ~1000 steps. If gradients have been small for a while, v is small; one large gradient then produces a gigantic update because it is divided by a stale, tiny √v. Lower β₂ to 0.95, or use an update-clipping variant.
5. 5

   5 — Check for architectural hot spots

   Unnormalised attention logits, a missing ε in a normalisation denominator, a log of something that can be zero, or exp of an unbounded value. Attention logit soft-capping (tanh clamping) is used in several recent models specifically to prevent this.

Dropout: the mechanism and the inference detail

Dropout zeroes each unit independently with probability $p$ during training. Two complementary readings: it prevents co-adaptation (a unit cannot rely on a specific partner always being present) and it approximately trains an exponentially large ensemble of subnetworks that share weights, then averages them at test time.

Label smoothing

With a hard one-hot target, cross-entropy is only minimised as the correct logit goes to $+\infty$ — so training pushes toward unbounded confidence, which produces overconfident, poorly-calibrated models. Label smoothing gives a finite optimum, improves calibration, and typically adds a small accuracy gain. The trade-off: it compresses the representation space and is known to slightly hurt when the model is used as a feature extractor for downstream retrieval or distillation.

• Augmentation is not regularisation in the same sense — it adds information rather than removing capacity. That is why it is strictly better when available.
• Mixup trains on $\lambda x_i+(1-\lambda )x_j$ with correspondingly mixed labels. It improves calibration and adversarial robustness beyond its accuracy gain.
• Stochastic depth is dropout over whole residual blocks; it makes very deep networks behave like an ensemble of shallower ones and cuts training time.
• Ensembling is the reliable last resort. Five seeds averaged beats almost any single-model regularisation tweak, at 5× inference cost — or 1× if you distil the ensemble back into one model.

The trade is entirely between range (exponent bits) and precision (mantissa bits). Deep learning turns out to care far more about range: gradients span many orders of magnitude, but you rarely need seven significant digits in a weight update that gets averaged over thousands of examples.

Memory accounting for a 7B model

For the frontier angle: fp8 training (E4M3 for forward, E5M2 for gradients) is production on H100 with per-tensor scaling factors, giving another ~2× on top of bf16. And for inference, weight-only quantisation to int8/int4 is standard — but that is a separate topic from training precision, and conflating the two is a common candidate mistake.