The master dial: what the learning rate actually does
Every optimizer update has the same skeleton: take the current parameters, compute a direction from the gradient, and step along it by some distance. The learning rate η is that distance. For plain SGD the update is θ ← θ − η · ∇L; for Adam it is θ ← θ − η · m̂ / (√v̂ + ε), where m̂ and v̂ are bias-corrected running estimates of the gradient’s mean and variance. In both cases η is the scalar that multiplies the whole step.
Because it multiplies everything, the learning rate dominates. Too large and the step overshoots the local minimum — the loss oscillates or diverges to NaN. Too small and each step barely moves; you burn compute crawling toward a solution you could have reached in a fraction of the steps. The catch that motivates the entire field of scheduling is that the right distance is not the same at every point in training. Early on, the parameters are random and the loss surface near them is treacherous; late on, the parameters are close to a good basin and want fine, careful steps. A single fixed number cannot be right for both regimes, which is exactly why we let η vary with the step count.
Why a constant learning rate fails
Imagine you must pick one η and live with it for the whole run. You face a dilemma with no good answer. Pick a value large enough to make fast progress in the bulk of training, and the very first updates — applied to a freshly initialized network whose gradients are large, noisy, and poorly conditioned — take huge, wild steps. The loss spikes, activations saturate, and in mixed precision you often see an outright divergence in the first few hundred iterations.
Now pick a value small enough to survive that dangerous opening. It works, but for the remaining 99% of training you are moving at a crawl. You have paid for the worst-case first step with the efficiency of every step after it. There is a second, subtler problem: the loss landscape a model sees changes as it learns. Near the start, almost any downhill direction helps and big steps are cheaply productive; near the end, the model sits in a narrow, sharp basin where a step that is too long bounces off the far wall and never settles. The ideal learning rate is therefore large in the middle and small at both ends — a profile a constant simply cannot express. The schedule exists to trace that profile deliberately: ramp up to dodge the unstable opening, run hot through the productive middle, cool down to settle into the minimum.
The canonical shape: warmup, then decay
Nearly every modern transformer training run uses a two-phase learning-rate curve. Phase one is warmup: starting from (near) zero, the rate climbs linearly over the first T_warmup steps until it reaches the peak learning rate η_max. Phase two is decay: from the peak, the rate falls smoothly — along a cosine, a line, or an inverse-square-root curve — toward a small final learning rate η_final at the last step T_max.
Visually it is a shark-fin: a short steep rise on the left, then a long, graceful descent. Three numbers pin the whole thing down: the peak η_max (how hot the run gets), the warmup length T_warmup (how long the ramp lasts), and the floor η_final (where the descent ends). The decay function — cosine versus linear versus inverse-sqrt — sets the exact curvature between peak and floor, but the skeleton is always the same. The rest of this article takes the skeleton apart: first why the warmup ramp is not optional, then the precise math of each decay curve, then how each of the three knobs behaves when you turn it. Once you see the shape as ‘dodge the unstable start, then cool into the minimum,’ the specific formulas become interchangeable tools for the same job.
Warmup: easing the learning rate in
Warmup is the simplest part of the schedule to write down. For the first T_warmup steps you scale the peak rate linearly by how far into warmup you are:
For step t ≤ T_warmup (linear warmup):
η(t) = η_max · ( t / T_warmup )
t = 0 → η = 0
t = T_warmup/2 → η = η_max / 2
t = T_warmup → η = η_max (peak reached)That is it: a straight line from zero (or a tiny floor) up to the peak. Some implementations warm up over a fixed step count like 2,000; others express it as a fraction of the run, say the first 2–5% of steps. A few use a mild nonlinear ramp, but linear is overwhelmingly the default because it is boringly robust. The intuition is that of easing a cold engine up to speed rather than flooring the accelerator: you give the optimizer’s internal statistics and the network’s activations a few thousand gentle steps to reach a sane operating regime before you subject them to full-size updates. The next section explains what those internal statistics are and why they are specifically broken at step zero — the real reason warmup earns its place, especially for Adam-family optimizers on transformers.
Why warmup stabilizes early Adam steps
Adam does not step along the raw gradient. It divides the gradient’s running mean m̂ by the square root of its running second moment v̂, so the effective step is roughly η · m̂ / √v̂. Both averages are seeded at zero and updated as exponential moving averages with decay rates like β₁ = 0.9 and β₂ = 0.999. Bias correction rescales them, but early on v̂ is estimated from only a handful of gradients and is therefore a high-variance, unreliable estimate of the true second moment.
That unreliability is the problem. When v̂ happens to be too small for some parameter, the ratio m̂ / √v̂ explodes and that coordinate takes an enormous step; a few such steps early in training can knock a randomly initialized transformer into a bad region it never recovers from. LayerNorm and residual connections make transformers especially touchy here. Warmup is the fix: by keeping η tiny while v̂ is still noisy, you cap the size of these erratic early updates. As the moving averages accumulate more gradients they converge to stable estimates, the adaptive denominator becomes trustworthy, and precisely then the schedule ramps η up to full size. In short: warmup buys the optimizer time to calibrate its own variance estimate before you let it take full-length steps.
Cosine decay: the formula
After warmup, the most popular way to bring the learning rate down is cosine annealing. Over the decay window — from the end of warmup at T_warmup to the final step T_max — you define a decay progress p that runs from 0 to 1, and shape the rate with half a cosine wave:
For step t > T_warmup (cosine decay):
p = (t − T_warmup) / (T_max − T_warmup) p ∈ [0, 1]
η(t) = η_final + ½(η_max − η_final)·(1 + cos(π·p))
p = 0 → cos(0) = 1 → η = η_max (start of decay)
p = ½ → cos(π/2)= 0 → η = midpoint
p = 1 → cos(π) = −1 → η = η_final (end)The (1 + cos(π·p)) term slides smoothly from 2 down to 0, so the bracket scales the peak-minus-floor gap from full to nothing and adds back the floor. The shape is the reason it is loved: the descent starts almost flat (the rate lingers near the peak while the model is still learning fast), steepens through the middle, then flattens again as it eases into η_final — a gentle landing rather than an abrupt stop. If you set η_final = 0 the curve decays all the way to zero; more often teams stop at roughly 0.1 × η_max so the model keeps making small refinements to the very end.
Worked example: cosine learning rate at a few steps
Numbers make the curve concrete. Take a run with peak η_max = 1×10⁻³, floor η_final = 0, warmup T_warmup = 1000 steps, and total T_max = 10000 steps. The warmup phase is the straight line η(t) = 10⁻³ · t/1000; the decay phase uses p = (t − 1000)/9000 and η(t) = ½·10⁻³·(1 + cos(πp)).
t phase computation η(t)
--------------------------------------------------------------
0 warmup 1e-3 · 0/1000 0
500 warmup 1e-3 · 500/1000 5.00e-4
1000 peak 1e-3 · 1000/1000 1.00e-3
3250 decay p=.25, ½e-3(1+cos45°=.707) 8.54e-4
5500 decay p=.50, ½e-3(1+cos90°=0) 5.00e-4
7750 decay p=.75, ½e-3(1+cos135°=-.707) 1.46e-4
10000 end p=1, ½e-3(1+cos180°=-1) 0Read the descent: the rate is still 8.54×10⁻⁴ a quarter of the way through decay — only 15% below peak — but has fallen to 1.46×10⁻⁴ three-quarters through. That is the cosine signature: hold high early, drop fast in the middle, and glide into the floor. Halfway through decay you are at exactly the midpoint learning rate, which is a handy sanity check when you plot your own schedule.
Linear decay
Cosine is popular but not special; a straight line from peak to floor works well too and is even simpler to reason about. Linear decay uses the same decay progress p and interpolates directly:
For step t > T_warmup (linear decay):
p = (t − T_warmup) / (T_max − T_warmup)
η(t) = η_max − p·(η_max − η_final)
= (1−p)·η_max + p·η_finalThe whole combined schedule — linear up during warmup, linear down after — is the classic warmup-linear-decay (sometimes ‘triangular’) profile, and it is a strong, no-surprises baseline used in many well-known fine-tuning recipes. Compared with cosine, linear decay spends less time near the peak: at p = 0.25 a linear schedule is already down to 75% of peak, whereas the cosine is still at about 85%. Correspondingly, linear reaches low rates sooner in the middle of the run. In practice the two often finish within a whisker of each other on final loss; cosine’s slightly-longer time at high rate and gentler final approach give it a small, consistent edge on many large-scale pretraining runs, which is why it became the default there, while linear remains common and perfectly respectable for fine-tuning.
Inverse square-root: the original Transformer schedule
The schedule that shipped with the transformer in Attention Is All You Need is neither cosine nor linear — it is inverse square root, and it folds warmup and decay into a single closed-form expression:
Original Transformer schedule (Vaswani et al., 2017):
η(t) = d_model^(−0.5) · min( t^(−0.5), t · T_warmup^(−1.5) )
• while t < T_warmup : the 2nd term wins → η ∝ t (linear rise)
• while t > T_warmup : the 1st term wins → η ∝ 1/√t (sqrt decay)
• the two terms are equal exactly at t = T_warmup (the peak)The min is the trick. For small t the term t · T_warmup^(−1.5) is smaller and grows linearly, so the rate ramps up — that is warmup, for free, with no separate branch. Once t passes T_warmup the t^(−0.5) term becomes smaller and takes over, decaying the rate as one over the square root of the step. The peak sits exactly at t = T_warmup, where both terms equal T_warmup^(−0.5), giving a peak of η_max = d_model^(−0.5) · T_warmup^(−0.5) = (d_model · T_warmup)^(−0.5). Notice the peak is not a free knob here — it is determined by the model width and warmup length. That coupling to d_model is deliberate: wider models get proportionally smaller steps, a crude but effective bit of scale-invariance baked into the formula.
Worked example: inverse-sqrt at a few steps
Use the paper’s own settings: d_model = 512 and T_warmup = 4000. Then d_model^(−0.5) = 1/√512 ≈ 0.04419 and T_warmup^(−1.5) = 4000^(−1.5) ≈ 3.953×10⁻⁶. Plug a few steps through the min:
t t^(-0.5) t·T_w^(-1.5) min η(t)=0.04419·min
---------------------------------------------------------------
100 0.10000 3.95e-4 3.95e-4 1.75e-5 (warmup)
1000 0.03162 3.95e-3 3.95e-3 1.75e-4 (warmup)
4000 0.01581 0.01581 0.01581 6.99e-4 (PEAK)
16000 0.00791 0.06324 0.00791 3.49e-4 (decay)The story the table tells: during warmup the rate climbs linearly — it is exactly ten times larger at t = 1000 than at t = 100, because the active term is proportional to t. It peaks at t = 4000 at about 7.0×10⁻⁴, which matches (512 · 4000)^(−0.5) exactly. After the peak the square-root decay is slow: quadrupling the step from 4000 to 16000 only halves the rate (since √4 = 2). That long, heavy tail — the rate never rushing to zero — is the defining behavior of inverse-sqrt and the main way it differs from cosine, which drives firmly toward its floor.
Constant-then-decay and warmup-stable-decay
Not every schedule is monotonic after the peak. A useful family holds the learning rate constant at the peak for a long stretch and only decays at the end. The simplest version — warmup, then a flat plateau, then a short decay — has a practical superpower: because the middle is constant, you do not have to commit to a total step count T_max in advance. Cosine and linear both need to know where the finish line is (the whole curve is defined relative to T_max), so changing your mind about how long to train means recomputing the schedule.
This is the idea behind warmup-stable-decay (WSD), which has become popular for training small language models on a compute budget you might want to extend. You warm up, hold a stable high rate for as long as you like (training a checkpoint you can keep going from), and then apply a rapid decay — often over just the final 10–20% of steps — to ‘anneal’ into a strong final model. The striking empirical finding is that the loss drops sharply during that short final decay, as if the constant phase had been accumulating potential the anneal then cashes in. For CPU-trained or budget SLM work, WSD is attractive precisely because it decouples ‘how long can I afford to train’ from ‘when must the schedule end.’
Peak learning rate: the knob that matters most
Of the three numbers, the peak learning rate η_max is the one to get right first, because it sets the scale of every step in the productive middle of training. Too high and even with warmup the run becomes unstable once it reaches the peak — loss spikes, gradient norms blow up, and you may diverge halfway through. Too low and the model trains stably but slowly, leaving loss on the table for the same compute.
There is no universal value; it depends on the optimizer, batch size, model size, and parameterization. For AdamW on transformers, peaks in the rough band of 1×10⁻⁴ to 3×10⁻³ are common, with larger models generally wanting smaller peaks. A robust way to find one is a short learning-rate range test: train briefly while sweeping η upward and watch where the loss stops improving and starts to climb — the peak lives just below that edge. Two couplings are worth remembering. First, peak learning rate and batch size move together: bigger batches give lower-variance gradients that tolerate (and often need) a larger peak, following square-root or linear scaling rules. Second, peak and warmup interact — a more aggressive peak generally wants a longer warmup to reach it safely. Tune the peak first; the other knobs are refinements around it.