Adam and AdamW

Combining momentum and RMSProp, bias correction, and why weight decay is not L2 regularization.

You are training a transformer. Loss starts at 4.2, decreases for a few thousand steps with SGD, then plateaus at 3.8 for 10,000 steps before barely moving again. The problem is not your data or architecture — it is that your transformer has attention matrices, embedding tables, and feedforward layers all updating simultaneously, each at wildly different gradient scales. A single global learning rate is completely wrong for all of them at once.

This is the problem Adam was built to solve. By 2014, practitioners had two partial solutions sitting separately on the shelf. RMSProp tracked each parameter's gradient magnitude via an exponential moving average $v_t = β_2 v_{t-1} + (1-β_2)g_t^2$, then divided by $\sqrt{v_t}$ to normalize steps — slow-updating parameters got large steps, fast-updating ones got small steps. SGD with momentum tracked gradient direction history $m_t = β_1 m_{t-1} + (1-β_1)g_t$, smoothing out oscillations and building velocity in consistent directions. Adam (Kingma & Ba, 2014) ran both simultaneously. The $m_t$ term provides direction stability. The $v_t$ term provides per-parameter scale adaptation. Dividing the smoothed direction by the smoothed magnitude gives a step that is both directionally stable and scale-normalized: $θ_t = θ_{t-1} - α \cdot \hat{m}_t / (\sqrt{\hat{v}_t} + ε)$.

There is a critical initialization trap. At step 1, both $m_0 = 0$ and $v_0 = 0$. So $m_1 = (1-β_1)g_1 = 0.1 g_1$ — the first moment is 10x too small. $v_1 = (1-β_2)g_1^2 = 0.001 g_1^2$ — the second moment is 1000x too small. Without correction, the ratio $m_1 / \sqrt{v_1}$ is inflated by a fixed, predictable factor — $(1-β_1)/\sqrt{1-β_2} ≈ 3.16$ for the default $β_1=0.9$, $β_2=0.999$ — independent of the gradient magnitude. Bias correction divides by the initialization factor: $\hat{m}_t = m_t / (1 - β_1^t)$, $\hat{v}_t = v_t / (1 - β_2^t)$. At $t=1$: $\hat{m}_1 = m_1 / 0.1 = g_1$, $\hat{v}_1 = v_1 / 0.001 = g_1^2$. Correct. Without this, early transformer training can corrupt embeddings in ways that are nearly impossible to recover from.

NOT this. Most people think Adam + L2 regularization in the loss = weight decay. They are not equivalent, and the difference is not small. When you add $λ||θ||^2$ to the loss, the gradient becomes $g + λθ$. Adam then divides this combined gradient by $\sqrt{\hat{v}}$. For a parameter with a large gradient history, $\sqrt{\hat{v}}$ is large — the regularization term $λθ$ gets divided down to almost nothing. The parameters that receive the most gradient (probably the most important ones) get the least regularization. This is backwards.

AdamW (Loshchilov & Hutter, 2019) fixes this surgically. Instead of modifying the gradient, weight decay is applied directly to the parameters before the gradient step: $θ_t ← (1-αλ)θ_{t-1} - α\hat{m}_t/(\sqrt{\hat{v}_t}+ε)$. The $(1-αλ)$ factor decays every parameter by the same fraction per step, completely independent of gradient history. This is true weight decay. Adam + L2 in the loss is not.

Every GPT, BERT, and Llama-class model is trained with AdamW, not Adam. The difference is real but the exact magnitude varies by model and run rather than a single fixed percentage — what matters mechanically is that AdamW decay is proportional to the parameter itself, not skewed by gradient history the way Adam-plus-L2 is. For your transformer that was plateauing at 3.8: switch to AdamW, set weight_decay=0.1, add warmup. The loss plateau disappears because the per-parameter adaptation of Adam finally has matching regularization.

Key points

• Use AdamW for transformers and most deep learning; SGD+momentum for CNNs where you have time to tune and need best generalization. The deployment signal: heterogeneous parameter scales (transformers, multi-modal models, anything with embeddings) → AdamW. Homogeneous architectures (ResNet, VGG) where SGD has a well-documented training recipe → SGD+momentum. This isn't just convention: on ResNet-50/ImageNet specifically, well-tuned SGD+momentum typically beats Adam's final top-1 accuracy by roughly 1-2 percentage points. The standard explanation is that Adam's adaptive per-parameter scaling reduces gradient noise and tends to converge toward sharper minima that generalize worse, while SGD+momentum's noisier updates are more likely to settle into flatter minima that generalize better. Default AdamW hyperparameters: $β_1=0.9$, $β_2=0.999$, $ε=1e$-$8$, $weight_decay=0.01$–$0.1$. PyTorch's AdamW is correct; Adam + manual L2 in the loss is not equivalent.
• The most common production trap: using Adam with L2 regularization in the loss and calling it weight decay. Symptom: model overfits despite high $λ$, or regularization seems to have no effect. Root cause: the adaptive denominator $\sqrt{\hat{v}}$ is dividing away your L2 penalty on the parameters that need it most. Fix: switch to AdamW and pass weight_decay directly to the optimizer. Never add L2 to the loss when using any adaptive optimizer — it does not do what you think.
• Diagnostic: if loss diverges in the first 100–500 steps with Adam, bias correction or missing warmup is the culprit. At step $t=1$ with $β_2=0.999$: $\hat{v}_1 = g_1^2$ after bias correction — fine. But if warmup is absent, the full $α$ is applied from step 1 before gradient statistics have stabilized across layers. Standard fix: linear warmup from $α_{min}=1e$-$7$ to $α=1e$-$4$ over 1%–4% of total training steps. If loss is stable but plateauing: check weight_decay is set on AdamW (not zero). If loss oscillates throughout training: $α$ is too high — reduce by 3x–10x.

Recap

• One global LR is wrong when a Transformer's attention, embeddings, and FFN all update at wildly different scales at once — no single rate suits all three, which is what motivates per-parameter adaptation.
• Adam = momentum + RMSProp combined: it divides a momentum term $m_t$ (direction stability) by $\sqrt{v_t}$ (per-parameter scale from squared gradients), giving updates that are both directionally stable and scale-normalized for every weight.
• Bias correction is essential, not optional: $m_0=v_0=0$, so for the first steps the running estimates are ~10× / 1000× too small — divide by $(1-β_1^t)$ and $(1-β_2^t)$ to inflate them back to unbiased size, or the early updates are tiny and training limps out of the gate.
• Adam + L2-in-the-loss is *not* proper weight decay: the $\sqrt{\hat v}$ denominator divides the $λθ$ penalty *down* on exactly the high-gradient parameters that most need shrinking — the regularization ends up uneven and backwards.
• AdamW fixes this by decaying the weights directly: $θ ← (1-αλ)θ - α\hat m/(\sqrt{\hat v}+ε)$ applies the decay uniformly and independently of the gradient scale — proper weight decay restored.
• Every serious language model (GPT, BERT, Llama) uses AdamW, not plain Adam — the exact size of the gain varies by model and run, but the mechanism is consistent: decay proportional to the parameter itself, not skewed by gradient history the way Adam-plus-L2 is.
• Defaults to memorize: $β_1=0.9$, $β_2=0.999$, $ε=1e$-$8$, weight_decay 0.01–0.1; drop $β_2$ (e.g. to 0.9, window ~10) for non-stationary objectives like RL or fine-tuning where the gradient distribution shifts.

Check your understanding