Learning-rate schedules

Problem

Every optimizer on the optimizers page multiplies its update by a learning rate η. Set η too high and the loss oscillates or explodes; set it too low and training takes many times longer than it should. Between those, a fixed η still leaves a gap: with noisy minibatch gradients the parameters never settle at the minimum but jitter around it, and the size of the jitter is proportional to η.

So the right learning rate depends on where training is. Early on the parameters are far from any minimum and a large step makes fast progress. Late in training the remaining error is mostly gradient noise, and only a small step averages it away. Adaptive optimizers add a third phase: in the first few hundred steps their statistics are estimated from almost no data, and a full-size step can push the model somewhere it never recovers from. A schedule is a function from the step count to η that serves all three phases.

Intuition

Picture a ball rolling into a valley while being shaken. A big step size crosses the valley quickly but the shaking keeps it bouncing near the bottom; a small step size bounces less but takes forever to get there. Starting big and shrinking the step does both: cover the distance while the step is large, then let the bouncing die down as the step shrinks. Every decay schedule (steps, an exponential, a cosine) is a different curve for that shrinking.

Warmup is the opposite move at the very beginning: start tiny and ramp up over a few hundred or thousand steps, so the first updates, made with a freshly initialised network and an optimizer that has seen almost no gradients, cannot do lasting damage. One-cycle combines both: warm up to a high peak, then anneal far below the starting value.

Mechanics

What a scheduler is in PyTorch

A scheduler in torch.optim.lr_scheduler owns an optimizer and rewrites each param group's "lr" entry. Its constructor stores the starting rate as group["initial_lr"], copies those into base_lrs, and immediately calls step() once, which sets last_epoch to 0. Each later step() increments last_epoch and writes the new rate. Despite the name, last_epoch counts calls to step(): call it once per batch and every schedule length below is in batches.

The order matters. Since PyTorch 1.1 the loop is optimizer.step() then scheduler.step(); the constructor wraps optimizer.step so the first scheduler.step() can warn when it is called before any optimizer step, because that skips the first value of the schedule.

Most schedulers compute the next rate recursively from the current group["lr"], so that several schedulers can modify the same rate in turn, and also have a _get_closed_form_lr used only when an explicit epoch is passed. The closed forms below are the ones this page computes with, and the tests check that they equal the recursions step for step to within 1e−12.

Decay shapes

StepLR: η_t = η₀ · γ^{⌊t / step_size⌋}    ExponentialLR: η_t = η₀ · γ^t
CosineAnnealingLR: η_t = η_min + ½(η₀ − η_min)(1 + cos(π t / T_max))

StepLR (default γ = 0.1) and its sibling MultiStepLR, which decays at a list of milestones, are the schedules classic image classifiers were trained with: a long plateau at a high rate, then drops by ten. ExponentialLR shrinks by the same factor every step. The cosine spends longer near the top than a straight line would, falls fastest in the middle and flattens out near η_min, without the discontinuities of step decay and with only one shape parameter, the length T_max. CosineAnnealingWarmRestarts restarts the cosine every T₀ steps, optionally growing the period by T_mult.

Linear warmup

LinearLR: η_t = η₀ · (s + (e − s) · min(t, T) / T), defaults s = 1/3, e = 1, T = 5

LinearLR multiplies the base rate by a factor that moves linearly from start_factor to end_factor over total_iters steps and then holds. Used with a small start factor it is a warmup. Warmup became standard for two reasons. Large-batch SGD (Goyal et al., 2017) scales the rate with the batch size, and the scaled rate is too large for the first few epochs. Adam-family optimizers divide by the square root of a second-moment estimate built from a handful of gradients, so their first updates have a step size near η in every coordinate regardless of how small or noisy the gradient is (see the worked example). Transformers trained with the original Post-LN layout are especially sensitive: Xiong et al. (2020) show the gradients near the output are large at initialisation, and Liu et al. (2020) trace the problem to the variance of Adam's adaptive rate in early steps, which RAdam corrects without a warmup.

One-cycle

OneCycleLR (Smith and Topin) starts at max_lr / div_factor (default 25), anneals up to max_lr over the first pct_start of total_steps (default 30%), then anneals down to initial_lr / final_div_factor (default 10⁴), so it ends far below where it started. Each phase uses the cosine interpolation end + (start − end)/2 · (cos(π·pct) + 1) or a straight line with anneal_strategy="linear". With cycle_momentum=True (the default) momentum, or Adam's β₁, moves the opposite way between 0.95 and 0.85, so the effective step is not doubly amplified at the peak. It must be told total_steps in advance and raises an error if stepped past it.

Plateaus and composition

ReduceLROnPlateau is the one scheduler that watches the training rather than the clock: step(metric) counts epochs without improvement, and after patience (default 10) of them multiplies the rate by factor (default 0.1), floored at min_lr. SequentialLR runs a list of schedulers one after another, switching at milestones; at each switch it calls the next scheduler's _update_lr(0), which restarts that scheduler from its base rate. ChainedScheduler instead steps every scheduler on every step, so their multiplicative effects stack. Warmup followed by cosine is SequentialLR([LinearLR, CosineAnnealingLR], milestones=[warmup]) with the cosine's T_max set to the steps that remain.

Worked example

Implementation

import torch

model = ...
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)

total_steps = epochs * len(train_loader)
warmup = int(0.02 * total_steps)
scheduler = torch.optim.lr_scheduler.SequentialLR(
    optimizer,
    schedulers=[
        torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=0.01, total_iters=warmup),
        torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps - warmup),
    ],
    milestones=[warmup],
)

for epoch in range(epochs):
    for x, y in train_loader:
        loss = loss_fn(model(x), y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        scheduler.step()   # after optimizer.step(), once per batch

checkpoint = {"model": model.state_dict(), "optimizer": optimizer.state_dict(),
              "scheduler": scheduler.state_dict()}

# One-cycle instead:
# scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=1e-3, total_steps=total_steps)

The schedule is part of the training state: resuming from a checkpoint without scheduler.state_dict() restarts the warmup and the cosine. The training loop page places this call in a full loop, and mixed precision adds a grad scaler that can skip an optimizer step while the scheduler still advances.

Tradeoffs

• Step decay is easy to reason about and resumes cleanly, but needs the milestones tuned to the run length, and each drop causes a sudden change in the loss that can be mistaken for an improvement in the model.
• Cosine has a single length parameter and no discontinuities, which is why it is a common default in modern training, but it commits to a total length up front: stopping early leaves the rate high, training longer needs a new schedule.
• Warmup costs a small fraction of training and removes a common source of early divergence for adaptive optimizers and large batches. Too short a warmup gives no protection; for SGD at moderate batch size it is often unnecessary.
• One-cycle reaches good results in short runs and fixes the momentum schedule too, but needs total_steps exactly and a sensible max_lr, usually from a range test.
• ReduceLROnPlateau needs no knowledge of the run length but depends on a noisy validation metric and on patience, and cannot be put inside SequentialLR or ChainedScheduler.
• The peak rate matters far more than the shape. A schedule tuned at one batch size usually needs its peak re-tuned at another.

Connections

The rate being scheduled is the step size of the update rules on the optimizers page, and the stability limit 2/h is the same bound as on derivatives and gradients. Warmup matters most for the transformer, whose layer norm placement decides how fragile the first steps are, and the training loop shows where scheduler.step() sits.

References

• Smith. Cyclical Learning Rates for Training Neural Networks. WACV, 2017 (arXiv 1506.01186; the LR range test).
• Loshchilov and Hutter. SGDR: Stochastic Gradient Descent with Warm Restarts. ICLR, 2017 (arXiv 1608.03983).
• Goyal, Dollár, Girshick, Noordhuis, Wesolowski, Kyrola, Tulloch, Jia and He. Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour. 2017 (arXiv 1706.02677; linear scaling and gradual warmup).
• Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser and Polosukhin. Attention Is All You Need. NeurIPS, 2017 (arXiv 1706.03762; warmup then inverse square root decay).
• Smith and Topin. Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates. 2017 (arXiv 1708.07120; one-cycle).
• Liu, Jiang, He, Chen, Liu, Gao and Han. On the Variance of the Adaptive Learning Rate and Beyond. ICLR, 2020 (arXiv 1908.03265; RAdam).
• Xiong, Yang, He, Zheng, Zheng, Xing, Zhang, Lan, Wang and Liu. On Layer Normalization in the Transformer Architecture. ICML, 2020 (arXiv 2002.04745).
• PyTorch v2.15.0: torch/optim/lr_scheduler.py.