Learning Rate Scheduling

Learning Objectives

By the end of this section, you will be able to:

1. Explain why a constant learning rate fails in noisy optimization and why decaying the LR over time enables tight convergence
2. Derive the cosine annealing formula and explain how a half-cosine curve maps training steps to learning rates
3. Explain the warmup problem and why starting with a small LR prevents early training instability in transformers
4. Implement step decay, cosine annealing, and warmup+cosine schedules from scratch in NumPy
5. Use PyTorch's lr_scheduler classes (StepLR, CosineAnnealingLR) to train a network with scheduled learning rates
6. Describe the modern transformer training recipe: warmup + cosine decay with AdamW, and why each piece is necessary

Where We Left Off

In Section 1, we learned that SGD with momentum accelerates convergence by accumulating past gradients into a velocity buffer. The velocity $\mathbf{v}_t = \beta \mathbf{v}_{t-1} + \nabla L$ smooths out noise and amplifies consistent directions, giving roughly a 3–5× speedup over vanilla SGD.

But we used the same learning rate $\eta = 0.05$ for all 30 training epochs. That raises a question: should the learning rate stay constant throughout training?

The answer is almost always no. A constant learning rate forces a painful tradeoff. Set it high enough to make fast progress early, and the optimizer oscillates wildly near the minimum. Set it low enough to converge precisely, and early training crawls. Learning rate scheduling resolves this tension by changing the learning rate during training — large steps in the beginning, small steps at the end.

Why Schedules Matter

Consider what happens when you train with SGD on noisy mini-batch gradients. At any step, the observed gradient is:

$\hat{g}_t = \nabla L(\mathbf{w}_t) + \varepsilon_t$ where $\varepsilon_t \sim \mathcal{N}(0, \sigma^2)$

The weight update is $\mathbf{w}_{t+1} = \mathbf{w}_t - \eta \hat{g}_t$. The noise in the gradient introduces a random walk component with variance proportional to $\eta^2 \sigma^2$. This means:

• With a large $\eta$: the random walk amplitude is large. The optimizer bounces around the minimum with big jumps. The loss oscillates even after many steps.
• With a small $\eta$: the random walk amplitude is small, so convergence is tight. But the signal component $\eta \nabla L$ is also tiny, so early progress is glacial.

The resolution is straightforward: start with a large $\eta$ for fast exploration, then decrease it for precise convergence. This is learning rate scheduling.

The oscillation bound: for SGD with constant learning rate $\eta$, the expected distance from the minimum is bounded by $\mathcal{O}(\eta \sigma^2)$. To halve the oscillation, you must halve the learning rate — but halving the LR also halves the convergence speed. Scheduling lets you escape this tradeoff by changing $\eta$ over time.

The Big Picture: Exploration vs Exploitation

There is a deeper reason why scheduling works so well. Think of training as having two phases:

Exploration (high LR): the loss landscape has many regions — some with good minima, some with bad ones. A high learning rate lets the optimizer jump over small bad minima and explore broadly. The noise from mini-batches is actually helpful here: it prevents the optimizer from committing too early to a suboptimal minimum.

Exploitation (low LR): once the optimizer has found a promising region, a low learning rate lets it settle into the best point within that region. The noise is now harmful — it prevents precise convergence. A small LR suppresses the noise and allows fine-tuning.

Every learning rate schedule is a particular strategy for transitioning from exploration to exploitation. The question is how to make that transition.

Step Decay

The simplest schedule: multiply the learning rate by a factor $\gamma < 1$ every $N$ epochs.

$\eta_t = \eta_0 \cdot \gamma^{\lfloor t / N \rfloor}$

For example, with $\eta_0 = 0.05$, $\gamma = 0.5$, and $N = 10$:

Step decay was the dominant schedule in early deep learning. The original ResNet paper (He et al., 2015) trained for 90 epochs, dividing the LR by 10 at epochs 30 and 60. It works well but requires choosing two hyperparameters ($\gamma$ and $N$) and introduces abrupt transitions that can destabilize training.

The Abrupt Transition Problem

When the LR drops suddenly, the optimizer's momentum buffer is still calibrated for the old LR. The velocity built up over 10 epochs at LR=0.05 is suddenly too large for LR=0.025. This causes a brief spike in loss right after each step. The model recovers quickly, but the transition is wasteful.

Cosine Annealing

Cosine annealing (Loshchilov & Hutter, 2016) replaces the staircase with a smooth curve. The learning rate follows a half-period of a cosine function:

$\eta_t = \eta_{\min} + \frac{1}{2}(\eta_0 - \eta_{\min})\left(1 + \cos\left(\frac{\pi t}{T}\right)\right)$

Let's unpack this formula piece by piece:

1. $\cos(\pi t / T)$ traces a half-cosine from $+1$ (at $t=0$) to $-1$ (at $t=T$)
2. Adding 1 shifts the range to $[0, 2]$
3. Multiplying by $\frac{1}{2}(\eta_0 - \eta_{\min})$ scales to $[0, \eta_0 - \eta_{\min}]$
4. Adding $\eta_{\min}$ shifts to $[\eta_{\min}, \eta_0]$

The result is a smooth S-shaped decay: slow at the start, fast in the middle, slow again near the end.

Why Cosine Beats Step Decay

• No abrupt transitions: the LR changes continuously, so the momentum buffer stays calibrated. No loss spikes.
• Fewer hyperparameters: just $\eta_0$, $\eta_{\min}$, and $T$ (total steps). No need to choose decay milestones.
• Slow decay near start and end: the cosine curve spends more time near $\eta_0$ (exploring) and near $\eta_{\min}$ (fine-tuning). The fast decay in the middle wastes the least time in the intermediate region.

Warmup: Stabilizing Early Training

All the schedules above assume that training starts at the peak LR. For simple models and SGD, this works fine. But for transformers and other large models with adaptive optimizers like Adam, starting at the full LR is dangerous.

Why Large Models Need Warmup

Adam (and similar adaptive optimizers) maintains a running estimate of each gradient's variance. In the first few steps, these estimates are based on very few samples and can be wildly inaccurate. If the learning rate is also high during these early steps, the combination of bad variance estimates + high LR produces huge, unstable weight updates that can derail training before it starts.

The solution: start with a tiny learning rate and linearly ramp up to the target LR over the first $W$ steps:

$\eta_t = \eta_0 \cdot \frac{t + 1}{W} \quad \text{for } t < W$

During warmup, the optimizer builds up accurate gradient statistics at a safe pace. By the time the LR reaches its peak, the variance estimates are reliable and training proceeds normally.

The Modern Standard: Warmup + Cosine Decay

The standard training recipe for transformers and most modern deep learning combines warmup with cosine decay:

$\eta_t = \begin{cases} \eta_0 \cdot \dfrac{t + 1}{W} & \text{if } t < W \\[8pt] \eta_{\min} + \dfrac{1}{2}(\eta_0 - \eta_{\min})\!\left(1 + \cos\!\left(\dfrac{\pi(t - W)}{T - W}\right)\right) & \text{if } t \geq W \end{cases}$

The two phases:

1. Phase 1 (Warmup): linear ramp from $\eta_0 / W$ to $\eta_0$ over steps $0$ to $W-1$
2. Phase 2 (Cosine Decay): half-cosine from $\eta_0$ to $\eta_{\min}$ over steps $W$ to $T-1$

This gives the optimizer the best of everything: safe startup (warmup prevents early instability), fast progress (high LR after warmup), and precise convergence (cosine decay settles into the minimum).

The recipe in one line: AdamW with warmup + cosine decay is to modern deep learning what SGD with momentum was to pre-2017 training. If you are starting a new project and need a default optimizer setup, this is it.

Interactive: LR Schedule Explorer

The visualization below lets you explore different learning rate schedules. Switch between LR Curves (the schedule shapes) and Training Loss (how they affect optimization of $f(w) = w^2$ with noisy gradients). Use the sliders to adjust parameters and see how each schedule responds.

Try these experiments:

1. LR Curves, default settings: press Play. Watch how cosine (blue) curves smoothly while step decay (purple) drops abruptly. Warmup+cosine (green) ramps up before decaying.
2. Increase warmup steps to 15: the green curve ramps for longer before peaking. With many warmup steps, the available decay time shrinks.
3. Switch to Training Loss mode: watch the constant LR (orange) oscillate wildly near the end while cosine and warmup+cosine settle down.
4. Set noise (σ) to 0: without noise, all schedules converge smoothly. Scheduling matters most when gradients are noisy (which they always are in practice).
5. Set noise (σ) to 5: extreme noise. The constant LR bounces violently. Scheduled LRs still converge because their tiny late-stage LR dampens the noise.

NumPy: Schedules from Scratch

Let's implement cosine annealing and warmup+cosine from scratch. We minimize $f(w) = w^2$ with noisy gradients ($\sigma = 2.0$) over 50 steps, comparing constant LR, cosine annealing, and warmup+cosine. The noise simulates mini-batch SGD: each gradient is the true gradient plus Gaussian noise.

1import numpy as np
2
3# ── Minimize f(w) = w² with noisy gradients ──
4# True gradient: df/dw = 2w.  We add noise to simulate
5# mini-batch SGD where each batch gives a noisy estimate.
6
7np.random.seed(42)
8eta_0 = 0.08       # peak learning rate
9eta_min = 0.001    # minimum learning rate
10T = 50             # total training steps
11sigma = 2.0        # gradient noise std dev
12
13# === Schedule 1: Cosine Annealing ===
14def cosine_lr(step):
15    return eta_min + 0.5 * (eta_0 - eta_min) * (
16        1 + np.cos(np.pi * step / T))
17
18# === Schedule 2: Linear Warmup + Cosine Decay ===
19def warmup_cosine_lr(step, warmup=5):
20    if step < warmup:
21        return eta_0 * (step + 1) / warmup
22    progress = (step - warmup) / (T - warmup)
23    return eta_min + 0.5 * (eta_0 - eta_min) * (
24        1 + np.cos(np.pi * progress))
25
26# === Compare three schedules ===
27for name, get_lr in [("Constant",      lambda s: eta_0),
28                     ("Cosine",         cosine_lr),
29                     ("Warmup+Cosine",  warmup_cosine_lr)]:
30    np.random.seed(42)
31    w = 5.0
32    for step in range(T):
33        noise = np.random.normal(0, sigma)
34        grad = 2 * w + noise
35        lr = get_lr(step)
36        w = w - lr * grad
37        if step % 10 == 0 or step == T - 1:
38            print(f"[{name:14s}] Step {step:2d}: "
39                  f"lr={lr:.5f}  w={w:7.4f}  loss={w**2:.4f}")
40    print()

The results reveal the key insight:

All three reach similar neighborhoods, but constant LR oscillates with a standard deviation of 0.118 while the scheduled approaches have std < 0.005 — a 25× reduction in oscillation. The schedules "park" the optimizer at a good point and stop bouncing.

PyTorch: Training with Schedulers

PyTorch provides built-in schedulers in $\texttt{torch.optim.lr\_scheduler}$. Let's train our FlipNet from Chapters 7–8 with three approaches: constant LR, StepLR, and CosineAnnealingLR. All three models start with identical weights. The only difference is how the learning rate evolves.

1import torch
2import torch.nn as nn
3
4torch.manual_seed(42)
5
6# ── Same FlipNet from Chapters 7-8 ──
7class FlipNet(nn.Module):
8    def __init__(self):
9        super().__init__()
10        self.layer1 = nn.Linear(4, 3)
11        self.layer2 = nn.Linear(3, 4)
12        self.relu = nn.ReLU()
13
14    def forward(self, x):
15        return self.layer2(self.relu(self.layer1(x)))
16
17# ── Dataset: all 16 binary 2×2 images ──
18X = torch.tensor([[int(b) for b in f"{i:04b}"]
19                   for i in range(16)], dtype=torch.float32)
20Y = X[:, [0, 2, 1, 3]]  # diagonal flip
21
22def train_epoch(model, optimizer):
23    total_loss = 0.0
24    for i in range(16):
25        pred = model(X[i])
26        loss = torch.mean((pred - Y[i]) ** 2)
27        optimizer.zero_grad()
28        loss.backward()
29        optimizer.step()
30        total_loss += loss.item()
31    return total_loss / 16
32
33# ── Model 1: Constant LR (no scheduler) ──
34torch.manual_seed(42)
35m_const = FlipNet()
36opt_const = torch.optim.SGD(m_const.parameters(),
37                            lr=0.05, momentum=0.9)
38
39# ── Model 2: StepLR (halve every 10 epochs) ──
40torch.manual_seed(42)
41m_step = FlipNet()
42opt_step = torch.optim.SGD(m_step.parameters(),
43                           lr=0.05, momentum=0.9)
44sched_step = torch.optim.lr_scheduler.StepLR(
45    opt_step, step_size=10, gamma=0.5)
46
47# ── Model 3: CosineAnnealingLR ──
48torch.manual_seed(42)
49m_cos = FlipNet()
50opt_cos = torch.optim.SGD(m_cos.parameters(),
51                          lr=0.05, momentum=0.9)
52sched_cos = torch.optim.lr_scheduler.CosineAnnealingLR(
53    opt_cos, T_max=30, eta_min=0.001)
54
55for epoch in range(30):
56    l1 = train_epoch(m_const, opt_const)
57    l2 = train_epoch(m_step, opt_step)
58    sched_step.step()
59    l3 = train_epoch(m_cos, opt_cos)
60    sched_cos.step()
61    if epoch % 5 == 0 or epoch == 29:
62        lr_s = opt_step.param_groups[0]['lr']
63        lr_c = opt_cos.param_groups[0]['lr']
64        print(f"Epoch {epoch:2d}:  Constant={l1:.4f}  "
65              f"Step(lr={lr_s:.4f})={l2:.4f}  "
66              f"Cosine(lr={lr_c:.5f})={l3:.4f}")

The final epoch losses tell the story:

Notice two things. First, the constant LR model has higher loss at epoch 29 (0.152) than at epoch 20 (0.135) — it is oscillating around the minimum. Second, cosine annealing achieves the lowest final loss (0.125) with the smoothest convergence trajectory.

Connection to Modern Training

Learning rate scheduling is not just a nice-to-have — it is a critical component of every modern training pipeline. Here is how it connects to the systems that power today's AI.

The Transformer Training Recipe

Nearly every large language model follows the same recipe:

1. Optimizer: AdamW (Adam with decoupled weight decay)
2. Schedule: Linear warmup + cosine decay to $\eta_{\min} = 0.1 \cdot \eta_0$
3. Warmup: 0.1–2% of total training steps
4. Weight decay: 0.1 (regularization, independent of LR)
5. Gradient clipping: max norm = 1.0 (prevents exploding gradients)

This recipe was established by BERT (2018), refined by GPT-2 (2019), and has been remarkably stable since. LLaMA-2 (2023) and Mistral (2023) use essentially the same setup.

Why Warmup Is Essential for Transformers

Transformers have a specific instability that makes warmup non-optional. The attention mechanism computes $\text{softmax}(QK^T / \sqrt{d_k})$. In the first few steps, the $Q$ and $K$ matrices are randomly initialized, so the attention weights are nearly uniform (each token attends equally to all others). The gradients from these uniform attention patterns are large and noisy.

If the learning rate is also large at this point, the weight updates are so big that the attention weights can become sharply peaked (one token dominates) before the model has learned which tokens should actually attend to which. This creates a self-reinforcing cycle where bad attention patterns produce bad gradients that reinforce the bad patterns. Warmup breaks this cycle by keeping updates small until the attention weights have time to organize.

Connection to KV-Cache and Inference

The learning rate schedule affects not just training but also the quality of the trained model. Models trained with proper scheduling (warmup + cosine decay) produce more stable attention patterns, which has two important implications for inference:

• KV-Cache efficiency: models with well-formed attention patterns have more predictable memory access patterns during autoregressive generation, making KV-cache utilization more efficient
• Quantization robustness: models that converged smoothly (via scheduling) have weight distributions with fewer outliers, making them more amenable to post-training quantization (INT8, INT4) without quality loss

Flash Attention and Training Speed

Flash Attention (Dao et al., 2022) accelerates the attention computation, reducing the wall-clock time per training step. This means more steps per hour, which makes the total number of training steps $T$ larger for a given time budget. A larger $T$ means the cosine decay stretches over more steps — slower, more gradual decay — which generally improves final model quality. In this sense, Flash Attention indirectly improves model quality by enabling longer, more gradual schedules.

Scaling Laws and Schedule Design

The Chinchilla scaling laws (Hoffmann et al., 2022) established that the optimal balance of model size and training data depends on the total compute budget. This directly affects scheduling: if you know you will train for $T$ steps, you set $T_{\max} = T$ in the cosine schedule. Training for more steps with a proportionally slower decay consistently improves performance — one reason why organizations invest in longer training runs rather than just larger models.

Summary

1. Constant learning rates force a tradeoff: high LR for fast progress causes oscillation near the minimum; low LR for precise convergence makes early training slow. Scheduling resolves this by decaying the LR over time.
2. Step decay ($\eta_t = \eta_0 \gamma^{\lfloor t/N \rfloor}$) multiplies the LR by $\gamma$ at fixed intervals. Simple but has abrupt transitions.
3. Cosine annealing ($\eta_t = \eta_{\min} + \frac{1}{2}(\eta_0 - \eta_{\min})(1 + \cos(\pi t / T))$) provides smooth, continuous decay. The modern default.
4. Warmup ramps the LR linearly from a small value to $\eta_0$ over the first $W$ steps. Essential for transformers and adaptive optimizers where variance estimates are unreliable early on.
5. Warmup + cosine decay is the standard recipe for modern training: safe startup, fast exploration at peak LR, smooth convergence at the end.
6. In PyTorch: use $\texttt{torch.optim.lr\_scheduler.CosineAnnealingLR}$ for cosine decay and $\texttt{StepLR}$ for step decay. Call $\texttt{scheduler.step()}$ once per epoch after $\texttt{optimizer.step()}$.

The complete optimizer toolbox: SGD with momentum (Section 1) gives every parameter the same smoothed learning rate. Adam (Section 2) adds per-parameter adaptive rates. Learning rate scheduling (this section) controls the global trajectory of the LR over training. Together, they form the complete optimization system used in every modern neural network: AdamW + warmup + cosine decay.