Learning rate schedules

At 3 a.m., a training run can look perfectly healthy: the loss falls for two hours, then starts bouncing between 1.8 and 2.1. The validation score has stopped improving. You try a larger model. You change the optimizer. You blame the data.

Often, the learning rate is simply still too large for the part of the landscape the model has reached.

The learning rate scales an optimizer’s update after it uses a gradient. The gradient points toward increasing loss; subtracting it, or an optimizer-specific transformed version of it, gives a local descent update. A fixed rate asks for the same movement from random initialization to final convergence. Learning-rate schedules replace that with a timetable: ease in, move quickly, then settle down.

The intuition: walking downhill

Imagine trying to reach the bottom of a foggy valley.

At the beginning, large steps cover useful ground. Near the bottom, the measurements are noisy and a large step can send you from one side of the valley to the other.

A training batch is one noisy measurement of the slope. Reducing the learning rate lowers the size of late corrections, allowing the weights to settle instead of jittering around a good region.

The beginning has a different problem: weights and optimizer state are not well calibrated. A large update can amplify an activation, make later gradients enormous, and send the run into NaN. Warmup increases the rate gradually while the run stabilizes.

A common shape therefore has three parts:

• Warmup: start small and ramp to the peak rate.
• Main training: use the peak rate while broad progress is valuable.
• Decay: reduce the rate as the run approaches its budget or useful improvements slow down.

This is a useful prior, not a law. A small classifier trained for 500 updates may work well with a constant rate; a transformer trained for billions of tokens usually needs a deliberate plan.

What the learning rate actually changes

For plain stochastic gradient descent:

weight_next = weight_now - learning_rate * gradient

AdamW adds moving averages and normalization, but its update is approximately:

weight_next = weight_now
              - learning_rate * normalized_momentum
              - learning_rate * weight_decay * weight_now

The exact implementation includes bias correction and numerical details. The important point is that the schedule still scales the adaptive update. Adam does not make the learning rate irrelevant; it makes the relationship between the displayed rate and an individual parameter’s movement less direct.

Decay helps because late gradient estimates contain more noise than useful direction, and because loss surfaces have narrow and wide directions at once. A rate that is safe in a wide direction can overshoot in a narrow one. Early aggressive movement may be worthwhile; later, the same movement can repeatedly cross the minimum.

Decay is not automatically beneficial. If it begins while the model is still underfitting, useful learning stops early. If it reaches zero before the budget is spent, the run becomes an expensive slideshow.

A schedule with numbers

Use one concrete run throughout this page:

• 10,000 optimizer updates in total
• 500 warmup updates
• peak learning rate 3e-4
• AdamW as the optimizer

An optimizer update means one call that changes the weights. It is not necessarily one batch: gradient accumulation may use several batches before one update.

With linear warmup:

learning_rate = peak_rate * (update + 1) / warmup_updates

The first update uses 3e-4 * 1 / 500, or 6e-7. Update 499 uses 3e-4. After warmup, cosine decay uses:

learning_rate = 0.5 * peak_rate
                * (1 + cos(pi * progress))

Here, progress goes from 0 at the end of warmup to 1 at the end of the schedule.

There are 9,500 updates after warmup. Update 5,250 is exactly halfway through that section, so the cosine factor is 0.5. The final value is almost zero because the final update is one index before the mathematical endpoint. The practical conclusion is that the last part of this schedule contributes almost no movement.

Why transformers often use warmup

Warmup is common in transformer training because unstable effects meet at the start of a run. Early representations may be poorly scaled. An update that changes an output projection alters attention scores, which changes how tokens exchange information and can reinforce the original change through many layers.

Adam’s moment estimates also begin with little history. Bias correction does not provide information the optimizer has not yet observed. An unlucky first batch or unusual gradient magnitude can make a full-size update too violent.

Warmup limits those early parameter changes while the model and optimizer gather history. It does not change the target loss or fix bad data. A common starting point is linear warmup for 1% to 5% of planned optimizer updates; 500 updates is 5% in the running example.

The peak matters as much as the ramp. Warm up to an excessive rate and you have only postponed the explosion; warm up to a tiny rate and the model may spend its budget taking timid steps.

Cosine decay: why the smooth curve is popular

Cosine decay starts at the peak and approaches zero:

rate(t) = 0.5 * peak_rate * (1 + cos(pi * t / T))

T is the number of decay updates and t is the current position. The curve changes smoothly, unlike step decay’s abrupt jumps, and spends substantial time at moderate rates. That makes it a reasonable default when the total training budget is known.

It is not magic. A cosine schedule with the wrong T is a confidently wrong plan: if training is extended far beyond the planned budget, the rate may already be near zero; if T is too large, the model remains noisy near the end.

Choosing among the schedules

Choose based on what you know before training starts:

Step decay multiplies the rate at fixed boundaries. It is easy to reason about and remains useful for established CNN recipes, but its discontinuities and dataset-specific boundaries make it less portable.

Decide whether the budget is measured in updates, examples, or tokens. If you double the batch size while keeping the update count fixed, the model sees twice as many examples. If you keep epochs fixed, it performs fewer updates. The relationship between batch size and rate deserves its own treatment in batch size and learning rate.

ReduceLROnPlateau: useful feedback, noisy control

ReduceLROnPlateau reduces the rate after a monitored metric stops improving. For example: validate each epoch, wait five unimproved validations, then multiply the rate by 0.1.

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode="min",
    factor=0.1,
    patience=5,
)

for epoch in range(epochs):
    train_one_epoch(...)
    val_loss = validate(...)
    scheduler.step(val_loss)

This fits projects with genuinely unknown stopping times, such as a small tabular model or some fine-tunes. It is a poor fit for large pretraining: validation is expensive and noisy, so two runs may reduce at different times. Do not call it after every batch; it expects the monitored metric, not one noisy mini-batch loss.

The production pattern in PyTorch

Define the unit once and make the loop use it. Here total_steps means optimizer updates, not epochs or data-loader iterations.

import math
import torch

total_steps = 10_000
warmup_steps = 500
peak_lr = 3e-4

def lr_multiplier(update):
    if update < warmup_steps:
        return (update + 1) / warmup_steps

    decay_steps = total_steps - warmup_steps
    progress = (update - warmup_steps) / decay_steps
    progress = min(progress, 1.0)
    return 0.5 * (1.0 + math.cos(math.pi * progress))

optimizer = torch.optim.AdamW(model.parameters(), lr=peak_lr)
scheduler = torch.optim.lr_scheduler.LambdaLR(
    optimizer,
    lr_lambda=lr_multiplier,
)

for update in range(total_steps):
    optimizer.zero_grad(set_to_none=True)
    loss = train_step(...)
    loss.backward()
    optimizer.step()
    scheduler.step()

Call optimizer.step() before scheduler.step() so the initial scheduled value is not skipped. Log the rate from the optimizer, especially when parameter groups have different rates.

With accumulation, the scheduler advances after the weight update. If four microbatches produce one optimizer update, four microbatches equal one schedule step. Mixed-precision scaling may skip an update after overflow; keep the schedule aligned with real updates, and log skipped updates.

Save optimizer and scheduler state in checkpoints. Restoring weights while restarting warmup, or restoring the optimizer with a newly high rate, can cause an immediate loss spike.

Failure modes you can see in the logs

The rate is still at its peak after thousands of batches. You probably called a step-based scheduler once per epoch while configuring total_steps in updates. Count optimizer calls and move scheduler.step() into that loop.

The loss becomes NaN during warmup. Shorten the peak rate, lengthen warmup, and inspect gradient norms, inputs, labels, initialization, and mixed-precision scaling. Clipping can contain an occasional large gradient; it should not hide a systematically excessive rate.

Training becomes flat while validation is still poor. The rate may already be near zero. If you step once per microbatch, accumulation makes it decay too early. If you configure total microbatches but step per optimizer update, it decays too slowly. Compute the budget and step frequency from actual optimizer updates.

The resumed run spikes or plateau reductions follow every validation wobble. Compare the learning-rate log before and after the checkpoint. Restore scheduler state and global update count; for noisy validation, increase patience or use a deterministic schedule.

Plot learning rate against optimizer update beside the loss curve. A schedule bug then looks like a schedule bug instead of an argument about architecture.

A sensible starting recipe

For transformer pretraining with a known token budget, choose a peak rate from a small pilot, warm up for roughly 1% to 5% of optimizer updates, and cosine-decay over the remaining budget. Match the schedule to the actual batch size, accumulation, and token count.

For fine-tuning, use a smaller peak rate and consider a short warmup. A constant rate can be better for a very short run. For a CNN with a trusted recipe, step decay remains a reproducible choice. For a small tabular model, try constant rate first, then plateau reduction if validation behavior calls for it.

Before spending the full budget, log the first, peak, midpoint, and final rates. This catches most schedule mistakes immediately.

Quick check