Learning Rate Scheduling
Introduction
The learning rate η is the single most important hyperparameter in gradient-based training. It scales every update: θ ← θ − η·∇L(θ). Pick it too large and training overshoots and diverges; too small and training crawls and may stall in a poor region. A learning rate schedule changes η over the course of training instead of holding it fixed — and a good schedule almost always beats the best constant value.
Core intuition: you want a large learning rate early to cover ground quickly and escape bad regions, and a small learning rate late to stop bouncing around the minimum and settle precisely into it. A fixed rate cannot do both.
Why a fixed learning rate is suboptimal
Imagine rolling toward the bottom of a valley. Far from the bottom, big steps make fast progress. But near the bottom, the same big steps overshoot, sending you up the opposite wall — the loss oscillates and never settles. Shrink the steps and you converge smoothly, but if you used those small steps from the start, you would waste enormous time crossing the flat early terrain.
Rate too high (fixed)
- Fast early progress
- Overshoots the minimum, loss bounces
- Can diverge entirely
Rate too low (fixed)
- Stable, settles precisely
- Painfully slow early on
- May get stuck in flat/saddle regions
A schedule resolves the tension: start high, then anneal toward a small value. The demo below lets you see exactly this trade-off play out on a real descent.
Common schedules and their formulas
Throughout, η (or η_max) is the initial learning rate, t is the step or epoch index, and T is the total number of steps. Here are the workhorse schedules.
Step decay
η(t) = η · γ^⌊t / s⌋
Multiply the rate by a factor γ < 1 (e.g. 0.1 or 0.5) every s epochs. The ⌊t/s⌋ floor produces a staircase: flat, then a sudden drop. Simple and very common in classic vision training (e.g. drop ×0.1 at epochs 30 and 60).
Exponential decay
η(t) = η · e⁻ᵏᵗ
A smooth continuous decay controlled by rate k > 0. Larger k decays faster. Unlike the staircase, it shrinks the rate every single step.
Cosine annealing
η(t) = η_min + ½(η_max − η_min)(1 + cos(π·t / T))
Follows the first half of a cosine wave from η_max down to η_min over T steps. It stays high for a while, then decays fast in the middle, then flattens out gently near the end — an empirically excellent profile that is now the default for many modern networks. At t = 0 the cosine term is +1 so η = η_max; at t = T it is −1 so η = η_min.
Linear warmup
η(t) = η · (t / w) for t < w
Ramp the rate linearly from 0 up to η over the first w steps, then hand off to a decay schedule. Warmup is crucial for Transformers and adaptive optimizers like Adam (see the dedicated section below).
Warmup + cosine
The standard recipe for large models: linearly warm up 0 → η_max over the first w steps, then cosine-anneal η_max → η_min over the remaining T − w steps. You get stable early training and a clean late settle.
Cyclical / one-cycle
Instead of monotonic decay, oscillate the rate between a low and high bound. Cyclical LR (Smith) repeatedly ramps up and down; one-cycle does a single big rise then a long fall to a value below the start. The high phases help escape saddle points and sharp minima, often training faster.
ReduceLROnPlateau
An adaptive, data-driven schedule rather than a fixed formula: monitor a validation metric and, when it stops improving for a "patience" number of epochs, multiply the rate by a factor (e.g. 0.1). It reacts to the actual training curve instead of a preset timetable.
Interactive schedule explorer
Pick a schedule and tune its hyperparameters. The top chart plots the learning-rate curve computed directly from the formula above. The bottom chart runs an actual gradient descent on the convex loss L(x) = x² using that scheduled rate at each step, and compares it against a constant rate fixed at the schedule's peak — so you can watch how the shape of the schedule changes convergence.
The descent starts at x = 4 on the bowl L(x) = x² (gradient ∇L = 2x) and takes one step per iteration using the learning rate read off the curve above. A constant learning rate set to the schedule's peak can keep overshooting the minimum and bounce, while a schedule that decays late settles in. Try cranking the initial LR very high to see the constant baseline diverge.
- Set peak η high (≈0.9) and watch the red constant baseline bounce or diverge.
- Switch to warmup+cosine and shrink warmup to 0 — early steps get jumpy.
- Compare step vs. cosine final loss for the same peak η.
Why warmup stabilizes large models
Warmup matters most at the very start of training, before the network has learned anything useful. There are two reinforcing reasons it helps:
1. Adaptive optimizers have unreliable early statistics
Adam scales each update by 1/√(v̂ + ε), where v̂ is a running estimate of the squared gradient. In the first few steps that estimate is based on almost no data, so it has high variance and can produce wildly large, badly-scaled steps. A small learning rate during warmup keeps those early, noisy updates from corrupting the weights; by the time η reaches its peak, the variance estimates have settled.
2. Deep/Transformer landscapes are fragile at init
With residual connections and layer norm, a large step at initialization can push activations into a regime the network never recovers from (loss spikes, NaNs). Ramping η up gently lets the weights organize before full-strength updates arrive. This is why the original Transformer used an explicit warmup-then-decay schedule, and why nearly every large language model trains with warmup today.
The original Transformer schedule scales the rate as η ∝ d_model⁻⁰·⁵ · min(t⁻⁰·⁵, t · w⁻¹·⁵): a linear warmup for the first w steps (the t · w⁻¹·⁵ term grows with t) followed by an inverse-square-root decay (the t⁻⁰·⁵ term shrinks). The two terms cross at exactly t = w, the warmup length.
Choosing a schedule
| Schedule | Shape | Best for | Watch out |
|---|---|---|---|
| Step decay | Staircase drops | Classic CNNs, simple recipes | Drop timing is a guess |
| Exponential | Smooth exponential | When you want continuous decay | Can decay too fast late |
| Cosine annealing | Half cosine to η_min | Modern vision & NLP default | Needs known total length T |
| Warmup + cosine | Ramp then anneal | Transformers, large models, Adam | Two extra knobs (w, η_min) |
| Cyclical / one-cycle | Oscillating | Fast training, escaping minima | Bounds need tuning (LR range test) |
| ReduceLROnPlateau | Data-driven drops | When you cannot preset a timetable | Reacts late; needs a good metric |
Common pitfalls
- Decaying too early or too aggressively starves the model of progress before it has converged.
- Forgetting warmup with Adam/Transformers — early loss spikes or NaNs are a classic symptom.
- Cosine annealing assumes you know T in advance; if you train longer you must re-plan it.
- Setting η_min to 0 with cosine means the very last steps make almost no progress — sometimes a small floor is better.
- Schedules interact with batch size: larger batches usually want a larger peak rate (and often longer warmup).
Key Takeaways
- A fixed learning rate forces a bad compromise: large enough to progress early means too large to settle late.
- Schedules start high and anneal low, getting fast early progress and a precise final convergence.
- Step decay (η·γ^⌊t/s⌋) and exponential decay (η·e⁻ᵏᵗ) are simple monotonic schedules; cosine annealing is the modern default.
- Linear warmup ramps η from 0 to its peak over the first w steps and is essential for Transformers and Adam, stabilizing noisy early updates.
- Warmup + cosine is the standard large-model recipe; cyclical/one-cycle can train faster by oscillating the rate.
- ReduceLROnPlateau adapts to the validation curve instead of following a preset formula.