Optimizers

Learning Objectives

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

1. Understand the optimizer landscape: Know the key optimizers (SGD, Momentum, RMSprop, Adam) and the problems each one solves
2. Derive optimizer update rules: Understand the mathematical formulation behind each optimizer and what each term contributes
3. Apply physical intuition: Connect momentum and adaptive learning rates to physics concepts like velocity, friction, and acceleration
4. Implement optimizers in PyTorch: Use PyTorch's optimizer API and understand when to choose each optimizer
5. Diagnose optimization problems: Recognize when training issues stem from optimizer choice and how to fix them

Why This Matters: The optimizer determines how your network learns. A well-chosen optimizer can mean the difference between a model that converges in hours versus days—or one that never converges at all. Understanding optimizers is essential for training deep networks effectively.

The Big Picture

The Problem We're Solving

In the previous section, we established the training loop: forward pass, compute loss, backward pass, update weights. The optimizer is the algorithm that performs that final step—deciding exactly how to adjust each weight based on its gradient.

The simplest approach is vanilla gradient descent: move each weight in the direction opposite to its gradient, scaled by a learning rate. But this simple approach has serious problems:

• Oscillation: In ravines (long, narrow valleys in the loss landscape), gradients point steeply across the valley but only gently along it. SGD oscillates back and forth across the ravine instead of moving efficiently toward the minimum.
• Uniform learning rates: All parameters get the same learning rate, but some features are sparse (appear rarely) while others are dense. Sparse features need larger updates when they appear.
• Local minima and saddle points: With no momentum, the optimizer can get stuck in suboptimal regions.

A Brief History

The evolution of optimizers is a story of progressively solving these problems:

Today, Adam is the default choice for most practitioners, though SGD with momentum remains competitive for some tasks (especially vision models). Let's understand why by building up from the basics.

Vanilla Stochastic Gradient Descent

The simplest optimizer updates each parameter by stepping in the direction opposite to the gradient:

Where:

• $\theta_t$ is the parameter value at step $t$
• $\eta$ (eta) is the learning rate
• $g_t = \nabla_\theta \mathcal{L}(\theta_t)$ is the gradient of the loss

Why "Stochastic"?

In true gradient descent, we would compute the gradient over the entire dataset. Stochastic gradient descent (SGD) instead computes gradients on small batches, introducing noise into the gradient estimate. This noise:

• Makes training faster (we update after each batch, not each epoch)
• Acts as regularization (noise helps escape local minima)
• Reduces memory requirements (we only need one batch in memory)

The Problem: Ravines

Consider a loss function shaped like a long, narrow valley. The gradient perpendicular to the valley (across it) is large, while the gradient along the valley (toward the minimum) is small.

What happens with vanilla SGD? The optimizer takes large steps across the valley and small steps along it. Result: it oscillates back and forth across the valley, making slow progress toward the minimum.

Momentum: Learning from Physics

Momentum solves the oscillation problem by borrowing an idea from physics: a ball rolling down a hill accumulates velocity. If the hill curves back and forth, the ball's inertia smooths out the oscillations.

The Physical Analogy

Imagine a ball rolling on the loss surface. Without friction, the ball would accelerate indefinitely. With friction, it reaches a terminal velocity where the gravitational force (gradient) balances the friction.

Momentum introduces a velocity term that accumulates across updates:

Where:

• $v_t$ is the velocity at step $t$
• $\mu$ (mu) is the momentum coefficient, typically 0.9
• $g_t$ is the current gradient

Why Momentum Works

In a ravine:

• Across the valley: Gradients alternate direction (left, right, left, right). The momentum term averages these out, canceling the oscillation.
• Along the valley: Gradients consistently point toward the minimum. Momentum accumulates, accelerating progress.

The velocity is an exponentially weighted moving average of past gradients. Recent gradients have more influence (weight $\mu^0 = 1$) than older ones (weight $\mu^t \to 0$ as $t \to \infty$).

Nesterov Accelerated Gradient (NAG)

A clever variant of momentum is Nesterov momentum, which "looks ahead" before computing the gradient:

Instead of computing the gradient at the current position, NAG computes it at the position we would reach if we followed momentum alone. This provides a correction: if momentum would overshoot, we notice sooner and slow down.

Adaptive Learning Rates

Momentum solves oscillation, but what about the uniform learning rate problem? Not all parameters should be updated at the same rate:

• Sparse features: Some features (like rare words in NLP) appear infrequently. When they do appear, we want to learn from them quickly.
• Different scales: Weights in different layers may have gradients of different magnitudes. A learning rate that works for one layer may be too large or too small for another.

AdaGrad: Accumulate Squared Gradients

AdaGrad (Adaptive Gradient) adapts the learning rate for each parameter based on how much that parameter has been updated:

Where:

• $r_t$ is the accumulated squared gradient (one value per parameter)
• $\odot$ denotes element-wise multiplication
• $\epsilon$ is a small constant (e.g., $10^{-8}$) to prevent division by zero

Intuition: Parameters with large accumulated gradients get smaller learning rates; parameters with small accumulated gradients get larger learning rates. This naturally handles sparse features.

RMSprop: Exponential Moving Average

RMSprop (Root Mean Square Propagation) fixes AdaGrad's problem by using an exponentially weighted moving average instead of a sum:

Where $\beta$ is typically 0.9 or 0.99.

Key difference: Unlike AdaGrad, RMSprop "forgets" old gradients. If a parameter stops receiving updates, its accumulated gradient decays, and its learning rate can recover.

Adam: The Best of Both Worlds

Adam (Adaptive Moment Estimation) combines the best ideas from momentum and RMSprop: it maintains both a first moment (momentum-like) and a second moment (RMSprop-like) of the gradients.

The Algorithm

Default hyperparameters (from the original paper):

• $\beta_1 = 0.9$ (momentum decay)
• $\beta_2 = 0.999$ (RMSprop decay)
• $\epsilon = 10^{-8}$
• $\eta = 0.001$

Understanding Each Component

First Moment (m): Momentum

The first moment $m_t$ is an exponentially weighted average of past gradients—exactly like momentum. It smooths out oscillations and accumulates velocity in consistent directions.

Second Moment (v): Adaptive Learning Rate

The second moment $v_t$ is an exponentially weighted average of squared gradients—exactly like RMSprop. It provides per-parameter adaptive learning rates.

Bias Correction

Why the correction terms $\hat{m}_t$ and $\hat{v}_t$? Both $m_t$ and $v_t$ are initialized to zero. Early in training, they're biased toward zero.

Consider $m_t$ after one step:

This is much smaller than $g_1$! The bias correction divides by $1 - \beta_1^1 = 0.1$, giving $\hat{m}_1 = g_1$. As $t \to \infty$, $\beta_1^t \to 0$, and the correction vanishes.

Interactive: Optimizer Comparison

Watch different optimizers navigate the same loss landscape. The Rosenbrock function creates a curved valley—a challenging test for optimizers. Observe how:

• SGD oscillates dramatically, struggling in the curved valley
• Momentum smooths the path but may overshoot
• RMSprop adapts to the landscape's shape
• Adam combines smooth movement with adaptation

Other Notable Optimizers

AdamW: Adam with Decoupled Weight Decay

The original Adam paper applied L2 regularization by adding the weight penalty to the loss, which means it gets scaled by the adaptive learning rate. AdamW decouples weight decay, applying it directly to the weights:

This seemingly small change improves generalization significantly. AdamW is now the default for training transformers and many other architectures.

NAdam: Nesterov + Adam

NAdam incorporates Nesterov momentum into Adam, computing the gradient at a "look-ahead" position. This can provide faster convergence in some cases.

Lion: Evolved Sign Momentum

Lion (2023) was discovered via automated search over optimizer algorithms. It uses only the sign of the momentum, not the magnitude:

Lion uses less memory than Adam (no second moment) and has shown strong performance on large language models.

Summary Table

PyTorch Implementation

PyTorch provides all major optimizers in torch.optim. Here's how to use them:

Basic Usage

Good for computer vision tasks when carefully tuned

1import torch
2import torch.nn as nn
3import torch.optim as optim
4
5# Create a simple model
6model = nn.Sequential(
7    nn.Linear(784, 256),
8    nn.ReLU(),
9    nn.Linear(256, 10)
10)
11
12# SGD with momentum
13optimizer_sgd = optim.SGD(
14    model.parameters(),
15    lr=0.01,
16    momentum=0.9,
17    nesterov=True  # Use Nesterov momentum
18)
19
20# Adam (the default choice)
21optimizer_adam = optim.Adam(
22    model.parameters(),
23    lr=0.001,
24    betas=(0.9, 0.999),  # (β₁, β₂)
25    eps=1e-8
26)
27
28# AdamW (Adam with decoupled weight decay)
29optimizer_adamw = optim.AdamW(
30    model.parameters(),
31    lr=0.001,
32    weight_decay=0.01  # L2 regularization strength
33)

Per-Parameter Learning Rates

Sometimes you want different learning rates for different parts of the model. PyTorch supports this with parameter groups:

1# Different learning rates for different layers
2optimizer = optim.Adam([
3    {'params': model.backbone.parameters(), 'lr': 1e-5},  # Pretrained: small LR
4    {'params': model.head.parameters(), 'lr': 1e-3}      # New layers: larger LR
5], lr=1e-4)  # Default LR for any params not in groups
6
7# Fine-tuning a pretrained model
8optimizer = optim.AdamW([
9    {'params': model.encoder.parameters(), 'lr': 1e-5, 'weight_decay': 0.01},
10    {'params': model.decoder.parameters(), 'lr': 1e-4, 'weight_decay': 0.01},
11    {'params': model.classifier.parameters(), 'lr': 1e-3, 'weight_decay': 0.0}
12])

Implementing Adam from Scratch

Understanding optimizers deeply means being able to implement them yourself. Here's Adam from scratch:

1class AdamFromScratch:
2    def __init__(self, params, lr=0.001, betas=(0.9, 0.999), eps=1e-8):
3        self.params = list(params)
4        self.lr = lr
5        self.beta1, self.beta2 = betas
6        self.eps = eps
7        self.t = 0
8
9        # Initialize moment estimates
10        self.m = [torch.zeros_like(p) for p in self.params]  # First moment
11        self.v = [torch.zeros_like(p) for p in self.params]  # Second moment
12
13    def zero_grad(self):
14        for p in self.params:
15            if p.grad is not None:
16                p.grad.zero_()
17
18    def step(self):
19        self.t += 1
20
21        for i, p in enumerate(self.params):
22            if p.grad is None:
23                continue
24
25            g = p.grad.data
26
27            # Update biased first moment estimate
28            self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * g
29
30            # Update biased second moment estimate
31            self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * g * g
32
33            # Compute bias-corrected estimates
34            m_hat = self.m[i] / (1 - self.beta1 ** self.t)
35            v_hat = self.v[i] / (1 - self.beta2 ** self.t)
36
37            # Update parameters
38            p.data -= self.lr * m_hat / (torch.sqrt(v_hat) + self.eps)

Choosing the Right Optimizer

With so many optimizers available, how do you choose? Here are practical guidelines:

Decision Tree

1. Starting a new project? Use Adam with default hyperparameters. It's robust and works well across domains.
2. Training a transformer or LLM? Use AdamW. The decoupled weight decay is important for these architectures.
3. Training a CNN for image classification? Consider SGD with momentum and a learning rate schedule. It often gives better final performance than Adam, though it requires more tuning.
4. Memory constrained? Consider Lion or SGD with momentum. They use less memory than Adam.
5. Training an RNN? RMSprop or Adam work well. Gradient clipping is often needed too.

Common Mistakes

Summary

Optimizers determine how neural networks learn from gradients. Let's review the key concepts:

Key Takeaways

• Momentum smooths oscillations by accumulating velocity across updates
• Adaptive learning rates handle parameters of different scales and sparsities
• Adam combines both ideas and works well out-of-the-box
• AdamW is preferred for transformers due to proper weight decay handling
• SGD + momentum can achieve better final performance with careful tuning

1# Quick reference: PyTorch optimizers
2import torch.optim as optim
3
4# For quick prototyping (most projects)
5opt = optim.Adam(model.parameters(), lr=1e-3)
6
7# For transformers and LLMs
8opt = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
9
10# For CNNs when tuning for best performance
11opt = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, nesterov=True)
12
13# For RNNs
14opt = optim.RMSprop(model.parameters(), lr=1e-3)

Exercises

Conceptual Questions

1. Explain why momentum helps with ravines in the loss landscape. What would happen if you used momentum coefficient $\mu = 0.99$ instead of $0.9$?
2. Why does Adam need bias correction but vanilla momentum doesn't? What would happen to Adam's first steps without bias correction?
3. You're training a model and notice that loss oscillates wildly. Which optimizer hyperparameters would you adjust, and how?
4. Explain why AdamW is preferred over Adam for transformers. What goes wrong with standard Adam + weight decay?

Coding Exercises

1. Implement SGD with momentum from scratch: Create a custom optimizer class that implements SGD with momentum (with optional Nesterov). Verify it produces the same results as PyTorch's optim.SGD.
2. Optimizer comparison experiment: Train a small CNN on MNIST with SGD, SGD+momentum, RMSprop, and Adam. Plot training and validation loss curves. Which converges fastest? Which achieves the best final accuracy?
3. Learning rate sensitivity analysis: Train the same model with Adam using learning rates [0.0001, 0.001, 0.01, 0.1]. Repeat with SGD. Compare how sensitive each optimizer is to learning rate choice.
4. Visualize optimizer trajectories: Create a 2D loss surface (like the Rosenbrock function) and animate the trajectories of different optimizers. Highlight where each optimizer struggles or excels.

In the next section, we'll explore learning rate strategies—techniques for changing the learning rate during training. You'll learn about warmup, cosine annealing, and other schedules that can significantly improve training.