Learning rate schedules

Overview

The learning rate sets how big a step the optimiser takes on each update. A single fixed value is rarely ideal for a whole training run. Early on, the model is far from a good solution and larger steps make quick progress, though the very first steps can be unstable. Later, the model is close to a good region and large steps bounce around it instead of settling in. A learning rate schedule changes the rate over time to suit each phase: often a short warm-up, a long period at a high rate, and a gradual decay to a small one.

An everyday picture is parking a car. You drive at normal speed along the street, slow down as you approach the space, and creep the last few centimetres. You also do not floor the accelerator the moment you start the engine. Driving the whole way at parking speed would take forever; parking at street speed would end in a dent. A schedule is the plan for when to speed up and when to slow down.

This page covers constant rates, step and exponential decay, cosine annealing and warm restarts, linear warm-up, the one-cycle policy, the inverse square root schedule used by the original Transformer, plateau-based reduction, and the learning-rate range test for finding good values. It ends with PyTorch code and practical advice.

What it is

A schedule is a function η(t) that gives the learning rate at step or epoch t. It works with any optimiser, from plain gradient descent with momentum to Adam and AdamW (see optimizers). The main families:

Schedules are one part of hyperparameter tuning. The peak learning rate, the warm-up length and the final value are the settings that matter most.

Where it came from

The idea that step sizes should shrink over time has a long theoretical history. Herbert Robbins and Sutton Monro's 1951 stochastic approximation paper showed that noisy iterative methods converge when step sizes decrease in a suitable way: their sum must be infinite but the sum of their squares finite. Schedules such as 1/t follow directly from that idea.

In deep learning, step decay became standard through image classification. The ResNet paper by Kaiming He and colleagues (2016) started at 0.1 and divided the rate by 10 when the error plateaued. In 2015 Leslie Smith's "Cyclical Learning Rates for Training Neural Networks" (published at WACV 2017) proposed letting the rate cycle between bounds and introduced the learning-rate range test. In 2016 Ilya Loshchilov and Frank Hutter published "SGDR: Stochastic Gradient Descent with Warm Restarts" (ICLR 2017), which popularised cosine annealing with periodic restarts.

In 2017 Ashish Vaswani and colleagues trained the original Transformer with a linear warm-up over 4,000 steps followed by inverse square root decay. The same year, Priya Goyal and colleagues' "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour" combined a linear scaling rule for large batches with a gradual warm-up over the first five epochs. Leslie Smith and Nicholay Topin's "Super-Convergence" paper (2017) described the one-cycle policy that PyTorch's OneCycleLR implements. Since then, warm-up followed by cosine or linear decay has become a near-standard recipe for training Transformers.

How it works

The formulas

step decay:     η(t) = η₀ · γ^floor(t / s)
exponential:    η(t) = η₀ · γ^t
cosine:         η(t) = η_min + ½ (η_max − η_min) · (1 + cos(π · t / T))
linear warm-up: η(t) = η_max · t / T_w                 for t < T_w
inverse sqrt:   η(t) = d^(−0.5) · min( t^(−0.5), t · T_w^(−1.5) )
one-cycle:      rise from η_max / D to η_max over p·T steps,
                then anneal to η_max / (D · F) by step T

Here t is the current step (or epoch), η₀ is the initial rate, γ (gamma) is the decay factor such as 0.1 or 0.95, s is the step size in epochs between drops, η_max and η_min are the highest and lowest rates, T is the total length of the schedule or cycle, T_w is the number of warm-up steps, d is the Transformer's model dimension, p is the fraction of training spent rising, D is the starting divisor and F the final divisor. In PyTorch's OneCycleLR, p defaults to 0.3, D (div_factor) to 25 and F (final_div_factor) to 10,000.

Constant

The simplest choice. It is fine for short experiments and for some fine-tuning runs with small rates. For full training it usually leaves accuracy on the table, because the final steps are too noisy for the model to settle.

Step and exponential decay

Step decay holds the rate flat and multiplies it by a factor, often 0.1, at chosen epochs. It is easy to reason about and works well, but the drop epochs are extra hyperparameters, and the sudden changes show as steps in the loss curve. Exponential decay multiplies the rate by a constant slightly below 1 every epoch or step. It is smooth but can decay too fast or too slowly depending on γ, which is sensitive to the run length.

Cosine annealing and warm restarts

Cosine annealing follows half a cosine wave from the peak to a minimum over the run. It stays fairly high for a long time, then falls faster in the middle and flattens again near the end. It has one main setting (the length) and tends to work well across many tasks. Warm restarts (SGDR) repeat the cosine shape several times, jumping back up to the peak at each restart, optionally with each cycle longer than the last (T_mult in PyTorch). The restarts can push the model out of one basin into another, and saving a model at the end of each cycle gives an inexpensive ensemble.

Linear warm-up

At the very start, weights are random, gradient statistics are unreliable, and adaptive optimizers such as Adam have barely estimated their running averages. A large step then can send the model somewhere bad from which it never recovers. Warm-up starts with a small rate and raises it linearly to the peak over a few hundred to a few thousand steps (or a few epochs). It matters most for Transformers, for large batch sizes and for Adam-family optimizers. It is almost always combined with a decay schedule afterwards.

One-cycle

The one-cycle policy raises the rate from a low value to a high peak for roughly the first third of training, then lowers it to far below the starting value for the rest. With SGD it also moves momentum in the opposite direction: low momentum when the rate is high, high momentum when it is low. The high middle phase acts as a regulariser and lets training use larger rates than would otherwise be stable, which can shorten training considerably.

Inverse square root

The original Transformer used a schedule that rises linearly during warm-up and then decays in proportion to one over the square root of the step number. It decays more slowly than step or cosine schedules and has no fixed end, so it suits runs where the total length is not known in advance. Many later models replaced it with warm-up plus cosine or linear decay to a fixed end point. A related pattern in some recent large-model training holds the rate flat after warm-up and decays only in a final phase, which makes it easy to extend a run.

Reduce on plateau

ReduceLROnPlateau watches a metric, usually validation loss, and multiplies the rate by a factor when the metric has not improved for a number of epochs. In PyTorch the defaults are factor=0.1 and patience=10. It adapts to the run rather than to a fixed plan, which is useful when you do not know how long training should be. Its drawbacks are that it reacts late and depends on a noisy metric.

The learning-rate range test

To find a good peak rate, train for a few hundred steps while increasing the rate exponentially from very small (say 10⁻⁷) to very large (say 10). Plot the smoothed loss against the rate on a log scale. The loss is flat at first, then falls, then explodes. A good maximum rate is somewhat below the point where the loss is lowest, often around where it is falling most steeply; many practitioners pick a value roughly a factor of ten below the minimum. The test takes minutes and removes most guesswork.

Using schedulers in PyTorch

import torch
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR

model = torch.nn.Linear(64, 1)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

total_steps, warmup = 10_000, 500
warm = LinearLR(opt, start_factor=0.01, end_factor=1.0, total_iters=warmup)
cosine = CosineAnnealingLR(opt, T_max=total_steps - warmup, eta_min=3e-6)
sched = SequentialLR(opt, schedulers=[warm, cosine], milestones=[warmup])

for step in range(total_steps):
    x, y = next(batches)                     # your data iterator
    opt.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    opt.step()
    sched.step()                             # after optimizer.step()
    if step % 1000 == 0:
        print(step, sched.get_last_lr())

Other built-in schedulers include StepLR, MultiStepLR, ExponentialLR, CosineAnnealingWarmRestarts, OneCycleLR, LambdaLR for any custom formula, and ReduceLROnPlateau, whose step() takes the monitored metric as an argument. Call the scheduler after optimizer.step(), and step it once per batch for step-based schedules or once per epoch for epoch-based ones.

Key concepts and terms

A worked example

An illustrative team trains a small Transformer for 10,000 steps with AdamW, peak rate 3 × 10⁻⁴, 500 warm-up steps and cosine decay, as in the code above. For simplicity this calculation takes the cosine floor as zero; the code's floor of 3 × 10⁻⁶ changes the numbers only slightly.

• Step 0. Warm-up starts at 1 per cent of the peak: 3 × 10⁻⁶.
• Step 250. Halfway through warm-up: about 1.5 × 10⁻⁴.
• Step 500. Peak: 3 × 10⁻⁴.
• Step 2,875. A quarter of the cosine phase: 0.5 × 3 × 10⁻⁴ × (1 + cos(π/4)) ≈ 2.56 × 10⁻⁴.
• Step 5,250. Halfway: 0.5 × 3 × 10⁻⁴ × (1 + 0) = 1.5 × 10⁻⁴.
• Step 7,625. Three quarters: 0.5 × 3 × 10⁻⁴ × (1 − 0.707) ≈ 4.4 × 10⁻⁵.
• Step 10,000. Near the floor.

Now compare with the original Transformer's inverse square root schedule with model dimension 512 and 4,000 warm-up steps. The peak at step 4,000 is 512^−0.5 × 4,000^−0.5 ≈ 0.0442 × 0.0158 ≈ 7.0 × 10⁻⁴. At step 16,000 it is half the peak, about 3.5 × 10⁻⁴, and at step 64,000 a quarter. It falls quickly at first and then very slowly, never reaching zero.

Finally, a one-cycle run with max_lr=0.01 and the PyTorch defaults starts at 0.01 ÷ 25 = 4 × 10⁻⁴, rises to 0.01 at 30 per cent of training, and ends at 4 × 10⁻⁴ ÷ 10,000 = 4 × 10⁻⁸.

Strengths and limitations

Strengths:

• Usually improves final accuracy over a constant rate at no extra compute cost.
• Warm-up stabilises the start of training, especially for Transformers and large batches.
• Decay lets the model settle into a minimum instead of bouncing around it.
• Cyclical schedules and restarts can speed up training and give cheap ensembles.
• Plateau-based schedules adapt to runs whose length is unknown.

Limitations:

• Adds hyperparameters: peak rate, warm-up length, decay shape, floor and total length.
• Fixed-length schedules (cosine, one-cycle) must know the total steps in advance; extending a finished run is awkward.
• Plateau schedules react late and depend on noisy validation metrics.
• Schedules interact with batch size, weight decay and optimiser choice, so a recipe may not transfer.
• A schedule cannot rescue a badly chosen peak learning rate.

How it compares with alternatives

Practical guidance and best practices

1. Pick the peak rate first, with a learning-rate range test or a small sweep; the schedule shape matters less than the peak.
2. Default to linear warm-up followed by cosine decay for Transformers and most new projects.
3. Use a warm-up of roughly 1 to 5 per cent of total steps, longer for large batches or unstable models.
4. Decay to a small floor (for example 1 to 10 per cent of the peak, or lower) rather than ending at the peak.
5. Use step-based scheduling for long runs so the rate changes smoothly within epochs.
6. When you scale the batch size up a lot, scale the peak rate and lengthen warm-up.
7. For fine-tuning pre-trained models (see fine-tuning), use a lower peak than for pre-training, a short warm-up and linear or cosine decay.
8. Log the learning rate alongside the loss; many training bugs are schedules that did not do what you thought.
9. If you may need to extend a run, prefer a schedule that holds the rate flat and decays only at the end, or ReduceLROnPlateau.

Common mistakes

• Calling scheduler.step() before optimizer.step(), which skips the first value and triggers a PyTorch warning.
• Stepping a per-batch schedule once per epoch (or the reverse), so the schedule runs far too slowly or too fast.
• Setting cosine T_max to the number of epochs while stepping every batch.
• Forgetting warm-up for Transformers and blaming the architecture when training diverges early.
• Resuming from a checkpoint without restoring the scheduler's state, so the rate jumps back to its starting value.
• Tuning the schedule shape endlessly while the peak rate is still wrong.

Frequently asked questions

Why use warm-up at all?

At the start of training, weights are random and gradient statistics are unreliable, especially for Adam-family optimizers, whose running averages need some steps to settle. A full-size step at that point can push the model into a bad region. Warm-up keeps the first steps small and then lets the rate rise once training is stable.

Is cosine decay better than step decay?

Neither is universally better. Cosine has fewer hyperparameters and changes smoothly, so it is a good default. Step decay remains strong for classic CNN training and is easier to reason about. With a well-tuned peak rate, the gap between them is often small.

Do I need a schedule with Adam?

Yes, usually. Adam adapts step sizes per parameter, but its overall learning rate still benefits from warm-up and decay. Most Transformer training uses AdamW with warm-up and cosine or linear decay.

How do I find a good maximum learning rate?

Run a learning-rate range test: increase the rate exponentially over a few hundred steps and plot the loss. Choose a peak below the point where the loss is lowest, typically where it is still falling steeply. Then confirm with one or two short runs.

What is the difference between OneCycleLR and cosine annealing?

Cosine annealing starts high and falls. One-cycle first rises from a low value to a high peak over about 30 per cent of training, then falls to a very low value, and with SGD it also cycles momentum in the opposite direction. One-cycle often allows a higher peak and faster training.

When should I use ReduceLROnPlateau?

When you do not know how long training should run, or when you want the schedule to react to validation results. It is common in smaller projects and in fine-tuning. For large planned runs, a fixed schedule is more predictable and easier to reproduce.

Glossary

• Learning rate schedule: a rule for changing the learning rate during training.
• Warm-up: a short initial period where the learning rate increases.
• Cosine annealing: decay following half a cosine wave.
• SGDR: stochastic gradient descent with warm restarts.
• One-cycle policy: a single rise-then-fall schedule, often with inverse momentum cycling.
• Step decay: dropping the rate by a factor at set epochs.
• Inverse square root schedule: decay proportional to one over the square root of the step.
• ReduceLROnPlateau: a scheduler that lowers the rate when a metric stalls.
• Range test: a quick sweep of learning rates to find a good maximum.
• Floor: the minimum learning rate a schedule decays to.

Summary

A learning rate schedule changes the optimiser's step size during training so that it can move fast when far from a solution and carefully when close to one. Step and exponential decay are the classic shapes; cosine annealing is a smooth modern default, and warm restarts repeat it to explore. Linear warm-up protects the unstable first steps and is essential for Transformers and large batches. The one-cycle policy rises and then falls to train quickly with high rates, the inverse square root schedule from the original Transformer decays slowly without a fixed end, and ReduceLROnPlateau reacts to validation metrics. Find the peak rate with a range test, use warm-up plus cosine unless you have a reason not to, log the learning rate, and remember that no schedule can fix a badly chosen peak.

Related

It belongs to Neural network basics in my stack, under Deep learning. Others in the same group: