Optimizers, Learning Rates, and Schedulers

Optimizers, learning rates, and schedulers control how a PyTorch model moves from current parameters toward lower loss. The purpose is plain: after autograd computes gradients, the optimizer decides how large a parameter update to apply, the learning rate sets the scale of that update, and the scheduler changes that scale over training. In this course section, those choices are the difference between a network that learns steadily, one that barely moves, and one that explodes into unusable weights.

What Actually Changes During Training

A neural network parameter is a tensor with stored values such as weights and biases. During the forward pass, PyTorch builds a computation graph connecting those tensors to a scalar loss. Calling loss.backward() fills each trainable parameter’s .grad field with the derivative of the loss with respect to that parameter. An optimizer does not rediscover calculus; it reads those gradient buffers and mutates parameter data according to an update rule.

The basic stochastic gradient descent update is parameter = parameter - lr * gradient. The word stochastic means the gradient is usually estimated from a mini-batch instead of the full dataset. Momentum adds a velocity buffer so repeated gradient directions accumulate. Adam and AdamW keep moving averages of gradients and squared gradients, then scale updates per parameter so noisy or consistently large gradients are handled differently from small ones. AdamW decouples weight decay from the adaptive gradient step, which usually makes regularization easier to reason about than classic Adam weight decay.

PyTorch stores optimizer state separately from the module. The model owns parameters; the optimizer owns per-parameter state such as momentum buffers, Adam moments, and step counts. This matters when checkpointing: saving only model.state_dict() lets you resume weights, but it does not resume the optimizer’s memory. For a long run, save the model, optimizer, scheduler, and epoch together.

API Anatomy

A typical training step has five ordered operations. First, call optimizer.zero_grad() so old gradients do not accumulate. Second, compute predictions and loss. Third, call loss.backward(). Fourth, optionally clip or inspect gradients. Fifth, call optimizer.step(). If a scheduler is used, call scheduler.step() at the cadence expected by that scheduler, commonly once per epoch for epoch schedulers and once per batch for warmup or one-cycle schedules.

The constructor torch.optim.SGD(model.parameters(), lr=...) receives iterable parameters and hyperparameters. Parameter groups allow different settings for different parts of the model: for example, a pretrained backbone may use a smaller learning rate than a newly initialized classifier head. Schedulers wrap an optimizer and update each parameter group’s lr. The current learning rate lives in optimizer.param_groups[i]["lr"], so it can be logged without guessing.

Example 1: One Weight, One Minimum

This smallest example optimizes the function (w - 3)^2. The gradient at w=0 is -6, so SGD with learning rate 0.1 increases w by 0.6 on the first step. Each later step is smaller because the gradient shrinks near the minimum.

import torch

w = torch.tensor([0.0], requires_grad=True)
optimizer = torch.optim.SGD([w], lr=0.1)

for step in range(5):
    optimizer.zero_grad()
    loss = (w - 3).pow(2).sum()
    loss.backward()
    optimizer.step()
    print(step, round(w.item(), 4), round(loss.item(), 4))

The deterministic output is approximately 0 0.6 9.0, 1 1.08 5.76, 2 1.464 3.6864, 3 1.7712 2.3593, and 4 2.017 1.5099. The loss falls because the learning rate is large enough to move but small enough not to jump past the minimum too aggressively.

Example 2: Learning Rate Size

The same loss shows why learning rate is not a cosmetic setting. Too small wastes steps; reasonable values converge quickly; too large can oscillate or diverge. This example prints the weight trajectory for three learning rates.

import torch

for lr in [0.01, 0.2, 1.1]:
    w = torch.tensor([0.0], requires_grad=True)
    optimizer = torch.optim.SGD([w], lr=lr)
    values = []
    for _ in range(4):
        optimizer.zero_grad()
        loss = (w - 3).pow(2).sum()
        loss.backward()
        optimizer.step()
        values.append(round(w.item(), 3))
    print(lr, values)

Expected behavior: 0.01 creeps from 0.06 toward 0.233 after four steps, 0.2 reaches about 2.611, and 1.1 overshoots repeatedly: about 6.6, -1.32, 8.184, -3.221. In real networks the loss surface is not a parabola, but the trade-off is the same. The learning rate is the main knob controlling update energy.

Example 3: AdamW and Step Decay

Schedulers change learning rates without recreating the optimizer. Here, StepLR halves the learning rate every two scheduler steps. The optimizer is AdamW, so each parameter also has adaptive moment state and decoupled weight decay.

import torch
from torch import nn

model = nn.Linear(2, 1)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.05, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.5)

for epoch in range(5):
    current_lr = optimizer.param_groups[0]["lr"]
    print(f"epoch={epoch} lr={current_lr:.4f}")
    optimizer.zero_grad(set_to_none=True)
    dummy_loss = sum(parameter.sum() * 0 for parameter in model.parameters())
    dummy_loss.backward()
    optimizer.step()
    scheduler.step()

The printed learning rates are 0.0500, 0.0500, 0.0250, 0.0250, and 0.0125. The scheduler is not training the model by itself; it only edits optimizer hyperparameters. In a full loop, you would normally call an epoch scheduler after the epoch’s training pass.

Example 4: A Tiny PyTorch Training Loop

This example fits y = 2x + 1 with a single linear layer. It uses set_to_none=True, which clears gradients by assigning None instead of zero tensors. That can reduce memory writes and makes accidental use of missing gradients easier to notice.

import torch
from torch import nn

torch.manual_seed(7)
x = torch.tensor([[-1.0], [0.0], [1.0], [2.0]])
y = 2 * x + 1
model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.2)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=6)

for epoch in range(6):
    optimizer.zero_grad(set_to_none=True)
    prediction = model(x)
    loss = nn.functional.mse_loss(prediction, y)
    loss.backward()
    optimizer.step()
    scheduler.step()
    print(epoch, round(loss.item(), 4), round(optimizer.param_groups[0]["lr"], 4))

The exact loss values depend on the seeded initialization and PyTorch numeric kernels, but the pattern should be a decreasing loss and a cosine learning rate that starts below 0.2 after the first scheduler step and reaches 0.0 at the final printed step. Verification is not just seeing numbers print; inspect that loss trends down and the scheduler changes the logged learning rate.

Design Choices and Trade-Offs

SGD with momentum is simple, memory-efficient, and often strong for vision models when tuned carefully. It may need more learning-rate tuning and warmup. AdamW is a common default for transformers and many tabular or language tasks because it adapts per-parameter steps and handles sparse or uneven gradients well. Its optimizer state is larger because it stores two moment tensors per parameter. RMSprop appears in some recurrent or reinforcement learning workloads, but it is less common as a modern default.

A constant learning rate is easy to debug but rarely ideal for long training. Decay schedules let training take large early steps and smaller later refinements. Warmup can prevent early instability when randomly initialized heads, mixed precision, or large batches produce volatile gradients. Cosine decay is smooth and requires a planned training horizon. Reduce-on-plateau reacts to validation metrics, but it is sensitive to noisy validation curves and needs patience settings.

Batch size interacts with learning rate. Larger batches usually produce less noisy gradient estimates and can often tolerate larger learning rates, but they also reduce the number of parameter updates per epoch and may generalize differently. Gradient accumulation imitates a larger batch for memory reasons, but the scheduler cadence must be chosen deliberately: stepping the scheduler every micro-batch is different from stepping after each accumulated optimizer step.

Failure Modes and Troubleshooting

Symptom: loss becomes nan or shoots upward. The common causes are too high a learning rate, bad input scaling, exploding gradients, or mixed-precision overflow. Diagnose by logging the learning rate, loss before backward, gradient norms after backward, and a small sample of inputs and targets. Correct by lowering the learning rate, normalizing inputs, using gradient clipping, checking labels, or temporarily disabling mixed precision to isolate the source.

Symptom: loss barely changes. Causes include learning rate too small, forgotten optimizer.step(), detached tensors, frozen parameters, or zero gradients. Diagnose by confirming each parameter has requires_grad=True, checking that p.grad is not None after backward, and printing one parameter before and after optimizer.step(). Correct the broken graph, unfreeze intended layers, or increase the learning rate carefully.

Symptom: resumed training behaves differently after a checkpoint. The likely cause is missing optimizer or scheduler state. AdamW without restored moments behaves like a fresh optimizer applied to old weights. Diagnose by checking whether the checkpoint includes optimizer state, scheduler state, and the last completed step or epoch. Correct by saving and restoring all state dictionaries together.

Symptom: a scheduler appears one step off. The cause is usually calling scheduler.step() at the wrong point or wrong frequency. Diagnose by logging learning rate at the start of each epoch and after the scheduler step. Correct the order and cadence to match the scheduler’s intended unit.

Performance and Reliability Implications

Optimizer choice affects memory. SGD with momentum stores one extra tensor per parameter; AdamW commonly stores two plus counters. For very large models, optimizer state can exceed model weight memory. Learning-rate instability also wastes compute because failed runs may consume hours before producing unusable checkpoints. Reliable training code logs loss, learning rate, gradient norm, batch size, accumulation factor, random seed, model commit, and checkpoint path so a result can be reproduced or rejected.

Security is mostly indirect but still real in shared training environments. Avoid printing raw private training examples while debugging failed optimization. Store checkpoints where only appropriate users can read them, because optimizer state can contain information derived from training data and exposes the exact model state.

Hands-On Lab: Compare Optimizer Behavior

Prerequisites: Python with PyTorch installed, a shell, and a clean working directory. No GPU is required. The lab trains the tiny linear example twice, once with SGD and once with AdamW, then verifies that both update parameters and that their learning rates are logged.

1. Create a new Python file or notebook cell and import torch plus torch.nn.
2. Build the four-point dataset x=[-1,0,1,2] and y=2x+1.
3. Write a function that creates a fresh nn.Linear(1, 1), optimizer, optional scheduler, and six-step training loop.
4. Run it with torch.optim.SGD at lr=0.2, then with torch.optim.AdamW at lr=0.1.
5. Print epoch, loss, learning rate, weight, and bias each step.
6. Verification: each run should show parameters changing after every optimizer step, finite loss values, and scheduler-controlled learning rates when a scheduler is attached.
7. Cleanup: remove the scratch file or notebook cell outputs if they contain local paths or data samples. No persistent model artifact is needed for this lab.

Assessment Exercises

1. A model’s training loss is flat for 500 steps. Describe three checks that distinguish a too-small learning rate from a broken gradient path.
2. You switch from SGD with momentum to AdamW and run out of memory. Explain why optimizer state can be the cause and name one mitigation.
3. Design parameter groups for fine-tuning a pretrained network with a new classifier head. Which group should usually receive the larger learning rate, and why?
4. A cosine scheduler is stepped every mini-batch in one experiment and every epoch in another. Explain how the effective schedule differs.
5. Write a checkpoint dictionary that can resume training without losing optimizer and scheduler progress.

Summary

PyTorch optimization is a stateful loop: gradients are accumulated by autograd, optimizers mutate parameters using those gradients and their own state, learning rates scale the updates, and schedulers change that scale over time. Good training practice means logging the actual learning rate, choosing optimizer memory and behavior deliberately, saving full training state, and debugging failures by inspecting gradients, parameters, and scheduler cadence rather than guessing.