🎯 What You'll Learn
- Identify the four core failure modes of vanilla SGD: oscillation, flat regions, saddle points, and unequal feature scales
- Understand Exponentially Weighted Moving Average (EWMA) — the mathematical building block of all modern optimizers
- Implement SGD with Momentum from scratch and explain why it accelerates convergence in consistent gradient directions
- Understand RMSProp's adaptive learning rate: parameters with large gradients get smaller steps, and vice versa
- Understand Adam: the combination of momentum (first moment) + adaptive LR (second moment) + bias correction
- Know the difference between Adam and AdamW, and why AdamW is the modern standard for LLMs and vision models
- Choose the right optimizer for your task from a practical decision guide
Imagine rolling a ball down a hilly landscape towards the lowest point. Plain gradient descent is like carefully placing the ball at each step — it only looks at the current slope and moves a tiny fixed distance. Momentum is like giving the ball actual physical momentum — it accumulates velocity from previous steps, rolls faster through flat regions, and can push over small bumps. RMSProp is like a smart cart that adjusts its wheel friction independently in each direction — softer where the terrain is flat, stiffer where it's steep. Adam combines both: the ball has physical momentum AND smart adaptive wheels. That's why Adam is the default optimizer for almost all deep learning — it works well without careful tuning.
1 Why Plain SGD Is Not Enough
The basic SGD update rule is deceptively simple: W ← W − α · g, where g is the gradient. On a simple convex loss function (like linear regression with one parameter), this works beautifully. But real neural network loss landscapes have four problems that plain SGD struggles with.
Problem 1: Oscillation in Narrow Valleys
If the loss landscape is much steeper in one direction than another (like a narrow valley running diagonally), SGD oscillates back and forth across the narrow dimension while making slow progress along the valley. The gradient perpendicular to the valley floor is large (causing wild swings) while the gradient along the floor is small (causing slow progress). Using a larger learning rate worsens the oscillation; using a smaller one just makes everything slow.
Problem 2: Same Learning Rate for All Parameters
SGD uses the same α for every parameter. But parameters receive very different gradient magnitudes. In a network with embedding layers (like word embeddings in NLP), most embeddings are barely updated (they appear rarely) while a few embeddings get large updates every batch. A single α cannot be right for both.
Problem 3: Flat Regions and Saddle Points
In high-dimensional loss landscapes, saddle points — points where the gradient is near zero but it's neither a minimum nor maximum — are far more common than local minima. At a saddle point, plain SGD essentially stalls because the gradient magnitude is tiny. Training appears to have converged when it hasn't.
Problem 4: Noisy Mini-Batch Gradients
Mini-batch gradients are estimates of the true gradient. With a small batch, the noise can be large. SGD takes every noisy step faithfully — zigzagging through the loss landscape instead of moving steadily toward the minimum.
import numpy as np
import torch
import torch.nn as nn
# Illustrate oscillation: the Beale function has narrow valleys
def beale(x, y):
"""A classic test function with a narrow valley — challenges plain SGD."""
return ((1.5 - x + x*y)**2 +
(2.25 - x + x*y**2)**2 +
(2.625 - x + x*y**3)**2)
# Plain SGD trajectory from a bad starting point
x, y = torch.tensor(-3.0, requires_grad=True), torch.tensor(-3.0, requires_grad=True)
lr = 0.001
trajectory = [(x.item(), y.item())]
for step in range(300):
loss = beale(x, y)
loss.backward()
with torch.no_grad():
x -= lr * x.grad
y -= lr * y.grad
x.grad.zero_()
y.grad.zero_()
if step % 50 == 0:
trajectory.append((x.item(), y.item()))
print(f"Step {step:4d} | loss={loss.item():.4f} | x={x.item():.3f} y={y.item():.3f}")
# Minimum is at x=3, y=0.5
print(f"\nFinal: x={x.item():.4f}, y={y.item():.4f}")
print(f"Optimum: x=3.0, y=0.5")Every improvement over vanilla SGD addresses one or more of these problems. Momentum smooths out oscillations and noisy gradients. Adaptive learning rates (RMSProp, Adagrad) solve the unequal-parameter problem. Adam combines both. Understanding this gives you a clear mental model for why each optimizer makes the choices it does.
2 Exponentially Weighted Moving Average (EWMA)
Before we can understand Momentum, RMSProp, or Adam, we need to understand Exponentially Weighted Moving Average (EWMA) — the mathematical primitive that underlies all of them.
The problem EWMA solves: you have a noisy signal (like gradient estimates from mini-batches), and you want a smooth version that represents the trend without reacting to every spike. A simple moving average (average of last N values) works but requires storing N values. EWMA is cleverer: it maintains a running single-number summary that gives more weight to recent values while still retaining memory of the past.
EWMA formula:
v_t = β · v_{t-1} + (1 − β) · x_t
Where:
- v_t is the smoothed value at step t (what we track)
- x_t is the new observation (e.g., a gradient)
- β is the decay factor (typically 0.9, 0.99)
- (1−β) is the weight given to the new observation
Connection to "averaging the last N steps": With β=0.9, the effective window is approximately 1/(1−β) = 10 steps. With β=0.99, it's approximately 100 steps. Recent steps have exponentially higher weight, older steps decay exponentially.
import numpy as np
def ewma(signal, beta=0.9):
"""
Compute Exponentially Weighted Moving Average of a signal.
Returns the smoothed version of the same length.
"""
smoothed = np.zeros_like(signal)
v = 0.0 # initialize to zero
for t, x in enumerate(signal):
v = beta * v + (1 - beta) * x
smoothed[t] = v
return smoothed
# Noisy temperature signal
np.random.seed(42)
days = np.arange(100)
true_temp = 20 + 5 * np.sin(days * 0.2) # underlying trend
noisy_temp = true_temp + np.random.randn(100) * 3 # add noise
smoothed_90 = ewma(noisy_temp, beta=0.9) # fast-moving average (last ~10 days)
smoothed_99 = ewma(noisy_temp, beta=0.99) # slow-moving average (last ~100 days)
# Compare how well each tracks the trend
for t in [0, 20, 50, 80, 99]:
print(f"Day {t:3d}: raw={noisy_temp[t]:6.2f}°C | "
f"EWMA(β=0.9)={smoothed_90[t]:6.2f}°C | "
f"EWMA(β=0.99)={smoothed_99[t]:6.2f}°C | "
f"true={true_temp[t]:6.2f}°C")Notice Day 0: EWMA(β=0.9) = 1.95°C, but the true value is ~20°C. That's the initialization bias — because v starts at 0, early estimates are pulled towards 0. This is exactly the bias that Adam's bias correction fixes (Section 5).
You've already seen EWMA without knowing it: BatchNorm's running mean and running variance (tracked across training batches for use at inference) are computed using EWMA. The momentum parameter in BatchNorm is actually β in the EWMA formula. It's also used in financial time-series analysis, signal processing, and web server load monitoring — anywhere you need a smooth trend from a noisy signal.
Drag the slider to change β and watch the smoothed curve react: low β tracks the noisy raw signal closely (barely smoothing at all), while high β produces a very slow, lagging trend line that averages over many more points.
β = 0.90 — effective window ≈ 10 days. Smooths noise while still tracking the seasonal trend closely.
3 SGD with Momentum: Gradient Smoothing
Now we apply EWMA to the gradients themselves. Instead of using the raw gradient g_t for each update, we use an EWMA of recent gradients — called the velocity v_t.
SGD with Momentum update rules:
v_t = β · v_{t-1} + g_t W ← W − α · v_t
With β=0.9, the velocity is a smoothed version of recent gradients. The physical analogy: a ball rolling downhill accumulates speed. In a region where gradients consistently point in the same direction (like along the floor of a valley), the velocity builds up — the ball accelerates. In a region where gradients oscillate (like across the narrow dimension of a valley), positive and negative contributions cancel out — the oscillations are damped.
Two key effects of momentum:
- Acceleration in consistent directions: If gradients have been pointing left for the past 10 steps, the velocity points strongly left. The effective step size in that direction is larger than α alone.
- Damping in oscillating directions: If gradients alternate +/-, they cancel in the velocity. The effective step size in the oscillating direction is much smaller than α.
import numpy as np
def sgd_momentum(X, y, lr=0.01, momentum=0.9, n_epochs=50):
"""
SGD with Momentum from scratch.
v_t = beta * v_{t-1} + g_t
W = W - lr * v_t
"""
n, p = X.shape
X_b = np.column_stack([np.ones(n), X])
w = np.zeros(p + 1)
v = np.zeros(p + 1) # velocity vector — one per parameter
loss_history = []
for epoch in range(n_epochs):
# Shuffle and iterate mini-batches
perm = np.random.permutation(n)
for start in range(0, n, 32):
batch = perm[start : start + 32]
X_b_b = X_b[batch]
y_b = y[batch]
# Gradient on this batch
y_hat = X_b_b @ w
g = (2 / len(y_b)) * X_b_b.T @ (y_hat - y_b)
# Momentum update: accumulate velocity
v = momentum * v + g
# Parameter update using velocity
w -= lr * v
y_pred_all = X_b @ w
loss = np.mean((y_pred_all - y) ** 2)
loss_history.append(loss)
return w, loss_history
def sgd_plain(X, y, lr=0.01, n_epochs=50):
"""Vanilla SGD for comparison."""
n, p = X.shape
X_b = np.column_stack([np.ones(n), X])
w = np.zeros(p + 1)
loss_history = []
for epoch in range(n_epochs):
perm = np.random.permutation(n)
for start in range(0, n, 32):
batch = perm[start : start + 32]
X_b_b, y_b = X_b[batch], y[batch]
g = (2 / len(y_b)) * X_b_b.T @ (X_b_b @ w - y_b)
w -= lr * g
loss_history.append(np.mean((X_b @ w - y) ** 2))
return w, loss_history
np.random.seed(42)
n = 1000
X = np.random.randn(n, 5)
y = np.array([1.5, -0.8, 2.1, 0.4, -1.2]) @ X.T + 3.0 + np.random.randn(n) * 0.3
_, losses_plain = sgd_plain(X, y, lr=0.01, n_epochs=30)
_, losses_momentum = sgd_momentum(X, y, lr=0.01, momentum=0.9, n_epochs=30)
print(f"{'Epoch':>5} | {'Plain SGD':>10} | {'Momentum':>10}")
for epoch in [0, 5, 10, 20, 29]:
print(f"{epoch+1:5d} | {losses_plain[epoch]:10.4f} | {losses_momentum[epoch]:10.4f}")
print(f"\nMomentum reaches {losses_plain[-1]:.4f} faster (plain SGD needs ~30 epochs)")
print(f"Momentum final loss: {losses_momentum[-1]:.4f}")Nesterov Momentum: Look Before You Leap
Standard momentum computes the gradient at the current position, then steps. Nesterov momentum computes the gradient at where you're going (the position after the momentum step) — a "lookahead" correction. This typically gives slightly faster convergence.
import torch.optim as optim
model = nn.Sequential(nn.Linear(5, 32), nn.ReLU(), nn.Linear(32, 1))
# Standard momentum
optimizer_momentum = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# Nesterov momentum (usually marginally better, especially on convex losses)
optimizer_nesterov = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, nesterov=True)
# Nesterov update (conceptually):
# Compute gradient at (w - lr * momentum * v) ← lookahead position
# v = momentum * v + gradient_at_lookahead
# w = w - lr * v4 RMSProp: Adaptive Learning Rates Per Parameter
Momentum solves the oscillation problem, but both momentum and plain SGD still use the same effective learning rate for all parameters. RMSProp solves this by giving each parameter its own adaptive learning rate — automatically larger for parameters with small gradients, smaller for parameters with large gradients.
The key idea: Track the EWMA of the squared gradient for each parameter. If a parameter has been receiving consistently large gradients, divide its update by a large number (small effective LR). If it's been receiving small gradients, divide by a small number (large effective LR). This self-normalization makes training much more robust to different scales across parameters.
RMSProp update rules:
s_t = β · s_{t-1} + (1−β) · g_t² W ← W − (α / √(s_t + ε)) · g_t
Where s_t is the EWMA of squared gradients and ε (typically 1e-8) prevents division by zero.
import numpy as np
def rmsprop(X, y, lr=0.01, beta=0.99, eps=1e-8, n_epochs=30):
"""
RMSProp from scratch.
s = beta * s + (1-beta) * g^2
W = W - lr * g / sqrt(s + eps)
"""
n, p = X.shape
X_b = np.column_stack([np.ones(n), X])
w = np.zeros(p + 1)
s = np.zeros(p + 1) # EMA of squared gradients — initialized to 0
loss_history = []
for epoch in range(n_epochs):
perm = np.random.permutation(n)
for start in range(0, n, 32):
batch = perm[start : start + 32]
X_b_b, y_b = X_b[batch], y[batch]
B = len(y_b)
# Gradient
g = (2 / B) * X_b_b.T @ (X_b_b @ w - y_b)
# EMA of squared gradient (element-wise squaring)
s = beta * s + (1 - beta) * g ** 2
# Adaptive update: divide by sqrt of EMA squared gradient
w -= lr * g / (np.sqrt(s) + eps)
loss_history.append(np.mean((X_b @ w - y) ** 2))
return w, loss_history
# Sparse gradient scenario: imagine features 0-4 appear at different rates
np.random.seed(42)
n = 2000
X = np.random.randn(n, 5)
X[:, 3] *= 10 # feature 3 has 10x larger scale → its gradient will be 10x larger
y = 1.0 * X[:, 0] + 2.0 * X[:, 1] + 0.5 * X[:, 2] + 0.3 * X[:, 3] + 1.8 * X[:, 4] + 2.0
_, losses_rms = rmsprop(X, y, lr=0.01, n_epochs=30)
_, losses_plain = sgd_plain(X, y, lr=0.001, n_epochs=30) # smaller lr for stability
print("RMSProp vs Plain SGD (with scaled feature):")
for epoch in [0, 5, 10, 20, 29]:
print(f"Epoch {epoch+1:3d}: RMSProp={losses_rms[epoch]:.4f} SGD={losses_plain[epoch]:.4f}")RMSProp converges much faster when features have different scales. For feature 3 (10x scale → 10x gradient), RMSProp automatically reduces the effective LR by √100 ≈ 10, preventing it from dominating the update. Plain SGD has to use a very small LR across the board to avoid overshooting feature 3, which slows down learning for all other features.
import torch.optim as optim
# PyTorch RMSprop
optimizer = optim.RMSprop(
model.parameters(),
lr=0.01, # base learning rate
alpha=0.99, # β — EMA decay for squared gradients
eps=1e-8, # ε — numerical stability
momentum=0.0, # optional: can combine with momentum
centered=False # if True, uses EMA of gradient instead of EMA of squared gradient
)
# RMSprop is particularly good for:
# - Recurrent neural networks (RNNs, LSTMs) — Geoff Hinton's original use case
# - Non-stationary objectives (gradients change significantly over time)
# - Online learning settingsThe ε in the denominator (√s + ε) has a dual purpose: (1) it prevents division by zero when s ≈ 0 (which happens for parameters that have received zero gradient — common with sparse inputs like word embeddings), and (2) it sets a floor on the effective learning rate. Even when s is very large, the effective LR can't go below α/√(s_max + ε). Too small an ε can cause numerical instability; too large an ε reduces the adaptive benefit. The value 1e-8 is the empirically determined sweet spot for float32 arithmetic.
5 Adam: The Best of Both Worlds
Adam (Adaptive Moment Estimation) combines the ideas from both Momentum and RMSProp. It maintains two running statistics per parameter:
- First moment m (EWMA of gradients) — like momentum: smooths the gradient direction
- Second moment v (EWMA of squared gradients) — like RMSProp: adapts the learning rate per parameter
Adam update rules:
m_t = β₁ · m_{t-1} + (1−β₁) · g_t [first moment: EWMA of gradient]
v_t = β₂ · v_{t-1} + (1−β₂) · g_t² [second moment: EWMA of squared gradient]
m̂_t = m_t / (1 − β₁ᵗ) [bias-corrected first moment]
v̂_t = v_t / (1 − β₂ᵗ) [bias-corrected second moment]
W ← W − α · m̂_t / (√v̂_t + ε) [parameter update]
Default hyperparameters: β₁ = 0.9, β₂ = 0.999, ε = 1e-8, α = 0.001
import numpy as np
def adam(X, y, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8, n_epochs=30):
"""
Adam optimizer from scratch.
Maintains per-parameter first and second moment estimates with bias correction.
"""
n, p = X.shape
X_b = np.column_stack([np.ones(n), X])
w = np.zeros(p + 1)
m = np.zeros(p + 1) # first moment (EMA of gradients)
v = np.zeros(p + 1) # second moment (EMA of squared gradients)
t = 0 # global step counter (needed for bias correction)
loss_history = []
for epoch in range(n_epochs):
perm = np.random.permutation(n)
for start in range(0, n, 32):
batch = perm[start : start + 32]
X_b_b, y_b = X_b[batch], y[batch]
B = len(y_b)
# Step counter — increments with EVERY gradient update
t += 1
# Compute gradient on this mini-batch
g = (2 / B) * X_b_b.T @ (X_b_b @ w - y_b)
# Update biased first moment estimate (EMA of gradient)
m = beta1 * m + (1 - beta1) * g
# Update biased second moment estimate (EMA of squared gradient)
v = beta2 * v + (1 - beta2) * g ** 2
# Compute bias-corrected estimates
# These corrections are large when t is small (early steps)
m_hat = m / (1 - beta1 ** t)
v_hat = v / (1 - beta2 ** t)
# Adam parameter update
w -= lr * m_hat / (np.sqrt(v_hat) + eps)
loss_history.append(np.mean((X_b @ w - y) ** 2))
return w, loss_history
# Compare all 4 optimizers
np.random.seed(42)
n = 1000
X = np.random.randn(n, 5)
y = (np.array([1.5, -0.8, 2.1, 0.4, -1.2]) @ X.T) + 3.0 + np.random.randn(n) * 0.3
_, l_sgd = sgd_plain(X, y, lr=0.01, n_epochs=30)
_, l_momentum = sgd_momentum(X, y, lr=0.01, momentum=0.9, n_epochs=30)
_, l_rms = rmsprop(X, y, lr=0.01, n_epochs=30)
_, l_adam = adam(X, y, lr=0.001, n_epochs=30)
print(f"{'Optimizer':>12} | Epoch 1 | Epoch 5 | Epoch 15 | Epoch 30")
print("-" * 57)
for name, losses in [("SGD", l_sgd), ("SGD+Momentum", l_momentum),
("RMSProp", l_rms), ("Adam", l_adam)]:
print(f"{name:>12} | {losses[0]:7.4f} | {losses[4]:7.4f} | {losses[14]:8.4f} | {losses[29]:8.4f}")Adam reaches near-optimal loss by epoch 5, where plain SGD is still at ~1.0. This is why Adam is the default choice when you start any new project — it works well without careful learning rate tuning.
import torch.optim as optim
# PyTorch Adam — the one line you'll write thousands of times
optimizer = optim.Adam(
model.parameters(),
lr=1e-3, # α — base learning rate (1e-3 is the default and a good starting point)
betas=(0.9, 0.999), # (β₁, β₂) — usually left at defaults
eps=1e-8, # ε — numerical stability
weight_decay=0 # L2 regularization coefficient (0 means no regularization)
)Watch All Four Optimizers Race on an Elongated Valley
Numbers in a table are useful, but nothing builds intuition like watching the actual paths. Below is a loss surface L(x, y) = 0.1x² + 2y² — a narrow, elongated valley (20× steeper in y than in x), rendered as a contour plot. All four optimizers start from the exact same point high on the valley wall (x₀, y₀) = (−9, 4.5) and take 60 steps toward the minimum at the origin, using the same update equations you just implemented from scratch above.
All four paths shown. Toggle a button to hide/show that optimizer's trajectory.
Notice the shapes: plain SGD (red) zig-zags sharply back and forth across the narrow y-direction for its first several steps and then creeps along the valley floor in tiny increments — after 60 steps it still hasn't reached the minimum. Momentum (amber) builds up velocity along the valley floor; it overshoots past the minimum once or twice (the "ball" has real inertia, so it doesn't stop exactly on target), but it covers far more ground along x in the same number of steps. RMSProp (teal) divides each dimension's step by its own running gradient magnitude, so it automatically shrinks the steep y-steps and stretches the shallow x-steps — it reaches the minimum the fastest and most smoothly of all four, with essentially no oscillation. Adam (violet) — combining both ideas plus bias correction — also converges quickly, with a touch of the same momentum-style overshoot before it settles.
6 Bias Correction: Why It Matters in Early Training
Both m and v in Adam are initialized to zero. This creates a systematic problem in the early steps: the EWMA is biased towards zero because it's averaging a tiny number of actual observations with a lot of implicit zeros from the initialization.
Let's see the magnitude of this effect numerically:
import numpy as np
# Simulate the first few steps with a constant gradient of 1.0
g = 1.0
beta1, beta2 = 0.9, 0.999
m, v = 0.0, 0.0
print(f"{'Step':>4} | {'Raw m':>8} | {'Raw v':>10} | {'m_hat':>8} | {'v_hat':>10} | {'Ratio m/v':>10}")
print("-" * 65)
for t in range(1, 11):
m = beta1 * m + (1 - beta1) * g
v = beta2 * v + (1 - beta2) * g ** 2
# Bias-corrected estimates
m_hat = m / (1 - beta1 ** t)
v_hat = v / (1 - beta2 ** t)
print(f"{t:4d} | {m:8.4f} | {v:10.6f} | {m_hat:8.4f} | {v_hat:10.6f} | {m_hat/np.sqrt(v_hat+1e-8):10.4f}")Without bias correction: at step 1, m = 0.1 and v = 0.001 — both are hugely underestimated. With bias correction (m_hat, v_hat): both equal 1.0, which is the true value (gradient has been 1.0 all along). The corrected ratio m_hat/√v_hat is exactly 1.0 — consistent with the true gradient. Without correction, it would be 0.1/√0.001 ≈ 3.16, which is a distorted update.
As t increases, β₁ᵗ → 0 and β₂ᵗ → 0, so the bias correction factors (1−β₁ᵗ) and (1−β₂ᵗ) both approach 1 — the correction disappears automatically in later training. The bias correction only has a strong effect in the first ~50 steps for β₁=0.9 and first ~5000 steps for β₂=0.999. This is why very short training runs are more sensitive to whether bias correction is applied.
7 AdamW: The Modern Standard
Adam has a subtle flaw in how it handles L2 regularization (weight decay). In standard gradient descent, L2 regularization adds a penalty λ·‖W‖² to the loss, which produces a gradient term λ·W that gets added to g. In Adam, this extra gradient term gets absorbed into the adaptive scaling — the second moment v adapts to include it, which incorrectly scales the weight decay contribution by the same adaptive factor as the gradient.
The consequence: Parameters with large gradients get their weight decay implicitly reduced (because v is large, dividing by √v reduces the effective weight decay). Parameters with small gradients get stronger weight decay. This is not what you want — weight decay should be applied uniformly.
AdamW fix: Apply weight decay directly to the weights, separately from the gradient-adaptive update. The weight decay is no longer affected by the adaptive scaling.
Adam: W ← W − α · m̂/(√v̂+ε) − α·λ·m̂_decay/(√v̂_decay+ε) [decay coupled to adaptive scaling]
AdamW: W ← W − α · m̂/(√v̂+ε) − α·λ·W [decay decoupled — cleaner]
import torch.optim as optim
# Adam with weight_decay — this is Adam, NOT AdamW behavior
optimizer_adam = optim.Adam(
model.parameters(),
lr=1e-3,
weight_decay=0.01 # L2 penalty, but absorbed into adaptive scaling — wrong
)
# AdamW — decoupled weight decay — the correct way
optimizer_adamw = optim.AdamW(
model.parameters(),
lr=1e-3,
weight_decay=0.01 # applied directly to weights, separate from gradient
)
# Show the difference in a simple experiment
import torch
import torch.nn as nn
torch.manual_seed(42)
model = nn.Sequential(nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 1))
X = torch.randn(100, 10)
y = torch.randn(100, 1)
def train_check(optimizer_class, **kwargs):
torch.manual_seed(42)
m = nn.Sequential(nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 1))
opt = optimizer_class(m.parameters(), **kwargs)
for _ in range(100):
opt.zero_grad()
loss = nn.MSELoss()(m(X), y)
loss.backward()
opt.step()
return loss.item(), sum(p.norm().item() for p in m.parameters())
loss_adam, norm_adam = train_check(optim.Adam, lr=1e-3, weight_decay=0.01)
loss_adamw, norm_adamw = train_check(optim.AdamW, lr=1e-3, weight_decay=0.01)
print(f"Adam final loss={loss_adam:.4f} param_norm={norm_adam:.4f}")
print(f"AdamW final loss={loss_adamw:.4f} param_norm={norm_adamw:.4f}")
# AdamW typically gives smaller param_norm — the weight decay is more effectiveAdamW is the optimizer used by GPT-2, GPT-3, BERT, and virtually every large language model published after 2019. The decoupled weight decay is essential when training billions of parameters — the regularization effect needs to be predictable and uniform across all parameter types, regardless of their gradient histories. Use torch.optim.AdamW as your default in any new project that requires regularization.
8 Comparing Optimizers: When to Use What
There's no single "best" optimizer — the right choice depends on your task, architecture, and resources. Here's a practical decision guide.
| Optimizer | Best For | Tuning Needed | Notes |
|---|---|---|---|
| Adam / AdamW | NLP, Transformers, any new project | Low — lr=1e-3 usually works | ⭐ Default choice. Use AdamW if using weight decay. |
| SGD + Momentum | Image classification (ResNets, CNNs) | High — lr, momentum, schedule | With careful tuning, often beats Adam on accuracy. |
| RMSprop | RNNs, LSTMs, online learning | Medium | Hinton's original use case for RNNs. Less common now. |
| Adagrad | Sparse features, NLP (older) | Low | LR monotonically decreases — learning stalls in long runs. |
The Adam vs SGD Convergence Speed vs Final Accuracy Trade-off
This is one of the most nuanced debates in deep learning. The empirical finding (especially for image classification) is:
- Adam converges faster in early training (lower loss after the same number of steps)
- SGD with momentum + careful LR schedule often achieves better final accuracy on fully-tuned runs
- For most practitioners: Adam's lower tuning requirement makes it the practical winner. The marginal accuracy improvement from tuned SGD is rarely worth the effort.
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# Compare optimizers on MNIST
device = 'cuda' if torch.cuda.is_available() else 'cpu'
def build_model():
return nn.Sequential(
nn.Flatten(), nn.Linear(784, 512), nn.ReLU(), nn.Dropout(0.2),
nn.Linear(512, 256), nn.ReLU(), nn.Dropout(0.2), nn.Linear(256, 10)
).to(device)
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_set = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_set = datasets.MNIST('./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=2)
test_loader = DataLoader(test_set, batch_size=512, shuffle=False)
criterion = nn.CrossEntropyLoss()
configs = [
("Plain SGD", lambda m: optim.SGD(m.parameters(), lr=0.01)),
("SGD+Momentum",lambda m: optim.SGD(m.parameters(), lr=0.01, momentum=0.9)),
("RMSprop", lambda m: optim.RMSprop(m.parameters(), lr=0.001)),
("Adam", lambda m: optim.Adam(m.parameters(), lr=0.001)),
("AdamW", lambda m: optim.AdamW(m.parameters(), lr=0.001, weight_decay=0.01)),
]
for name, opt_fn in configs:
model = build_model()
optimizer = opt_fn(model)
for epoch in range(5):
model.train()
for X_b, y_b in train_loader:
X_b, y_b = X_b.to(device), y_b.to(device)
optimizer.zero_grad()
loss = criterion(model(X_b), y_b)
loss.backward()
optimizer.step()
model.eval()
correct = sum((model(X.to(device)).argmax(1) == y.to(device)).sum().item()
for X, y in test_loader)
acc = correct / len(test_set) * 100
print(f"{name:15s}: Test Accuracy = {acc:.2f}%")Real-World Spotlight: 4 Optimizers, 1 Network, MNIST
The practical takeaway: use AdamW as your default. The time you save not tuning learning rates and schedules is worth more than the marginal accuracy improvement you might get from perfectly tuned SGD+momentum. When you're near the end of a project and optimizing the last 0.5% of accuracy, then revisit SGD+momentum with a tuned schedule.
Quick Check
✍️ Practice Exercises
- Implement Adagrad from scratch. Its update rule is: s = s + g², W = W − lr * g / (√s + ε). Note: unlike RMSProp, Adagrad sums (not averages) the squared gradients — this means the effective LR monotonically decreases. Show empirically that Adagrad's learning stalls after many steps while Adam's doesn't.
- Visualize the convergence paths of SGD, Momentum, RMSProp, and Adam on the Rosenbrock function f(x,y) = (1−x)² + 100(y−x²)² using matplotlib. Start all from the same point (-1, 1) and run for 500 steps. Compare the paths.
- Implement the bias correction effect demonstration: train a tiny model for just 10 steps with Adam, once with bias correction enabled (standard) and once with bias correction disabled (set m_hat=m and v_hat=v). Plot the loss curves and show that bias correction helps in early steps.
- Run a grid search over Adam learning rates [1e-4, 3e-4, 1e-3, 3e-3, 1e-2] on a dataset of your choice. Plot test accuracy at epoch 10 for each LR. Show that Adam is relatively robust to LR choice compared to SGD (which has a much narrower "good" LR range).
📚 Primary Sources for This Lesson
Adam: A Method for Stochastic Optimization (Kingma & Ba, 2014) — the original Adam paper. Short, readable, with clear derivations.
Decoupled Weight Decay Regularization (Loshchilov & Hutter, 2017) — the AdamW paper explaining why decoupled weight decay is more principled.
CS231n: Neural Networks Part 3 — Learning and Evaluation — Stanford's excellent notes covering all optimizers with visualisations of their convergence paths.