Optimizers: SGD, momentum, RMSProp, Adam, AdamW

The big idea

Backpropagation tells you which way is downhill for every weight. The optimizerThe rule that turns gradients into parameter updates, such as SGD, momentum, RMSProp, Adam or AdamW. decides how far to step, and that turns out to matter a lot. Real loss surfaces are not round bowls. They have long, narrow valleys (ravines) where the slope is steep across the valley and gentle along it.

Picture a ball rolling down such a valley. Plain gradient descent behaves like a ball with no mass: at every point it heads straight down the steepest direction, which is mostly across the valley, so it bounces from wall to wall and creeps along the floor. A heavy ball keeps its speed: the side-to-side pushes cancel out, and the steady push along the floor builds up. That is momentum. Adaptive methods go further and give each weight its own step size, so steep directions take small steps and flat directions take large ones.

Adam combines both ideas and is the default for most deep learning. AdamW fixes how Adam handles weight decay, and is the standard for training transformers.

How it works

Step 1: SGD

Stochastic gradient descentStochastic gradient descent: update each parameter by minus the learning rate times its gradient, computed on a mini-batch. (SGD) is the rule from Gradient descent, the engine of ML , applied to mini-batch gradients:

One learning rate eta for every weight. In a ravine that is a bad compromise: a step size small enough not to bounce out of the steep direction is far too small for the gentle one. Push eta up and it diverges. In the code, SGD with eta = 0.03 on a curved ravine goes from loss 21.6 to 143 to 31,456 in three steps.

Step 2: Momentum

MomentumAn optimizer that keeps a running velocity, a decaying sum of past gradients, and steps along it, which damps oscillations and speeds up progress along consistent directions. keeps a velocity v, a decaying sum of past gradients, and steps along it:

With beta = 0.9, the velocity is roughly the sum of the last 10 gradients. Gradients that flip sign every step (across the ravine) cancel; gradients that keep the same sign (along it) add up, up to 10 times larger. The ball builds speed on long gentle slopes and does not overreact to bumps.

Step 3: RMSProp, a step size per weight

RMSPropAn adaptive optimizer that divides each parameter's gradient by a running root-mean-square of its recent gradients, giving every parameter its own effective step size. tracks, for each weight, a running average of its squared gradients, and divides by its square root:

A weight whose gradients are consistently large gets its steps shrunk; a weight with tiny gradients gets them enlarged. The update size becomes roughly eta for everyone, whatever the scale of their gradients. That is exactly what a ravine needs, and it also helps with weights that receive rare or small gradients, such as embeddings of rare words.

Step 4: Adam

AdamAdaptive moment estimation: an optimizer that combines momentum (a running mean of gradients) with RMSProp-style scaling (a running mean of squared gradients), plus a bias correction for the first steps. combines both: a momentum-like running mean m of gradients and an RMSProp-like running mean v of squared gradients, each with its own decay (beta1 = 0.9, beta2 = 0.999 by default):

The hats are a bias correctionIn Adam, dividing the running averages by (1 - beta^t) so they are not biased towards zero in the first steps, when they start from zero.: m and v start at zero, so early on they underestimate the true averages; dividing by (1 - beta^t) fixes that. Adam's defaults work across a wide range of problems, which is a big reason it became the default. A learning rate around 1e-3 to 3e-4 is a common starting point.

Step 5: Weight decay, and why AdamW

Weight decayShrinking every weight slightly towards zero at each step, a form of L2 regularization that discourages large weights. shrinks every weight a little each step, a regularizer that prevents weights from growing large ( Regularization for deep nets ). With SGD, adding an L2 penalty to the loss and shrinking the weights directly are the same thing. With Adam they are not: the L2 term's gradient goes through Adam's per-weight scaling, so weights with large gradient history are barely decayed and weights with tiny gradients are decayed hard.

AdamWAdam with decoupled weight decay: the weights are shrunk directly (w -= lr * wd * w) instead of adding an L2 term to the gradient, so every weight decays at the intended rate. decouples the two: Adam handles the loss gradient, and a separate step multiplies each weight by (1 - eta x wd). Every weight decays at the intended rate. AdamW is the standard optimizer for transformers and most large models.

Try it

Try this:

1. At multiplier 1, move "steps shown" from 0 to 300. Which optimizer reaches the floor of the ravine first, and which reaches the minimum first?
2. Raise the multiplier to 2 and then 3. Which optimizer diverges first, and which barely changes? Why would dividing by the running RMS of gradients make an optimizer less sensitive to the learning rate?
3. Lower the multiplier to 0.25. Which ones still reach (1, 1) within 300 steps?

Going deeper

Code: the race, and Adam versus AdamW

import numpy as np

# A curved ravine (Rosenbrock style): minimum at (1, 1), steep walls, flat floor
f = lambda p: (1 - p[0]) ** 2 + 10 * (p[1] - p[0] ** 2) ** 2
grad = lambda p: np.array([-2 * (1 - p[0]) - 40 * p[0] * (p[1] - p[0] ** 2), 20 * (p[1] - p[0] ** 2)])

def run(kind, lr, steps, b1=0.9, b2=0.999, eps=1e-8):
    p, m, v = np.array([-1.5, 2.0]), np.zeros(2), np.zeros(2)
    for t in range(1, steps + 1):
        g = grad(p)
        if kind == "SGD":
            p = p - lr * g
        elif kind == "Momentum":                       # velocity remembers past gradients
            m = b1 * m + g; p = p - lr * m
        elif kind == "RMSProp":                        # divide by a running RMS of gradients
            v = 0.9 * v + 0.1 * g ** 2; p = p - lr * g / (np.sqrt(v) + eps)
        else:                                          # Adam: momentum + RMS scaling + bias correction
            m = b1 * m + (1 - b1) * g; v = b2 * v + (1 - b2) * g ** 2
            p = p - lr * (m / (1 - b1 ** t)) / (np.sqrt(v / (1 - b2 ** t)) + eps)
    return p

for kind, lr in [("SGD", 0.01), ("Momentum", 0.002), ("RMSProp", 0.01), ("Adam", 0.05)]:
    row = [f"{f(run(kind, lr, n)):.4f}" for n in (25, 100, 300)]
    print(f"{kind:<9} lr {lr:<5} loss after 25 / 100 / 300 steps: {' / '.join(row)}")
print(f"SGD with lr 0.03: loss after 1 / 2 / 3 steps: {f(run('SGD', 0.03, 1)):.3g} / {f(run('SGD', 0.03, 2)):.3g} / {f(run('SGD', 0.03, 3)):.3g} (diverging)")

# Weight decay: Adam with L2 in the gradient vs AdamW's decoupled decay
rng = np.random.default_rng(0)
def adam_decay(decoupled, lr=0.01, wd=0.1, steps=2000):
    p, m, v = np.ones(2), np.zeros(2), np.zeros(2)
    for t in range(1, steps + 1):
        g = rng.normal(0, 1, 2) * np.array([100.0, 0.01])  # noisy gradients, very different scales
        if not decoupled: g = g + wd * p                     # L2: decay enters the adaptive scaling
        m = 0.9 * m + 0.1 * g; v = 0.999 * v + 0.001 * g ** 2
        p = p - lr * (m / (1 - 0.9 ** t)) / (np.sqrt(v / (1 - 0.999 ** t)) + 1e-8)
        if decoupled: p = p - lr * wd * p                     # AdamW: shrink every weight equally
    return np.abs(p)
print("\n|weights| after 2000 steps, starting from 1 (large-gradient, small-gradient):")
print("  Adam + L2:", np.round(adam_decay(False), 3), "  AdamW:", np.round(adam_decay(True), 3))

SGD       lr 0.01  loss after 25 / 100 / 300 steps: 5.1262 / 2.6481 / 0.0485
Momentum  lr 0.002 loss after 25 / 100 / 300 steps: 4.8812 / 0.2462 / 0.0016
RMSProp   lr 0.01  loss after 25 / 100 / 300 steps: 5.5480 / 4.0852 / 0.0445
Adam      lr 0.05  loss after 25 / 100 / 300 steps: 5.0688 / 0.4684 / 0.0001
SGD with lr 0.03: loss after 1 / 2 / 3 steps: 21.6 / 143 / 3.15e+04 (diverging)

|weights| after 2000 steps, starting from 1 (large-gradient, small-gradient):
  Adam + L2: [1.448 0.009]   AdamW: [0.053 0.104]

On this ravine after 300 steps, Adam reaches loss 0.0001, momentum 0.0016, RMSProp 0.0445 and SGD 0.0485, each with a learning rate tuned to be reasonable for it. Do not read this as a universal ranking: on some problems well-tuned SGD with momentum generalizes better than Adam, and it remains common for training image models.

The second block isolates weight decay. Two weights receive pure noise gradients, one 10,000 times larger than the other, and should only shrink from decay. With Adam plus L2, the large-gradient weight does not shrink at all (it drifts to 1.448), while the small-gradient one collapses to 0.009. With AdamW both shrink at a comparable rate (0.053 and 0.104).

Memory and state

Optimizers carry state: momentum keeps one extra number per weight, Adam two (m and v). For a model with a billion parameters stored as 32-bit floats, Adam's state alone takes 8 GB, on top of the weights and gradients. This is why memory-saving optimizers and sharding optimizer state across GPUs matter for large models.

Trade-offs

Common mistakes

• Using the same learning rate when switching from SGD to Adam; Adam's typical rates are 10 to 100 times smaller.
• Adding L2 regularization to the loss with Adam and expecting it to act like weight decay; use AdamW.
• Applying weight decay to biases and normalization parameters, which usually should be excluded.
• Concluding that an optimizer "does not work" from one learning rate without a short sweep.

In the wild

AdamW is the default optimizer for training transformer language and vision models, usually with a warmup and cosine learning rate schedule ( Learning rate schedules and warmup ). SGD with momentum still trains many convolutional image classifiers, where its tuned schedules are well understood. Research continues on optimizers that need less memory or fewer tuning runs, but almost all of them are variations on the two ideas here: averaging gradients over time and scaling steps per weight.

Check yourself

Recap