“What is the difference between SGD and Adam?” and “why AdamW instead of Adam?” are standard questions in deep learning interviews, and they get harder quickly: why does Adam need bias correction, why do transformers use warmup, why might SGD with momentum still generalize better on image models. A good answer shows that you can write each update rule from memory, explain the problem it solves, and say what it does not solve. Interviewers also listen for practical judgment, such as knowing that fine-tuning learning rates are one or two orders of magnitude smaller than training-from-scratch rates.
Before you start
You need gradient descent (w = w - lr * grad) and the idea that a mini-batch gives a noisy estimate of the full gradient. The experiments use a two-parameter loss so every step can be printed, and run on Python 3.12 or newer with the standard library only. PyTorch snippets are illustrative (PyTorch 2.x); the defaults quoted come from the torch.optim documentation.
The short answer
SGD moves each weight against its mini-batch gradient by a single global learning rate. Momentum keeps a running sum of past gradients, which smooths noise and speeds progress along consistent directions. Adam keeps two running averages per parameter, the mean gradient and the mean squared gradient, and divides one by the square root of the other, so every parameter gets a step of roughly the learning rate regardless of its gradient’s scale. AdamW is Adam with weight decay applied directly to the weights rather than added to the gradient, because inside Adam an L2 term gets rescaled and stops acting like real decay. A learning-rate schedule, typically linear warmup followed by cosine or linear decay, sits on top of any of them.
How it works
The update rules, in the form PyTorch implements them:
SGD: w -= lr * g
Momentum: v = mu * v + g; w -= lr * v
Adam: m = b1 * m + (1 - b1) * g
s = b2 * s + (1 - b2) * g**2
m_hat = m / (1 - b1**t); s_hat = s / (1 - b2**t)
w -= lr * m_hat / (sqrt(s_hat) + eps)
AdamW: w -= lr * wd * w, then the Adam step on g alonePyTorch’s defaults for Adam and AdamW are lr=1e-3, betas=(0.9, 0.999) and eps=1e-8. The difference is weight_decay: 0 for Adam, 0.01 for AdamW. Passing weight_decay to Adam adds wd * w to the gradient (classic L2), which is not the same thing, as Step 4 shows.
The bias correction exists because m and s start at zero. After one step, m = 0.1 * g, which badly underestimates the gradient; dividing by 1 - 0.9**1 = 0.1 restores it. As t grows the correction approaches 1 and stops mattering.
Step-by-step walkthrough
The test problem is an elongated bowl, f(x, y) = 0.5 * (x**2 + 50 * y**2). It is steep in y and flat in x, the same shape as the ill-conditioned valleys real loss surfaces have.
Step 1: Implement the optimizers
import math
def loss(p):
return 0.5 * (p[0] ** 2 + 50 * p[1] ** 2)
def grad(p):
return [p[0], 50 * p[1]]
class SGD:
def __init__(self, lr, momentum=0.0):
self.lr, self.mu, self.v = lr, momentum, None
def step(self, p, g):
if self.v is None:
self.v = [0.0] * len(p)
self.v = [self.mu * v + gi for v, gi in zip(self.v, g)] # momentum buffer
return [pi - self.lr * v for pi, v in zip(p, self.v)]
class Adam:
def __init__(self, lr, betas=(0.9, 0.999), eps=1e-8):
self.lr, (self.b1, self.b2), self.eps = lr, betas, eps
self.m = self.v = None
self.t = 0
def step(self, p, g):
if self.m is None:
self.m, self.v = [0.0] * len(p), [0.0] * len(p)
self.t += 1
self.m = [self.b1 * m + (1 - self.b1) * gi for m, gi in zip(self.m, g)]
self.v = [self.b2 * v + (1 - self.b2) * gi * gi for v, gi in zip(self.v, g)]
mhat = [m / (1 - self.b1 ** self.t) for m in self.m] # bias correction
vhat = [v / (1 - self.b2 ** self.t) for v in self.v]
return [pi - self.lr * mh / (math.sqrt(vh) + self.eps) for pi, mh, vh in zip(p, mhat, vhat)]
def steps_to(opt, target=1e-3, start=(10.0, 1.0), limit=5000):
p = list(start)
for t in range(1, limit + 1):
p = opt.step(p, grad(p))
if not math.isfinite(loss(p)) or loss(p) > 1e12:
return "diverged"
if loss(p) < target:
return t
return f">{limit}"Step 2: Watch SGD hit its ceiling and momentum break through it
for name, opt in [
("SGD lr=0.041", SGD(0.041)),
("SGD lr=0.039", SGD(0.039)),
("momentum 0.6, lr=0.06", SGD(0.06, momentum=0.6)),
("Adam lr=0.5", Adam(0.5)),
]:
print(f"{name:24s} {steps_to(opt)}")
# SGD lr=0.041 diverged
# SGD lr=0.039 137
# momentum 0.6, lr=0.06 21
# Adam lr=0.5 80For gradient descent on a quadratic, any learning rate above 2 / curvature diverges. The steep direction has curvature 50, so the limit is 0.04, and 0.041 blows up. At 0.039 the flat x direction shrinks by only about 4 percent per step, so it takes 137 steps. Momentum raises the stable limit (to 2 * (1 + mu) / 50, or 0.064 for mu = 0.6) and accumulates velocity along the consistent x direction while oscillations in y partly cancel: 21 steps. The momentum settings came from a small grid search, which is the honest part of the story: tuned momentum is very hard to beat.
Adam takes 80 steps here. It is not the fastest on a problem where SGD has been tuned, because its steps stay close to lr in size and it overshoots near the minimum. Its strength is elsewhere.
Step 3: Change the gradient scale and see what Adam buys you
Multiply the loss by 1000, as happens when you change a loss reduction from mean to sum or rescale a target:
def steps_scaled(opt, scale, limit=5000):
p = [10.0, 1.0]
for t in range(1, limit + 1):
p = opt.step(p, [scale * gi for gi in grad(p)])
if not math.isfinite(loss(p)) or loss(p) > 1e12:
return "diverged"
if loss(p) < 1e-3:
return t
return f">{limit}"
for scale in (1, 1000):
print(scale, steps_scaled(SGD(0.039), scale), steps_scaled(Adam(0.5), scale))
# 1 137 80
# 1000 diverged 80SGD’s learning rate is tied to the gradient’s units; Adam’s is not, because the ratio m_hat / sqrt(s_hat) cancels the scale. The same property applies per parameter: embeddings for rare tokens, layers with very different gradient magnitudes and sparse features all get sensible step sizes without hand tuning. That robustness is why Adam-family optimizers are the default for transformers.
The flip side shows up on the very first step. With bias correction, step 1 is about lr * sign(g) for every parameter, however small or large its gradient:
print([round(v, 6) for v in Adam(0.001).step([0.0, 0.0, 0.0], [1e-6, 3.0, -500.0])])
# [-0.00099, -0.001, 0.001]Early on, s_hat is estimated from very few samples, so Adam’s steps are large relative to how much it knows. Warmup exists largely for this reason.
Step 4: Compare Adam plus L2 with AdamW
Give a weight a data gradient that is pure noise, so only weight decay should shrink it, and try two noise levels:
import random
def run(decoupled, noise_std, wd=0.1, lr=0.01, steps=200, seed=1):
rng = random.Random(seed)
opt, w = Adam(lr), [1.0]
for _ in range(steps):
g = [rng.gauss(0, noise_std)]
if decoupled:
w = [w[0] - lr * wd * w[0]] # AdamW: shrink the weight directly
w = opt.step(w, g)
else:
w = opt.step(w, [g[0] + wd * w[0]]) # Adam + L2: decay passes through 1/sqrt(s)
return round(w[0], 3)
for noise in (0.01, 10.0):
print(f"noise {noise:5}: Adam+L2 w={run(False, noise)} AdamW w={run(True, noise)}")
# noise 0.01: Adam+L2 w=0.005 AdamW w=0.708
# noise 10.0: Adam+L2 w=0.862 AdamW w=0.708With L2 inside Adam, the decay term is normalized along with the gradient. A weight with tiny gradients is crushed to almost zero, while a weight with large gradients is barely decayed at all. AdamW gives both the same treatment (the two runs match because Adam’s normalized noise steps are identical at both scales). Loshchilov and Hutter’s “Decoupled Weight Decay Regularization” showed this decoupling generalizes better, and AdamW became the standard.
Step 5: Add warmup and cosine decay
def lr_at(step, peak=3e-4, warmup=1000, total=10000, floor=3e-5):
if step < warmup:
return peak * (step + 1) / warmup
progress = (step - warmup) / (total - warmup)
return floor + 0.5 * (peak - floor) * (1 + math.cos(math.pi * progress))
print([f"{lr_at(s):.2e}" for s in (0, 499, 999, 1000, 5500, 9999)])
# ['3.00e-07', '1.50e-04', '3.00e-04', '3.00e-04', '1.65e-04', '3.00e-05']Warmup ramps the rate while Adam’s statistics and the network’s activations settle. Decay then lowers it so the optimizer can settle into a minimum instead of bouncing around it.
Worked scenario
A team fine-tunes a pretrained BERT-style classifier with torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=0.01) and no schedule. The loss spikes in the first hundred steps, and the final accuracy is below a simple logistic regression on the same features.
Three settings combine. 1e-3 is a from-scratch learning rate; with Adam’s roughly lr-sized first steps, every weight (typical magnitude a few hundredths) moves by a large fraction of its value at once, erasing what pretraining learned. Without warmup, nothing protects those early steps. And weight_decay on Adam is L2, not decoupled decay, so it also hits LayerNorm gains and biases unevenly. The illustrative PyTorch 2.x fix:
# Illustrative PyTorch 2.x
decay, no_decay = [], []
for name, p in model.named_parameters():
(no_decay if p.ndim < 2 else decay).append(p) # biases and norm weights skip decay
optimizer = torch.optim.AdamW(
[{"params": decay, "weight_decay": 0.01}, {"params": no_decay, "weight_decay": 0.0}],
lr=2e-5,
)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda s: lr_at(s, peak=1.0, floor=0.0))
# per step: loss.backward(); clip_grad_norm_(model.parameters(), 1.0); optimizer.step(); scheduler.step(); optimizer.zero_grad()LambdaLR multiplies the base lr by the returned factor, which is why peak=1.0 here. With a learning rate around 2e-5, warmup over the first few percent of steps and decoupled decay, fine-tuning is stable and beats the baseline.
Common mistake
- “Adam always converges faster.” It is more forgiving of scale. Tuned SGD with momentum matches or beats it on many vision tasks and sometimes generalizes better.
- “Adam with
weight_decayis AdamW.” It is L2 regularization, which Adam’s normalization distorts. - “Adam has a per-parameter learning rate, so the global one does not matter.” The global
lrstill sets the step size; Adam only normalizes the gradient’s scale. - Applying weight decay to everything. Decaying biases and LayerNorm gains usually hurts or does nothing useful; exclude them with parameter groups.
- Forgetting that optimizer state costs memory. Adam stores two extra float32 values per parameter, so for a 7B-parameter model that is 56 GB on top of weights and gradients.
Verify the behavior
def test_sgd_limit_is_two_over_curvature():
assert steps_to(SGD(0.039)) != "diverged" and steps_to(SGD(0.041)) == "diverged"
def test_adam_is_scale_invariant():
assert steps_scaled(Adam(0.5), 1) == steps_scaled(Adam(0.5), 1000)
def test_adamw_decay_ignores_gradient_scale():
assert run(True, 0.01) == run(True, 10.0)
assert run(False, 0.01) < 0.1 < 0.8 < run(False, 10.0)
test_sgd_limit_is_two_over_curvature(); test_adam_is_scale_invariant(); test_adamw_decay_ignores_gradient_scale()
print("ok")Follow-up questions
What does Nesterov momentum change? Conceptually, it evaluates the gradient at the look-ahead point w - lr * mu * v instead of the current point, which corrects overshoot sooner. In PyTorch it is SGD(..., momentum=0.9, nesterov=True).
Why does batch size interact with the learning rate? Larger batches give less noisy gradients, so a larger rate is usually safe; a linear scaling rule with warmup is a common starting point, though it breaks down at very large batches.
What is the role of eps in Adam? It prevents division by zero for parameters whose gradients are near zero. Raising it makes Adam behave more like momentum SGD, and with low-precision training it sometimes needs to be larger than 1e-8.
When would you pick SGD with momentum today? For convolutional vision models trained from scratch with well-known recipes, or when memory for Adam’s two state tensors is tight.
Interview exercise
Training a transformer from scratch with AdamW at lr=6e-4, the loss falls for 300 steps, then spikes and turns into NaN. Gradient norms climb steadily before the spike. There is no warmup and no gradient clipping. What do you change, and in what order?
Answer and reasoning
Rising gradient norms before the spike point to an instability, not bad data in one batch. First add a linear warmup (for example over the first 1 to 5 percent of steps) so the early steps, which Adam sizes at roughly the learning rate with poorly estimated statistics, are small. Second, clip the global gradient norm (1.0 is a common value) and log the pre-clipping norm so you can see whether spikes persist. Third, if it still diverges, lower the peak rate, for example to 3e-4, and check the setup: pre-LayerNorm placement, bfloat16 rather than float16 without loss scaling, and initialization of residual branches. Changing one thing at a time, with a fixed seed, is what makes the diagnosis credible. Switching to SGD is the wrong first move; it hides the problem and usually trains transformers worse.