Learning-Rate Schedules

The learning rate is the one hyperparameter that must move during training. Warmup exists because the optimiser starts blind; decay exists because the loss surface narrows.

[optimiser][training]Current standard

Standing

Current standardPart of the default recipe. Present in most frontier models trained today, and the thing a new design departs from rather than argues for.

Warmup plus decay is universal; the shape of the decay is where the disagreement is.

judged as of 2026-09 · what the labels mean

Theory

Warmup

Adam’s second moment is a running average with an effective window of $1/(1-\beta_2)$ steps. Before that many steps have passed, $\sqrt{\hat v}$ is noisy, and a noisy denominator on a full-size learning rate produces steps in arbitrary directions. Linear warmup holds the rate small until the estimate settles.

\eta_t = \eta_{\max} \cdot \min\!\left( 1, \frac{t}{W} \right)

Cosine decay

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

Cosine has one serious drawback, and it is structural rather than empirical: $T$ appears in the formula. The schedule must know the total step count in advance, and a run stopped early sits at a high learning rate with a correspondingly poor loss. Every intermediate checkpoint is off the compute-optimal frontier.

Warmup–stable–decay

WSD holds $\eta_{\max}$ constant for most of the run and decays sharply over the last 10–20%. Because the stable phase does not depend on $T$ , a run can be extended indefinitely, and any point on it can be turned into a finished model by decaying from there.

\eta_t = \begin{cases} \eta_{\max} \cdot t/W & t < W \\ \eta_{\max} & W \le t < (1-f)T \\ \eta_{\max} \cdot \left( \dfrac{T-t}{fT} \right) & t \ge (1-f)T \end{cases}

The loss curve under WSD looks worse than cosine for most of training and then drops abruptly during the decay, ending in the same place. Judging a WSD run by its mid-training loss is judging it before it has been asked the question.

Implementation

python · torch ≥ 2.1

import math
from torch.optim.lr_scheduler import LambdaLR


def cosine_with_warmup(optimizer, warmup: int, total: int, floor: float = 0.1):
    def factor(step: int) -> float:
        if step < warmup:
            return step / max(1, warmup)
        progress = (step - warmup) / max(1, total - warmup)
        cosine = 0.5 * (1 + math.cos(math.pi * min(1.0, progress)))
        return floor + (1 - floor) * cosine

    return LambdaLR(optimizer, factor)


def wsd(optimizer, warmup: int, total: int, decay_frac: float = 0.1):
    """Stable until the final decay_frac, so the run can be extended."""
    decay_start = int(total * (1 - decay_frac))

    def factor(step: int) -> float:
        if step < warmup:
            return step / max(1, warmup)
        if step < decay_start:
            return 1.0
        return max(0.0, (total - step) / (total - decay_start))

    return LambdaLR(optimizer, factor)

Decaying to zero rather than a floor is common and mostly harmless, but a floor of $0.1\,\eta_{\max}$ leaves the model able to keep moving if the run is continued. Note that both schedules step per optimiser step, not per epoch — stepping per epoch on a single-epoch pretraining run means never stepping at all.

adamw chinchilla-scaling-laws gradient-checkpointing

References

[1]Loshchilov & Hutter — SGDR: Stochastic Gradient Descent with Warm Restarts (2016)arXiv:1608.03983

[2]Hu et al. — MiniCPM: Unveiling the Potential of Small Language Models (2024)arXiv:2404.06395

[3]Hägele et al. — Scaling Laws and Compute-Optimal Training Beyond Fixed Durations (2024)arXiv:2405.18392