Deep Learning
Advanced 16 min read Deep Learning

Learning Rate Schedules & Training Loop

Shaping the descent over the whole run

One learning rate over the whole run is rarely optimal: warm-up, decay and restarts shape a trajectory that converges faster and better. This concept covers schedules and the loop that evaluates them.

learning rate schedulewarm-upcosine decayone-cycleconvergencetraining loop
Prerequisites: Loss Functions & Optimizers
Before we start

What you will learn

  • Explain why a static LR is a poor plan for most runs
  • Choose warm-up, decay and restart strategies per schedule
  • Inspect loss/weights trajectory to tune schedule knobs
Introduction

What is this?

A learning-rate schedule changes the optimizer's step size over training. Modern recipes combine a short warm-up (low LR rising to peak), a long decay (step/linear/cosine toward near zero), and occasional restarts that re-explores. The schedule must match the optimizer — Adam and SGD respond differently — and is usually chosen along with batch size in a few guided runs rather than by feel.

Motivation

Why do you need it?

  • Early steps with too-large LR diverge; with too-small LR you waste the period where gradients are strongest.
  • Decaying LR toward the end lets the optimizer settle into a sharp-but-accurate valley.
  • Warm-up protects Adam's extreme early updates on huge, poorly-scaled batches.
  • Restarts escape local flats and re-explore the landscape when loss plateaus.
Think about it

Mental model

Imagine

Landing an aircraft

Think of

LR schedule

You approach fast early (a high LR), then flare and bleed speed at the threshold (warm-up is unfewise not to mis-begin; decay is the final glide path, smoothing into the touch-down so you don't bounce off the runway).

Mechanics

How it works

  • Warm-up: raise LR from a tiny value to a peak over 1-10% of steps — lets the optimizer build stable moments/statistics.
  • Step decay: multiply LR by γ every k epochs — coarse approximation of a smooth decay.
  • Cosine decay: follow the first half of a cosine from peak to ~0 — smooth annealing that lands softly.
  • One-cycle: warm-up to a high peak, then cosine down below start; excellent for vision.
  • Restarts: periodically raise LR again (cosine with restarts) to jump out of shallow minima.

LR over training with warm-up and cosine decay.

Warm-up
Peak
Cosine decay
Landing
Real world

In production

  • ImageNet classifiers train with one-cycle/multi-step cosine schedules, not constant LR.
  • LLM fine-tuning uses linear warm-up + cosine decay over the final token budget.
  • RL agents anneal LR and explore by episode — schedule is part of the policy training recipe.
Guidelines

Best practices

  • Run a short LR range test to find the max stable LR; base your peak on it.
  • Pair warm-up length with total steps (10% warm-up is a common default).
  • Log LR alongside loss so 'improved after epoch 20' can be attributed to the schedule.
  • Keep schedule tied to step counts, not epoch counts, when data size changes.
Watch out

Common mistakes

  • Warming up a perfectly fine constant-LR run and only adding instability.
  • Decaying LR to 0 too fast — model stops sharpening and validation stagnates with slack.
  • Using SGD's decay pattern on Adam (primary learning rate) without re-checking the peak.
  • Relying on loss alone, when weight-norm and val-history would separate signal from noise.
Root cause analysis

Common issues

Error
Loss spikes exactly when warm-up ends
Root cause

Jump to a peak LR that's too high for the weights' current state

Fix

Lower the peak or make warm-up longer; verify with the LR range test.

Error
Validation gets worse in the last 10% of training
Root cause

LR too small late → fits training noise at fine granularity

Fix

Stop decay earlier, or early-stop at the best-val checkpoint regardless of schedule.

Error
Scheduler shows lr unchanged every step
Root cause

Scheduler not invoked per optimizer step (forgotten .step() call)

Fix

Call scheduler.step() each batch/epoch consistently and print lr in logs.

When things break

Debugging tips

  • If loss peaks during warm-up, your peak LR is too high for the optimizer — lower by half.
  • If late-training loss plateaus at a slightly-too-high value, extend/steepen the decay.
  • Graph lr.tolist() against step to confirm the scheduler is actually firing (a classic silent bug).
  • On restarts, val loss spikes are normal — judge by the troughs, not the peaks.
Interviews

Interview checkpoint

Test yourself

Quick quiz

0/2
1

What is the primary job of warm-up in an LR schedule?

2

A single constant learning rate is optimal for modern deep training.

Go deeper

Official documentation

PyTorch schedulers
docs