AdaGrad, RMSProp, Adam, AdamW

Before this lesson: momentum as a running sum of gradients (lesson 04), minibatch noise (lesson 05), and diagonal preconditioning by feature scaling (lesson 02). Time: about 70 minutes, plus the notebook.

In plain words

Gradient descent uses one step size for all parameters. But some parameters have large gradients and some have tiny ones, so one step size is too large for some and too small for others. Adaptive methods give every parameter its own step size. They keep a running average of each parameter's squared gradient and divide the step by its square root. A parameter with large gradients gets a small step; a parameter with small gradients gets a large step. AdaGrad sums all past squared gradients, RMSProp uses a moving average, and Adam adds momentum and a correction for the zero start. This fixes bad scaling along the coordinate axes, but not along tilted directions. Finally, weight decay behaves differently inside Adam: AdamW applies it directly to the weights, and this is the standard optimizer for training transformers today.

Why this matters #

Adam (Kingma and Ba 2015) and AdamW (Loshchilov and Hutter 2019) train almost every large neural network: language models (GPT-3 and LLaMA used AdamW with $\beta_2 = 0.95$ ), vision transformers, diffusion models and weather models such as GraphCast. Small details of the update decide whether a run is stable: the value of $\varepsilon$ , the bias correction, and whether weight decay goes through the gradient. Interviewers ask exactly these details, and you need them to debug an unstable run or to resume one correctly. Lesson 07 adds the learning-rate schedules, warmup and clipping that complete the recipe.

Notation. In this lesson and the next one we count steps with $t = 1, 2, \dots$ , as in the Adam paper and the PyTorch documentation. $\theta$ are the parameters, $g_t$ the (minibatch) gradient at step $t$ , $\eta$ the learning rate (lr), $\lambda$ the weight-decay coefficient. Squares, square roots and divisions of vectors act per coordinate.

1. Predict first #

2. The picture #

Moving averages remember the recent past #

Momentum (lesson 04) and Adam both use an exponential moving average (EMA):

m_t = \beta\, m_{t-1} + (1 - \beta)\, g_t, \qquad m_0 = 0 .

Unroll it: $m_t = (1-\beta)\big(g_t + \beta g_{t-1} + \beta^2 g_{t-2} + \dots\big)$ . The weight of a gradient decays by the factor $\beta$ per step. The weights $(1-\beta)\beta^j$ add up to almost 1, and most of the weight sits on the last $1/(1-\beta)$ steps: about 10 steps for $\beta = 0.9$ , 1000 steps for $\beta = 0.999$ . Components that keep their sign add up; components that flip sign cancel.

The zero start is a problem. At $t = 1$ , $m_1 = (1-\beta)g_1$ : for $\beta = 0.9$ this is only a tenth of the gradient. Dividing by $1 - \beta^t$ fixes this: $\hat m_t = m_t/(1 - \beta^t)$ is the bias-corrected average, and $\hat m_1 = g_1$ .

Slide $\beta$ and watch the two averages of the noisy gradients (grey dots). The orange EMA without correction starts near 0 and needs about $1/(1-\beta)$ steps to catch up; the aqua corrected average is right from the first step. With jump at t = 50, a large $\beta$ reacts slowly: long memory means smooth but late. The bars show the weight of each past gradient in the last average; the dashed line marks $1/(1-\beta)$ steps back.

Quick check

$\beta = 0.9$ , $m_0 = 0$ and the first gradient is $g_1 = 2$ . What are $m_1$ and the bias-corrected $\hat m_1$ ?

One step size per coordinate #

Take $f(x) = \tfrac12(x_1^2 + 100x_2^2)$ . The gradient is $(x_1, 100x_2)$ . Gradient descent is stable only for $\eta < 2/100$ (lesson 01), so $x_1$ moves at most 2% of its distance per step. Now divide each coordinate of the step by the size of its own gradient, as Adam does. The factor 100 cancels, and both coordinates move at the same speed. In matrix form the step is $-\eta D\nabla f$ with a diagonal matrix $D$ : a diagonal preconditioner, like feature scaling in lesson 02, but learned from the gradients during training.

A diagonal preconditioner can only rescale the coordinate axes. Rotate the same valley by $45°$ : the steep and the flat direction are now diagonal. Both coordinates of the gradient mix the two directions and have similar sizes, so $D$ is close to a multiple of the identity, and nothing cancels.

Race four optimizers. Gradient descent (blue), momentum (orange), RMSProp (aqua, dashed) and Adam (aqua, solid) start from the same point (drag it). On the axis-aligned quadratic, Adam and RMSProp go almost straight to the minimum. Switch to rotated 45°: with the same learning rates they lose their advantage and stall, while gradient descent is unchanged. Then try the saddle (Adam escapes in a few steps, because it rescales the tiny gradient along $x_2$ ), Rosenbrock and the double well, and turn on gradient noise. Each optimizer has its own learning-rate slider; Default learning rates resets them.

Quick check

Adam, with a well-tuned learning rate, quickly solves $f(x) = \tfrac12(x_1^2 + 100x_2^2)$ . You rotate the problem by $45°$ . What happens?

3. Definitions and main results #

The exponential moving average #

Definition. For $\beta \in [0, 1)$ , the EMA of a sequence $g_1, g_2, \dots$ is $m_t = \beta m_{t-1} + (1-\beta) g_t$ with $m_0 = 0$ . Unrolled:

m_t = (1-\beta)\sum_{i=1}^{t} \beta^{t-i} g_i .

The weights $(1-\beta)\beta^{t-i}$ sum to $(1-\beta)(1 + \beta + \dots + \beta^{t-1}) = 1 - \beta^t$ . Their "average age" is $\sum_{j \ge 0} j(1-\beta)\beta^j = \beta/(1-\beta)$ , so the memory is about $1/(1-\beta)$ steps.

Heavy-ball momentum in PyTorch form, $b_t = \beta b_{t-1} + g_t$ (lesson 04), is the same average without the factor $1-\beta$ : $b_t = m_t/(1-\beta)$ . With $\beta = 0.9$ the momentum buffer is 10 times the EMA, which is why a constant gradient gives the effective step $\eta g/(1-\beta)$ .

AdaGrad, RMSProp and Adam #

All state starts at 0. Here are the PyTorch rules (without weight decay) and their defaults.

AdaGrad (Duchi, Hazan and Singer 2011; $\eta = 10^{-2}$ , $\varepsilon = 10^{-10}$ ):

s_t = s_{t-1} + g_t^2, \qquad \theta_t = \theta_{t-1} - \eta\,\frac{g_t}{\sqrt{s_t} + \varepsilon} .

$s_t$ only grows. With gradients of steady size $|g|$ , $\sqrt{s_t} = |g|\sqrt t$ , so the step decays like $\eta/\sqrt t$ . This is good for convex problems with sparse features (rare features keep large steps), but too fast a decay for long nonconvex training.

RMSProp (Hinton's lecture notes, 2012; $\eta = 10^{-2}$ , $\alpha = 0.99$ , $\varepsilon = 10^{-8}$ ):

v_t = \alpha v_{t-1} + (1-\alpha)\, g_t^2, \qquad \theta_t = \theta_{t-1} - \eta\,\frac{g_t}{\sqrt{v_t} + \varepsilon} .

The sum of AdaGrad becomes an EMA, so old gradients are forgotten and the step does not decay to zero. RMSProp has no bias correction: $v_1 = 0.01\,g_1^2$ , so $\sqrt{v_1} = 0.1|g_1|$ and the first step is $10\,\eta$ in every coordinate.

Adam (Kingma and Ba 2015; $\eta = 10^{-3}$ , $\beta_1 = 0.9$ , $\beta_2 = 0.999$ , $\varepsilon = 10^{-8}$ ):

\begin{aligned} m_t &= \beta_1 m_{t-1} + (1-\beta_1)\, g_t, & \hat m_t &= m_t/(1-\beta_1^t),\\ v_t &= \beta_2 v_{t-1} + (1-\beta_2)\, g_t^2, & \hat v_t &= v_t/(1-\beta_2^t),\\ \theta_t &= \theta_{t-1} - \eta\,\frac{\hat m_t}{\sqrt{\hat v_t} + \varepsilon} . && \end{aligned}

Adam is RMSProp plus momentum ( $m$ ) plus bias correction. $\hat m_t$ estimates the mean of the gradient, $\hat v_t$ the mean of its square (the "uncentered variance"), so the ratio $\hat m_t/\sqrt{\hat v_t}$ is like a signal-to-noise ratio per coordinate.

Method	State per parameter	Bias correction	Effective step for a steady gradient
SGD	none	–	$\eta\,\lvert g\rvert$
AdaGrad	$s$	–	$\eta/\sqrt t$
RMSProp	$v$	no	$\eta$ (first step $10\eta$ )
Adam	$m$ , $v$	yes	$\eta$

Proposition 1 (bias correction) #

If the gradients are stationary, $\mathbb E[g_i] = \bar g$ and $\mathbb E[g_i^2] = \bar s$ for all $i$ , then $\mathbb E[m_t] = (1-\beta_1^t)\,\bar g$ and $\mathbb E[v_t] = (1-\beta_2^t)\,\bar s$ . So $\hat m_t$ and $\hat v_t$ are unbiased.

Proof.

Unroll from $m_0 = 0$ : $m_t = (1-\beta_1)\sum_{i=1}^{t}\beta_1^{t-i} g_i$ .
Take expectations: $\mathbb E[m_t] = (1-\beta_1)\,\bar g\sum_{j=0}^{t-1}\beta_1^{j}$ .
Geometric sum: $(1-\beta_1)\sum_{j=0}^{t-1}\beta_1^j = 1 - \beta_1^t$ .
The same steps with $g_i^2$ and $\beta_2$ give $\mathbb E[v_t] = (1-\beta_2^t)\,\bar s$ . Dividing by $1 - \beta^t$ removes the bias that the zero start creates. $\blacksquare$

(Kingma and Ba add a small error term when $\mathbb E[g^2]$ drifts slowly.)

Why the correction matters most for $v$ . With $\beta_2 = 0.999$ , $1 - \beta_2^t \approx t/1000$ for small $t$ , so $v_t$ is far too small for hundreds of steps. Without any correction, at $t = 1$ : $m_1 = 0.1\,g_1$ and $\sqrt{v_1} = \sqrt{0.001}\,|g_1| \approx 0.0316\,|g_1|$ , so the step is about $3.16\,\eta$ . With the correction, $\hat m_1 = g_1$ and $\hat v_1 = g_1^2$ , so the first step is

\eta\,\frac{g_1}{|g_1| + \varepsilon} \approx \eta\,\operatorname{sign}(g_1):

size $\eta$ in every coordinate, whatever the scale of the gradient.

Two consequences (Kingma and Ba §2.1).

Scale invariance. Multiply all gradients by $c > 0$ : $\hat m$ and $\sqrt{\hat v}$ both grow by $c$ , and the step does not change (as long as $\varepsilon$ is negligible).
Bounded step. By Jensen's inequality, $|\mathbb E g| \le \sqrt{\mathbb E g^2}$ , so typically $|\Delta\theta_t| \lesssim \eta$ . When the gradient sign is noisy, $|\hat m|$ is much smaller than $\sqrt{\hat v}$ , and the step shrinks automatically.

Type a gradient (or pick an example) and read the first three steps of each method, in units of $\eta$ . SGD copies the gradient: steps of $1000\eta$ and $0.001\eta$ . RMSProp (light aqua) starts with $10\eta$ ; Adam without correction (grey) starts with about $3.16\eta$ ; Adam (aqua) moves by $\eta$ in both coordinates. With (100, 0.01), then flip, the coordinate whose sign flips moves back by only 5% of $\eta$ . Set $\varepsilon = 1$ to see Adam lose its adaptivity.

Quick check

At the first step $g_1 = (1000, 0.001)$ . Adam with $\eta = 10^{-3}$ and the default $\varepsilon = 10^{-8}$ : what is the step $-\Delta\theta_1$ ?

The roles of $\beta_2$ and $\varepsilon$ #

$\beta_2$ . $v$ remembers about $1/(1-\beta_2)$ steps: 1000 for $0.999$ . A long memory gives a smooth estimate, but a slow one: after a sudden rise of the gradient scale, $\hat v$ lags behind and the steps exceed $\eta$ for a while. This is one reason for loss spikes in large runs. Many LLM and weather-model recipes use $\beta_2 = 0.95$ (GPT-3, LLaMA, GraphCast), with a memory of about 20 steps.

$\varepsilon$ . If $\sqrt{\hat v} \gg \varepsilon$ , it does nothing. If $\sqrt{\hat v} \ll \varepsilon$ , the update is about $(\eta/\varepsilon)\,\hat m$ : momentum SGD with step $\eta/\varepsilon$ . So a large $\varepsilon$ makes Adam less adaptive, and coordinates with gradients below $\varepsilon$ are not rescaled.

Quick check

You train with Adam and set $\varepsilon = 1$ , while the gradients have size about $10^{-3}$ . What does the optimizer behave like?

Weight decay and L2 regularization #

Weight decay shrinks the weights by a fixed factor at every step: $\theta \leftarrow (1 - \eta\lambda)\theta - \eta\, d_t$ , where $d_t$ is the step direction that the optimizer computes from the loss. L2 regularization adds $\tfrac{\lambda}{2}\|\theta\|^2$ to the loss, so its gradient $\lambda\theta$ is added to $g_t$ .

Proposition 2 (L2 equals weight decay for SGD, not for Adam) #

Plain SGD. The L2 step is $\theta - \eta(\nabla f + \lambda\theta) = (1 - \eta\lambda)\theta - \eta\nabla f$ : exactly weight decay with the same $\lambda$ (PyTorch form). In Loshchilov and Hutter's form, $\theta \leftarrow (1 - \lambda')\theta - \eta\nabla f$ , the two agree for $\lambda' = \eta\lambda$ (their Proposition 1).
Adam. The L2 gradient $\lambda\theta$ enters $m$ and $v$ and is divided by $\sqrt{\hat v}$ . When $\lambda\theta$ is small next to the gradient noise, $\sqrt{\hat v}$ is about the RMS of the gradient, and the per-step shrinkage is about $\eta\lambda/(\sqrt{\hat v} + \varepsilon)$ . Weights with large gradients are decayed less. For any preconditioner that is not a multiple of $I$ , no L2 coefficient reproduces weight decay (their Proposition 2).

(With momentum, SGD + L2 and SGD + weight decay are also no longer identical, because PyTorch's L2 term passes through the momentum buffer. The difference is small for small $\eta\lambda$ .)

AdamW takes the decay out of the normalization (PyTorch form):

\theta_t = \theta_{t-1}(1 - \eta\lambda) - \eta\,\frac{\hat m_t}{\sqrt{\hat v_t} + \varepsilon},

where $m$ and $v$ are computed from the raw $g_t$ . Every weight shrinks by the same factor $1 - \eta\lambda$ per step. PyTorch's default is $\lambda = 10^{-2}$ . The paper multiplies its decay by the schedule multiplier, not by the base learning rate; PyTorch multiplies it by lr. In both, the decay follows the learning-rate schedule.

Two groups of weights, no reason to move. All weights start at 1 and see zero-mean gradient noise: group A with RMS 10 (solid lines), group B with a smaller RMS (dashed). Only the decay moves them on average. With AdamW (aqua) both groups shrink alike, like $(1 - \eta\lambda)^t$ . With Adam + L2 (magenta), group A barely shrinks and group B shrinks fast, like $\exp(-\eta\lambda t/\text{RMS})$ . Change the RMS of group B and the decay $\lambda$ .

Quick check

For which optimizer is adding $\tfrac{\lambda}{2}\|\theta\|^2$ to the loss exactly the same as decaying the weights by $1 - \eta\lambda$ per step?

Adam as sign descent, and its memory cost #

With $\beta_1 = \beta_2 = 0$ and $\varepsilon = 0$ , Adam's step is $\eta\,g_t/|g_t| = \eta\operatorname{sign}(g_t)$ : sign descent. With the usual $\beta$ 's Adam is a smoothed version of it: each coordinate moves by about $\eta$ in the direction of its recent average sign. This view explains its robustness to gradient scale and connects it to Lion (Chen et al. 2023), which uses the sign of a momentum average directly, and to the matrix optimizers of lesson 12.

Adam stores two extra numbers per parameter ( $m$ and $v$ ). With mixed precision a parameter costs about 16 bytes: fp16 weight and gradient (2 + 2), and fp32 master weight, $m$ and $v$ (4 + 4 + 4) (Rajbhandari et al. 2020). This optimizer state is what ZeRO and FSDP shard across GPUs.

Adam does not always converge

Reddi, Kale and Kumar (2018, Theorem 3) show that for any constant $\beta_1 < \sqrt{\beta_2}$ there is a stochastic convex problem on which Adam does not converge to the optimum. In their example a rare, large gradient points the right way, but it also raises $v$ and so is damped; the frequent small gradients point the wrong way and get large steps once $v$ has forgotten the large one. Their fix, AMSGrad, uses $\max(\hat v_t, \hat v_{t-1})$ . In practice, a large $\beta_2$ and a decaying learning rate avoid the problem, and Adam with these settings is the default. Zhang et al. (2022) show that Adam converges when $\beta_2$ is large enough for the problem.

Answers to section 1. (1) Both move by about $10^{-3}$ : Adam's first step is $\eta\operatorname{sign}(g_1)$ in every coordinate ( $0.001$ and $0.00099999$ with $\varepsilon = 10^{-8}$ ). (2) The weight with small gradients. Inside Adam the L2 gradient is divided by $\sqrt{\hat v}$ , which is large for the weight with large gradients, so that weight is barely regularized. AdamW decays both by the same factor.

4. Worked example #

(a) Two Adam steps by hand (defaults, $\varepsilon$ negligible). Gradients $g_1 = (100, 0.01)$ , then $g_2 = (-100, 0.01)$ .

$t = 1$ : $m_1 = 0.1\,g_1 = (10, 0.001)$ and $v_1 = 0.001\,g_1^2 = (10, 10^{-7})$ . Corrected: $\hat m_1 = m_1/0.1 = (100, 0.01)$ and $\hat v_1 = v_1/0.001 = (10^4, 10^{-4})$ . Step $= \eta\,(100/100,\ 0.01/0.01) = \eta\,(1, 1)$ . Without correction: $m_1/\sqrt{v_1} = (10/3.162,\ 0.001/0.000316) = (3.16, 3.16)$ .
$t = 2$ : $m_2 = 0.9\,m_1 + 0.1\,g_2 = (9 - 10,\ 0.0009 + 0.001) = (-1, 0.0019)$ and $v_2 = 0.999\,v_1 + 0.001\,g_2^2 = (9.99 + 10,\ 0.999\cdot 10^{-7} + 10^{-7}) = (19.99,\ 1.999\cdot10^{-7})$ . With $1 - \beta_1^2 = 0.19$ and $1 - \beta_2^2 = 0.001999$ : $\hat m_2 = (-5.263, 0.01)$ and $\hat v_2 = (10^4, 10^{-4})$ . Step $= \eta\,(-5.263/100,\ 0.01/0.01) = \eta\,(-0.0526, 1)$ .

Gradients that differ by a factor $10^4$ give the same step. The coordinate whose sign flips moves back by only 5% of $\eta$ : the momentum still remembers the first gradient.

(b) RMSProp's first step. $v_1 = 0.01\,g_1^2$ , so $\sqrt{v_1} = 0.1|g_1|$ and the step is $\eta\,g_1/(0.1|g_1|) = 10\,\eta\operatorname{sign}(g_1)$ . With $\eta = 10^{-2}$ (the default) this is $0.1$ per coordinate, whatever the gradient.

(c) L2 against decoupled decay, one step. $\theta_0 = 2$ , $g = 1$ , $\eta = 0.1$ , $\lambda = 0.1$ .

SGD + L2: $\theta_1 = 2 - 0.1\,(1 + 0.1 \cdot 2) = 2 - 0.12 = 1.88$ .
SGD + weight decay: $\theta_1 = 2\,(1 - 0.01) - 0.1 = 1.98 - 0.1 = 1.88$ . Same.
Adam + L2: the gradient becomes $g' = 1 + 0.2 = 1.2$ ; the first Adam step is $\eta\operatorname{sign}(g') = 0.1$ , so $\theta_1 = 1.90$ . The decay disappeared in the normalization.
AdamW: $\theta_1 = 2\,(1 - 0.01) - 0.1 \cdot 1 = 1.88$ .

(d) Many steps. $\eta = 10^{-3}$ , $\lambda = 10^{-2}$ , zero-mean gradient noise, 20,000 steps.

AdamW shrinks every weight by $1 - \eta\lambda = 1 - 10^{-5}$ per step: $(1 - 10^{-5})^{20000} \approx e^{-0.2} \approx 0.82$ .
Adam + L2 shrinks by about $\eta\lambda/\text{RMS}$ per step: $10^{-6}$ for noise RMS 10, so $e^{-0.02} \approx 0.98$ ; and $10^{-4}$ for RMS 0.1, so $e^{-2} \approx 0.14$ .

A simulation with 500 weights per group (notebook, exercise 4) gives mean weights of about $0.98$ and $0.14$ for Adam + L2, and about $0.82$ for both groups with AdamW.

5. In code #

The NumPy function below follows the PyTorch formulas exactly. It reproduces the hand computations of section 4.

import numpy as np

def adam_steps(grads, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8,
               wd=0.0, decoupled=False, theta0=0.0, correct=True):
    """Adam / AdamW exactly as in torch.optim; returns the parameters after each step."""
    theta = np.full(len(grads[0]), theta0, dtype=float)
    m, v, out = np.zeros_like(theta), np.zeros_like(theta), []
    for t, g in enumerate(grads, start=1):
        g = np.asarray(g, dtype=float)
        if decoupled:
            theta = theta * (1 - lr * wd)          # AdamW: decay outside the normalization
        else:
            g = g + wd * theta                     # Adam(weight_decay=...): L2 via the gradient
        m = b1 * m + (1 - b1) * g
        v = b2 * v + (1 - b2) * g ** 2
        m_hat = m / (1 - b1 ** t) if correct else m
        v_hat = v / (1 - b2 ** t) if correct else v
        theta = theta - lr * m_hat / (np.sqrt(v_hat) + eps)
        out.append(theta.copy())
    return out

# (a) two Adam steps, in units of lr
th = adam_steps([[100.0, 0.01], [-100.0, 0.01]])
print(np.round(-th[0] / 1e-3, 6), np.round(-(th[1] - th[0]) / 1e-3, 6))
# [1.       0.999999] [-0.052632  0.999999]        (0.999999: the effect of eps = 1e-8)
print(np.round(-adam_steps([[100.0, 0.01]], correct=False)[0] / 1e-3, 4))
# [3.1623 3.1622]                                   without bias correction: about 3.16 lr

# (b) RMSProp's first step: v_1 = 0.01 g^2
g = np.array([100.0, 0.01]); v1 = 0.01 * g ** 2
print(g / np.sqrt(v1))                     # [10. 10.]

# (c) one step from theta = 2 with g = 1, lr = 0.1, weight decay 0.1
print(2 - 0.1 * (1 + 0.1 * 2))                                            # SGD + L2: 1.88
print(adam_steps([[1.0]], lr=0.1, wd=0.1, theta0=2.0)[0])                  # Adam + L2: [1.9]
print(adam_steps([[1.0]], lr=0.1, wd=0.1, decoupled=True, theta0=2.0)[0])  # AdamW: [1.88]

The same check with PyTorch, where these optimizers are used in practice:

import torch

for name, opt_class in [("SGD + L2", torch.optim.SGD), ("Adam + L2", torch.optim.Adam),
                        ("AdamW", torch.optim.AdamW)]:
    p = torch.tensor([2.0], requires_grad=True)
    opt = opt_class([p], lr=0.1, weight_decay=0.1)
    p.grad = torch.tensor([1.0])
    opt.step()
    print(f"{name:9s} theta_1 = {p.item():.4f}")   # 1.8800, 1.9000, 1.8800

6. Where it shows up #

Large language models. AdamW with a linear warmup and a cosine or warmup–stable–decay schedule (lesson 07) is the standard pretraining recipe. GPT-3 used $\beta_1 = 0.9$ , $\beta_2 = 0.95$ ; LLaMA used AdamW with $\beta_2 = 0.95$ and weight decay 0.1. Weight decay is usually applied to weight matrices only, not to biases and normalization gains.
ML for science. GraphCast trained with AdamW, $\beta_2 = 0.95$ and weight decay 0.1 (Lam et al. 2023). Diffusion models usually train with AdamW; DiT used a constant learning rate of $10^{-4}$ , no weight decay, and an EMA of the weights for sampling. Physics-informed networks often run Adam first and then switch to L-BFGS (lesson 10).
Distributed training. Adam's state ( $m$ , $v$ ) is twice the size of the model. ZeRO and FSDP shard it across GPUs; 8-bit optimizers store it in low precision.
Sparse features. AdaGrad's per-feature steps suit problems with rare features, such as embeddings and linear models for click prediction.
Beyond diagonal. Adam fixes only axis-aligned scaling. Shampoo, SOAP and Muon (lesson 12) precondition whole weight matrices.

7. Common mistakes #

torch.optim.Adam(weight_decay=...) is not AdamW. It adds an L2 term to the gradient. Use torch.optim.AdamW (recent PyTorch versions also accept decoupled_weight_decay=True in Adam).
Forgetting that PyTorch's AdamW decay is lr * weight_decay per step. Halving the learning rate also halves the decay.
Bias correction with the wrong $t$ . If $m$ and $v$ restart at 0 but $t$ keeps counting (a global step counter), the correction is about 1 and the first steps are several times too large: $3.16\,\eta$ at the first step, and up to about $6.6\,\eta$ near step 12 for a constant gradient. Counting from $t = 0$ divides by $1 - \beta^0 = 0$ .
Resuming a run without the optimizer state. $m$ , $v$ and $t$ restart, $\hat v$ again comes from a few gradients, and the first steps are about $\eta$ in every coordinate, as at a start without warmup.
$\varepsilon$ in a different place. The paper's "efficient" form puts $\hat\varepsilon$ next to the uncorrected $\sqrt{v_t}$ , so the same number means a different $\varepsilon$ in libraries that use it.
Expecting Adam to fix any ill-conditioning. A rotated valley stays ill-conditioned for a diagonal preconditioner.
Comparing optimizers at one learning rate. Each method has its own good range (Adam around $10^{-4}$ to $10^{-3}$ , SGD often 10 to 100 times larger); compare each at its tuned value.

8. Check yourself #

What is the size of RMSProp's first step ( $\alpha = 0.99$ , $\varepsilon \approx 0$ )? And of AdaGrad's first step?
Answer
RMSProp: $v_1 = 0.01\,g_1^2$ , so $\sqrt{v_1} = 0.1\,|g_1|$ and the step is $10\,\eta\operatorname{sign}(g_1)$ . AdaGrad: $s_1 = g_1^2$ , so the step is $\eta\operatorname{sign}(g_1)$ .
Adam sees the same gradient $g = (3, -0.5)$ at every step. What is its step at step $t$ ?
Answer
$m_t = (1-\beta_1^t)g$ and $v_t = (1-\beta_2^t)g^2$ exactly, so $\hat m_t = g$ , $\hat v_t = g^2$ , and the step is $\eta\operatorname{sign}(g) = \eta\,(1, -1)$ at every $t$ (up to $\varepsilon$ ).
With constant gradients of size $|g|$ , how does AdaGrad's step change with $t$ ? Why is this a problem for long training runs?
Answer
$s_t = t\,g^2$ , so the step is $\eta|g|/\sqrt{t g^2} = \eta/\sqrt t$ . It keeps shrinking even when the problem would allow large steps, so long nonconvex runs slow down. RMSProp's EMA forgets old gradients and keeps the step near $\eta$ .
Plain SGD with $\eta = 0.05$ and an L2 coefficient $10^{-4}$ . Which weight decay gives the same updates, in PyTorch form and in Loshchilov–Hutter form?
Answer
The per-step factor is $1 - \eta\lambda = 1 - 5\cdot10^{-6}$ . PyTorch form $\theta(1 - \eta\lambda)$ : $\lambda = 10^{-4}$ . Loshchilov–Hutter form $\theta(1 - \lambda')$ : $\lambda' = \eta\lambda = 5\cdot10^{-6}$ .
On $f(x) = \tfrac12(x_1^2 + 100x_2^2)$ from $x_0 = (1, 1)$ , Adam (exact gradients) keeps $x_1 = x_2$ at every step. Why?
Answer
Scale invariance per coordinate. While $x_1 = x_2$ , the gradient of coordinate 2 is 100 times that of coordinate 1 at every step, so its $m$ and its $\sqrt v$ are both 100 times larger, and the two steps are equal (up to $\varepsilon$ ). By induction this holds at every step.
A colleague restarts training from a checkpoint but forgets to load the optimizer state. What happens in the first steps?
Answer
$m$ , $v$ and $t$ restart at 0. After bias correction the first step is about $\eta\operatorname{sign}(g)$ in every coordinate, even for coordinates that had tiny steps before. This is like a start without warmup and can cause a loss spike. Load the state, or warm up again.

Key takeaways

An EMA $m_t = \beta m_{t-1} + (1-\beta)g_t$ remembers about $1/(1-\beta)$ steps; with $m_0 = 0$ it is biased by the factor $1 - \beta^t$ .
AdaGrad divides by the root of the sum of squared gradients (steps decay like $1/\sqrt t$ ); RMSProp uses an EMA instead, without bias correction (first step $10\eta$ ).
Adam: $m_t$ and $v_t$ are EMAs of $g$ and $g^2$ ; $\hat m_t = m_t/(1-\beta_1^t)$ , $\hat v_t = v_t/(1-\beta_2^t)$ ; $\theta_t = \theta_{t-1} - \eta\,\hat m_t/(\sqrt{\hat v_t} + \varepsilon)$ . Defaults $10^{-3}$ , $(0.9, 0.999)$ , $10^{-8}$ .
Adam's first step is about $\eta\operatorname{sign}(g_1)$ in every coordinate; the step is invariant to the gradient scale and is typically at most about $\eta$ .
Adam is a diagonal preconditioner: it removes bad scaling along the coordinate axes, not along rotated directions.
A large $\varepsilon$ turns Adam into momentum SGD with step $\eta/\varepsilon$ ; $\beta_2 = 0.95$ gives a shorter memory than $0.999$ and reacts faster to changes in the gradient scale.
L2 regularization equals weight decay for plain SGD, but inside Adam it is divided by $\sqrt{\hat v}$ , so weights with large gradients are decayed less. AdamW decays every weight by $1 - \eta\lambda$ per step.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

How many steps does an EMA with coefficient $\beta$ remember?
Write the Adam update with bias correction.
Why does Adam need bias correction?
What is the size of Adam's first step in each coordinate?
What is the size of RMSProp's first step with $\alpha = 0.99$ ?
Why does Adam not fix a valley rotated by $45°$ ?
What does a large $\varepsilon$ do to Adam?
Why is L2 regularization not weight decay inside Adam?
Write the AdamW update (PyTorch form).

Open the review →

9. Interview angle #

"Adam or AdamW?" — Adam's weight_decay is an L2 term through the gradient and is divided by $\sqrt{\hat v}$ ; AdamW decays every weight by $1 - \eta\lambda$ , outside the normalization. Use AdamW.
"Why bias correction?" — $m$ and $v$ start at 0 and underestimate by the factor $1 - \beta^t$ ; without correction the first step is about $3.16\,\eta$ with the default betas.
"What does Adam actually adapt to?" — The per-coordinate scale of the gradient: it is a diagonal preconditioner $\operatorname{diag}(1/(\sqrt{\hat v} + \varepsilon))$ , invariant to rescaling each coordinate. It does not adapt to rotations or correlations.
"Why $\beta_2 = 0.95$ for LLMs?" — A shorter memory (about 20 steps) lets $\hat v$ follow sudden changes of the gradient scale, so the steps do not exceed $\eta$ for long; this reduces instabilities.
"What is Adam's memory cost?" — Two extra tensors per parameter; about 16 bytes per parameter with mixed precision, which is what ZeRO and FSDP shard.

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution:

Implement AdaGrad, RMSProp, Adam and AdamW steps in NumPy, exactly as the PyTorch formulas, and test them against the hand computations of section 4.
Bias correction on a constant gradient: show that $\hat m_t = g$ for every $t$ , and measure the step without correction.
The diagonal preconditioner on an axis-aligned and on a rotated quadratic: compute the condition number after rescaling, and run Adam on both.
Adam + L2 against AdamW: simulate two groups of weights and compare their mean weights after 20,000 steps.
Train a softmax regression on the scikit-learn digits with SGD, momentum and Adam for the same number of steps, and compare them at several learning rates.