Home / Fundamentals / Pre-training & Optimization

Optimizers, learning rates and batch size

Adam and AdamW with equations, why decoupled weight decay, typical LLM hyperparameters, warmup, cosine and WSD schedules, batch size and the gradient noise scale, clipping and initialization.

15 min read

Adam

Adam keeps exponential moving averages of the gradient and the squared gradient, and normalizes each parameter's step by the latter:

m_t = β1 · m_{t−1} + (1 − β1) · g_t
v_t = β2 · v_{t−1} + (1 − β2) · g_t²
m̂_t = m_t / (1 − β1^t)        v̂_t = v_t / (1 − β2^t)      # bias correction
θ_t = θ_{t−1} − η · m̂_t / (√v̂_t + ε)

• m provides momentum; v gives a per-parameter scale, so parameters with consistently large gradients take smaller steps.
• Bias correction compensates for initializing m and v at zero; without it, early steps would be too small (for m) or too large (for the ratio).
• Memory: two extra fp32 values per parameter (8 bytes), which dominates optimizer memory in large-model training.

AdamW: decoupled weight decay

L2 regularization adds λθ to the gradient. Inside Adam, that term then gets divided by √v̂, so parameters with large historical gradients are regularized less, which is not what we want. AdamW applies decay directly to the weights instead:

θ_t = θ_{t−1} − η · ( m̂_t / (√v̂_t + ε) + λ · θ_{t−1} )

Decay is now uniform across parameters and independent of gradient history. It's the default for transformers.

Typical LLM settings: β1 = 0.9, β2 = 0.95 (lower than the classic 0.999, so the second moment adapts faster and training is more stable), ε around 1e-8 (some recipes use 1e-5), weight decay ≈ 0.1, gradient clipping at a global norm of 1.0. Biases and normalization gains are usually excluded from weight decay. Peak learning rates are typically around 3e-4 for models of a few billion parameters and lower for larger ones.

Learning-rate schedules

Warmup. Start near zero and ramp linearly over the first ~1–2k steps. Early on, Adam's second-moment estimates are based on few samples and can be badly off, producing oversized updates; activations and gradients also haven't settled. Warmup keeps early steps small until the statistics are reliable. Post-norm architectures need much more warmup than pre-norm.

Cosine decay. After warmup, decay along a half cosine to ~10% of the peak by the end of training. The decay phase is where much of the loss improvement is realized. The schedule length must match the training length: decaying too early or too late wastes compute.

Warmup–stable–decay (WSD). Hold the peak learning rate constant for most of training, then decay quickly over the last ~10–20% of steps. Results are competitive with cosine, and it's practical: you can branch a decay phase from any stable-phase checkpoint, extend training without restarting the schedule, and run scaling studies cheaply.

Batch size

• Larger batches give less noisy gradients, but each doubling helps less once noise is no longer the bottleneck.
• The gradient noise scale B_noise ≈ tr(Σ) / |G|² (gradient variance over squared gradient norm) estimates the critical batch size: below it, doubling the batch nearly halves the number of steps needed; above it, returns diminish and compute is wasted.
• The critical batch size grows as the loss falls, so large runs often ramp the batch size up during training.
• Batch sizes for large LLM pre-training are typically a few million tokens per step.
• When changing batch size, the learning rate usually needs adjusting; for Adam, scaling roughly with the square root of the batch-size ratio is a common heuristic, verified with small sweeps.

Gradient clipping

Clip the global gradient norm (across all parameters) to a threshold such as 1.0. It prevents single bad batches from producing destructive updates. A clipped-step rate that rises over time is an early warning of instability.

Initialization

• Weights are typically drawn with standard deviation ~0.02, or scaled as 1/√d_in.
• Output projections of residual blocks are often scaled down further, e.g. by 1/√(2L), so the residual stream doesn't grow with depth at initialization.
• μP (maximal update parametrization) scales initialization and per-layer learning rates with width so that optimal hyperparameters found on a small model transfer to a large one. It reduces expensive sweeps at scale.

Beyond Adam

Adafactor (factored second moments, less memory), Lion (sign-based updates), and second-order-inspired methods such as Shampoo, SOAP and Muon (orthogonalized momentum updates for matrix parameters) are active areas. Several report better tokens-to-quality than AdamW at some scales; AdamW remains the robust default.

Interview angles

• Write Adam's update with bias correction, and explain each term.
• Explain why L2 inside Adam differs from decoupled weight decay.
• Justify warmup, and compare cosine with WSD.
• Explain the critical batch size and what happens beyond it.