The learning rate is the single most important hyperparameter in training a transformer, and almost nobody holds it fixed. Instead they run a schedule: a rule that changes the learning rate at every step — usually ramping it up gently from near zero (warmup), holding a peak, then decaying it smoothly back down (cosine, linear, or inverse-square-root decay). That shape is not decoration. A constant rate that is safe early is far too timid late, and a rate that is efficient late will blow up an untrained network in the first hundred steps. This article builds the schedule from first principles: why constant fails, what warmup actually fixes in Adam’s early updates, the exact formulas for cosine and inverse-sqrt decay with worked numbers you can check by hand, how linear decay and constant-then-decay compare, and how the three knobs that matter — peak learning rate, warmup steps, and final learning rate — each change the run. Get the schedule wrong and no amount of data or compute saves you; get it right and the same model trains faster and lands lower.
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.
Warmup steps: how long to ramp
The number of warmup steps T_warmup controls how long you stay in the safe, low-rate regime before hitting full speed. The tension is simple. Too few warmup steps and you expose a barely-calibrated optimizer to large updates — the early-instability failure mode returns, sometimes as an immediate divergence, sometimes as a subtle early damage that shows up only as a worse final loss. Too many and you waste steps crawling at low learning rates when the optimizer was ready to move faster, slightly slowing the whole run.
In practice the safe direction to err is more warmup: the downside of over-warming is a mild efficiency loss, while the downside of under-warming can be a wrecked run. Typical values land anywhere from a few hundred to a few thousand steps, or expressed proportionally, the first 1–5% of total training. Larger and less stable configurations — big models, large batches, aggressive peak rates, low-precision training — generally want longer warmup, because they have more early instability to absorb. A pragmatic default many recipes reach for is a fixed count like 2,000 steps, or 2–4% of the schedule, then adjust upward if you see loss spikes shortly after the rate reaches its peak — the classic fingerprint of warmup that ended too soon.
The final learning rate, and why it rarely hits zero
The last knob is where the decay ends: the final learning rate η_final. It is tempting to decay all the way to zero — and cosine-to-zero is a perfectly common choice — but many large runs stop at a small nonzero floor instead, often around 0.1 × η_max. The reasoning is that a model near the end of training is still making useful, tiny refinements; cutting the rate to exactly zero freezes it in place, whereas a small floor lets it keep polishing the solution right up to the final step.
The floor also interacts with what happens after the planned run. If you might resume training, continue on more data, or fine-tune from the final checkpoint, a model whose learning rate collapsed to zero can be hard to get moving again, and its optimizer state reflects a fully-annealed regime. Leaving a small residual rate keeps the model in a state that resumes gracefully. The counter-consideration is that a nonzero floor means you never fully settle into the sharpest minimum, so for a run you are certain is final — no resumption, no continuation — decaying to (or very near) zero can shave a little extra loss. Like the other knobs, there is no universally correct answer; the useful instinct is that η_final trades final sharpness against resumability, and 0.1×η_max is a widely used compromise.
Practical implications and common pitfalls
A few failure modes recur often enough to name. Scheduling on the wrong unit: schedules are defined over optimizer steps, not epochs or raw samples — if you change batch size or gradient-accumulation and forget to recompute T_max and T_warmup, your carefully tuned curve silently deforms. Warmup that is too short shows up as a loss spike right after the rate peaks; lengthen it before you touch anything else. Mismatched T_max: because cosine and linear are defined relative to the final step, setting T_max to the wrong value — longer than you actually train — means you stop while the rate is still high and never get the benefit of the anneal; setting it too short decays to the floor early and then crawls.
For CPU and small-language-model work the schedule matters just as much, and arguably more, because compute is scarce and every wasted step is felt directly. Two habits pay off: prefer schedules that do not force an early commitment to the total budget — constant-then-decay or WSD — so you can train as long as your machine allows and anneal at the end; and always warm up, even on a tiny model, because Adam’s early-variance problem is about the optimizer, not the model size. Get the peak roughly right with a quick range test, use a generous warmup, pick a decay shape (cosine is a safe default), and decay to a small floor. That recipe is boring, and boring is exactly what you want from the most powerful hyperparameter you own.
v̂ is still noisy; decay lets the model settle into its minimum. Cosine holds high early and lands gently; inverse-sqrt (the original Transformer schedule) folds warmup and a slow one-over-root-t decay into one min formula; constant-then-decay and WSD let you postpone committing to a total step count. Three knobs steer it all: the peak sets the scale of every step and is the first thing to tune, the warmup length buys stability and should err long, and the final rate trades final sharpness against resumability. Get these right and the same model and data train faster and land lower.