Learning-Rate Schedules

A learning-rate schedule is a deterministic function $\eta: \mathbb{N} \to \mathbb{R}_{>0}$ mapping the current training step (or epoch) $t$ to a learning rate $\eta(t)$, used in the update rule

where $g_t$ is the (possibly momentum-adjusted) gradient estimate at step $t$. The schedule is fixed before training starts — it does not adapt to the observed loss (that would be an adaptive method like ReduceLROnPlateau, a related but distinct idea) — although some schedules are parameterized by quantities like total number of steps $T$, which must be known or estimated in advance. The reason a schedule helps is rooted in the shape of the loss landscape and the statistics of stochastic gradients. Early in training, the model is typically far from any minimum, gradients tend to be large and informative, and a larger step size lets the model make rapid progress across broad, low-curvature regions. Late in training, the model is near a minimum (or a flat basin containing one), the loss surface is locally more like a bowl with possibly high curvature in some directions, and stochastic gradient noise (from mini-batching) has variance that does not shrink as training proceeds. A large learning rate at this stage causes the parameters to bounce around the bottom of the bowl instead of settling into it — the update variance, roughly $\eta^2 \,\mathrm{Var}(g_t)$, must shrink for the iterate to converge rather than oscillate around a noise floor. Formally, in stochastic approximation theory (Robbins–Monro conditions), a decaying schedule satisfying

is sufficient (together with standard smoothness and noise assumptions) for SGD iterates to converge almost surely to a stationary point of the objective. The first condition ensures the schedule does not decay so fast that the algorithm cannot travel arbitrarily far (it must be able to escape a bad starting point and reach any region of the space). The second ensures the accumulated noise injected by stochastic gradients is summable, so the random fluctuations eventually die out and the sequence stabilizes. A constant learning rate satisfies neither condition properly: it can travel infinitely far ($\sum \eta = \infty$, fine) but the noise term $\sum \eta^2$ diverges too, meaning the iterate never truly settles — it converges only to a noise ball around the minimum whose radius is proportional to $\eta$. Schedules like $1/t$ decay or step decay satisfy both conditions and hence give a much stronger convergence guarantee, at the cost of needing to travel less late in training. Deep learning practice departs from this classical picture in three important ways, motivating the specific schedules covered here. First, training runs for a finite, known horizon $T$, so schedules can be shaped to hit a target minimum learning rate exactly at $T$ (as cosine annealing does) rather than decaying forever. Second, the very first few hundred to thousand iterations are often the most fragile part of training — activations and gradient scales have not yet equilibrated, and a full-sized learning rate can cause instability or even divergence — motivating warmup, a short initial phase where $\eta(t)$ ramps up rather than down. Third, non-convexity means the "right" trajectory of $\eta(t)$ is partly empirical: cosine schedules, for instance, were popularized less by a convergence proof and more by consistently strong empirical results across architectures.

We now specify each schedule precisely as a function of the training step $t$ (or epoch $e$), assuming a base learning rate $\eta_0$ and, where relevant, a total training horizon $T$.

Step decay multiplies $\eta$ by a fixed factor $\gamma \in (0,1)$ (commonly $0.1$ or $0.5$) every time training crosses one of a fixed set of milestone epochs $\{m_1, m_2, \dots\}$:

So $\eta$ is a piecewise-constant staircase: full rate until $m_1$, then $\eta_0\gamma$ until $m_2$, then $\eta_0\gamma^2$, and so on.

Exponential decay shrinks the learning rate smoothly and continuously at every step, governed by a decay rate $k>0$:

for a decay factor $\gamma$ applied every $s$ steps. Unlike step decay, there are no discontinuities — the rate glides down monotonically at every single update.

Cosine annealing follows one half-period of a cosine curve from $\eta_0$ down to a floor $\eta_{\min}$ (often $0$) over the horizon $T$:

At $t=0$ this gives $\eta(0)=\eta_0$; at $t=T$ it gives $\eta(T)=\eta_{\min}$. The curve decays slowly at first, drops fastest through the middle of training, and flattens out again as it approaches $\eta_{\min}$ — an S-shaped decay that spends more time near both the top and the bottom than a linear ramp would. A common variant, cosine with warm restarts (SGDR), resets $t$ to $0$ periodically, creating a sawtooth of cosine cycles used to encourage the optimizer to explore multiple basins and produce ensembles of snapshots.

Warmup is not a full schedule on its own but a short prefix, typically the first $T_w$ steps, during which $\eta$ ramps linearly (or sometimes with a low-order polynomial) from a small value (often $0$) up to $\eta_0$:

After $t = T_w$, control passes to a main schedule (step, exponential, or cosine) that takes over with $\eta_0$ as its own starting value. Warmup is essential for very deep networks, large batch sizes, and Transformer architectures, where the raw gradient signal at initialization can be poorly scaled and a full-size step causes the loss to spike or the optimizer's momentum/adaptive-moment estimates (as in Adam) to lock onto a bad early direction. ```

# Generic training loop with a pluggable learning-rate schedule.

# schedule(t) returns eta_t given the current global step t. function TRAIN(model, data, T, schedule, momentum_beta): velocity = zeros_like(model.parameters) t = 0 for epoch in 1..num_epochs: for batch in data: g = compute_gradient(model, batch) # stochastic gradient estimate eta_t = schedule(t) # <-- the only thing that changes velocity = momentum_beta velocity + g model.parameters -= eta_t velocity t = t + 1 return model.parameters

# --- Schedule implementations --- function STEP_DECAY(t, eta0, gamma, milestones): k = count(m in milestones such that m <= t) return eta0 gamma^k function EXPONENTIAL_DECAY(t, eta0, gamma, step_size): return eta0 gamma^(t / step_size) function COSINE_ANNEALING(t, eta0, eta_min, T): t_clamped = min(t, T) # hold at eta_min after t = T cos_term = cos(pi t_clamped / T) return eta_min + 0.5 (eta0 - eta_min) (1 + cos_term) function WARMUP_THEN(t, T_warmup, eta0, main_schedule): if t < T_warmup: return eta0 (t + 1) / T_warmup # linear ramp, avoids eta=0 at t=0 else: return main_schedule(t - T_warmup) # main schedule restarts its clock at 0 ``` Each `schedule(t)` function is a pure function of the step counter, queried once per update; it adds negligible overhead ($O(1)$ time, no extra memory) compared to the gradient computation itself.

Time and space complexity. Every schedule discussed is $O(1)$ time and $O(1)$ extra space per step: computing $\eta(t)$ requires a handful of arithmetic operations (a power, an exponential, or a cosine evaluation) independent of model size, dwarfed by the $O(\text{number of parameters})$ cost of the gradient step itself. Storing the schedule requires only a few scalars ($\eta_0$, $\gamma$, milestones, or $T$), not per-parameter state — this is in contrast to adaptive optimizers like Adam, which maintain $O(\text{parameters})$ auxiliary state. Consequently, a schedule is essentially "free": it changes what number multiplies the gradient, not the cost of computing that gradient. The only indirect cost is that a poorly chosen schedule can lengthen the number of steps needed to converge, which is a statistical/optimization cost, not a computational one per step.

Correctness argument for decaying schedules. Consider SGD on a smooth, possibly non-convex objective $f$ with $L$-Lipschitz-continuous gradient, using unbiased stochastic gradients $g_t$ with bounded variance $\mathrm{Var}(g_t) \le \sigma^2$. A standard descent-lemma argument gives, for one step with learning rate $\eta_t$,

Summing over $t=1,\dots,T$ and rearranging, the average squared gradient norm is bounded by a telescoping term (which is finite because $f$ is bounded below) plus a noise term proportional to $\sum_t \eta_t^2$. For the right-hand side to vanish as $T \to \infty$, we need $\sum_t \eta_t \to \infty$ (so the useful first term's denominator grows, letting us divide it out and still get a shrinking bound) while $\sum_t \eta_t^2 < \infty$ (so the accumulated noise stays bounded). This is exactly the Robbins–Monro condition stated earlier. Step decay, exponential decay (over a finite horizon, effectively), and cosine annealing (whose learning rate integrates to a finite but large value and whose squared values are summable since it is bounded and eventually near zero) all satisfy versions of this condition in the finite-horizon setting used in practice, giving a principled reason they outperform a constant rate on the same total noise budget.

Why warmup does not break this. Warmup only affects a short prefix of $O(T_w)$ steps out of a much longer horizon $T$; since $T_w \ll T$ in practice, its contribution to $\sum \eta_t$ and $\sum \eta_t^2$ is negligible, and it does not change the asymptotic convergence argument. Its benefit is a separate, more practical one: at initialization, the loss landscape's local curvature and gradient scale are poorly conditioned (weights have not adapted to the data distribution), and using the full $\eta_0$ immediately can cause an update large enough to move parameters into a region where gradients blow up (particularly acute for deep networks, large batch sizes, or Transformer layer-normalization interactions). Ramping $\eta$ up avoids taking a large step until the local geometry has stabilized somewhat, after which the full learning rate is safe.

Honest tradeoffs.

• Step decay is simple and interpretable but the milestone epochs are extra hyperparameters that must be tuned per dataset/architecture; a badly placed milestone wastes epochs at a rate that is already too small, or drops too late and never lets the model settle.
• Exponential decay has just one decay-rate hyperparameter but can decay the rate to near-zero well before training ends if $k$ is even slightly too aggressive, effectively freezing learning early.
• Cosine annealing needs the total horizon $T$ known in advance; extending training (e.g., resuming with more epochs later) requires recomputing the schedule, unlike step or exponential decay which naturally continue.
• Warmup adds one more hyperparameter ($T_w$) and, for very short training runs, can consume a nontrivial fraction of the whole budget in what is initially near-zero learning rate.

Consider training a network with base learning rate $\eta_0 = 0.1$, total horizon $T = 100$ epochs, warmup length $T_w = 5$ epochs, and each schedule's own decay parameters, tracked epoch by epoch.

Reading this table: step decay stays completely flat at $0.1$ until epoch $30$, then instantaneously drops to $0.01$ — a $10\times$ shock that typically causes a visible bump in training loss curves (the model was still making large steps and suddenly its exploration radius shrinks). Exponential decay is already down to about $0.078$ by epoch $5$ and continues shrinking every single epoch without any jump, reaching $0.0018$ (below step decay's final value) by epoch $99$ — for this $\gamma$, it decays faster in the long run than the $3$-stage step schedule with $\gamma=0.1$. Cosine annealing starts identically to the others at $\eta_0$, decays slowly for the first fifth of training, drops through its steepest region in the middle third, and flattens smoothly to nearly zero by epoch $99$, with no discontinuities anywhere. The warmup+cosine column shows the ramp from $0.02$ at epoch $0$ to the full $0.1$ at epoch $5$ (note $\eta(0)$ is not exactly zero — using $(t{+}1)/T_w$ avoids a literal zero-rate first step which would produce no update at all — a common off-by-one detail), after which cosine annealing takes over with its clock reset to zero at that point, following the same S-shape but over the remaining $95$ epochs. A second, smaller trace makes the discontinuity of step decay concrete on a per-batch loss curve. Suppose at epoch $29$, just before the first milestone, the running loss is $0.42$ and decreasing steadily by about $0.001$ per epoch. At epoch $30$, the learning rate drops $10\times$; empirically the loss often increases slightly for one or two epochs (say to $0.44$) because the optimizer's momentum buffer, built for the old step size, briefly overshoots relative to the new tiny step size before re-equilibrating, then resumes a steeper, more stable decrease than before the drop (because a smaller step size lets it settle into finer local structure). This transient bump-then-improve pattern is the classic visual signature of a step-decay milestone in a training log, and recognizing it is a key diagnostic skill developed further in Debugging the Training Process.

The theoretical foundation for decaying step sizes predates deep learning by decades, rooted in stochastic approximation theory developed for sequential statistical estimation.

Learning-rate schedules are embedded in essentially every major deep learning framework and training recipe used in industry and research. Computer vision training pipelines (ResNet, EfficientNet, and successors) historically relied on step decay with milestones tied to a fixed epoch budget, while more recent vision and self-supervised recipes default to cosine annealing, often with a linear warmup of a few epochs. Large language model training — GPT-style and encoder-only Transformer pretraining alike — universally uses a warmup phase (typically thousands of steps) followed by either linear or cosine decay to a small floor value, a recipe now considered essential rather than optional for stable large-batch, large-model training. Reinforcement learning and speech recognition pipelines frequently borrow the same cosine-with-warmup recipe, and hyperparameter-search tools (as covered in Hyperparameter Tuning) often treat the schedule family and its parameters — not just a single learning rate — as part of the search space. Cloud training frameworks (e.g., learning-rate scheduler modules in major deep learning libraries) implement all four schedules described here as first-class, composable objects, reflecting how central this simple idea has become to reliably training deep networks at scale.

Ask the tutor: How would you design a learning-rate schedule for a training run whose total number of steps isn't known in advance?