Adam and AdamW

Adam keeps two running averages of the gradient, so each parameter gets its own step size. AdamW applies weight decay after that scaling, which is why it became the default.

One gradient feeds two running averages, both get bias corrected, and their ratio sets a per parameter step near 0.001, while AdamW weight decay joins after the division

Why Does This Exist?

One learning rate cannot serve every layer. Picture fine-tuning a text classifier: the embedding gradients can run a thousand times smaller than the final head gradients. Set the single step for the head and the embeddings crawl. Set it for the embeddings and the head blows up. That mismatch is normal in deep nets, not a bug in your data.

Gradient descent uses one step size for every parameter, and stochastic gradient descent keeps that property while adding sampling noise. Adam answers with per-parameter scaling: each weight builds its own effective step from its own gradient history. This page builds that mechanism from its two averages, then the AdamW fix. The family around it lives in adaptive optimizers, and the penalty it rescues lives in weight decay.

You need one prerequisite and you already half know it. A gradient is one slope per parameter, pointing uphill, so training steps the other way: weights move against the gradient by a small amount. Everything below is about choosing that amount separately for each parameter instead of once for all of them.

Think of It Like This

Two hoses, one tap setting

You are watering with two hoses off one tap. The first blasts. The second trickles. One shared twist of the tap cannot be right for both: what soaks the bed with the first barely wets it with the second.

So you fit each hose with its own nozzle, set from how that hose ran over the last few minutes. You close the blast hose down and open the trickle hose up. Now one twist changes both flows by a useful amount. You also smooth the sputter: instead of jerking the nozzle with every spit of water, you average the last few seconds and turn steadily.

The twist is the global learning rate. Each nozzle is the second moment. Smoothing the sputter is momentum.

Where it stops: nozzles do not change the water pressure, but your steps change the loss surface you stand on. When the terrain turns a corner, last minute flow misleads. Adam trusts recent sizes blindly. It cannot tell a steep useful slope from loud noise, and near cliffs or flat plains that trust breaks.

How It Actually Works

Write $g_t$ for the gradient at step $t$ , one number per parameter. Adam feeds that same $g_t$ into two running averages. Every symbol below gets its plain meaning on first use, and the numbers in the worked example match the diagram above.

The two running averages

m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t \qquad \text{direction}, \qquad v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 \qquad \text{size}

$m_t$ is the first moment, an exponential moving average of the gradient with decay $\beta_1$ , usually 0.9. That is momentum: it points where gradients have consistently pointed and damps single-step noise. $v_t$ is the second moment, the same average of the squared gradient with decay $\beta_2$ , usually 0.999. Squaring drops the sign, so $v_t$ carries only size: how large this parameter's gradients have been running.

Dividing by $\sqrt{v_t}$ is where per-parameter scaling comes from. A parameter whose gradients ran tiny has a tiny denominator, so its step grows. One whose gradients ran large gets scaled down. A single global rate $\eta$ (eta, the learning rate you set) becomes workable because each parameter rescales it with its own history.

Bias correction

Both averages start at zero, so early steps drag toward zero. At step one, $m_1$ is only a tenth of the gradient and $v_1$ a thousandth of the squared gradient. Adam divides that startup drag away with the decay raised to the step count $t$ :

\hat{m}_t = \frac{m_t}{1 - \beta_1^{\,t}}, \qquad \hat{v}_t = \frac{v_t}{1 - \beta_2^{\,t}}, \qquad \theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}

$\hat{m}_t$ and $\hat{v}_t$ are the corrected averages. $\theta_t$ (theta) is the parameter vector at step $t$ . $\epsilon$ (epsilon), typically $10^{-8}$ , is a guard that stops division by something near zero when a parameter's gradients ran flat. It is not a tuning knob. The numerator recovers in tens of steps while the denominator needs thousands, since $\beta_1$ and $\beta_2$ differ. Skip this and the first few hundred steps run several times too large, and those steps decide where training goes.

Why AdamW decouples the decay

Plain Adam puts weight decay into the loss as a penalty, so it arrives inside $g_t$ and gets divided by $\sqrt{\hat{v}_t}$ like everything else. Trace that through: a parameter with large gradients owns a large denominator, so its penalty shrinks, while a quiet parameter gets punished hard. The strength $\lambda$ (lambda) you set is not the pressure any parameter feels.

AdamW applies the decay straight to the weights, after the division, as its own term:

\theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} - \eta \lambda \theta_t \qquad \text{AdamW}

Every weight now shrinks by the same fraction $\eta \lambda$ per step whatever its gradient history. That one-line move is why AdamW replaced Adam for transformers and most deep nets.

The price is memory: two extra numbers per parameter, about twice the weights in optimizer state. The assumption is that recent gradient sizes predict a good scale for the next step. That holds on steady slopes and breaks at cliffs, on noisy tiny batches, and for sparse features that wake up suddenly. For long vision schedules, SGD with momentum still matches or wins, and when memory hurts, leaner cousins on the adaptive optimizers page trade scale for state.

Worked example: two steps by hand

One parameter. Defaults $\beta_1 = 0.9$ , $\beta_2 = 0.999$ , $\eta = 0.001$ , $\epsilon = 10^{-8}$ . Gradients $g_1 = 0.5$ , then $g_2 = 0.3$ .

Step one:

m_1 = 0.9 \times 0 + 0.1 \times 0.5 = 0.05, \qquad v_1 = 0.999 \times 0 + 0.001 \times 0.5^2 = 0.00025

\hat{m}_1 = \frac{0.05}{1 - 0.9} = 0.5, \qquad \hat{v}_1 = \frac{0.00025}{1 - 0.999} = 0.25, \qquad \text{step} = 0.001 \times \frac{0.5}{\sqrt{0.25}} = 0.001

Without correction the step would be $0.001 \times 0.05 / \sqrt{0.00025} \approx 0.00316$ , about 3.16 times too large. That is the startup drag, removed exactly.

Step two folds in $g_2 = 0.3$ : $m_2 = 0.9 \times 0.05 + 0.1 \times 0.3 = 0.075$ and $v_2 = 0.999 \times 0.00025 + 0.001 \times 0.09 = 0.00033975$ . Correct with $1 - 0.9^2 = 0.19$ and $1 - 0.999^2 = 0.001999$ : $\hat{m}_2 \approx 0.3947$ , $\hat{v}_2 \approx 0.1700$ , step $\approx 0.000957$ . Same scale as step one. That steadiness across changing gradients is the product working.

Code

This snippet replays the worked example and prints the corrected step beside the uncorrected one, so you can see the startup drag vanish.

1import numpy as np2
3
4def first_steps(grads: list[float], lr: float = 1e-3) -> None:5    m, v = 0.0, 0.06    b1, b2 = 0.9, 0.9997    for t, g in enumerate(grads, start=1):8        m = b1 * m + (1 - b1) * g  # first moment: smoothed direction9        v = b2 * v + (1 - b2) * g * g  # second moment: this parameter's scale10        raw = lr * m / (np.sqrt(v) + 1e-8)  # still dragged toward zero11        m_hat = m / (1 - b1**t)  # divide out the startup drag12        v_hat = v / (1 - b2**t)13        fixed = lr * m_hat / (np.sqrt(v_hat) + 1e-8)14        print(f"step {t} g={g:.1f} raw {raw:.2e} fixed {fixed:.2e}")15
16
17first_steps([0.5, 0.3])18# -> step 1 g=0.5 raw 3.16e-03 fixed 1.00e-0319# -> step 2 g=0.3 raw 4.07e-03 fixed 9.57e-04

Uncorrected steps overshoot three to four times. Corrected ones sit near 0.001, matching the hand computation above.

Watch Out For

An L2 term in the loss, expected to behave like weight decay

Symptom: you sweep the penalty coefficient across two orders of magnitude under plain Adam and validation barely moves. The usual conclusion is that this model needs no regularization. Wrong conclusion.

The penalty entered through the gradient, so the second moment rescaled it per parameter, and the number you swept is not the pressure any weight felt. Switch to AdamW, where decay bypasses the division, and rerun the sweep. The curve usually gains a shape. Check your optimizer class too: some libraries couple under a weight_decay argument and keep the decoupled version in a separate class.

The default learning rate on a fine-tuning run

Symptom: fine-tuning pretrained weights, validation drops below the zero-shot score within a few hundred steps, then claws back to something mediocre. Lowering nothing else helps first.

0.001 is Adam's default because it suits training from scratch, where weights start random and large moves help. Pretrained weights already encode something, and per-parameter scaling makes effective steps larger than raw gradients suggest. A few hundred steps at 0.001 wash that away. Start near $10^{-5}$ , warm up over the first few hundred steps, and check validation before the first epoch ends.

The Quick Version

Adam tracks two running averages per parameter: the gradient mean for direction and the squared gradient mean for scale.
Dividing the direction by the square root of the scale hands each parameter its own effective step size from one global rate.
Both averages start at zero, so Adam divides out the startup drag with one minus beta to the step count. Skipping that overshoots the first few hundred steps severalfold.
Plain Adam folds weight decay into the gradient, where the scaling distorts it. AdamW applies the same decay after the division so every weight shrinks equally.
The price is memory for two extra numbers per parameter, about twice the weights in optimizer state.
Use AdamW for transformers and most deep nets. Keep SGD with momentum for long vision schedules where it still matches or wins.