Optimizers — SGD, Momentum, Adam, RMSprop

Why this session exists

A gradient is a local derivative, not an instruction to take a particular-sized step. SGD, momentum, RMSprop and Adam turn the same derivatives into different updates by retaining different state. Understanding that state explains learning-rate sensitivity, checkpoint requirements and why an optimizer comparison needs controlled initialization, batches and selection rules.

You will implement all four methods, verify them against PyTorch on an exact gradient sequence, derive Adam's startup correction, analyze a quadratic stability bound and choose a schedule with an explicit clock. No optimizer is declared the winner on 90% of problems.

Prerequisites

Know S098 gradients, the S097 MLP and exponential moving averages. Learning rates and optimizer settings must be selected using development data, not final test scores.

(a) Intuition · 5 min

Imagine moving down a narrow valley. The steep cross-valley direction can oscillate while the long direction changes slowly. Momentum filters the gradient history; adaptive methods additionally divide each coordinate by a recent scale estimate. This ball-on-a-hillside picture is an analogy: the optimizer is not a physical simulation, and its second moment is not a measurement of terrain curvature.

The three problems optimisers solve

Momentum can accumulate persistent directions, damp some alternating components and reduce some minibatch noise. RMSprop rescales coordinates with different recent gradient magnitudes and forgets distant history, unlike AdaGrad's accumulating squared-gradient sum. Adam combines a first-moment average with that second-moment scale and startup corrections. These mechanisms can help but do not remove the global learning rate, conditioning issues or stochastic noise.

Sparse values in a dense gradient are different from a sparse tensor API. PyTorch Adam/AdamW do not support sparse gradient tensors in this version; use a compatible optimizer such as SparseAdam or an appropriate SGD configuration where required. A blanket “AdamW is mandatory for embeddings” rule is wrong.

The evolution

AdaGrad keeps $s_t=s_{t-1}+g_t^2$ and divides each coordinate by its accumulated scale; the denominator never decreases. RMSprop discounts older squared gradients instead, while Adam adds a first-moment average and bias correction. Nesterov-style momentum evaluates a gradient at a lookahead location rather than only the current parameters. In one constant-rate convention, $g_t=\nabla f(\theta_{t-1}-\eta\beta v_{t-1})$, followed by $v_t=\beta v_{t-1}+g_t$ and $\theta_t=\theta_{t-1}-\eta v_t$. Libraries can use equivalent reparameterized forms with different velocity definitions. Convex acceleration theorems require specific assumptions and schedules; they do not prove every Nesterov run beats heavy-ball momentum on a neural network. AdamW separates parameter shrinkage from those adaptive moments. These are distinct algorithms rather than a sequence in which every older method became deprecated. The original Transformer paper explicitly used Adam with $(\beta_1,\beta_2)=(0.9,0.98)$ and a warmup/inverse-square-root schedule, not a universal modern cosine recipe.

The update rules, side by side

Let $g_t$ be the current gradient and $\eta_t$ the base rate. All operations on vectors below are coordinatewise except the loss used to obtain $g_t$.

SGD: $\theta_t=\theta_{t-1}-\eta_tg_t$. With full-data gradients this is deterministic gradient descent; the algebra is identical but there is no sampling noise.

Momentum convention used here:

An alternative EMA convention uses $(1-\beta)g_t$. For constant beta, zero initial state and matched rate scaling, these can describe the same trajectory; with the same numeric learning rate they do not. In our unnormalized convention a constant gradient yields $v_t=g(1-\beta^t)/(1-\beta)$, approaching $10g$ for beta 0.9. A normalized EMA instead approaches $g$. Keep that distinction in both explanation and code.

RMSprop (uncentered, no momentum):

This matches PyTorch's epsilon placement outside the square root; $\sqrt{s+\epsilon}$ is a different formula. The history length is approximately $1/(1-\rho)$ under a simple decay interpretation. A small recent gradient does not mean the loss is flat in the Hessian sense.

Adam:

Bias correction: work out the first step

Starting at zero, a stationary gradient mean gives $E[m_t]=(1-\beta_1^t)E[g]$; likewise for the raw second moment. The correction removes this initialization effect under that assumption; it does not estimate future curvature or make training globally unbiased.

At the first step with $g=2$, beta1 0.9 and beta2 0.999, $m_1=0.2$, $v_1=0.004$, $\hat m_1=2$, $\hat v_1=4$. Ignoring tiny epsilon, corrected update magnitude is $\eta$, whereas the uncorrected ratio is $0.2/\sqrt{0.004}\approx3.1623$. Omitting both corrections makes the first step larger, not 10× smaller. More generally the uncorrected/corrected ratio for a constant gradient is

Changing beta values changes this ratio. There is no universal claim that removing correction delays convergence by a fixed factor. Bias correction and warmup address different aspects of the update; one does not render the other unnecessary.

The four optimisers on the same loss surface

For $f(x,y)=\tfrac12(x^2+20y^2)$, gradient descent satisfies $x_{t+1}=(1-\eta)x_t$, $y_{t+1}=(1-20\eta)y_t$. Convergence requires $0<\eta<0.1$. At 0.05, the $y$ coordinate becomes zero after one step: this particular setup does not zigzag across the steep axis. At 0.09 it changes sign while contracting; above 0.1 its magnitude grows. These exact recurrences are a better explanation than invented optimizer race timings.

Momentum can accelerate persistent directions but can overshoot or destabilize with excessive beta/rate. RMSprop and Adam are coordinatewise, not invariant to arbitrary rotations of the problem. Compare finite-budget trajectories, not a claim that an optimizer generally converges in a fixed number of steps.

Choosing an optimiser

Plain SGD is a valid training method, not merely a teaching artifact. SGD with momentum has one velocity tensor per parameter; Adam has first and second moments, plus counters and optional variant state. At equal state dtype the two moments cost twice the parameter bytes, but total training memory also includes parameters, gradients and activations.

AdamW is a reasonable candidate when adaptive updates and decoupled decay fit the model recipe; SGD-momentum is a useful comparison. RMSprop remains supported. Switching Adam to SGD mid-run introduces a switch point and state/rate transition, so justify it with validation rather than assuming both methods' best properties combine automatically. Match data, initialization, tuning budget and stopping policy before comparing methods. Equal numeric learning rates are neither required nor forbidden; their effective updates differ, so tune fairly.

Weight decay is a separate decision

Adding $\lambda\|\theta\|^2/2$ to the loss adds $\lambda\theta$ to its gradient. For plain SGD, this is $(1-\eta\lambda)\theta-\eta g$. With coupled Adam, that term enters both moment estimates. AdamW instead applies shrinkage to parameters outside the adaptive moment calculation. This is a different regularization choice, not proof that an implementation of L2-regularized Adam is mathematically broken.

PyTorch 2.9's Adam(..., decoupled_weight_decay=True) explicitly provides AdamW-equivalent behavior; default False couples the penalty. Version-specific options matter. Parameter groups can give biases and normalization parameters zero decay as a chosen policy. SGD momentum with an L2 gradient remains coupled into velocity and is not generally equivalent to a separate per-step shrink with the same coefficient.

Learning rate schedules

A constant rate is a valid baseline. Step decay multiplies it at specified milestones; plateau schedules respond to a monitored validation metric. Cosine decay over a declared interval is

Warmup can linearly increase the rate before decay. OneCycle varies rate up and down within a budget and may vary momentum inversely. None universally improves validation quality. The original Transformer's documented schedule grows linearly for 4,000 updates then decays with inverse square root; the paper does not establish that every transformer without warmup diverges.

Specify whether the clock counts optimizer updates, epochs or tokens. When accumulating microbatches, an update-based schedule should advance only when parameters step. Save optimizer and scheduler state and counters for resumption. In PyTorch's usual scheduler pattern call optimizer.step() before scheduler.step(); ReduceLROnPlateau instead consumes its monitored metric at the chosen cadence.

(c) Hands-on · 25 min

This exact CPU lab implements the four recurrences and AdamW, matches PyTorch on a shared two-step gradient sequence, then runs a quadratic and a small MLP with controlled initialization/batch order. The MLP numbers are training diagnostics, not held-out generalization or optimizer benchmarks.

import numpy as np
import torch
from scipy.special import logsumexp
from sklearn.datasets import make_classification
torch.set_num_threads(1)
 
class Optimizer:
    def __init__(self,p,kind,lr,beta=.9,b1=.9,b2=.999,eps=1e-8,wd=0.):
        self.p,self.kind,self.lr=p,kind,lr
        self.beta,self.b1,self.b2,self.eps,self.wd=beta,b1,b2,eps,wd
        self.m={k:np.zeros_like(v) for k,v in p.items()}
        self.v={k:np.zeros_like(v) for k,v in p.items()}; self.t=0
    def step(self,grads):
        self.t+=1
        for k in self.p:
            g=grads[k]
            assert np.isfinite(g).all()
            if self.kind=='sgd': update=g
            elif self.kind=='momentum':
                self.m[k]=self.beta*self.m[k]+g; update=self.m[k]
            elif self.kind=='rmsprop':
                self.v[k]=self.beta*self.v[k]+(1-self.beta)*g*g
                update=g/(np.sqrt(self.v[k])+self.eps)
            else:
                assert self.kind in ['adam','adamw']
                if self.kind=='adam': g=g+self.wd*self.p[k]
                else: self.p[k]*=1-self.lr*self.wd
                self.m[k]=self.b1*self.m[k]+(1-self.b1)*g
                self.v[k]=self.b2*self.v[k]+(1-self.b2)*g*g
                mh=self.m[k]/(1-self.b1**self.t); vh=self.v[k]/(1-self.b2**self.t)
                update=mh/(np.sqrt(vh)+self.eps)
            self.p[k]-=self.lr*update
 
for kind in ['sgd','momentum','rmsprop','adam','adamw']:
    p={'w':np.array([1.,-2.])}; a=Optimizer(p,kind,.01,wd=.03 if kind in ['adam','adamw'] else 0.)
    w=torch.tensor([1.,-2.],dtype=torch.float64,requires_grad=True)
    if kind=='sgd': ref=torch.optim.SGD([w],lr=.01)
    elif kind=='momentum': ref=torch.optim.SGD([w],lr=.01,momentum=.9)
    elif kind=='rmsprop': ref=torch.optim.RMSprop([w],lr=.01,alpha=.9,eps=1e-8)
    elif kind=='adam': ref=torch.optim.Adam([w],lr=.01,weight_decay=.03)
    else: ref=torch.optim.AdamW([w],lr=.01,weight_decay=.03)
    for g in [np.array([.2,-.3]),np.array([-.4,.1])]:
        a.step({'w':g}); w.grad=torch.tensor(g); ref.step()
        assert np.allclose(p['w'],w.detach().numpy(),rtol=1e-11,atol=1e-12)
    print('matched',kind,p['w'])
ratio=(1-.9)/np.sqrt(1-.999)
assert np.isclose(ratio,np.sqrt(10))
print('first-step uncorrected/corrected ratio',ratio)
for kind,lr in [('sgd',.05),('momentum',.05),('rmsprop',.1),('adam',.1)]:
    p={'w':np.array([-5.,-1.])}; opt=Optimizer(p,kind,lr)
    start=.5*np.sum(np.array([1.,20.])*p['w']**2)
    for _ in range(200): opt.step({'w':np.array([1.,20.])*p['w']})
    end=.5*np.sum(np.array([1.,20.])*p['w']**2)
    assert np.isfinite(end) and end<start
    print('quadratic',kind,'loss after 200 updates',end)
 
X,y=make_classification(n_samples=240,n_features=6,n_informative=4,
                        n_redundant=0,n_classes=3,random_state=2)
Y=np.eye(3)[y]; rng=np.random.default_rng(3)
base={'W1':rng.normal(0,.2,(6,12)),'b1':np.zeros(12),
      'W2':rng.normal(0,.2,(12,3)),'b2':np.zeros(3)}
orders=[rng.permutation(len(y)) for _ in range(20)]
def fb(p,ids):
    x=X[ids]; target=Y[ids]; z=x@p['W1']+p['b1']; a=np.maximum(z,0)
    logits=a@p['W2']+p['b2']; lp=logits-logsumexp(logits,axis=1,keepdims=True)
    dz=(np.exp(lp)-target)/len(ids); dh=(dz@p['W2'].T)*(z>0)
    return float(-np.mean(np.sum(target*lp,axis=1))),{
        'W2':a.T@dz,'b2':dz.sum(0),'W1':x.T@dh,'b1':dh.sum(0)}
for kind,lr in [('sgd',.05),('momentum',.02),('rmsprop',.003),('adam',.01)]:
    p={k:v.copy() for k,v in base.items()}; opt=Optimizer(p,kind,lr)
    initial=fb(p,np.arange(len(y)))[0]
    for rows in orders:
        for i in range(0,len(rows),32): opt.step(fb(p,rows[i:i+32])[1])
    final=fb(p,np.arange(len(y)))[0]
    assert final<initial and np.isfinite(final)
    print('training-only MLP',kind,initial,final)
# A specified update-clock warmup/cosine contract, endpoints included.
def schedule(t,warm=10,total=100,peak=.01,floor=.001):
    if t<=warm: return peak*t/warm
    u=min((t-warm)/(total-warm),1.)
    return floor+.5*(peak-floor)*(1+np.cos(np.pi*u))
assert schedule(0)==0 and schedule(10)==.01 and np.isclose(schedule(100),.001)
print('schedule endpoints passed')

Anatomy of the script

Every method receives an independent copy of the same initial arrays. The gradient-sequence test isolates the update equations from any model derivative. PyTorch momentum is compared with zero dampening and no Nesterov adjustment, matching this implementation. RMSprop uses alpha 0.9 explicitly rather than relying on PyTorch's 0.99 default. The tiny epsilon is outside the root in both implementations.

The MLP comparison reuses exact minibatch orders; the earlier pattern advanced one global RNG separately for each method and confounded optimizer with initialization and shuffle. Learning rates here are declared illustrative settings, not an equal-budget hyperparameter search. Training-loss differences cannot establish a generalization winner.

Bias-correction and AdamW practice with answers

Exercise 1: Remove both Adam corrections. Predict the first update before running anything. Answer: At the stated betas it is approximately $\sqrt{10}$ times the corrected step for a nonzero constant gradient, not smaller. Compare complete trajectories separately; startup arithmetic does not determine the final winner.

Exercise 2: Change coupled L2 to decoupled shrinkage. Answer: Compute moments using the data gradient only and multiply the old parameter by $1-\eta\lambda$ outside the adaptive calculation. The lab checks both variants against PyTorch. To compare predictive quality, tune rate and decay for each on validation and evaluate the selected policies on a clean test. AdamW need not achieve lower training loss, and lower training loss is not itself better regularization.

(d) Production reality · 15 min

A flat loss may indicate too small a step, disconnected gradients, frozen parameters, wrong targets or already-converged behavior. Oscillation may reflect large steps, noisy batches or changing objectives. An early NaN may originate in invalid inputs, overflow, normalization or the optimizer state. Inspect each rather than mapping one curve shape to one cause.

Log the data objective, gradient norm, parameter norm, update-to-parameter ratio and current rate. Checkpoint moments as well as weights; restarting Adam from zero moments is a different procedure. With gradient accumulation, scale by the total example/token count, include the last incomplete group and advance the schedule on real updates. If using FP16 scaling elsewhere, unscale once after accumulation and before clipping; this lesson makes no accelerator execution claim.

The original Transformer schedule is a documented example of warmup plus inverse-square-root decay. It does not support a story about a company's first thousand divergent steps. Likewise, published adaptive-versus-SGD comparisons are conditional on their datasets and recipes, not a universal one-to-two-point vision advantage.

Where this shows up in the rest of the plan

S100 wraps these updates in torch.optim; S101 studies decay together with dropout and normalization. Fine-tuning rates depend on parameterization, data and which parameters train; there is no universal “LoRA uses 100× smaller LR” rule. Monitoring and resumption must retain optimizer-specific state.

(e) Recall + stretch · 10 min

1. SGD? $\theta\leftarrow\theta-\eta g$.
2. Momentum's benefits? Persistent components accumulate and alternating components can cancel; excessive momentum can also overshoot.
3. RMSprop denominator? $\sqrt{s}+\epsilon$ here; it rescales by recent raw second moment, not curvature.
4. Adam? Bias-correct first/second moments, divide the first by square root of the second plus epsilon, then apply the global rate.
5. Correction at step one? With the default betas, removing both produces an approximately 3.162-fold larger nonzero-gradient step.
6. When SGD-momentum? When a controlled recipe or validation comparison supports it, or its state/behavior suits the task—not solely because the input is an image.
7. AdamW? Decay bypasses moment accumulation. PyTorch 2.9 also exposes the decoupled option in Adam.
8. Warmup required? It is a training-design choice with successful precedents, not a universal guarantee or model-size threshold.
9. Oscillation? Examine rates, gradients, noise and objective changes; reducing rate is a useful test, not proof of the cause.
10. Stretch winner? Coupled and decoupled decay follow different equations. Test under a clean selection protocol; no final-loss or generalization ordering is guaranteed.

Explain-out-loud test

Derive the quadratic stability interval, distinguish normalized EMA from accumulated momentum, calculate Adam's first step and state the scheduler clock you would checkpoint.

What comes next

PyTorch fundamentals separates registered parameters, gradient buffers, optimizer state and module mode.

Sources and execution boundary

PyTorch 2.9 optimizers, tagged SGD, RMSprop and Adam implementations, AdamW primary paper, and Transformer optimizer section. The executed comparisons are CPU equation checks and toy training diagnostics, not wall-clock, hardware or generalization benchmarks.