Learning Rate Schedules

Chapter 1: Warmup — Stabilizing the Start

You've decided on a target learning rate of 0.001. Naively, you set it to 0.001 from step 0. The model has random weights. Gradients are wild — pointing in random directions with huge magnitudes. Multiplying wild gradients by a non-trivial learning rate: chaos. The first few hundred steps are a car crash.

Learning rate warmup is a simple fix: instead of starting at the target LR, start at zero (or near-zero) and linearly increase to the target over the first N steps. This gives the optimizer time to get its bearings before you ask it to take real steps.

Why Random Initialization Makes Early Gradients Terrible

At step 0, every weight in your network is random. The model is essentially computing garbage — its output has nothing to do with the input. Gradients computed on garbage outputs are noisy, enormous, and point in nearly random directions.

Now multiply those wild gradients by a non-trivial learning rate. The result is large, random weight updates. The model lurches in an arbitrary direction. This lurch produces even wilder gradients at step 1. A positive feedback loop forms: bad predictions → wild gradients → large updates → worse predictions → wilder gradients. Training can diverge in the first hundred steps.

Why Adam Makes It Worse (Without Warmup)

Adam maintains two running averages for each parameter: the first moment m (exponential moving average of gradients) and the second moment v (exponential moving average of squared gradients). The actual update divides the first moment by the square root of the second moment: Δw = m / √v.

At step 0, both m and v are initialized to zero. Adam uses bias correction to account for this — dividing by (1 - β^t) — but the corrected estimates are still based on just one or two gradient samples. They're unreliable.

Here's the dangerous scenario: suppose the first gradient for some parameter is accidentally small. Then v (squared gradient) is tiny. The update m / √v divides by a tiny number, producing an enormous step. One unlucky gradient can send a weight into orbit.

Hand Calculation: Linear Warmup

Linear warmup is the simplest and most common form. Given a target LR of η_target and a warmup duration of T_w steps, the learning rate at step t is:

Let's trace it. Target LR = 0.001, warmup = 1000 steps:

Notice the LR at step 0 is exactly zero — no update happens at all. By step 100, the LR is just 10% of the target. This means even if the gradient is 10× too large (because Adam's estimates are bad), the effective update is only as big as it would be at the target LR with a normal gradient.

How Much Warmup Do You Need?

The standard recipe is 1-5% of total training steps. This is not a magic number — it's roughly how long Adam needs to build reliable running averages.

GPT-3 used 375 warmup steps out of 300,000 total (0.125%). BERT used 10,000 out of 1,000,000 (1%). LLaMA used 2,000 out of roughly 1,500,000 (0.13%). The trend in large-scale training is toward shorter warmup — just enough to stabilize, no more.

See It: Warmup vs. No Warmup

The simulation below shows two training runs side by side. The left curve trains without warmup — the LR starts at the target from step 0. The right curve uses linear warmup. Watch the loss in the first few hundred steps: without warmup, the loss often spikes or diverges before recovering (if it recovers at all). With warmup, the loss decreases smoothly from the start.

The Code

python
# Linear warmup from scratch
def linear_warmup_lr(step, target_lr, warmup_steps):
    """Returns the learning rate at the given step."""
    if step < warmup_steps:
        return target_lr * (step / warmup_steps)
    return target_lr

# Usage in a training loop
target_lr = 1e-3
warmup_steps = 1000
optimizer = torch.optim.Adam(model.parameters(), lr=target_lr)

for step in range(total_steps):
    lr = linear_warmup_lr(step, target_lr, warmup_steps)
    for pg in optimizer.param_groups:
        pg['lr'] = lr
    loss = model(batch)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

python
# PyTorch built-in: LambdaLR with warmup
import torch
from torch.optim.lr_scheduler import LambdaLR

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
warmup_steps = 1000

scheduler = LambdaLR(
    optimizer,
    lr_lambda=lambda step: min(1.0, step / warmup_steps)
)

for step in range(total_steps):
    loss = model(batch)
    loss.backward()
    optimizer.step()
    scheduler.step()   # updates LR after each step
    optimizer.zero_grad()

Chapter 2: Step Decay — The Staircase

Before fancy schedules, researchers used a simple recipe: train at LR = 0.1 for 30 epochs, then drop to 0.01, then to 0.001. Each drop lets the optimizer settle into a minimum that the previous LR was oscillating around. ResNet's original training — the paper that put 152-layer networks on the map — used exactly this recipe.

Step decay (also called staircase decay or piecewise constant schedule) is the simplest schedule beyond a constant LR. You pick a set of milestones (specific epochs or steps) and multiply the LR by a decay factor γ at each milestone.

The Formula

Given an initial learning rate η₀, a decay factor γ (typically 0.1 or 0.5), and a step size S (the number of epochs between drops), the learning rate at epoch e is:

The floor function ⌊e / S⌋ counts how many decay events have happened. At epoch 0, zero events have happened, so γ⁰ = 1 and the LR is η₀. At epoch S, one event has happened, so the LR drops to η₀ × γ. At epoch 2S, it drops again to η₀ × γ².

Hand Calculation: ResNet-Style Step Decay

The original ResNet paper (He et al., 2015) trained for 120 epochs on ImageNet with η₀ = 0.1 and γ = 0.1 at epochs 30, 60, and 90. Let's trace through:

Now let's work a gentler example. Initial LR = 0.1, γ = 0.5 (halve the LR each time), drop every 10 epochs, 50 epochs total.

The final LR is 0.1 × 0.5⁴ = 0.00625 — that's 16× smaller than the initial LR. Over 50 epochs, the total "learning budget" (sum of LRs across all epochs, a rough proxy for total learning) is:

The Plateau-Then-Drop Pattern

Step decay produces a distinctive loss curve shape that you'll recognize instantly once you know what to look for.

Between LR drops, the loss plateaus. The optimizer has found a minimum at the current LR scale and is oscillating around it. Each gradient step bounces the weights back and forth across the minimum — the LR is too large to settle, but small enough to stay in the neighborhood.

When the LR drops, the optimizer suddenly has smaller steps. It can now settle deeper into that minimum, causing a sudden decrease in loss. The loss drops sharply, then plateaus again at the new LR scale.

This plateau → drop → plateau pattern is the visual signature of step decay. If you see staircase-shaped loss curves in someone's training logs, they're using step decay.

Interactive: Build Your Staircase

The simulation below lets you design your own step decay schedule. The top plot shows the LR staircase. The bottom plot shows a simulated training loss curve — watch for the plateau-then-drop pattern at each LR reduction.

The Code

python
# Step decay from scratch
def step_decay_lr(epoch, init_lr, gamma, step_size):
    """LR at the given epoch under step decay."""
    return init_lr * (gamma ** (epoch // step_size))

# Example: ResNet schedule
for epoch in [0, 29, 30, 59, 60, 89, 90, 119]:
    lr = step_decay_lr(epoch, 0.1, 0.1, 30)
    print(f"Epoch {epoch:3d}: LR = {lr:.4f}")
# Epoch   0: LR = 0.1000
# Epoch  29: LR = 0.1000
# Epoch  30: LR = 0.0100
# Epoch  59: LR = 0.0100
# Epoch  60: LR = 0.0010
# Epoch  89: LR = 0.0010
# Epoch  90: LR = 0.0001
# Epoch 119: LR = 0.0001

python
# PyTorch built-in: StepLR
import torch

optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer, step_size=30, gamma=0.1
)

for epoch in range(120):
    train_one_epoch(model, optimizer)
    scheduler.step()  # call AFTER optimizer.step()
    print(f"Epoch {epoch}: LR = {scheduler.get_last_lr()[0]:.4f}")

# For custom milestones (not evenly spaced):
scheduler = torch.optim.lr_scheduler.MultiStepLR(
    optimizer, milestones=[30, 60, 90], gamma=0.1
)

Chapter 3: Cosine Annealing

Step decay requires knowing when to drop the LR — and getting it wrong costs you. What if we didn't drop at all, but smoothly lowered the LR over the entire run? Cosine annealing does exactly this, and the cosine shape turns out to be nearly optimal.

Proposed by Loshchilov and Hutter in 2016, cosine annealing has become the default schedule for most of modern deep learning. GPT-3, LLaMA, PaLM, Chinchilla — they all use cosine annealing (with warmup). Its appeal is simple: one fewer hyperparameter to tune (no milestone epochs), smooth gradient flow, and a shape that naturally spends more time at productive learning rates.

Deriving the Formula

We want a function that starts at η_max and smoothly decays to η_min over T total steps. Let's build it from the cosine function.

Step 1: The cosine function cos(θ) goes from 1 (at θ = 0) to -1 (at θ = π). We want something that goes from 1 to 0. So we take:

At t = 0: ½(1 + cos(0)) = ½(1 + 1) = 1. At t = T: ½(1 + cos(π)) = ½(1 - 1) = 0. It goes from 1 to 0 over T steps, following a smooth cosine curve.

Step 2: Scale it to our LR range. We want to go from η_max down to η_min:

Verify: At t = 0: η_min + (η_max - η_min) × 1 = η_max. At t = T: η_min + (η_max - η_min) × 0 = η_min. Correct.

Hand Calculation: Tracing the Curve

Let η_max = 0.001, η_min = 0, T = 100 steps. We'll compute the LR at key points along the schedule.

Look at the rate of change. From step 0 to step 25 (first quarter), the LR drops from 0.001 to 0.000854 — only a 15% decrease. From step 75 to step 100 (last quarter), it drops from 0.000146 to 0 — also slow in absolute terms. But from step 25 to step 75 (the middle half), it drops from 0.000854 to 0.000146 — a massive 83% of the total decrease happens in the middle half.

Why Cosine Beats Linear Decay

A linear decay drops the LR at a constant rate: η(t) = η_max(1 - t/T). It spends equal time at every LR level. But not all LR levels are equally productive.

The middle of training is the most productive phase. The model has already learned the broad strokes (it's past the random-initialization chaos) but hasn't yet settled into a minimum (it's still making meaningful improvements). This is when the optimizer is doing real work — refining features, sharpening decision boundaries, tuning internal representations.

Cosine annealing keeps the LR higher during this productive phase and drops it sharply only toward the end. Linear decay, by contrast, has already wasted the LR budget — by the midpoint, it's at half the initial LR, reducing the effective learning rate during the most productive phase.

Another way to see this: imagine a histogram of "time spent at each LR level." Linear decay gives a flat histogram — equal time at every level. Cosine gives a U-shaped histogram that's weighted toward the middle values. Since middle LR values are where the optimizer does its best work, cosine is more efficient.

Interactive: Three Schedules Compared

The simulation below plots cosine annealing, linear decay, and step decay on the same axes. Watch how the three schedules distribute the LR budget differently. The shaded areas represent total learning budget (integral of LR over time).

Cosine Restarts (SGDR)

Loshchilov and Hutter's original paper (SGDR: Stochastic Gradient Descent with Warm Restarts, 2016) proposed an elegant extension: instead of one cosine decay, use repeated cosine cycles. When the LR reaches η_min, restart from η_max.

Why would you want to increase the LR after carefully decreasing it? Because the optimizer might have settled into a local minimum that's not the best one. A sudden LR increase — a "warm restart" — gives the optimizer enough energy to escape and explore nearby regions. It might find a better minimum on the next cycle.

The resulting schedule looks like a sawtooth wave with cosine-shaped teeth: each tooth starts at η_max and decays to η_min. A common variation doubles the period after each restart (T, 2T, 4T, ...), giving the optimizer progressively longer settling phases.

The Code

python
# Cosine annealing from scratch — just 3 lines
import math

def cosine_lr(step, total_steps, lr_max, lr_min=0):
    """Cosine annealing LR at the given step."""
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * step / total_steps))

# Trace it
for t in [0, 25, 50, 75, 100]:
    print(f"Step {t:3d}: LR = {cosine_lr(t, 100, 0.001):.6f}")
# Step   0: LR = 0.001000
# Step  25: LR = 0.000854
# Step  50: LR = 0.000500
# Step  75: LR = 0.000146
# Step 100: LR = 0.000000

python
# PyTorch: CosineAnnealingLR
import torch

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=100,       # total steps
    eta_min=1e-5    # minimum LR (default 0)
)

for step in range(100):
    loss = model(batch)
    loss.backward()
    optimizer.step()
    scheduler.step()
    optimizer.zero_grad()

python
# PyTorch: CosineAnnealingWarmRestarts (SGDR)
scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
    optimizer,
    T_0=50,        # steps in first cosine cycle
    T_mult=2,      # double the period after each restart
    eta_min=1e-5
)
# Cycle 1: steps 0-49 (period 50)
# Cycle 2: steps 50-149 (period 100)
# Cycle 3: steps 150-349 (period 200)

python
# The complete recipe: warmup + cosine decay
import math

def warmup_cosine_lr(step, total_steps, warmup_steps, lr_max, lr_min=0):
    """Linear warmup then cosine decay — the modern standard."""
    if step < warmup_steps:
        return lr_max * (step / warmup_steps)
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * progress))

# Typical usage: 2000 warmup, 300K total, peak LR 3e-4
for s in [0, 1000, 2000, 50000, 150000, 299999]:
    lr = warmup_cosine_lr(s, 300000, 2000, 3e-4)
    print(f"Step {s:6d}: LR = {lr:.6f}")
# Step      0: LR = 0.000000  (warmup)
# Step   1000: LR = 0.000150  (mid-warmup)
# Step   2000: LR = 0.000300  (peak)
# Step  50000: LR = 0.000290  (slowly decaying)
# Step 150000: LR = 0.000150  (midpoint)
# Step 299999: LR = 0.000000  (near zero)

Chapter 4: Cyclical Learning Rates

What if lowering the learning rate isn't always the right move? What if raising it periodically — kicking the optimizer out of whatever valley it settled into — actually leads to better minima?

That's the radical idea behind Cyclical Learning Rates (CLR), proposed by Leslie Smith in 2017. Instead of a monotonic descent, the LR oscillates between a base_lr (minimum) and a max_lr (maximum) over repeating cycles. Each cycle has two halves: ramp up, then ramp down. The ramp up pushes the optimizer out of the current basin. The ramp down lets it settle into a new one — often a better one.

The key insight is about the loss landscape. Neural networks have many local minima. Some are sharp and narrow (bad — they overfit), some are wide and flat (good — they generalize). A low LR gets stuck in whatever minimum it finds first. A periodically high LR has enough energy to escape narrow minima but not enough to escape wide ones. Over multiple cycles, the optimizer naturally migrates toward wider, flatter minima.

The Three Policies

Smith proposed three variants, all built on the same triangle wave:

All three share the same base formula. The only difference is how the peak shrinks (or doesn't) over successive cycles.

The CLR Formula

Given a training step t, a step_size (half the cycle length in steps), base_lr, and max_lr:

The variable x is a normalized position within the current cycle. When x = 0, we're at the peak (LR = max_lr). When x = 1, we're at the trough (LR = base_lr). The absolute value creates the triangle wave — x decreases from 1 to 0 during the ramp-up half, then increases from 0 to 1 during the ramp-down half.

Hand Calculation: Triangular Policy

Step 1: Which cycle are we in?

• cycle = floor(3500 / (2 × 2000)) = floor(3500 / 4000) = floor(0.875) = 0

Still in the first cycle.

Step 2: Where within the cycle?

• x = | 3500 / 2000 − 2 × 0 − 1 | = | 1.75 − 0 − 1 | = | 0.75 | = 0.75

We're 75% of the way through the ramp-down half (x > 0 means we're descending).

Step 3: Compute the LR.

• LR = 0.001 + (0.01 − 0.001) × max(0, 1 − 0.75)
• LR = 0.001 + 0.009 × 0.25
• LR = 0.001 + 0.00225 = 0.00325

Makes sense: step 3500 is three-quarters through a 4000-step cycle, deep into the ramp-down, so the LR is close to (but above) base_lr.

Verification checkpoints:

• At t = 0: x = |0 − 1| = 1, LR = 0.001 + 0.009 × 0 = 0.001 (base)
• At t = 2000: x = |1 − 1| = 0, LR = 0.001 + 0.009 × 1 = 0.01 (peak)
• At t = 4000: x = |2 − 2×1 − 1| = 1, LR = 0.001 (back to base, new cycle starts)

Simulation: Cyclical LR Policies

From Scratch: Cyclical LR in Code

python
import numpy as np

def cyclical_lr(step, base_lr=0.001, max_lr=0.01,
                 step_size=2000, policy='triangular', gamma=0.99994):
    """Compute LR at a given step using cyclical schedule."""
    cycle = np.floor(step / (2 * step_size))
    x = np.abs(step / step_size - 2 * cycle - 1)

    if policy == 'triangular':
        scale = 1.0
    elif policy == 'triangular2':
        scale = 1 / (2 ** cycle)      # Halve max each cycle
    elif policy == 'exp_range':
        scale = gamma ** step             # Exponential decay
    else:
        raise ValueError(f"Unknown policy: {policy}")

    lr = base_lr + (max_lr - base_lr) * max(0, 1 - x) * scale
    return lr

# Verify our hand calculation
lr_3500 = cyclical_lr(3500)
print(f"LR at step 3500: {lr_3500:.5f}")
# Output: 0.00325 ✓

# Full schedule
steps = np.arange(10000)
lrs = [cyclical_lr(s) for s in steps]
print(f"Min LR: {min(lrs):.4f}, Max LR: {max(lrs):.4f}")
# Output: Min LR: 0.0010, Max LR: 0.0100

With PyTorch:

python
import torch

model = torch.nn.Linear(10, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

scheduler = torch.optim.lr_scheduler.CyclicLR(
    optimizer,
    base_lr=0.001,
    max_lr=0.01,
    step_size_up=2000,    # Half-cycle (ramp up)
    step_size_down=2000,  # Half-cycle (ramp down)
    mode='triangular',     # or 'triangular2', 'exp_range'
    gamma=0.99994,         # Only used for exp_range
    cycle_momentum=True,  # Inversely cycle momentum too!
)

# Usage: call scheduler.step() after each BATCH, not each epoch
for step in range(10000):
    optimizer.zero_grad()
    loss = model(torch.randn(32, 10)).sum()
    loss.backward()
    optimizer.step()
    scheduler.step()  # Step per batch, not per epoch!

Chapter 5: The 1Cycle Policy

What if you could find the optimal learning rate AND train to completion in the same run? The 1cycle policy does both. It's the schedule that achieves super-convergence — reaching the same accuracy in 10 times fewer steps.

Leslie Smith (the same researcher behind CLR) teamed with Nicholay Topin in 2019 to propose something elegant: instead of repeating many small cycles, use one single large cycle that spans the entire training run. The LR starts low, climbs to a high peak in the middle of training, descends back to the starting point, and then plummets even further in a brief annihilation phase at the very end.

The Three Phases

1cycle divides training into three phases with very different purposes:

The phase fractions are controlled by pct_start (typically 0.3 to 0.45 — the fraction of training spent ramping up). Phase 2 mirrors phase 1 in length. Phase 3 gets whatever remains.

The LR Range Test

Before using 1cycle, you need to know the right max_lr. Smith proposed a brilliant diagnostic: the LR range test. Here's how it works:

1. Train for one epoch (or a few hundred steps) while linearly increasing the LR from very small (say 1e-7) to very large (say 10).
2. Record the loss at each step.
3. Plot loss versus LR (log scale on x-axis).

The resulting curve has a characteristic shape: the loss stays flat (LR too small to matter), then drops steeply (the sweet spot), then explodes upward (LR too large, training diverges). Your max_lr should be slightly below the explosion point — where the loss is still decreasing but the curve starts to bend upward. The min_lr is typically 1/10th to 1/25th of max_lr.

Why Super-Convergence Works

The magic of 1cycle is that the high LR in the middle of training acts as implicit regularization. At high learning rates, the optimizer can't settle into sharp, narrow minima — the update steps are too large to fit inside them. It can only stay in wide, flat minima. And wide, flat minima generalize better because nearby points in weight space produce similar loss values.

This is the same mechanism that makes large batch training with high LR generalize well — it's an implicit bias toward flat regions of the loss landscape. The 1cycle policy exploits this deliberately: spend half of training at a high enough LR that only wide minima are reachable, then reduce the LR to fine-tune within that wide basin.

Hand Calculation: 1Cycle Schedule

Phase boundaries:

• Phase 1 (ramp up): steps 0 to 4,500 (45% of 10,000)
• Phase 2 (ramp down): steps 4,500 to 9,000 (another 45%)
• Phase 3 (annihilation): steps 9,000 to 10,000 (final 10%)

At step 2,250 (mid phase 1):

• Progress through phase 1: 2250 / 4500 = 0.5
• LR = min_lr + (max_lr − min_lr) × 0.5
• LR = 0.0001 + (0.001 − 0.0001) × 0.5 = 0.0001 + 0.00045 = 0.00055

At step 4,500 (peak):

• LR = 0.001 (max_lr)

At step 6,750 (mid phase 2):

• Progress through phase 2: (6750 − 4500) / 4500 = 0.5
• LR = max_lr − (max_lr − min_lr) × 0.5
• LR = 0.001 − 0.0009 × 0.5 = 0.001 − 0.00045 = 0.00055

Notice the symmetry: mid phase 1 and mid phase 2 give the same LR.

At step 9,500 (mid phase 3 — annihilation):

• Progress through phase 3: (9500 − 9000) / 1000 = 0.5
• LR = min_lr − (min_lr − final_lr) × 0.5
• LR = 0.0001 − (0.0001 − 0.000001) × 0.5 = 0.0001 − 0.0000495 = 0.0000505

Simulation: 1Cycle Policy

From Scratch: 1Cycle + LR Range Test

python
import numpy as np

def one_cycle_lr(step, total_steps, max_lr=0.001,
                   div_factor=10, final_div_factor=100,
                   pct_start=0.45):
    """Compute LR at a given step using 1cycle schedule."""
    min_lr = max_lr / div_factor
    final_lr = min_lr / final_div_factor

    phase1_end = int(total_steps * pct_start)
    phase2_end = int(total_steps * (2 * pct_start))
    # If pct_start=0.45, phase2 ends at 90%, leaving 10% for annihilation

    if step < phase1_end:
        # Phase 1: ramp up
        t = step / phase1_end
        lr = min_lr + (max_lr - min_lr) * t
    elif step < phase2_end:
        # Phase 2: ramp down
        t = (step - phase1_end) / (phase2_end - phase1_end)
        lr = max_lr - (max_lr - min_lr) * t
    else:
        # Phase 3: annihilation
        t = (step - phase2_end) / (total_steps - phase2_end)
        lr = min_lr - (min_lr - final_lr) * t
    return lr

# Verify hand calculation
print(f"Step 2250: {one_cycle_lr(2250, 10000):.5f}")   # 0.00055
print(f"Step 4500: {one_cycle_lr(4500, 10000):.5f}")   # 0.00100
print(f"Step 6750: {one_cycle_lr(6750, 10000):.5f}")   # 0.00055
print(f"Step 9500: {one_cycle_lr(9500, 10000):.5f}")   # 0.00005

The LR range test:

python
def lr_range_test(model, train_loader, optimizer,
                   start_lr=1e-7, end_lr=10, num_steps=200):
    """Run the LR range test. Returns (lrs, losses)."""
    mult = (end_lr / start_lr) ** (1 / num_steps)
    lr = start_lr
    lrs, losses = [], []
    best_loss = float('inf')

    for step, (x, y) in zip(range(num_steps), train_loader):
        # Set LR
        for pg in optimizer.param_groups:
            pg['lr'] = lr

        loss = model.train_step(x, y)
        lrs.append(lr)
        losses.append(loss)

        # Smoothed loss tracking
        if loss < best_loss:
            best_loss = loss
        if loss > 4 * best_loss:
            break  # Loss exploded — stop here

        lr *= mult  # Geometric increase

    return lrs, losses

# After running: plot lrs (log scale) vs losses
# Pick max_lr just before the loss starts rising

With PyTorch:

python
import torch

model = torch.nn.Linear(10, 1)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.0001)

scheduler = torch.optim.lr_scheduler.OneCycleLR(
    optimizer,
    max_lr=0.001,          # Peak LR (from range test)
    total_steps=10000,     # Total training steps
    pct_start=0.3,         # 30% warmup, 30% decay, 10% annihilation
    div_factor=10,         # initial_lr = max_lr / 10
    final_div_factor=100, # final_lr = initial_lr / 100
    anneal_strategy='cos', # Cosine (smoother) or 'linear'
)
# Note: PyTorch's OneCycleLR uses cosine annealing within each phase by default
# Our hand calc used linear for simplicity — both work well

Chapter 6: Warmup-Stable-Decay — The LLM Standard

You've seen warmup. You've seen cosine annealing. You've seen cyclical tricks. Now forget all of them — because the schedule actually running inside GPT-4, Llama 3, Chinchilla, and Gemini is something much simpler. It's three straight lines:

1. Warmup: linear ramp from near-zero to peak LR (typically 0.5-2% of training).
2. Stable: hold at the peak LR for the vast majority of training (70-90%).
3. Decay: cosine or linear decay to a small fraction of peak (final 5-25%).

This is Warmup-Stable-Decay (WSD), also called the "trapezoidal" schedule. It's the workhorse of large language model training, and understanding why it works reveals deep truths about how LLMs learn.

Why Three Phases?

Each phase serves a distinct purpose tied to the dynamics of LLM training:

Warmup (Phase 1): The optimizer state (momentum, second moments in Adam) starts from zero. During the first few hundred steps, these running averages are unreliable — they're dominated by the very first gradients, which are computed on random weights. The warmup phase keeps the LR low while the optimizer state "fills up" with reliable statistics. Without warmup, the first few large updates can push the model into a region of weight space from which it never recovers.

Stable (Phase 2): This is where the model learns. At a constant, high learning rate, the optimizer has enough energy to explore the loss landscape broadly. For an LLM, this means developing syntax, semantics, world knowledge, reasoning patterns — all the capabilities that emerge during pre-training. The constant LR is not optimal for any one point in training, but it's good enough across the entire phase. And "good enough with simplicity" beats "theoretically optimal with complexity" when you're training for weeks on thousands of GPUs.

Decay (Phase 3): The final annealing phase reduces the LR to squeeze out the last bits of performance. During stable-phase training, the model oscillates around a good solution — the constant LR is high enough that the loss fluctuates. Decaying the LR lets the optimizer settle into the bottom of the basin. This phase typically improves benchmark scores by 1-3%, which doesn't sound like much but represents billions of dollars in capability at frontier scale.

The WSD Formula

Let η_peak be the peak LR, T_total the total training steps, w the warmup fraction, and s the stable fraction. Then d = 1 − w − s is the decay fraction.

During warmup (t < t_warmup):

During stable (t_warmup ≤ t < t_stable):

During cosine decay (t ≥ t_stable):

Where η_min is typically η_peak / 10 (a 10x ratio between peak and minimum LR).

Real-World Numbers

Here's what the big labs actually use:

Notice the pattern: warmup is always tiny (under 1%), stable is always dominant (around 90%), and the decay phase gets roughly 10%. The 10x ratio between peak and min LR is nearly universal. This isn't coincidence — it's the result of extensive scaling law research at Google DeepMind and Meta.

SHOWCASE: The WSD Simulator

Key Insights the Simulator Reveals

From Scratch: WSD Scheduler

python
import math

def wsd_lr(step, total_steps, peak_lr=3e-4, min_lr=3e-5,
            warmup_frac=0.02, stable_frac=0.88,
            decay='cosine'):
    """Warmup-Stable-Decay schedule used by LLaMA, Chinchilla, etc."""
    warmup_steps = int(total_steps * warmup_frac)
    stable_end = int(total_steps * (warmup_frac + stable_frac))

    if step < warmup_steps:
        # Phase 1: linear warmup
        return peak_lr * (step / warmup_steps)

    elif step < stable_end:
        # Phase 2: constant
        return peak_lr

    else:
        # Phase 3: decay
        decay_steps = total_steps - stable_end
        t = (step - stable_end) / decay_steps  # 0 → 1

        if decay == 'cosine':
            return min_lr + (peak_lr - min_lr) * 0.5 * (1 + math.cos(math.pi * t))
        elif decay == 'linear':
            return peak_lr - (peak_lr - min_lr) * t
        elif decay == 'sqrt':
            return min_lr + (peak_lr - min_lr) * (1 - math.sqrt(t))
        else:
            raise ValueError(f"Unknown decay: {decay}")

# Chinchilla config
for s in [0, 1000, 2000, 150000, 272000, 290000, 300000]:
    lr = wsd_lr(s, 300000)
    print(f"Step {s:>7d}: LR = {lr:.6f}")
# Step       0: LR = 0.000000  (start of warmup)
# Step    1000: LR = 0.000050  (mid warmup)
# Step    2000: LR = 0.000100  (... still warming)
# Step  150000: LR = 0.000300  (stable phase)
# Step  272000: LR = 0.000300  (still stable!)
# Step  290000: LR = 0.000082  (decaying)
# Step  300000: LR = 0.000030  (final min_lr)

How LLaMA sets it up in its training config:

python
# LLaMA-style training config (simplified from fairscale/metaseq)
config = {
    "optimizer": {
        "type": "AdamW",
        "lr": 3e-4,          # peak_lr
        "betas": (0.9, 0.95),
        "weight_decay": 0.1,
    },
    "scheduler": {
        "type": "WSD",
        "warmup_steps": 2000,
        "total_steps": 2000000,
        "min_lr": 3e-5,      # 10x ratio
        "decay": "cosine",
    },
}

# The actual training loop (pseudocode):
for step in range(config["scheduler"]["total_steps"]):
    lr = wsd_lr(step, **config["scheduler"])
    for pg in optimizer.param_groups:
        pg["lr"] = lr
    # ... forward, backward, step

Chapter 7: Polynomial Decay & Other Schedules

Cosine isn't the only smooth decay shape. Polynomial decay lets you control the curvature of the schedule with a single parameter. Want to decay fast early and slow late? Set the power above 1. Want to decay slow early and fast late? Set the power below 1.

And before cosine took over the world, the original Transformer paper used inverse square root decay — a specific shape that held as the standard for three years. Let's derive and compare them all.

The Polynomial Decay Formula

Polynomial decay interpolates from η_init to η_end over T steps, with a power parameter p that controls the shape:

At t = 0: (1 − 0)^p = 1, so η(0) = η_init. At t = T: (1 − 1)^p = 0, so η(T) = η_end. The power p controls the path between these endpoints:

• p = 1 (linear): The LR drops at a constant rate. Equal decay per step.
• p = 2 (quadratic): Fast decay early, slow decay late. The LR drops quickly at first, then lingers near η_end.
• p = 0.5 (square root): Slow decay early, fast decay late. The LR holds high for longer, then drops steeply at the end.
• p = 3+: Extremely aggressive early decay. The LR practically jumps to η_end within the first quarter of training.

The physical intuition: high power = front-loaded decay (spend most of your time at a low LR), low power = back-loaded decay (spend most of your time at a high LR). The choice depends on your task: vision models often benefit from p = 2 (quadratic), while NLP tasks tend to prefer p = 1 (linear) or cosine.

Hand Calculation: Comparing Powers

Linear (p = 1):

• η(50) = 0.001 × (1 − 50/100)¹ = 0.001 × 0.5 = 0.0005

Exactly half. That's what linear means.

Quadratic (p = 2):

• η(50) = 0.001 × (0.5)² = 0.001 × 0.25 = 0.00025

Already at 25% of the initial LR halfway through. Quadratic decays faster early.

Square root (p = 0.5):

• η(50) = 0.001 × (0.5)^0.5 = 0.001 × 0.7071 = 0.000707

Still at 70% of the initial LR at the halfway point. Sqrt preserves high LR longer.

Cubic (p = 3):

• η(50) = 0.001 × (0.5)³ = 0.001 × 0.125 = 0.000125

Only 12.5% left at halfway. Aggressive.

Compare all four at the same moment: sqrt preserves 2.8x more LR than quadratic at the midpoint. The choice of power fundamentally changes where the optimizer spends its "learning budget."

The Inverse Square Root Schedule (Original Transformer)

Vaswani et al.'s "Attention Is All You Need" (2017) used a distinctive schedule that combines linear warmup with inverse square root decay:

The min selects between two curves: the linear warmup term (t × warmup^-1.5) dominates when t is small, and the inverse sqrt term (t^-0.5) takes over once t exceeds warmup_steps. The crossover happens exactly at t = warmup_steps.

Hand calculation: d_model = 512, warmup_steps = 4000.

At the peak (t = 4000):

• d_model^-0.5 = 512^-0.5 = 1 / 22.63 = 0.04419
• t^-0.5 = 4000^-0.5 = 1 / 63.25 = 0.01581
• t × warmup^-1.5 = 4000 × 4000^-1.5 = 4000 / 253,000 = 0.01581
• Both terms are equal (crossover point)
• Peak LR = 0.04419 × 0.01581 = 0.000699

At t = 16000 (4x past warmup):

• t^-0.5 = 16000^-0.5 = 0.00791
• η = 0.04419 × 0.00791 = 0.000349

Half the peak LR. Under inverse sqrt decay, you need 4x the steps to halve the LR (since √4 = 2). This makes it a very gentle decay — much slower than cosine or linear.

The Complete Schedule Zoo

Simulation: The Schedule Comparison

From Scratch: Polynomial + Inverse Sqrt in Code

python
import math

def polynomial_lr(step, total_steps, init_lr=0.001,
                    end_lr=0.0, power=2.0):
    """Polynomial decay: fast early for p>1, slow early for p<1."""
    t = min(step / total_steps, 1.0)
    return (init_lr - end_lr) * (1 - t) ** power + end_lr

def inverse_sqrt_lr(step, d_model=512, warmup_steps=4000):
    """Original Transformer schedule (Vaswani et al. 2017)."""
    if step == 0:
        step = 1
    return d_model ** (-0.5) * min(step ** (-0.5),
                                       step * warmup_steps ** (-1.5))

def exponential_lr(step, init_lr=0.001, gamma=0.9999):
    """Exponential decay: LR = init * gamma^step."""
    return init_lr * gamma ** step

# Compare at midpoint (step 50 out of 100)
print("Polynomial comparisons at t=50/100:")
for p in [0.5, 1.0, 2.0, 3.0]:
    lr = polynomial_lr(50, 100, power=p)
    print(f"  p={p:.1f}: {lr:.6f}")
# p=0.5: 0.000707  (sqrt — gentle early)
# p=1.0: 0.000500  (linear)
# p=2.0: 0.000250  (quadratic — aggressive early)
# p=3.0: 0.000125  (cubic — very aggressive early)

# Inverse sqrt peak
peak = inverse_sqrt_lr(4000)
print(f"Transformer peak LR: {peak:.6f}")  # 0.000699

With PyTorch:

python
import torch

model = torch.nn.Linear(10, 1)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)

# Polynomial decay (built-in since PyTorch 1.13)
poly_sched = torch.optim.lr_scheduler.PolynomialLR(
    optimizer,
    total_iters=100,
    power=2.0,  # Quadratic decay
)

# Exponential decay
exp_sched = torch.optim.lr_scheduler.ExponentialLR(
    optimizer,
    gamma=0.9999,  # Multiply LR by 0.9999 each step
)

# Inverse sqrt (use LambdaLR)
d_model = 512
warmup = 4000
inv_sqrt_sched = torch.optim.lr_scheduler.LambdaLR(
    optimizer,
    lr_lambda=lambda step: (
        d_model ** (-0.5) * min(
            (step + 1) ** (-0.5),
            (step + 1) * warmup ** (-1.5)
        )
    ),
)

Chapter 8: The Schedule Racing Arena

Everything we've learned, in one simulation. Drop multiple schedules onto the same training run and watch them race. Each one runs its actual formula — the real schedule equations, not an approximation.

This is the payoff. You've learned what each schedule does; now you'll see the differences in real time. Which schedule reaches the lowest loss fastest? Which one oscillates? Which one wastes compute at the wrong learning rates?

What to Notice

Cosine vs WSD: Cosine starts decaying from step 0, meaning it operates at a lower effective LR throughout mid-training. WSD holds the peak LR for ~90% of training, giving it a higher average LR and more total learning. The WSD loss curve drops faster during the stable phase.

1Cycle's high-LR hump: Watch the 1cycle loss — it actually increases during the ramp-up phase as the LR climbs. This is normal! The optimizer is being pushed out of its current basin. When the LR drops in the second half, the loss falls to a new low. The temporary spike buys a better final result.

Cyclical's oscillation: The CLR loss curve has a sawtooth pattern — periodic spikes at each LR peak followed by drops at each trough. This is the exploration/exploitation cycle in action. Over time, the spikes get smaller as the optimizer settles into progressively wider minima.

Step decay's plateaus: Look for the characteristic flat segments in the step decay loss curve. The loss stops improving between drops (the optimizer oscillates around a minimum at the current LR scale), then drops sharply when the LR is reduced.

Chapter 9: Cheat Sheet & What to Use When

You now understand the full schedule toolkit — from raw constant LR to the WSD schedule running inside every frontier LLM. This chapter is your practical reference. No new concepts. Just the decision guide you'll actually use.

The Decision Flowchart

Hyperparameter Cheat Sheet

Tuning Priority

When you have limited time for schedule tuning, this is the order of importance:

1. Peak learning rate. Sweep 5-10 values on a log scale (1e-5 to 1e-1). This alone determines 80% of training success. Use the LR range test (Chapter 5) for a principled starting point.
2. Schedule shape. Cosine decay is the safe default. If training LLMs at scale, use WSD. For fast experiments, try 1cycle.
3. Warmup duration. 1-5% of total steps. Shorter for large LRs, longer for very large batch sizes. Almost never needs tuning beyond "use 2000 steps."
4. Min LR / decay ratio. Peak/10 is the universal starting point. Peak/100 for aggressive decay; Peak/3 for gentle decay.
5. Specific schedule parameters. γ for step decay, power for polynomial, pct_start for 1cycle. Only tune these after the above are set.

Summary of Everything

The Scaling Law Connection

The Chinchilla paper (Hoffmann et al., 2022) showed that peak LR scales as roughly 1 / √(model_size). Larger models need smaller learning rates. The schedule shape matters less than getting the peak right — which is why WSD's simplicity (one flat plateau at peak) is so appealing for large-scale training.