Learning Rate Schedules and Warmup in Deep Learning

A learning-rate schedule changes the step-size coefficient over training. It can help a run move quickly early and make smaller adjustments later, but constant rates also work for some tasks. A useful schedule depends on the optimizer, model, batch size, and training budget. This article starts with a short range test, then develops warmup and decay with explicit step counts.

Finding the starting point: the LR range test

A learning-rate range test increases the rate during one short training run, often over a few hundred updates, while recording loss against log learning rate. Choose a broad but plausible range for the model. This is a trajectory through changing parameters and optimizer state, not an independent trial at every learning rate.

The loss may stay flat, fall, and then rise sharply, but these regions need not all appear. Data variation and earlier updates affect the curve. A rate below a sustained deterioration point can be a candidate for a fresh run; choosing one tenth of that rate is a heuristic, not an optimum. Confirm candidates with short runs from the same starting state and then compare validation performance.

import numpy as np

def lr_range_test(train_step, num_iters=200, lr_min=1e-7, lr_max=10.0):
    """train_step(lr) updates once and returns a nonnegative post-update loss."""
    if num_iters < 2 or not (0 < lr_min < lr_max < np.inf):
        raise ValueError("Use at least two iterations and finite positive LR bounds")
    history, best = [], np.inf
    for lr in np.geomspace(lr_min, lr_max, num_iters):
        loss = float(train_step(lr))
        history.append((lr, loss))
        if not np.isfinite(loss):
            break
        if loss < 0:
            raise ValueError("This multiplicative stopping heuristic requires nonnegative loss")
        if loss > 4*max(best, 1e-12):
            break
        best = min(best, loss)
    return history

# A fresh scalar model for the demonstration: J(theta) = theta**2.
def make_train_step():
    theta = 1.0
    def train_step(lr):
        nonlocal theta
        theta -= lr*2*theta
        return theta**2
    return train_step

history = lr_range_test(make_train_step(), num_iters=12, lr_min=1e-3, lr_max=3.0)
print("recorded updates", len(history))
print("first rate", round(history[0][0], 6))
print("last rate", round(history[-1][0], 6))
print("last loss", round(history[-1][1], 6))
# recorded updates 12
# first rate 0.001
# last rate 3.0
# last loss 0.36047

The geometric grid includes both requested endpoints if all iterations run; early stopping may prevent reaching the upper one. The four-times-best rule is only a guard for this nonnegative-loss example. It depends on loss scale and offset and does not prove divergence; a noisy batch can trigger it. Other losses need a different stopping rule. The scalar example checks the mechanics, not a universally useful rate range.

A range test can narrow a search but does not replace validation across candidate schedules. Use a disposable model or restore the model, optimizer, scheduler, and relevant random/data-loader state afterwards. Continuing from sweep-modified weights would confound the next trial. Reconsider the range when the architecture, batch size, optimizer, or loss scale changes.

The optimizers that turn these gradients into parameter updates are compared in Deep Learning Optimizers: Mini-Batch, Momentum, RMSprop, and Adam.

Warmup: why the first steps are different

Early training can have different activation and update scales from later training. A large initial rate may destabilize some configurations, but initial gradients are not necessarily large or uninformative. Adam’s moment estimates also have little history at first; bias correction normalizes their weights without eliminating their variability. These are reasons to test a gradual start, not a diagnosis that applies to every model.

Linear warmup raises the rate from a small initial value to a chosen peak over a specified number of optimizer updates. It is common in many Transformer recipes, but its necessity and length depend on initialization, normalization placement, optimizer, batch size, and peak rate. Compare plausible warmup lengths, including no warmup when appropriate, under the intended training budget.

The main decay shapes

Step decay multiplies the rate by \(0<\gamma<1\) at selected milestones, for example by 0.1 twice. The learning-rate curve jumps there; the loss need not immediately fall or show a visible discontinuity. Specify whether milestones count epochs or optimizer updates.

Cosine annealing decays smoothly from \(\alpha_{\max}\) to \(\alpha_{\min}\) over \(T\) steps:

\[\alpha_t=\alpha_{\min}+\tfrac{1}{2}(\alpha_{\max}-\alpha_{\min})\left(1+\cos\frac{\pi t}{T}\right)\]

For \(0\le t\le T\), this cosine phase begins at \(\alpha_{\max}\), ends at \(\alpha_{\min}\), and falls fastest near its midpoint. It requires choosing a decay horizon. Extending a run after reaching the floor requires an explicit continuation rule, such as holding that floor or starting a new phase.

Inverse square root with warmup can be written in terms of peak rate \(\alpha_{\rm peak}\) and warmup duration \(w>0\) as \(\alpha_t=\alpha_{\rm peak}\min(t/w,\sqrt{w/t})\) for update numbers \(t\ge1\). It rises linearly until \(t=w\), then decays like \(t^{-1/2}\), without needing a final step count. The original Transformer schedule uses \(d_{\rm model}^{-1/2}\min(t^{-1/2},tw^{-3/2})\), where \(d_{\rm model}\) is the model width. A formula based only on \(1/\sqrt{\max(t,w)}\) would be flat before \(w\), not a linear warmup.

One-cycle raises the rate to a peak and then lowers it, often to below the starting value. The rising fraction is configurable; 30% is one choice. Some versions vary momentum inversely with the rate. The larger-rate phase can alter regularization and optimization behavior, but a high peak can also destabilize training. Specify the phase lengths, minimum rate, and any momentum schedule.

def lr_at(step, total_steps, base_lr, warmup_steps=0, min_lr=0.0):
    if total_steps < 2 or not 0 <= warmup_steps <= total_steps-2:
        raise ValueError("Leave at least two decay samples")
    if not 0 <= step < total_steps:
        raise ValueError("step must index an actual optimizer update")
    if not np.isfinite([base_lr, min_lr]).all() or not 0 <= min_lr <= base_lr:
        raise ValueError("Require finite 0 <= min_lr <= base_lr")
    if step < warmup_steps:
        return base_lr*(step+1)/warmup_steps
    progress = (step-warmup_steps)/(total_steps-warmup_steps-1)
    return min_lr+(base_lr-min_lr)*0.5*(1+np.cos(np.pi*progress))

total, warm = 10000, 500
for s in [0, 250, 499, 500, 2500, 5000, 9999]:
    print(s, round(lr_at(s, total, 1e-3, warm), 8))
# 0 2e-06
# 250 0.000502
# 499 0.001
# 500 0.001
# 2500 0.00089455
# 5000 0.00054121
# 9999 0.0

This function uses zero-based update indices \(0,\ldots,N-1\), unlike the one-based inverse-square-root formula. With \(w\) warmup updates, indices \(w-1\) and \(w\) both use the peak; the last update uses exactly min_lr. The denominator \(N-w-1\) counts intervals between decay samples. With no warmup, the first update uses the peak. A zero final rate means the final update makes no gradient-driven parameter change; choose a positive floor if that is undesirable. The floor describes the decay endpoint and need not bound the early warmup rates.

Batch size and the linear scaling rule

At fixed parameters, the mean of \(B\) independent per-example gradients with finite covariance has covariance proportional to \(1/B\). This is a sampling statement, not permission to double the step size; shuffling without replacement, batch-dependent layers, and correlated examples modify the setup. Linear scaling, \(\alpha\propto B\), approximates combining several small-batch SGD steps when their gradients change little. It has worked in particular large-batch vision experiments, but there is no universal batch-size cutoff or square-root-rule replacement. Optimizer and problem details matter.

Effective batch size is the number of examples contributing to one optimizer update. In a typical synchronized run it equals per-device batch times device count times gradient-accumulation steps. Moving to more devices while preserving that effective batch does not itself demand a new rate. If it changes, distinguish equal-update and equal-example budgets and retune as needed; the hardware count alone does not determine the schedule.

Reading the loss curve

How to allocate a tuning budget across these choices is the subject of A Hyperparameter Tuning Strategy That Fits Your Budget.

Practical defaults

• Log the learning rate as a series alongside the loss. This exposes counter and phase mistakes, including stepping per epoch when the schedule expects optimizer updates.
• For the function above, set the returned rate before that indexed optimizer update. For a stateful library scheduler, follow its call-order contract; many PyTorch schedules advance after the optimizer step. Advance only at the configured interval. With gradient accumulation, count optimizer updates, not microbatches, and handle skipped optimizer steps consistently.
• For fine-tuning, start with a model-specific recipe and compare rates. Full-model tuning and training newly added adapters or heads can require different settings; no fixed ratio to the pretraining rate covers both.
• Record peak rate, warmup duration, decay horizon, and minimum rate separately. Test these choices under the actual model and batch configuration.
• State the comparison budget: optimizer updates, processed examples or tokens, or wall-clock time. Evaluate at the budget relevant to the intended run; schedules need not finish at zero to be compared.

Log the rate actually used by each parameter group, not merely the value requested for the next update. Save optimizer and scheduler state when checkpointing so that resuming does not restart warmup or shift the decay phase.

How these models are pretrained at scale is covered in LLM Pretraining: Objectives, Data, and Scaling Laws.

Exercises

1. Investigate three loss curves. For each description, give a plausible explanation, an alternative, and the first check: (a) loss becomes nan at step 40; (b) loss falls for 500 steps then plateaus above zero; (c) loss decreases slowly and is still falling when training ends.

You should get: candidate explanations; the curve alone does not establish a unique cause.

2. Schedule boundaries. Using lr_at, verify the warmup handover and final rate. Deliberately lower only the final warmup sample by one increment. Calculate the resulting jump and explain why it need not cause a visible jump in loss.

You should get: matching rates at the intended handover, an exact decay endpoint, and a measurable rate jump in the modified schedule.

total, warm, base = 10000, 500, 1e-3
left = lr_at(warm-1, total, base, warm)
right = lr_at(warm, total, base, warm)
bad_left = base*(warm-1)/warm
print("handover", round(left, 8), round(right, 8))
print("final", lr_at(total-1, total, base, warm))
print("modified jump", round(right-bad_left, 8))
# handover 0.001 0.001
# final 0.0
# modified jump 2e-06

3. Linear scaling in practice. A one-device recipe uses batch 256, peak rate 0.001, 500 warmup updates, and 10,000 total updates. Move to eight devices with the same per-device batch and no accumulation. Assuming full batches, preserve the total number of processed examples and the warmup example budget. Compute the effective batch, candidate linearly scaled rate, total updates, and approximate warmup updates.

You should get: a larger batch and candidate rate, but fewer total and warmup updates for the same example budgets.

References

• Smith (2017). Cyclical Learning Rates for Training Neural Networks.
• Loshchilov and Hutter (2017). SGDR: Stochastic Gradient Descent with Warm Restarts. ICLR.
• Goyal et al. (2017). Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour. arXiv preprint.
• Vaswani et al. (2017). Attention Is All You Need.
• Smith and Topin (2019). Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates.