Module 1 · Reading 1.6 · read before S06: Optimizers

The PDF posted in Canvas is the assigned version of this reading. It is built from this page, so the text is identical; anything here marked as deep reading is optional and is not in the PDF.

Gradient descent, as introduced in loss functions and gradient descent, steps downhill, proportional to the gradient. Deep-learning losses can curve much more steeply in some directions than others, and minibatch gradients add noise. A learning rate small enough for a steep direction may make little progress in a shallow one. The sequence SGD → momentum → RMSProp → Adam → AdamW organizes different responses: momentum remembers direction, RMSProp rescales coordinates, Adam combines those two ideas, and AdamW separates weight decay from the adaptive update. RMSProp does not automatically retain momentum; these are related ideas, not a sequence in which every method contains all earlier methods.

Optimizer choice is not a detail. The same network on the same data can train to convergence or sit at chance depending on the optimizer and learning rate, and papers you will read throughout this course quietly assume you know why AdamW plus a warmup-cosine schedule is a common Transformer recipe. This section builds each method from its update rule, then races them on two problems small enough that you can rerun every experiment in seconds.

How to read this

What goes wrong with plain gradient descent

Write the update as

\[w_{t+1} = w_t - \eta g_t\]

where \(g_t\) is the gradient of the current minibatch loss and \(\eta\) is the learning rate (lr in code). Two separate problems limit lr.

Curvature (ill-conditioning). Consider the positive-definite quadratic \(L(w)=\tfrac12(w_0^2+25w_1^2)\). Along a direction with curvature \(c\gt0\), a gradient-descent step multiplies the displacement from the minimum by \(1-\eta c\). Convergence therefore requires \(\lvert1-\eta c\rvert\lt1\), or \(0\lt\eta\lt2/c\). Our steep direction imposes \(\eta\lt2/25=0.08\), whereas the shallow direction alone would permit \(\eta\lt2\). This is a 25-fold restriction on the allowable learning rate, not a claim that every error curve takes exactly 25 times as many steps. Real networks need not have positive curvature everywhere, so this bound is a local quadratic guide, not a global guarantee.

Noise. With minibatches, g is a random estimate of the true gradient. This is stochastic gradient descent, and its convergence theory dates to Robbins and Monro’s 1951 work on stochastic approximation — decades before deep learning. Noise means individual steps partially cancel, and it is one reason decreasing the learning rate over time helps: early on, noise is small relative to the signal; near a minimum, it dominates.

For a positive-definite quadratic, the ratio of largest to smallest curvature is the condition number \(\kappa\). With the optimal fixed learning rate \(2/(c_{\min}+c_{\max})\), ordinary gradient descent contracts the parameter-error norm by at most \((\kappa-1)/(\kappa+1)\) per step. For \(\kappa=25\) this is \(24/26\approx0.923\), so about 9 steps halve that error bound; the loss-gap bound contracts by its square. Optimally tuned heavy-ball momentum has asymptotic spectral factor \((\sqrt\kappa-1)/(\sqrt\kappa+1)=2/3\) on this quadratic class, giving the familiar square-root improvement in conditioning. It is not a bound on each individual step: transient overshoot and polynomial factors can occur. The experiment below uses a shared learning rate and momentum 0.9, not these theoretically optimal settings.

With a stable constant learning rate and gradient noise that remains nonzero near a minimum, SGD typically fluctuates around the minimum instead of settling exactly at it. This is the noise-ball picture; its scale depends on the learning rate, noise, and curvature. It is not universal: noiseless gradient descent, or problems whose minibatch gradients all vanish at a shared solution, can converge with a constant rate. Classical stochastic-approximation convergence results use conditions such as \(\sum_t\eta_t=\infty\) and \(\sum_t\eta_t^2\lt\infty\), together with assumptions on the objective and noise; \(\eta_t\propto1/t\) is one example. Practical schedules negotiate a similar tradeoff between making progress and reducing late fluctuations.

The lineage that follows attacks curvature with memory (momentum), attacks anisotropy with per-parameter scaling (AdaGrad, RMSProp), and finally combines both (Adam).

Momentum: remembering where you were going

Momentum, associated with Boris Polyak’s 1964 “heavy ball” method, replaces the raw gradient with an exponentially weighted sum of past gradients:

\[v_t=\beta v_{t-1}+g_t,\qquad w_{t+1}=w_t-\eta v_t,\qquad v_0=0.\]

Here beta is typically 0.9. With a constant gradient \(g\) and zero initial velocity, \(v_t=g(1-\beta^t)/(1-\beta)\) approaches \(10g\). This is a steady-gradient result, not an upper bound relative to the current gradient when gradients change. Alternating gradient signs can cancel in the velocity, but a large learning rate or momentum can still cause overshoot. This convention is an unnormalized sum; an exponential average with a factor \(1-\beta\) on the new gradient needs a corresponding learning-rate conversion. Sutskever and colleagues (2013) demonstrated the importance of carefully chosen initialization and momentum schedules for deep and recurrent networks. In PyTorch, the basic heavy-ball variant is torch.optim.SGD(..., momentum=0.9).

Nesterov momentum: looking before you leap

Running the experiments. Run all cells in order in a fresh kernel. Each experiment resets its seeds and creates fresh parameters and optimizer state. Printed numbers are reference outputs; training trajectories can vary with software and hardware. Compare methods within the same environment. The bowl uses full gradients; the first ring comparison uses full-batch training; the warmup and schedule experiments use minibatches. All ring metrics are training metrics.

Heavy-ball momentum can overshoot near a minimum because its accumulated velocity continues moving before the gradient corrects it. Nesterov momentum evaluates the gradient at a lookahead point. In one common convention, form \(q_t=w_t-\eta\beta v_{t-1}\), evaluate \(g_t=\nabla L(q_t)\), then update \(v_t=\beta v_{t-1}+g_t\) and \(w_{t+1}=w_t-\eta v_t\). The gradient can therefore respond to where momentum is taking the iterate. This still needs one gradient evaluation per step. Nesterov acceleration has optimal-rate results in specified smooth convex settings; the fixed-momentum recipe used here is not a general guarantee for deep networks.

PyTorch uses the lookahead iterate as its stored parameter. With constant learning rate and momentum, zero dampening, and no weight decay, its rule is \(v_t=\beta v_{t-1}+g_t\) followed by \(q_{t+1}=q_t-\eta(g_t+\beta v_t)\), where \(g_t=\nabla L(q_t)\). This follows by substituting \(q_t=w_t-\eta\beta v_{t-1}\) into the lookahead rule above. To match PyTorch step for step, compare the same stored iterate and initialize its buffer consistently. The following experiment compares nesterov=True with heavy-ball momentum at the same learning rate. 

import torch

torch.manual_seed(0)

# The same bent bowl as the from-scratch race, now through torch.optim.SGD,
# comparing heavy-ball momentum against its Nesterov variant.
curv = torch.tensor([1.0, 25.0])

def loss_fn(w):
    return 0.5 * (curv * w * w).sum()

def run(nesterov):
    w = torch.nn.Parameter(torch.tensor([2.0, 2.0]))
    opt = torch.optim.SGD([w], lr=0.03, momentum=0.9, nesterov=nesterov)
    out = {}
    for t in range(1, 201):
        opt.zero_grad()
        loss = loss_fn(w)
        loss.backward()
        opt.step()
        if t in (10, 30, 50, 100, 200):
            out[t] = loss_fn(w).item()
    return out

checkpoints = (10, 30, 50, 100, 200)
print(f"{'loss after step:':>18s}" + "".join(f"{t:>12d}" for t in checkpoints))
for name, nesterov in [("heavy-ball", False), ("Nesterov", True)]:
    h = run(nesterov)
    print(f"{name:>18s}" + "".join(f"{h[t]:12.6f}" for t in checkpoints))

Running this prints:

  loss after step:          10          30          50         100         200
        heavy-ball   20.022793    2.325771    0.170028    0.000155    0.000000
          Nesterov    0.001276    0.000577    0.000210    0.000001    0.000000

The heavy-ball row reproduces the early overshoot in the NumPy race: its loss is about 20 at step 10. Nesterov reaches about 0.001 at that checkpoint with the same learning rate and momentum. This demonstrates damping on this particular quadratic and setting; it does not establish that Nesterov always converges faster or remains stable at every learning rate.

Adaptive scaling: AdaGrad, RMSProp, Adam

Momentum still applies one global learning rate. Adaptive methods scale each coordinate by its gradient history. AdaGrad (Duchi, Hazan, and Singer, 2011) uses \(r_t=r_{t-1}+g_t^2\) and divides the update by \(\sqrt{r_t}+\epsilon\). The coefficient cannot increase and tends to zero if the accumulated squared gradients grow without bound, which can make long runs progress slowly. RMSProp replaces the cumulative sum with an exponential average (all squares, roots, and divisions below are coordinatewise):

\[r_t=\beta_2r_{t-1}+(1-\beta_2)g_t^2,\qquad w_{t+1}=w_t-\eta\frac{g_t}{\sqrt{r_t}+\epsilon}.\]

For a steady nonzero gradient, after the average has settled and when epsilon is negligible, the magnitude of \(g_t/\sqrt{r_t}\) is near 1, so the step magnitude is near lr. This is not true at every step: from \(r_0=0\) with \(\beta_2=0.99\), the first RMSProp update has magnitude approximately \(10\eta\). Coordinatewise scaling also cannot in general remove correlations in a rotated loss surface.

Adam (Kingma and Ba, ICLR 2015) combines a first-moment average with a second-moment average:

\[m_t=\beta_1m_{t-1}+(1-\beta_1)g_t,\qquad v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2,\] \[\hat m_t=\frac{m_t}{1-\beta_1^t},\qquad \hat v_t=\frac{v_t}{1-\beta_2^t},\qquad w_{t+1}=w_t-\eta\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}.\]

The defaults are beta1 = 0.9, beta2 = 0.999, and eps = 1e-8, with both buffers initialized to zero and \(t\) starting at 1. Bias correction compensates for that zero initialization; under stationary gradient moments, it removes their initialization bias. It does not mean that the uncorrected update is necessarily too small. At the first step, \(m_1=0.1g_1\) and \(v_1=0.001g_1^2\). Ignoring epsilon, the uncorrected ratio has magnitude \(0.1/\sqrt{0.001}\approx3.16\), whereas the corrected ratio has magnitude 1. Correcting both moments makes the first step approximately lr, not 3.16 times lr. For a nearly constant nonzero gradient, the corrected ratio stays near unit magnitude; changing gradients can produce other values, so this is not a step-size cap.

Historical note. RMSProp was never published in a paper. Geoffrey Hinton introduced it in Lecture 6 of his 2012 Coursera course “Neural Networks for Machine Learning,” and for years the standard citation was literally a lecture slide (usually written “Tieleman & Hinton, 2012”). It became a workhorse optimizer (used, for example, in early deep reinforcement learning work) while existing only as course material, a reminder of how informally this field’s tooling spread.

Racing the update rules on a bent bowl

The four update rules above fit in a few lines each, so we implement them from scratch and race them on the ill-conditioned quadratic from earlier: curvature 1 in w[0], 25 in w[1], minimum at the origin, everyone starting at (2, 2) with the same learning rate. 

import numpy as np

np.random.seed(0)

# An ill-conditioned quadratic bowl: steep in w[1], shallow in w[0].
# L(w) = 0.5 * (1 * w[0]^2 + 25 * w[1]^2), minimum at (0, 0).
curv = np.array([1.0, 25.0])

def loss(w):
    return 0.5 * np.sum(curv * w * w)

def grad(w):
    return curv * w

w0 = np.array([2.0, 2.0])
steps = 200
checkpoints = (10, 50, 100, 200)

def run(update, state_init):
    w, state = w0.copy(), state_init()
    history = {}
    for t in range(1, steps + 1):
        g = grad(w)
        w, state = update(w, g, state, t)
        if t in checkpoints:
            history[t] = loss(w)
    return history

lr = 0.03

def sgd(w, g, s, t):
    return w - lr * g, s

def momentum(w, g, s, t):           # heavy-ball, beta = 0.9
    v = 0.9 * s["v"] + g
    return w - lr * v, {"v": v}

def rmsprop(w, g, s, t):            # beta2 = 0.99
    r = 0.99 * s["r"] + 0.01 * g * g
    return w - lr * g / (np.sqrt(r) + 1e-8), {"r": r}

def adam(w, g, s, t):               # beta1 = 0.9, beta2 = 0.999
    m = 0.9 * s["m"] + 0.1 * g
    v = 0.999 * s["v"] + 0.001 * g * g
    m_hat = m / (1 - 0.9 ** t)      # bias correction
    v_hat = v / (1 - 0.999 ** t)
    return w - lr * m_hat / (np.sqrt(v_hat) + 1e-8), {"m": m, "v": v}

print(f"start loss = {loss(w0):.4f},  lr = {lr}")
print(f"{'loss after step:':>16s}" + "".join(f"{t:>12d}" for t in checkpoints))
for name, update, init in [
    ("SGD", sgd, dict),
    ("momentum", momentum, lambda: {"v": 0.0}),
    ("RMSProp", rmsprop, lambda: {"r": 0.0}),
    ("Adam", adam, lambda: {"m": 0.0, "v": 0.0}),
]:
    h = run(update, init)
    print(f"{name:>16s}" + "".join(f"{h[t]:12.6f}" for t in checkpoints))

Running this prints:

start loss = 52.0000,  lr = 0.03
loss after step:          10          50         100         200
             SGD    1.087589    0.095105    0.004522    0.000010
        momentum   20.022798    0.170029    0.000155    0.000000
         RMSProp    7.321635    0.014481    0.000001    0.000000
            Adam   37.642532    6.077811    0.120019    0.000002

Read this table carefully, because it undercuts the folklore “Adam is always fastest.”

  • SGD kills the steep direction quickly, then spends the rest of the run inching along the shallow one — the ill-conditioning tax.
  • Momentum overshoots badly at first (loss 20 at step 10 — the heavy ball swings past the minimum) but its accumulated velocity then wins decisively.
  • RMSProp rescales both directions to similar effective steps and is the fastest to reach 1e-6 territory.
  • Adam is slowest early in this run. With a consistent gradient its normalized per-coordinate steps are roughly lr = 0.03, so covering a distance of 2 takes on the order of 67 such steps; its momentum then causes some orbiting around the minimum before it settles. This is an explanation of the observed trajectory, not a universal upper bound on Adam’s update.

The ranking depends on when you stop and on the shared learning rate, which is the honest, general lesson: optimizer comparisons are always joint statements about the method and its hyperparameters and the budget.

SGD versus Adam on a real network

Quadratic bowls have no nonlinearity and no data. The second experiment trains an actual MLP on a synthetic “ring” classification problem (label 1 if a point lies in an annulus), comparing plain SGD, SGD with momentum, and Adam at its default learning rate. Every run starts from identical initial weights. Loss and accuracy are measured on the training set; there is no held-out test set, so this experiment compares optimization at these settings, not generalization. 

import torch
import torch.nn as nn

torch.manual_seed(0)

# Synthetic "ring" classification: class 1 if the point lies in an annulus.
n = 600
X = torch.randn(n, 2) * 1.5
r = X.norm(dim=1)
y = ((r > 1.0) & (r < 2.5)).long()

def make_model():
    torch.manual_seed(1)                     # identical init for every optimizer
    return nn.Sequential(
        nn.Linear(2, 32), nn.ReLU(),
        nn.Linear(32, 32), nn.ReLU(),
        nn.Linear(32, 2),
    )

def train(opt_name):
    model = make_model()
    if opt_name == "SGD":
        opt = torch.optim.SGD(model.parameters(), lr=0.1)
    elif opt_name == "SGD+mom":
        opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
    elif opt_name == "Adam":
        opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    else:
        raise ValueError(f"unknown optimizer: {opt_name}")
    loss_fn = nn.CrossEntropyLoss()
    checkpoints = {}
    for epoch in range(1, 401):
        opt.zero_grad()
        loss = loss_fn(model(X), y)
        loss.backward()
        opt.step()
        if epoch in (50, 100, 200, 400):
            with torch.no_grad():
                checkpoints[epoch] = loss_fn(model(X), y).item()  # AFTER update
    with torch.no_grad():
        acc = (model(X).argmax(dim=1) == y).float().mean().item()
    return checkpoints, acc

print(f"{'train loss after epoch:':>24s}{50:>10d}{100:>10d}{200:>10d}{400:>10d}{'train acc':>10s}")
for name in ("SGD", "SGD+mom", "Adam"):
    cp, acc = train(name)
    row = "".join(f"{cp[e]:10.4f}" for e in (50, 100, 200, 400))
    print(f"{name:>24s}{row}{acc:10.3f}")

Running this prints:

 train loss after epoch:        50       100       200       400 train acc
                     SGD    0.6645    0.6258    0.5427    0.2681     0.947
                 SGD+mom    0.4618    0.0879    0.0342    0.0146     1.000
                    Adam    0.6471    0.5664    0.2873    0.0899     0.992

Plain SGD at lr = 0.1 is still making progress at epoch 400. Adding momentum reaches lower training loss and perfect training accuracy in this run. Adam at 1e-3 lies between them. We did not tune each method over an equal search budget, so this table does not establish the best achievable result for any optimizer. Wilson and colleagues (2017) studied a separate question, held-out generalization, and found cases where SGD generalized better than adaptive methods even when the latter fit training data well. Our training-only table cannot test that claim. Adam and AdamW remain useful starting points across varied gradient scales, but their learning rates still need tuning.

People. Diederik P. Kingma and Jimmy Ba introduced Adam at ICLR 2015. Kingma is also a co-inventor of the variational autoencoder studied later in this course. Ilya Loshchilov and Frank Hutter introduced decoupled weight decay in a 2017 preprint, published at ICLR 2019. Their paper makes a precise distinction between an L2 penalty in the loss and a direct shrinkage step on the parameters.

AdamW: weight decay done correctly

L2 regularization and multiplicative weight decay are equivalent for plain SGD without momentum: adding \(\tfrac\lambda2\lVert w\rVert^2\) to the loss adds \(\lambda w\) to the gradient, yielding \(w_{t+1}=(1-\eta\lambda)w_t-\eta g_t\). With Adam, an L2 penalty enters both the first- and second-moment buffers; its effect depends on the data-gradient history and is not a separate uniform shrinkage. AdamW instead applies \(w\leftarrow(1-\eta\lambda)w\) separately from the adaptive data-gradient update. torch.optim.Adam uses coupled L2 by default; current versions also expose decoupled_weight_decay=True. torch.optim.AdamW states the decoupled choice explicitly.

The experiment below uses two weights initialized to 1.0, with alternating data gradients of magnitude 30 and 0.3. These gradients have zero arithmetic mean, but momentum and normalization can still create net parameter motion. We therefore include a zero-decay control. The comparison isolates how adding the same nominal decay changes the two coordinates; it does not assume the data-gradient updates cancel. 

import torch

torch.manual_seed(0)

# Two weights, both starting at 1.0.
# Each step, the "data gradient" is +/-30 for w_a but +/-0.3 for w_b
# (zero raw mean, but the adaptive updates need not cancel).
# Compare against zero decay to expose that baseline motion.

def run(optimizer_cls, weight_decay):
    w = torch.nn.Parameter(torch.tensor([1.0, 1.0]))
    opt = optimizer_cls([w], lr=0.01, weight_decay=weight_decay)
    for t in range(500):
        sign = 1.0 if t % 2 == 0 else -1.0
        w.grad = torch.tensor([30.0 * sign, 0.3 * sign])  # loss gradient only
        opt.step()
        opt.zero_grad()
    return w.detach()

control = run(torch.optim.Adam, weight_decay=0.0)
control_w = run(torch.optim.AdamW, weight_decay=0.0)
assert torch.allclose(control, control_w)
adam = run(torch.optim.Adam, weight_decay=0.1)     # coupled L2
adamw = run(torch.optim.AdamW, weight_decay=0.1)   # decoupled decay

print("after 500 steps (started at 1.0):")
print(f"  no decay (both):         w_a = {control[0]:.4f}   w_b = {control[1]:.4f}")
print(f"  Adam  (L2-in-gradient):  w_a = {adam[0]:.4f}   w_b = {adam[1]:.4f}")
print(f"  AdamW (decoupled decay): w_a = {adamw[0]:.4f}   w_b = {adamw[1]:.4f}")
print(f"  pure decay, no data updates: {(1 - 0.01 * 0.1)**500:.4f}")

Running this prints:

after 500 steps (started at 1.0):
  no decay (both):         w_a = 0.9824   w_b = 0.9824
  Adam  (L2-in-gradient):  w_a = 0.9662   w_b = 0.1859
  AdamW (decoupled decay): w_a = 0.5958   w_b = 0.5958
  pure decay, no data updates: 0.6064

The zero-decay control moves away from 1.0, confirming that alternating raw gradients do not imply zero net adaptive motion. With coupled L2, the large-gradient weight ends near 0.97, whereas the small-gradient weight ends near 0.19. With AdamW both coordinates end near 0.60: their normalized data updates are nearly identical and each receives the same multiplicative decay. The pure-decay prediction is about 0.61; the difference reflects the nonzero data-update trajectory, not a failure of decoupling. A paper saying “Adam with weight decay” is ambiguous: check its equations or implementation before deciding whether it means coupled L2 or AdamW.

That \((1-\eta\lambda)^{500}\) formula contains one more coupling worth internalizing: even in AdamW, the learning rate still multiplies the decay. PyTorch’s implementation shrinks each weight by lr * wd * w per step, so the per-step multiplicative shrinkage is controlled by lr * wd, not wd alone. The effect on the full training trajectory also depends on the adaptive updates and number of steps. Halve the learning rate at fixed wd and you halve the per-step decay; run a cosine schedule and the decay decays with it. This has two practical consequences. First, when you sweep the learning rate, you are also sweeping the regularization — if you want the comparison to isolate lr, some practitioners adjust wd inversely to hold lr * wd fixed. Second, when you copy a weight_decay value from a paper, check what learning rate and schedule it was paired with; wd = 0.1 at lr = 3e-4 and wd = 0.01 at lr = 3e-3 are closer cousins than they look.

One more habit belongs in this section because many Transformer recipes use it: not everything is necessarily decayed. Bias terms and normalization scale/shift parameters are often placed in a zero-decay group. This is a recipe-level heuristic, not a theorem that decay can never help those parameters; follow the reference implementation when reproducing a model and ablate the grouping when it matters. PyTorch exposes the choice through optimizer parameter groups.

Learning-rate schedules

A fixed learning rate can work, but decreasing it often reduces late fluctuations. One common recipe is linear warmup, which gradually increases the rate during early training, followed by cosine decay. Warmup can reduce early instability; it does not guarantee stability or help every task. Loshchilov and Hutter’s SGDR paper (ICLR 2017) studied cosine schedules with warm restarts; here we use a single cosine decay without restarts. The following code demonstrates the scheduler only, without training a model: 

import math
import torch

torch.manual_seed(0)

model = torch.nn.Linear(4, 1)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

total_steps, warmup_steps = 1000, 100

def lr_lambda(step):
    if step < warmup_steps:                       # linear warmup: 0 -> 1
        return step / warmup_steps
    progress = min(1.0, (step - warmup_steps) / (total_steps - warmup_steps))
    return 0.5 * (1.0 + math.cos(math.pi * progress))   # cosine: 1 -> 0

sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)

for step in range(total_steps):
    # Training would run zero_grad(), forward, backward() before opt.step().
    # Here gradients stay None: opt.step() does not change the parameters.
    if step in (0, 50, 100, 500, 900, 999):
        print(f"step {step:4d}: lr = {sched.get_last_lr()[0]:.3e}")
    opt.step()          # optimizer first, scheduler second
    sched.step()
print(f"next step {total_steps}: lr = {sched.get_last_lr()[0]:.3e}")

Running this prints:

step    0: lr = 0.000e+00
step   50: lr = 1.500e-04
step  100: lr = 3.000e-04
step  500: lr = 1.760e-04
step  900: lr = 9.046e-06
step  999: lr = 9.139e-10
next step 1000: lr = 0.000e+00

The scheduler multiplies the base lr = 3e-4. Here step is the zero-based update index: LambdaLR applies multiplier 0 during construction, so the first update uses zero learning rate. Step 100 uses the full base rate; step 999 is very small but positive. The rate reaches exactly zero for the next update, index 1000, after the demonstration ends. We clamp the cosine progress at 1 so extra scheduler calls do not make the rate rise again. Call sched.step() once per optimizer update, after opt.step(). A schedule defined in epochs needs epoch-level calls instead; confusing these units changes the decay speed by the number of batches per epoch.

Does warmup actually do anything?

Warmup is easy to cargo-cult, so it is worth seeing the thing it prevents. The standard account: in the first steps, Adam’s moment estimates are built from a handful of noisy gradients, and the network is in whatever ill-conditioned region initialization put it, so full-size steps early are the most dangerous steps of the run. Goyal and colleagues’ 2017 large-batch ImageNet work made gradual warmup a headline technique — their recipe scaled the learning rate linearly with batch size and reported that a few epochs of warmup were what kept the enormous resulting rates from destabilizing early training. The experiment below reproduces the phenomenon at toy scale: minibatch Adam on the ring problem, at three base learning rates, with and without 150 steps of linear warmup, tracking the worst loss seen in the first 50 steps. 

import torch
import torch.nn as nn

torch.manual_seed(0)

# Ring data, as in the SGD-versus-Adam experiment.
n = 600
X = torch.randn(n, 2) * 1.5
r = X.norm(dim=1)
y = ((r > 1.0) & (r < 2.5)).long()

def make_model():
    torch.manual_seed(1)                     # identical init for every run
    return nn.Sequential(
        nn.Linear(2, 64), nn.ReLU(),
        nn.Linear(64, 64), nn.ReLU(),
        nn.Linear(64, 2),
    )

def train(base_lr, warmup_steps, total_steps=600, batch=32):
    """Minibatch Adam; returns the worst loss seen in the first 50 steps
    and the final training-set accuracy."""
    model = make_model()
    opt = torch.optim.Adam(model.parameters(), lr=base_lr)
    loss_fn = nn.CrossEntropyLoss()
    g = torch.Generator().manual_seed(2)     # identical batch order for every run
    step, worst_early = 0, 0.0
    while step < total_steps:
        for idx in torch.randperm(n, generator=g).split(batch):
            if step >= total_steps:
                break
            step += 1
            lr = base_lr * min(1.0, step / warmup_steps) if warmup_steps else base_lr
            for group in opt.param_groups:
                group["lr"] = lr
            opt.zero_grad()
            loss = loss_fn(model(X[idx]), y[idx])
            loss.backward()
            opt.step()
            worst_early = max(worst_early, loss.item()) if step <= 50 else worst_early
    model.eval()
    with torch.no_grad():
        acc = (model(X).argmax(dim=1) == y).float().mean().item()
    return worst_early, acc

print(f"{'base lr':>8s} | {'worst loss, steps 1-50':>28s} | {'final train accuracy':>20s}")
print(f"{'':>8s} | {'no warmup':>13s}{'warmup 150':>13s}  | {'no warmup':>9s}{'warmup 150':>11s}")
for base_lr in (0.01, 0.05, 0.1):
    w0, a0 = train(base_lr, warmup_steps=0)
    w1, a1 = train(base_lr, warmup_steps=150)
    print(f"{base_lr:8.2f} | {w0:13.4f}{w1:13.4f}  | {a0:9.3f}{a1:11.3f}")

Running this prints:

 base lr |       worst loss, steps 1-50 | final train accuracy
         |     no warmup   warmup 150  | no warmup warmup 150
    0.01 |        0.7669       0.7224  |     0.957      0.963
    0.05 |        1.4401       0.7218  |     0.985      0.982
    0.10 |        3.0751       0.7436  |     0.965      0.962

Read the middle column pair: in this run, warmup substantially reduces the worst minibatch loss observed during the first 50 updates, especially at larger base rates. That maximum does not show the whole trajectory, so it cannot establish that the loss stayed flat or that later spikes never occurred. Final training accuracies are similar here, with small changes in either direction. This warmup uses one-based update counting: its first rate is base_lr / 150, reaching base_lr at update 150, whereas the preceding scheduler demonstration deliberately starts at zero. Both are conventions; state which one an experiment uses. Warmup changes the early optimization trajectory without adding gradient evaluations, but may delay useful progress and still needs a suitable base rate.

A schedule bake-off

Which decay shape should you use? The candidates you will meet most often: constant (no schedule), step decay (drop by 10x at fixed milestones, the classic ResNet-era recipe), cosine (smooth descent to zero), and warmup+cosine (a common Transformer recipe). The following run races all four on the ring problem with minibatch AdamW, deliberately using a base learning rate on the aggressive side (1e-2) so that late-training gradient noise matters, and evaluates the full-data loss at three checkpoints: 

import math
import torch
import torch.nn as nn
torch.manual_seed(0)
# Ring data again; four learning-rate schedules, same model, same base lr.
n = 600
X = torch.randn(n, 2) * 1.5
r = X.norm(dim=1)
y = ((r > 1.0) & (r < 2.5)).long()
total_steps, base_lr, warmup = 800, 1e-2, 80

def schedule(name, step):
    """Multiplier on base_lr at a given step (0-indexed)."""
    if name == "constant":
        return 1.0
    if name == "step decay":            # x0.1 at 50% and again at 75%
        if step < total_steps // 2:
            return 1.0
        return 0.1 if step < 3 * total_steps // 4 else 0.01
    if name == "cosine":
        return 0.5 * (1 + math.cos(math.pi * step / total_steps))
    if name == "warmup+cosine":
        if step < warmup:
            return step / warmup
        prog = (step - warmup) / (total_steps - warmup)
        return 0.5 * (1 + math.cos(math.pi * prog))
    raise ValueError(f"unknown schedule: {name}")

def train(name, batch=32):
    torch.manual_seed(1)
    model = nn.Sequential(
        nn.Linear(2, 64), nn.ReLU(),
        nn.Linear(64, 64), nn.ReLU(),
        nn.Linear(64, 2),
    )
    opt = torch.optim.AdamW(model.parameters(), lr=base_lr, weight_decay=1e-4)
    loss_fn = nn.CrossEntropyLoss()
    g = torch.Generator().manual_seed(2)
    step, checkpoints = 0, {}
    while step < total_steps:
        for idx in torch.randperm(n, generator=g).split(batch):
            if step >= total_steps:
                break
            for group in opt.param_groups:
                group["lr"] = base_lr * schedule(name, step)
            opt.zero_grad()
            loss = loss_fn(model(X[idx]), y[idx])
            loss.backward()
            opt.step()
            step += 1
            if step in (200, 400, 800):
                model.eval()
                with torch.no_grad():
                    checkpoints[step] = loss_fn(model(X), y).item()
                model.train()
    return checkpoints
print(f"{'full-data loss at step:':>24s}{200:>10d}{400:>10d}{800:>10d}")
for name in ("constant", "step decay", "cosine", "warmup+cosine"):
    cp = train(name)
    print(f"{name:>24s}" + "".join(f"{cp[s]:10.4f}" for s in (200, 400, 800)))

Running this prints:

 full-data loss at step:       200       400       800
                constant    0.1485    0.0619    0.0789
              step decay    0.1485    0.0619    0.0328
                  cosine    0.1287    0.0493    0.0281
           warmup+cosine    0.1302    0.0511    0.0243

In this reference run the constant-rate loss rises between steps 400 and 800, while the decaying schedules finish lower. This is consistent with late fluctuations at an aggressive constant rate, but three checkpoints do not prove a stationary noise ball or identify a unique cause. Constant and step decay match at checkpoint 400 because the first reduced-rate update occurs immediately afterward, at zero-based index 400. Warmup+cosine finishes lowest in this run; compare multiple seeds and tuned base rates before choosing a winner. Keep the update budget fixed: cosine’s shape depends on its planned endpoint. All reported losses use the training set.

Pitfall. Retune the learning rate when changing optimizers. For a constant gradient, the momentum convention here approaches 1/(1-beta) times the raw gradient; that is not a universal bound on updates. The SGD and Adam rates in our table are starting points for this model, not portable ranges. Compare matched budgets and initialization, sweep rates in factors of roughly 3, and inspect both optimization and held-out performance before drawing conclusions.

Going deeper

The large-batch story. Everything above quietly assumed batches small enough that gradient noise is a feature. Industrial training pushes the other way: more accelerators are most easily used by making the batch bigger, and that changes the optimization problem. The empirical picture, mapped out between roughly 2017 and 2020, has three parts. First, the linear scaling rule (Goyal et al., 2017): multiply the batch size by \(k\), multiply the learning rate by \(k\), and (for SGD on their ImageNet setup, up to batches of 8192) training curves approximately match. If successive small-batch gradients were evaluated at nearly unchanged parameters, their \(k\) SGD updates would approximately equal one update on the combined batch with \(k\) times the rate. This approximation can fail early or at very large batch sizes. For independent samples, averaging \(k\) times as many examples divides gradient-noise variance by \(k\), and standard deviation by \(\sqrt{k}\). Second, warmup becomes load-bearing at those rates, as discussed above. Third, scaling eventually breaks anyway: past a problem-dependent point, bigger batches stop buying faster wall-clock convergence (each step is better, but you take proportionally fewer of them for the same data budget), and naive scaling harms final accuracy.

Layer-adaptive methods pushed the breaking point further out. LARS (You, Gitman, and Ginsburg, 2017) observed that in large-batch SGD, some layers’ updates are enormous relative to their weight norms while others’ are tiny, and normalized each layer’s step by the ratio \(\lVert w\rVert / \lVert g\rVert\): one trust-ratio per layer on top of momentum. The paper reports scaling ResNet-50 training to batch 32K without accuracy loss. LAMB (You et al., ICLR 2020) applied the same layerwise trust ratio on top of Adam and reports pretraining BERT with batch sizes up to 32K, cutting wall-clock training from days to 76 minutes on a TPU pod. You are unlikely to need either in this course; you should recognize the pattern (per-layer step normalization as the third rung after global SGD and per-parameter Adam) and the papers, because they define the large-batch regime the field now takes for granted.

Why compare optimizer families? Coordinatewise scaling, momentum, layerwise normalization, and curvature approximations respond to different features of the optimization problem. Their usefulness depends on the architecture, parameter groups, gradient statistics, tuning budget, and implementation cost. Treat a reported gain as a result for the experiments actually measured; include a tuned baseline and evaluate both convergence and held-out performance. The comparison discipline from the bowl and ring experiments remains useful as optimizer recipes change.

Where the theory stands. Smooth convex and strongly convex optimization have classical rates and matching lower bounds under specified oracle assumptions. This does not make every optimizer optimal, nor transfer those guarantees to arbitrary deep networks. Nonconvex analyses often bound progress toward stationary points under assumptions on smoothness, noise, and step sizes. Stationarity alone does not establish a global minimum or good generalization. Claims that SGD noise favors flatter, better-generalizing solutions need additional assumptions and empirical evidence; treat them as hypotheses to investigate.

Common pitfalls

  • Symptom: loss decreases, then explodes to NaN mid-training. Possible causes: an excessive learning rate, momentum amplifying an update, or non-finite inputs or loss operations. Check: find the first non-finite value before adjusting lr, warmup, or gradient clipping. Clipping does not repair a NaN gradient.
  • Symptom: weight decay behaves differently across coordinates in Adam. Possible cause: coupled L2 enters both moment buffers, so it is not uniform multiplicative shrinkage. Check: inspect the optimizer configuration; use AdamW when decoupled decay is intended, and retune its strength.
  • Symptom: switching from Adam to SGD “breaks” a model that trained fine. Possible cause: the old learning rate is unsuitable, or the model benefits from adaptive scaling. Check: retune lr for the chosen optimizer; the useful ratio of rates is problem-dependent.
  • Symptom: training loss plateaus early and stays flat. Cause: effective learning rate too small — decay schedule stepped per-batch instead of per-epoch, or warmup never ended due to a miscounted total_steps. Fix: print sched.get_last_lr() alongside the loss; the learning-rate curve should look like the one in the schedule example.
  • Symptom: two optimizer comparisons in your experiments disagree with each other run to run. Cause: different random initializations or data order across runs. Fix: control seeds so all optimizers start from identical weights (as make_model does above) and average over several seeds before concluding anything — this discipline returns in experiments and reproducibility.

Try it yourself

  1. In the bowl experiment, try lr = 0.075 and lr = 0.085. For plain SGD on this deterministic quadratic, predict the steep-coordinate factor \(1-25\eta\) before running: when does it oscillate with shrinking amplitude, and when does it diverge? Compare the other optimizers empirically, but do not apply SGD’s \(2/25\) bound to their different update rules.
  2. Add PyTorch’s stored-iterate Nesterov rule to the NumPy race: with g = grad(w), compute v_new = beta * v + g, then w_new = w - lr * (g + beta * v_new). Initialize v to zero; use zero dampening and no weight decay. Compare every update with torch.optim.SGD(..., nesterov=True) in float64. Then implement the explicit lookahead form and verify the change of variables given in the Nesterov section; comparing its w directly with PyTorch’s stored lookahead point compares different iterates.
  3. In the ring experiment, sweep Adam’s learning rate over {1e-4, 3e-4, 1e-3, 3e-3, 1e-2} and check whether any tested rate beats SGD+momentum at lr = 0.1. What does this tell you about the earlier table?
  4. Reproduce the AdamW demonstration with beta1 = 0 (no momentum). Predict the final w_a and w_b for both optimizers before running.

Further reading

  • Kingma and Ba, “Adam: A Method for Stochastic Optimization,” ICLR 2015. arxiv.org/abs/1412.6980 — the original Adam paper; section 3 derives the moment bias correction used above.
  • Loshchilov and Hutter, “Decoupled Weight Decay Regularization,” ICLR 2019. arxiv.org/abs/1711.05101 — the AdamW paper; short, readable, and the source of the coupling argument this section demonstrated numerically.
  • Loshchilov and Hutter, “SGDR: Stochastic Gradient Descent with Warm Restarts,” ICLR 2017. arxiv.org/abs/1608.03983 — where cosine-shaped schedules entered mainstream practice (the restarts are less used than the shape).
  • Sutskever, Martens, Dahl, and Hinton, “On the importance of initialization and momentum in deep learning,” ICML 2013. proceedings.mlr.press/v28/sutskever13.html — the paper that established momentum (heavy-ball and Nesterov) as central to training deep networks, and the source of the formulation PyTorch implements.
  • Duchi, Hazan, and Singer, “Adaptive Subgradient Methods for Online Learning and Stochastic Optimization,” JMLR 2011. jmlr.org/papers/v12/duchi11a.html — AdaGrad: where per-parameter adaptive scaling started, with the theory for sparse problems.
  • Wilson, Roelofs, Stern, Srebro, and Recht, “The Marginal Value of Adaptive Gradient Methods in Machine Learning,” NeurIPS 2017. arxiv.org/abs/1705.08292 — the influential skeptical take: tuned SGD often matches adaptive methods; a model for how to compare optimizers honestly.
  • Goyal et al., “Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour,” 2017. arxiv.org/abs/1706.02677 — the linear scaling rule and gradual warmup, from the report that opened the large-batch era.
  • You et al., “Large Batch Optimization for Deep Learning: Training BERT in 76 minutes,” ICLR 2020. arxiv.org/abs/1904.00962 — LAMB, with LARS (You, Gitman, and Ginsburg, 2017, arxiv.org/abs/1708.03888) as its layerwise-scaling predecessor.
  • Dive into Deep Learning, chapter 12, “Optimization Algorithms.” d2l.ai/chapter_optimization/ — every method in this section rebuilt slowly with code and convergence visualizations.
  • PyTorch documentation, torch.optim. docs.pytorch.org/docs/stable/optim.html — the authoritative update-rule pseudocode for every optimizer named here, including the exact Nesterov reformulation.

Before class you should be able to

  • Write the update rules for SGD, momentum, RMSProp, and Adam (with bias correction) from memory.
  • Explain the plain-GD stability bound on a positive-definite quadratic and distinguish its per-step contraction from heavy-ball momentum’s asymptotic rate.
  • Predict Adam’s approximate steady-gradient step size, and explain why robustness across gradient scales, not raw speed, is its selling point.
  • Explain the difference between Adam’s L2-in-the-gradient and AdamW’s decoupled decay, and why coupled L2 does not give uniform multiplicative shrinkage.
  • State why the per-step shrinkage depends on lr * wd, and what that implies when copying hyperparameters from a paper.
  • Read a loss curve and recognize a noise-ball plateau, a post-warmup spike, and a too-small learning rate.
  • Describe the warmup+cosine schedule and the per-step versus per-epoch sched.step() convention trap.

Connects to

  • This session is covered by Quiz 1 (Module 1: S01–S08), given at S09.
  • Core preparation for every training run you will do from here on: the NumPy race exposes the SGD and Adam update rules, and the Nesterov exercise asks you to match PyTorch step for step. The moment initialization and time index deserve special attention.
  • The training-loop mechanics (zero_grad, step, schedulers) reappear in 1.8 (S08); weight decay returns as a regularizer in 1.7 (S07).

← S05 · Backpropagation II · All readings · S07 · Regularization, Initialization, Normalization →