Optimizers — SGD, Adam, AdamW

After loss.backward(), every parameter has a .grad — the direction that would increase the loss. The optimizer’s job is to convert those gradients into actual weight updates. Different optimizers make different bets about how to do that.

TryOptimizers · the race

Race four optimizers down the same hill

Same surface, same start, same learning rate — only the update rule differs. Hit Race and watch SGD zig-zag down the ravine while Momentum and Adam cut straight to the bottom. Click the plot to move the start; slide the learning rate up until SGD diverges.

surfaceRavineSaddle-ish bowl

runnersSGDMomentumRMSPropAdam

Click to move the start. minimum

leaderboardstep 80 / 80

1. 1Momentum0.001
2. 2SGD1.047
3. 3RMSProp1.252
4. 4Adam1.780

learning rate0.060

0.002log scale0.6

Race againReset

Momentum smooths the zig-zag; RMSProp and Adam scale the step per axis, so they handle the steep wall and the flat floor at once.

SGD — the baseline

Vanilla stochastic gradient descent does the obvious thing:

w ← w − lr · grad

Take a step opposite the gradient, scaled by the learning rate. Simple, unbiased, and surprisingly hard to beat with the right learning rate schedule. SGD is still the gold standard for training ResNets on ImageNet — but it’s slow to tune and brittle to scale.

Momentum adds a velocity that accumulates past gradients:

v ← β · v + grad
w ← w − lr · v

With β ≈ 0.9, the velocity is a weighted average of all past gradients, with recent ones weighted most. In a consistent direction (flat valley floor) gradients always point the same way, so they accumulate — the optimizer accelerates. In an oscillating direction (steep walls of a ravine) gradients cancel each other out, so velocity stays small — oscillations are damped. SGD-with-momentum is what “SGD” usually means in papers.

Adam — the default that took over

Adam keeps two running estimates per parameter:

• m — EMA of the gradient (like momentum)
• v — EMA of the squared gradient

The update divides m by √v. Each parameter effectively gets its own adaptive learning rate, scaled down where gradients are consistently large.

Let’s implement Adam in NumPy on a tiny problem.

import numpy as np

# Toy loss: f(w) = (w1 - 3)^2 + 10*(w2 + 1)^2
# Gradient: [2(w1-3), 20(w2+1)]
def grad(w):
    return np.array([2 * (w[0] - 3), 20 * (w[1] + 1)])

# Adam state
w = np.array([0.0, 0.0])
m = np.zeros_like(w)         # 1st moment (gradient EMA)
v = np.zeros_like(w)         # 2nd moment (squared-gradient EMA)
beta1, beta2, eps = 0.9, 0.999, 1e-8
lr = 0.5

print(f"{'step':>5} {'w1':>8} {'w2':>8} {'loss':>10}")
for t in range(1, 41):
    g = grad(w)
    # 1. Update biased moments
    m = beta1 * m + (1 - beta1) * g
    v = beta2 * v + (1 - beta2) * g * g
    # 2. Bias correction (matters early when m,v are still close to 0)
    m_hat = m / (1 - beta1 ** t)
    v_hat = v / (1 - beta2 ** t)
    # 3. Update
    w = w - lr * m_hat / (np.sqrt(v_hat) + eps)
    if t % 5 == 0:
        loss = (w[0] - 3)**2 + 10 * (w[1] + 1)**2
        print(f"{t:5d} {w[0]:8.3f} {w[1]:8.3f} {loss:10.4f}")

print("target: w1 = 3.0, w2 = -1.0")

 step       w1       w2       loss
    5    2.403   -1.585     3.7841
   10    3.953   -0.820     1.2343
   15    3.966   -0.778     1.4256
   20    3.164   -1.239     0.5960
   25    2.534   -0.992     0.2183
   30    2.565   -0.864     0.3736
   35    2.992   -1.101     0.1031
   40    3.263   -1.001     0.0690
target: w1 = 3.0, w2 = -1.0

Notice how Adam handles the scale difference in the gradients (w2’s gradient is 10x larger) automatically — it normalizes each parameter by its own typical magnitude. With plain SGD you’d need to pick a learning rate small enough for w2 (which crawls w1) or large enough for w1 (which oscillates w2).

AdamW — what you actually use

Adam plus “weight decay” is what most people think they’re doing when they pass weight_decay=0.01 to Adam — but the original implementation mixes weight decay into the gradient before the adaptive scaling, which isn’t the same as L2 regularization for Adam. With SGD the two are mathematically equivalent; with Adam they are not, because dividing by √v distorts the effective penalty per parameter. The fix is AdamW.

AdamW (Adam with decoupled Weight Decay) applies weight decay separately from the gradient step:

w ← w − lr · update − lr · wd · w

The lr · wd · w term is just “pull each weight a little toward zero.” It does what you intended when you read “L2 regularization” in a paper.

RMSProp (briefly)

RMSProp is “Adam minus the first moment” — only the squared-gradient EMA. Mostly historical (Adam dominates), but still shows up in RL where adaptive lr without momentum is sometimes more stable.

Now that all four are on the table, picture them on the same surface from the same start. Plain SGD bounces off the steep walls of the ravine; Momentum smooths that out; RMSProp and Adam rescale each axis and glide to the bottom:

What’s new: Muon, Lion, Sophia

AdamW is still the safe default in 2026, but it’s no longer unchallenged. A few newer optimizers are worth knowing by name:

• Muon — the most credible challenger. It orthogonalizes the momentum update for 2D weight matrices (the hidden layers), and is typically paired with AdamW for the embeddings and output head. It was used to train the Moonlight 16B MoE model and reported roughly 2× the compute efficiency of AdamW at large batch sizes. A 2026 optimizer lesson that ignores Muon is already dated.
• Lion — discovered by symbolic search; uses only the sign of a momentum term, so it stores one state per parameter instead of Adam’s two (less memory). Often competitive with AdamW at a lower learning rate.
• Sophia — a light second-order method that uses a cheap Hessian estimate; shown to speed up LLM pretraining in some settings.

The honest takeaway: start with AdamW. Reach for Muon if you’re training at scale and want the efficiency win, and treat Lion/Sophia as experiments to benchmark, not defaults.

How to pick

Situation	Optimizer	Typical lr
Transformer / LLM	AdamW	1e-4 to 3e-4
Vision transformer	AdamW	1e-3
ResNet on ImageNet	SGD + momentum	0.1 (with cosine decay)
Fine-tuning	AdamW	1e-5 to 5e-5
Reinforcement learning	Adam or RMSProp	3e-4
Large-scale LLM pretraining	AdamW (or Muon for the matrices)	1e-4 to 3e-4
Anything else	AdamW	start at 1e-3

Putting it together in PyTorch

from torch.optim import AdamW

optimizer = AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

for x, y in loader:
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

That’s the default. Change the optimizer only when you have a specific reason — replicating a paper or you’ve profiled an alternative.

In one breath

• The optimizer turns each parameter’s gradient into an actual weight update; SGD just steps opposite the gradient, scaled by the learning rate.
• Momentum accumulates a velocity (an EMA of past gradients) — it accelerates in consistent directions and damps oscillations in a ravine.
• Adam keeps an EMA of the gradient and of the squared gradient and divides one by the other, giving each parameter its own adaptive step — robust to badly-scaled gradients.
• Use AdamW, not Adam, whenever weight_decay > 0: it decouples weight decay from the adaptive step, matching the L2 effect you intended.
• AdamW is the 2026 default (Muon is the credible challenger at scale), and the learning rate is the king hyperparameter — tune it first, on a log scale.

Quick check

Quick check

0/2

Q1Why does Adam often converge faster than SGD on poorly-scaled problems?

Adam uses a higher learning rate by defaultAdam normalizes each parameter by its own gradient magnitude (the v term)Adam runs on GPU; SGD doesn'tAdam adds noise that helps escape local minima

Check answer

Q2You're training a transformer with `weight_decay=0.01`. Should you use Adam or AdamW?

Adam — they're the same when weight_decay > 0AdamW — it applies weight decay decoupled from the adaptive update, which is what you actually wantSGD — Adam can't do weight decayDoesn't matter as long as `weight_decay > 0`

Check answer