KBKnowledge Base
Machine Learning · 2.2.3

Accelerated & Adaptive Optimizers

Momentum, Nesterov, AdaGrad, RMSProp, Adam & AdamW — derived, not just named.

On this page
In plain English — beginner to advanced

Beginner: the previous lesson showed plain gradient descent zig-zagging across a narrow, stretched valley — bouncing back and forth across the steep direction while crawling painfully slowly along the shallow one. Momentum fixes this the way a heavy ball rolling downhill would: give the step some inertia, so it doesn't reverse direction every time the gradient flips sign, and it keeps building speed in whichever direction has stayed consistent. The methods in this lesson are all variations on that one idea — remembering something about past gradients instead of only ever looking at the current one.

Intermediate: each method changes the plain update rule in one specific, nameable way:

  • Momentum keeps a running, exponentially-weighted average of past gradients and steps in that averaged direction instead of the raw current gradient.
  • Nesterov momentum makes one subtle but real improvement: it computes the gradient at where momentum is about to carry the point, not at the current point — a small "look-ahead" correction that measurably speeds up convergence on smooth problems.
  • AdaGrad gives every parameter its own learning rate, shrinking it for parameters whose gradients have historically been large — so a frequently-updated, steep direction automatically slows down, letting a rarely-updated, shallow direction catch up.
  • RMSProp fixes AdaGrad's one real flaw: AdaGrad's per-parameter rate is built from an ever-growing sum of squared gradients, so it eventually shrinks to essentially zero and training stalls. RMSProp replaces that sum with a decaying average, so old gradients eventually stop mattering instead of permanently weighing the rate down.
  • Adam combines both ideas at once: momentum (a first-moment running average of the gradient itself) and RMSProp-style adaptive per-parameter scaling (a second-moment running average of the squared gradient), plus a bias-correction term that fixes both averages being unreliable in the first few steps, when they've barely started accumulating anything.

Advanced: AdamW fixes a specific, subtle bug in how Adam is normally combined with weight decay (an L2-style penalty pulling parameters toward zero): if you just add the penalty's gradient into the raw gradient before Adam's adaptive scaling divides by the accumulated squared-gradient average, the effective decay strength ends up different for every parameter — parameters with historically large gradients get less decay than intended, exactly backwards from what a uniform regularization strength is supposed to do. AdamW decouples the two: apply Adam's adaptive update to the loss gradient alone, then apply the weight decay as a separate, uniform shrinkage step afterward. Despite this, no single optimizer in this lesson strictly dominates the others — Adam/AdamW are the default for training transformers and most modern deep nets, but plain SGD with momentum still wins in some carefully-tuned computer vision training recipes, generalizing very slightly better even though it converges a bit more slowly.

Formula
vk=βvk1+f(θk),θk+1=θkαvkv_k = \beta v_{k-1} + \nabla f(\theta_k), \qquad \theta_{k+1} = \theta_k - \alpha v_k

Momentum's update: a running average of gradients, stepped by a fixed learning rate.

mk=β1mk1+(1β1)gk,vk=β2vk1+(1β2)gk2,m^k=mk1β1k,v^k=vk1β2km_k = \beta_1 m_{k-1} + (1-\beta_1)g_k, \quad v_k = \beta_2 v_{k-1} + (1-\beta_2)g_k^2, \quad \hat m_k = \frac{m_k}{1-\beta_1^k}, \quad \hat v_k = \frac{v_k}{1-\beta_2^k}
θk+1=θkαm^kv^k+ϵ\theta_{k+1} = \theta_k - \alpha \frac{\hat m_k}{\sqrt{\hat v_k} + \epsilon}

Adam's full update: a bias-corrected first moment (momentum) divided by a bias-corrected second moment's square root (the adaptive, per-parameter scaling), plus a tiny ε to avoid dividing by zero early on.

Derivation: why momentum specifically fixes the elongated-valley problem

Set up the same toy quadratic from the previous lesson, f(x,y) = ½(c₁x² + c₂y²) with c₁ ≫ c₂ (a bowl much steeper in x than in y). Plain gradient descent updates each coordinate independently: x_{k+1} = (1-αc₁)x_k and y_{k+1} = (1-αc₂)y_k. Stability in the steep x direction requires |1-αc₁| < 1, i.e. α < 2/c₁ — this caps how large α can be. But that same small α, applied to the y update, gives a convergence factor (1-αc₂) that's very close to 1 whenever c₂ ≪ c₁ — meaning progress along the shallow direction is agonizingly slow, precisely because the step size had to be kept small by the unrelated steep direction.

Now trace what momentum's running average v_k = βv_{k-1} + ∇f(θ_k) does to each coordinate separately. Along the steep x direction, gradient descent overshoots the minimum every step or two once α is anywhere near its stability limit — the per-step gradient c₁x_k keeps flipping sign as x_k oscillates around zero. Averaging a sequence of gradients that keeps flipping sign causes substantial cancellation: the running average v_k's x-component ends up smaller in magnitude than the raw gradient, damping the oscillation. Along the shallow y direction, by contrast, the gradient c₂y_k keeps the same sign step after step (the point is still consistently far from the minimum in that direction) — averaging a sequence of same-signed numbers causes no cancellation at all; instead the terms accumulate, so v_k's y-component grows larger than any single gradient, in effect taking a bigger, more confident step exactly where a bigger step was safe all along.

This is the precise mechanism, not just an empirical observation: momentum damps oscillating (steep-direction) components via cancellation in the running average, and amplifies consistent (shallow-direction) components via accumulation in that same average, using one and the same update rule.

Where this is used: this exact elongated-valley failure mode is extremely common in real loss landscapes whenever input features are correlated or poorly scaled — which is nearly always — and it's precisely why momentum (or Adam, which includes it) is a default choice in virtually every deep learning training loop, not an optional nicety.

Three optimizers race toward the minimum of the same stretched bowl

Blue = plain gradient descent, amber = momentum, violet = Adam — all three compute their real update rule live and are played back step by step. Drag the learning-rate slider and re-run to see GD destabilize far sooner than the other two.

Practical example — plain GD, momentum, and Adam on the same elongated bowl
python
import numpy as np

KX, KY = 1.0, 6.0  # same elongated bowl as the diagram

def grad(p):
    return np.array([KX * p[0], KY * p[1]])

def run_gd(lr, steps, p0):
    p = np.array(p0, dtype=float)
    for _ in range(steps):
        p = p - lr * grad(p)
    return p

def run_momentum(lr, beta, steps, p0):
    p = np.array(p0, dtype=float)
    v = np.zeros(2)
    for _ in range(steps):
        v = beta * v + grad(p)
        p = p - lr * v
    return p

def run_adam(lr, steps, p0, b1=0.9, b2=0.999, eps=1e-8):
    p = np.array(p0, dtype=float)
    m = np.zeros(2)
    v = np.zeros(2)
    for k in range(1, steps + 1):
        g = grad(p)
        m = b1 * m + (1 - b1) * g
        v = b2 * v + (1 - b2) * g * g
        m_hat = m / (1 - b1 ** k)
        v_hat = v / (1 - b2 ** k)
        p = p - lr * m_hat / (np.sqrt(v_hat) + eps)
    return p

start = [-3.6, 1.9]
print("GD final:      ", run_gd(0.12, 60, start))
print("Momentum final:", run_momentum(0.12, 0.9, 60, start))
print("Adam final:    ", run_adam(0.25, 60, start))
# All converge toward (0, 0); momentum and Adam typically get there in far fewer
# effective steps for a given stable learning rate than plain GD does.
Real-world examples
  • Adam/AdamW is the default optimizer for training transformers and almost every modern large deep learning model — it's forgiving of imperfect learning rate tuning and handles wildly different gradient scales across a huge model's layers.
  • SGD with momentum is still preferred in some well-tuned computer-vision training recipes (certain ResNet-style pipelines), where it has been observed to generalize very slightly better than Adam despite converging somewhat more slowly.
  • RMSProp originated to fix an instability in training recurrent networks and remains a common choice in some reinforcement-learning training setups.
  • AdaGrad is a genuinely good fit for sparse-gradient settings — e.g. large NLP embedding tables where most rows are updated only rarely — since its learning-rate-decays-to-zero flaw matters far less when most parameters barely get touched anyway.
  • Learning-rate warmup is routinely paired with Adam in large-model training specifically because Adam's bias-corrected second-moment estimate is unreliable during the first handful of steps, before it has accumulated enough gradient history to be trustworthy.
  • Production ML frameworks (PyTorch, TensorFlow, JAX) all expose every optimizer in this lesson as a one-line drop-in swap, which is exactly why practitioners can afford to treat the optimizer as a hyperparameter to try a few of, rather than committing to one up front.
Common mistakes
  • Assuming Adam is tuning-free just because it's more forgiving than plain SGD — the learning rate still matters, and a badly chosen one still causes slow convergence or instability with Adam, just usually less dramatically than with plain gradient descent.
  • Porting code between frameworks/versions and getting a naive L2 penalty added inside Adam's gradient instead of true decoupled AdamW weight decay — a genuinely common, subtle bug, since the two look similar in code but behave differently per-parameter.
  • Using a high momentum coefficient together with too large a learning rate and seeing oscillation or divergence, then wrongly concluding "momentum doesn't work" instead of recognizing an unstable combination of hyperparameters.
Going deeper

Nesterov's look-ahead correction has a clean interpretation as a more accurate approximation of the continuous-time trajectory a truly frictionless, accelerating ball would follow — evaluating the gradient at the point momentum is about to reach, rather than where the point currently sits, corrects for a lag that plain momentum otherwise introduces. This isn't just folklore: for convex, smooth (Lipschitz-gradient) functions, Nesterov's method provably achieves an O(1/k²) convergence rate, compared to plain gradient descent's O(1/k) — a real, provable speedup, not merely an empirical tendency.

Check yourself
On the elongated-bowl toy problem, why does momentum eventually take a LARGER effective step along the shallow axis than plain gradient descent ever safely could, using the exact same per-step learning rate α?

Because momentum's update isn't the raw current gradient — it's a running average of many past gradients. Along the shallow axis, the gradient keeps the same sign every step (the point is still consistently far from the minimum there), so those same-signed terms accumulate in the average, producing an effective step larger than any single gradient. Along the steep axis, by contrast, the gradient keeps flipping sign as the point oscillates, so those terms cancel in the average instead of accumulating — momentum gets to be aggressive exactly where it's safe and cautious exactly where it needs to be, without α itself ever changing.

Key takeaway

Momentum, Nesterov, AdaGrad, RMSProp, and Adam are not five unrelated tricks — they're five different, precise answers to "what should I remember about past gradients before taking the next step," layered on top of the exact same generic update rule from this module's overview. With gradient-based methods (this lesson) and second-order methods (the next lesson) both in hand, the toolbox for smooth, unconstrained optimization is essentially complete.

Newsletter

Stay in the loop

Subscribe to get new docs, diagrams, and engineering write-ups by Dharaneesh Boobalan delivered to your inbox.

  • Deep-dive write-ups on ML, inference, and systems.
  • New Draw.io diagrams & interactive canvases.
  • Agentic patterns and rocket-science notes.
  • No spam. One tasteful email when there's something new.

Crafted by Dharaneesh Boobalan

Newsletter

Get new docs, diagrams, and write-ups in your inbox.

We never share your details. Unsubscribe anytime.