Learning Rate Scheduling and Warmup Strategies

In June 2017, a team at Facebook AI Research led by Priya Goyal, with Kaiming He among the co-authors, published a result that looked almost like a trick: they trained ResNet-50 on the full ImageNet dataset in one hour, using 256 GPUs, without losing accuracy compared to the standard multi-day recipe. The trick was not more compute alone. The standard recipe trains with a mini-batch of 256 images and a base learning rate of 0.1. To keep 256 GPUs busy, Goyal's team used a mini-batch of 8192 — 32 times larger. Their "linear scaling rule" says that when you multiply the batch size by k, you should multiply the learning rate by k as well, to keep the expected size of each weight update roughly constant. That gives a target learning rate of 0.1 × (8192 / 256) = 0.1 × 32 = 3.2.

Set the learning rate to 3.2 from step zero, on a network whose weights are still random, and training diverges within the first few hundred iterations — the loss spikes to NaN before the model has learned anything. Goyal's fix was a five-epoch warmup: start the learning rate near zero and ramp it linearly up to 3.2 over those first epochs, only then switching to the normal decay schedule. That one addition is what made the 32x batch scale-up trainable at all. This chapter is about the two ideas that made it work: how a learning rate should change shape over the course of training (scheduling), and why it often needs to start low and climb before it does anything else (warmup).

Why a constant learning rate fails

Gradient descent updates a weight vector θ by θ ← θ - η · ∇L(θ), where η is the learning rate. A single fixed η for the entire run has to serve two incompatible phases of training.

Early in training, the loss surface around a randomly initialized network is comparatively coarse: large gradients, large useful steps are available, and a bigger η lets the model cover ground fast. Late in training, the parameters sit near a minimum (or a flat, wide basin around one), and the useful step size shrinks. A learning rate that was well-tuned for the start of training is now too large: instead of settling into the minimum, the optimizer overshoots it every step, and the loss oscillates or even climbs back up. Conversely, a learning rate tuned for careful late-stage convergence is far too small at the start — training crawls, and on a compute budget of a fixed number of GPU-hours, most of that budget is wasted taking tiny steps across a landscape that would have tolerated large ones.

A schedule resolves the conflict by making η a function of the training step, η(t), large during the phase that benefits from large steps and small during the phase that needs precision. Warmup adds a third, short phase before either of those: a controlled ramp-up at the very start, for a reason that turns out to be specific to how modern optimizers estimate their own step size.

Why warmup: what actually goes wrong at step one

Almost all Grade-12-relevant training (transformers, CNNs at scale) uses an adaptive optimizer — Adam or one of its variants — rather than plain SGD. Adam keeps a running estimate of the second moment (uncentered variance) of each gradient, v_t = β2 v_{t-1} + (1-β2)g_t², and divides the update by √v_t. This is what makes Adam "adaptive": parameters with historically large gradients get smaller effective steps, and vice versa.

The problem is that at step 1, v_t has been estimated from a single gradient sample. It is not yet a reliable estimate of anything — it has enormous variance. Liu et al. (Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han, "On the Variance of the Adaptive Learning Rate and Beyond," ICLR 2020) showed this rigorously: they derived the variance of Adam's adaptive learning rate as a function of the number of samples seen so far, and demonstrated that this variance is largest in exactly the first few hundred steps, then falls off as more gradient samples accumulate. A large, unreliable adaptive step size, applied to a randomly initialized network with no prior training signal to correct it, is what produces the divergence or the "bad local minimum" that unwarmed-up Adam training is notorious for. Their paper's headline contribution, Rectified Adam (RAdam), computes this variance analytically at each step and rescales the update automatically — it is, in their own framing, a principled substitute for the empirical trick of warmup.

Correcting a common misconception

A plausible-sounding but imprecise explanation students often reach for is: "we warm up because the loss landscape is very steep or chaotic right at random initialization, so small steps avoid overshooting." This is not quite wrong as an intuition for plain SGD, but it is not the mechanism that Liu et al.'s research identifies as the dominant cause for adaptive optimizers, and it does not explain several things the steepness story alone cannot: why the problem is specific to Adam-family optimizers and largely absent for plain SGD; why the fix that works is a ramp measured in optimizer steps taken, not a fixed number tied to landscape geometry; and why RAdam — which does nothing about landscape steepness, only about the variance of the second-moment estimate — removes the need for warmup almost entirely in their experiments. The precise mechanism is statistical, not geometric: the adaptive learning rate itself is a noisy estimate early on, because it is built from too few samples, and warmup's real job is to keep the step size small until that internal estimate has enough data to be trustworthy.

Schedule families

Once past warmup, several standard shapes for η(t) are in wide use. Each has a defensible reason to exist, and each shows up in a real published training recipe.

Step decay is the oldest and crudest: hold η flat, then chop it by 10x at fixed milestones. It works but leaves a visible "staircase" in the loss curve — a sudden drop right after each cut, because the optimizer had been oscillating around the minimum at the old, too-large step size and only settles once the step shrinks. Cosine annealing decays smoothly and continuously to (near) zero following one arch of a cosine curve, which avoids the staircase and, empirically, tends to reach a slightly better final loss for the same total step budget — this is why it is the default in most modern large-model pretraining. The Transformer's own inverse square-root schedule, examined numerically below, is a warmup-plus-decay schedule folded into a single closed-form expression.

Worked example 1: the Transformer's inverse square-root schedule

Vaswani et al. (2017) define the learning rate at optimizer step t as

lr(t) = d_model^-0.5 · min( t^-0.5 , t · warmup_steps^-1.5 )

with d_model = 512 and warmup_steps = 4000 in the base model. Two branches compete inside the min: a linearly rising term t · warmup_steps^-1.5, and a decaying term t^-0.5. At t = warmup_steps the two branches are exactly equal (substituting t = w into the linear branch gives w · w^-1.5 = w^-0.5, identical to the decay branch), so this is the schedule's peak: before it, the smaller, rising linear term is selected; after it, the smaller, falling t^-0.5 term is selected. The formula is warmup and decay expressed as a single continuous function of t, with no explicit "if" branch needed to switch phases.

Computing the constant: 512^-0.5 = 1/√512 = 1/22.6274 = 0.044194.

def transformer_lr(step, d_model=512, warmup_steps=4000):
    step = max(step, 1)
    return d_model ** -0.5 * min(step ** -0.5, step * warmup_steps ** -1.5)

for step in (1000, 4000, 8000, 16000):
    print(step, round(transformer_lr(step), 6))

Tracing each call by hand: at step=1000 (inside warmup, since 1000 < 4000), the linear branch is smaller: 4000^-1.5 = 1/(4000·√4000) = 1/252982.2 = 3.9528×10^-6, so 1000 · 3.9528×10^-6 = 0.0039528, versus 1000^-0.5 = 0.031623 for the other branch — the minimum is 0.0039528. Multiplying by 0.044194 gives 0.0001747, which rounds to 0.000175. At step=4000, the peak, both branches equal 4000^-0.5 = 0.015811, giving 0.044194 × 0.015811 = 0.0006988, rounding to 0.000699. At step=8000 (past the peak, decay branch smaller): 8000^-0.5 = 1/√8000 = 0.011180, giving 0.044194 × 0.011180 = 0.0004941, rounding to 0.000494. At step=16000: 16000^-0.5 = 1/126.491 = 0.0079057, giving 0.044194 × 0.0079057 = 0.0003494, rounding to 0.000349. The printed output is therefore:

1000 0.000175
4000 0.000699
8000 0.000494
16000 0.000349

Notice the shape: rising from step 1000 to step 4000, then falling at every later step, exactly the warmup-then-decay behavior described above, produced by one formula with no explicit phase switch.

Worked example 2: linear warmup with cosine decay to zero

The schedule used by most current large-model pretraining runs is simpler to reason about directly: ramp the learning rate linearly from 0 up to a peak over warmup_steps, then decay it along a cosine curve from that peak down to 0 by total_steps.

import math

def lr_lambda(step, warmup_steps=1000, total_steps=10000):
    if step < warmup_steps:
        return step / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    return 0.5 * (1 + math.cos(math.pi * progress))

for step in (500, 1000, 5500, 9100, 10000):
    print(step, round(lr_lambda(step), 6))

This function returns a multiplier on the peak learning rate, not the learning rate itself, so it can be reused with any peak_lr by multiplying afterward. Tracing it with warmup_steps=1000, total_steps=10000: at step=500, still inside warmup, the function returns 500/1000 = 0.5. At step=1000, the warmup condition 1000 < 1000 is false, so control falls to the cosine branch with progress = (1000-1000)/9000 = 0, and 0.5·(1+cos(0)) = 0.5·2 = 1.0 — the two branches meet exactly at the peak, so the curve is continuous there. At step=5500, the midpoint of the decay window, progress = 4500/9000 = 0.5, and cos(0.5π) = 0, giving 0.5·(1+0) = 0.5. At step=9100, progress = 8100/9000 = 0.9, and cos(0.9π) = cos(162°) = -cos(18°) = -0.951057, giving 0.5·(1-0.951057) = 0.024472. At step=10000, progress = 1, cos(π) = -1, giving 0.5·(1-1) = 0.0. Printed output:

500 0.5
1000 1.0
5500 0.5
9100 0.024472
10000 0.0

With a peak learning rate of 6×10^-4, these multipliers correspond to actual learning rates of 3×10^-4 at step 500, 6×10^-4 at the step-1000 peak, back down to 3×10^-4 at the step-5500 midpoint, roughly 1.47×10^-5 near the end at step 9100, and exactly 0 at step 10000. This is exactly the schedule diagrammed below.

Diagram: the combined schedule

The blue line is the exact curve traced by worked example 2: a straight rise across the shaded warmup region (steps 0–1000) to the peak, then a smooth cosine fall across the shaded decay region (steps 1000–10000), passing through the midpoint at step 5500 exactly halfway down, and reaching zero at step 10000. The dashed vertical line marks the warmup boundary, the single value where the two schedule phases must agree so the curve has no discontinuity.

Choosing warmup length and interaction with batch size

Warmup length is not a free parameter to guess at. The Transformer paper's warmup_steps=4000 was tuned for its specific batch size and model width; Goyal et al.'s 5-epoch warmup was tuned for a specific 32x batch-size jump. The general pattern behind both: the larger the effective batch size, the more steps of warmup are needed, because a larger batch means each optimizer step is a larger, more confident jump — and the variance argument from Liu et al. says the adaptive learning-rate estimate needs enough steps, not enough wall-clock time, to stabilize. Doubling the batch size while holding warmup_steps fixed effectively halves the number of independent gradient samples seen by the same optimizer step, so the standard practice is to scale warmup length together with the linear-scaling-rule learning rate increase, not leave it fixed. A second practical rule: cosine decay's minimum, ηmin, is rarely set to exactly zero in very long runs — a small floor (often around 10% of peak) is kept so the last steps of training still make some progress rather than freezing entirely, though the pure-to-zero form used in example 2 above is common for shorter runs and is what most textbook implementations show first.

Active recall

Attempt each question before reading its answer.

1. Why does a single fixed learning rate for an entire training run create two distinct, opposite failure modes? Name both.
2. Using transformer_lr(step, d_model=512, warmup_steps=4000) from worked example 1, compute the learning rate at step=8000 by hand, showing which branch of the min is selected and why.
3. A classmate says: "we use warmup because the loss landscape is very steep near random initialization, so small steps avoid overshooting." What does this explanation fail to account for, and what does Liu et al.'s (2020) RAdam research identify as the more precise mechanism?
4. You are scaling ResNet-50 training from batch size 256 (base learning rate 0.1) to batch size 4096, using the linear scaling rule from Goyal et al. (2017). What is the new peak learning rate, and why is it unsafe to start training at that value from step 1 without a warmup ramp?
5. In worked example 2's lr_lambda function, what multiplier and what actual learning rate (given peak_lr = 6×10^-4) does step=9100 return, and why is that value so close to (but not exactly) zero?
6. Suppose you change warmup_steps from 1000 to 2000 in worked example 2, keeping total_steps=10000 and peak_lr=6×10^-4 unchanged. Recompute the learning rate at step=1000 and at step=5500, and state what happens to the length of the decay phase.

Answers

1. Early in training, the loss surface tolerates and benefits from large steps — a fixed rate tuned for this phase is too small later, wasting compute on tiny, over-cautious updates near the minimum where curvature is sharper. Late in training, the parameters are near a minimum where large steps overshoot it every iteration, causing the loss to oscillate or plateau instead of converging — a fixed rate tuned for this phase would have made early training glacially slow. A single constant cannot serve both regimes.

2. Since step=8000 > warmup_steps=4000, the decay branch is smaller. Check: 8000^-0.5 = 1/√8000 = 0.011180; the linear branch would give 8000 × 4000^-1.5 = 8000 × 3.9528×10^-6 = 0.031623, clearly larger, so min selects 0.011180. Multiplying by the constant 512^-0.5 = 0.044194: 0.044194 × 0.011180 = 0.0004941, rounding to 0.000494 — matching the code trace shown earlier.

3. The steepness story does not explain why the effect is specific to adaptive optimizers (Adam-family) and largely absent for plain SGD, why the useful warmup length is measured in optimizer steps rather than a landscape property, or why Rectified Adam removes the need for manual warmup by correcting the variance of the second-moment estimate alone, without touching landscape geometry at all. Liu et al. (2020) show the real cause is statistical: Adam's adaptive step size is computed from a running estimate of gradient variance (v_t) that is built from very few samples in the first steps, so it is itself a high-variance, unreliable number early on. Warmup works by keeping the actual step size small until that internal estimate has accumulated enough samples to be trustworthy — not by working around a steep landscape.

4. New batch size / old batch size = 4096 / 256 = 16. Linear scaling rule: new LR = 0.1 × 16 = 1.6. Starting at 1.6 immediately is unsafe because the network's weights are still random and its Adam-style (or momentum) second-moment / running statistics have not accumulated any samples yet; a step of that size applied to an untrained, high-variance gradient estimate produces an update large enough to push weights into a region with exploding activations or gradients, typically manifesting as the loss jumping to NaN within the first few hundred iterations — exactly the failure Goyal et al.'s 5-epoch warmup was built to avoid.

5. lr_lambda(9100) with warmup_steps=1000, total_steps=10000: progress = (9100-1000)/9000 = 0.9, multiplier = 0.5·(1+cos(162°)) = 0.5·(1-0.951057) = 0.024472. Actual LR = 6×10^-4 × 0.024472 ≈ 1.468×10^-5. It is close to zero because step 9100 is 90% of the way through the decay window, and the cosine shape spends most of its second half falling steeply toward zero, only truly reaching it at the very last step (step=10000, progress=1, cos(180°)=-1, multiplier exactly 0).

6. With warmup_steps=2000, step=1000 is now inside warmup (1000 < 2000), not at the old peak: multiplier = 1000/2000 = 0.5, giving LR = 6×10^-4 × 0.5 = 3×10^-4 — half of its old value of 6×10^-4, because step 1000 is no longer the peak but the ramp's midpoint. At step=5500, now in the decay branch: progress = (5500-2000)/(10000-2000) = 3500/8000 = 0.4375; cos(0.4375π) = cos(78.75°) = 0.19509; multiplier = 0.5·(1+0.19509) = 0.59755; LR = 6×10^-4 × 0.59755 ≈ 3.585×10^-4 — slightly higher than the old value of 3×10^-4, not lower, because the decay window itself has changed. The full ripple: the peak shifts later (from step 1000 to step 2000), and the decay window shrinks from 9000 steps (1000–10000) to 8000 steps (2000–10000) — the same cosine shape is now compressed into fewer absolute steps, so the schedule descends more steeply per step during decay even though the fraction-of-window value at any given progress is identical to before.

Think About It

Think about this: How would you explain learning rate scheduling and warmup strategies to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind learning rate scheduling and warmup strategies, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.