Warmup, cosine and WSD schedules

Why warmup is needed, mechanically#

Adam's update is

$$ \theta_{t+1} = \theta_t - \eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} $$

with $\hat m_t$ and $\hat v_t$ the bias-corrected first and second moment estimates of the gradient, and $\eta$ the learning rate. Because the update is normalised by $\sqrt{\hat v_t}$, its magnitude is roughly $\eta$ regardless of gradient scale.

Early in training $\hat v_t$ is estimated from a handful of samples and has enormous variance. Liu et al., 2020 (On the Variance of the Adaptive Learning Rate and Beyond, RAdam) showed the variance of the adaptive term is unbounded at $t$ small, and that linear warmup is a variance-reduction device. RAdam derives a rectification term that makes the warmup implicit; in practice explicit warmup is still preferred because it is one line.

A second, architecture-specific reason applies to post-norm transformers. Xiong et al., 2020 (On Layer Normalization in the Transformer Architecture) proved that post-LN transformers have gradients at initialisation that scale as $O(\sqrt{d}\,\ln L)$ near the output layer, and that warmup is necessary for stability. Pre-LN removes the requirement in theory, and every large model still uses warmup, because the Adam-variance argument survives the architecture change.

Cosine#

$$ \eta_t = \eta_{\min} + \tfrac{1}{2}(\eta_{\max}-\eta_{\min})\left(1 + \cos!\left(\pi \frac{t - t_w}{T - t_w}\right)\right) $$

$t$ is the step, $t_w$ the warmup length, $T$ the total steps, $\eta_{\max}$ the peak and $\eta_{\min}$ the floor (typically $0.1\,\eta_{\max}$).

Introduced as SGDR (Loshchilov and Hutter, 2017) for the restarts, and adopted for LLMs without them. Chinchilla (Hoffmann et al., 2022) documented the critical constraint: the cosine cycle length must match the number of training steps. Setting $T$ to 1× the run gives the best loss; a mismatch of 10× costs a substantial fraction of the run's value. This single fact makes every cosine run a fixed-budget commitment, and it corrupts scaling-law experiments — every point on the curve needs its own full run.

WSD, and why it took over#

The Warmup–Stable–Decay schedule (Hu et al., 2024, MiniCPM):

$$ \eta_t = \begin{cases} \eta_{\max}\, t/t_w & t < t_w \ \eta_{\max} & t_w \le t < T - t_d \ f!\left(\frac{t - (T - t_d)}{t_d}\right)\eta_{\max} & t \geq T - t_d \end{cases} $$

$t_d$ is the cooldown length, typically $0.1$ to $0.2\,T$, and $f$ decays from 1 to a small floor. MiniCPM used an exponential $f$; linear and $1-\sqrt{\cdot}$ variants are also used, and Hägele et al., 2024 found $1-\sqrt{\cdot}$ marginally best.

Three properties matter:

1. Budget-agnostic. The stable phase is a valid checkpoint at any point. One run yields a family of models at different token counts, each finished with a short cooldown. Hägele et al. showed this reduces the compute for a scaling-law study by roughly an order of magnitude.
2. Loss matches cosine. Both papers report constant-LR-plus-cooldown reaching cosine's final loss, sometimes slightly better.
3. The cooldown produces a sharp, reproducible loss drop. Loss is flat and mediocre during the stable phase, then falls steeply during the decay. Watching a WSD run mid-training and concluding it has plateaued is a standard misreading.

Wen et al., 2024 (Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective) give the mechanism: the loss surface resembles a river in a steep valley. A high constant learning rate makes fast progress along the river while bouncing between the steep walls; the cooldown stops the bouncing and drops the iterate to the river bed. It also predicts, correctly, that a cooled-down checkpoint is a poor place to resume high-LR training from — resume from the stable branch instead.

The competing correction: µP#

Scaling the peak learning rate down as models grow is a workaround for a parameterisation problem, not a law of nature. Yang et al., 2022 (Tensor Programs V, µTransfer) show that under Maximal Update Parametrisation, the optimal learning rate is invariant to width. You tune on a 40M-parameter proxy and transfer the hyperparameters to a 6.7B model directly. Cerebras-GPT and several frontier labs use it; it removes the most expensive hyperparameter search in pretraining.

Weight averaging as an alternative to decay#

Hägele et al., 2024 and Sanyal et al., 2023 both report that stochastic weight averaging over the constant-LR phase recovers most of the cooldown's benefit with no schedule change. LAWA (latest weight averaging) keeps a rolling mean of recent checkpoints. Liu et al., 2025 (WSM) push this further, arguing checkpoint merging can replace the decay phase outright. This is an active area; treat it as promising rather than settled.

Papers#

• Goyal et al., Accurate, Large Minibatch SGD, 2017 — arxiv.org/abs/1706.02677 (origin of gradual warmup)
• Loshchilov and Hutter, SGDR: Stochastic Gradient Descent with Warm Restarts, 2017 — arxiv.org/abs/1608.03983
• Liu et al., On the Variance of the Adaptive Learning Rate and Beyond (RAdam), 2020 — arxiv.org/abs/1908.03265
• Xiong et al., On Layer Normalization in the Transformer Architecture, 2020 — arxiv.org/abs/2002.04745
• Yang et al., Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, 2022 — arxiv.org/abs/2203.03466
• Hoffmann et al., Training Compute-Optimal Large Language Models, 2022 — arxiv.org/abs/2203.15556
• Hu et al., MiniCPM, 2024 — arxiv.org/abs/2404.06395
• Hägele et al., Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations, NeurIPS 2024 — arxiv.org/abs/2405.18392
• Wen et al., Understanding Warmup-Stable-Decay Learning Rates, 2024 — arxiv.org/abs/2410.05192

What to learn next#

• Loss spikes and gradient clipping — what happens when the schedule is not enough.
• Finding a learning rate — how to pick the peak in the first place.
• Scaling laws and compute-optimal training — where the total step count comes from.