Deep Learning

Training Deep Networks Effectively

Learning Rate Schedules & Warmup

Chapter 2's optimizers still require choosing a base learning rate, held constant throughout training in the simplest setup. In practice, the ideal learning rat

JrCodex·6 min read

Jr Codex Deep Learning Notes

Level: Intermediate Prerequisites: Chapter 4: Regularization Time to complete: ~20 minutes

Table of Contents

  1. Why a Fixed Learning Rate Is Suboptimal
  2. Step Decay
  3. Cosine Annealing
  4. Learning Rate Warmup
  5. Combining Warmup with Decay
  6. Learning Rate Schedules in PyTorch
  7. Finding a Good Initial Learning Rate
  8. Summary & Next Steps

1. Why a Fixed Learning Rate Is Suboptimal

Chapter 2's optimizers still require choosing a base learning rate, held constant throughout training in the simplest setup. In practice, the ideal learning rate typically changes over the course of training:

The Intuition
─────────────────────────────────────────
  EARLY in training:      weights are far from any good
                             solution — LARGE steps make fast
                             progress and are relatively safe
  LATE in training:           weights are CLOSE to a good
                                 solution — LARGE steps risk
                                 overshooting and oscillating
                                 around the minimum instead of
                                 settling into it precisely
─────────────────────────────────────────

A learning rate schedule adjusts the learning rate over time, typically starting higher and decreasing — directly generalizing the ML Notes' hyperparameter tuning discipline into something applied during a single training run rather than only between runs.

2. Step Decay

The simplest schedule: multiply the learning rate by a fixed factor every fixed number of epochs.

def step_decay(initial_lr, epoch, drop_factor=0.5, epochs_per_drop=10):
    return initial_lr * (drop_factor ** (epoch // epochs_per_drop))
 
for epoch in [0, 5, 10, 15, 20, 25, 30]:
    lr = step_decay(initial_lr=0.1, epoch=epoch)
    print(f"Epoch {epoch}: learning rate = {lr:.5f}")
 
# Epoch 0:  0.10000
# Epoch 10: 0.05000  (dropped after 10 epochs)
# Epoch 20: 0.02500  (dropped again after 20 epochs)
# Epoch 30: 0.01250

3. Cosine Annealing

Cosine annealing smoothly decreases the learning rate following a cosine curve, from an initial value down to (near) zero — a popular, smoother alternative to step decay's abrupt drops.

import numpy as np
 
def cosine_annealing(initial_lr, epoch, total_epochs):
    return initial_lr * 0.5 * (1 + np.cos(np.pi * epoch / total_epochs))
 
for epoch in [0, 10, 20, 30, 40, 50]:
    lr = cosine_annealing(initial_lr=0.1, epoch=epoch, total_epochs=50)
    print(f"Epoch {epoch}: learning rate = {lr:.5f}")
Step Decay vs Cosine Annealing
─────────────────────────────────────────
  Learning Rate
   │████                    Step decay — abrupt DROPS at
   │    ████                  fixed intervals
   │        ████
   │            ████
   │
   │╲                        Cosine annealing — smooth,
   │ ╲__                       continuous decrease
   │    ╲___
   │        ╲______
   └──────────────────────────► Epoch
─────────────────────────────────────────

4. Learning Rate Warmup

Warmup does the opposite at the very start of training: it begins with a small learning rate and gradually increases it over the first several hundred or thousand steps, before switching to a decay schedule.

def warmup_schedule(step, warmup_steps, target_lr):
    if step < warmup_steps:
        return target_lr * (step / warmup_steps)      # LINEAR ramp-up
    return target_lr
 
for step in [0, 250, 500, 750, 1000, 1500]:
    lr = warmup_schedule(step, warmup_steps=1000, target_lr=0.001)
    print(f"Step {step}: learning rate = {lr:.6f}")
Why Warmup Helps, Especially for Transformers (Module 6)
─────────────────────────────────────────
  At the VERY START of training, weights are randomly
  initialized (Chapter 1) and the running statistics used by
  Adam (Chapter 2) and BatchNorm/LayerNorm (Chapter 3) haven't
  stabilized yet. Taking large steps immediately, before these
  statistics have settled, can push the network into a poor,
  hard-to-escape region early on. A brief warmup period lets
  these statistics stabilize gently before applying full-sized
  updates.

  Warmup is now considered close to MANDATORY for training
  Transformers (Module 6) — most published training recipes
  for BERT, GPT, and similar architectures explicitly specify
  a warmup period.
─────────────────────────────────────────

5. Combining Warmup with Decay

The standard modern recipe combines both: warmup for the first portion of training, followed by a decay schedule (often cosine annealing) for the remainder.

def warmup_then_cosine_decay(step, warmup_steps, total_steps, target_lr):
    if step < warmup_steps:
        return target_lr * (step / warmup_steps)      # Section 4's warmup
 
    # Section 3's cosine decay, over the REMAINING steps after warmup
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    return target_lr * 0.5 * (1 + np.cos(np.pi * progress))
Learning Rate Over the FULL Training Run
─────────────────────────────────────────
  Learning Rate
   │      ╱‾‾╲
   │    ╱      ╲___
   │  ╱              ╲______
   │╱                        ╲__________
   └──────────────────────────────────────► Step
      ↑warmup↑          cosine decay
─────────────────────────────────────────

6. Learning Rate Schedules in PyTorch

import torch
import torch.nn as nn
 
model = nn.Linear(10, 1)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
# Step decay (Section 2)
step_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
 
# Cosine annealing (Section 3)
cosine_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
 
# A typical training loop, applying the scheduler AFTER each epoch
for epoch in range(50):
    # ... training steps for this epoch ...
    cosine_scheduler.step()      # update the learning rate for the NEXT epoch
    current_lr = optimizer.param_groups[0]["lr"]
    print(f"Epoch {epoch}: lr = {current_lr:.6f}")

7. Finding a Good Initial Learning Rate

A Practical Technique: the Learning Rate Range Test
─────────────────────────────────────────
  1. Start training with a VERY small learning rate
  2. GRADUALLY increase it every few steps (e.g. exponentially)
  3. Plot loss against learning rate
  4. Pick a learning rate from JUST BEFORE the point where
       loss starts increasing sharply — this is typically a
       strong INITIAL learning rate to use with a proper
       schedule (Sections 2-5)
─────────────────────────────────────────

This range test is a fast, empirical alternative to the ML Notes' systematic hyperparameter search methods — specifically adapted for the learning rate, given how large its search range typically spans (often several orders of magnitude).

8. Summary & Next Steps

Key Takeaways

  • A fixed learning rate is a compromise; ideal learning rates tend to be larger early in training and smaller as training approaches a good solution.
  • Step decay drops the learning rate at fixed intervals; cosine annealing decreases it smoothly and continuously — both are common decay schedules.
  • Learning rate warmup starts small and ramps up, letting optimizer and normalization statistics stabilize before large updates begin — close to mandatory for training Transformers.
  • The learning rate range test provides a fast, empirical way to choose a strong initial learning rate before applying a full schedule.

Concept Check

  1. Why might a learning rate that works well early in training cause problems later in training?
  2. What specific training instability does warmup help prevent, and why is it especially important for Transformers?
  3. How does the learning rate range test help you pick an initial learning rate?

Next Chapter

→ Chapter 6: Vanishing/Exploding Gradients & Diagnosing Training Problems

Jr Codex — 1-on-1 Personalized Coaching | Back to Module Index | Back to DL Index