33  Adaptive Optimizers: Adam & AdamW

← Back to: Optimization — GD, SGD, Momentum

TipTL;DR

The optimization note ended at SGD + momentum. Adam (Kingma & Ba, 2014) adds a per-parameter adaptive step size — a running estimate of each gradient’s mean (momentum) and variance — so every weight gets its own effective learning rate. AdamW (Loshchilov & Hutter, 2019) fixes how weight decay interacts with that scaling, and is the optimizer that trains essentially every modern Transformer.

33.1 The problem momentum doesn’t solve

Momentum smooths the direction of the update but uses one global learning rate \(\eta\) for all parameters. When gradients differ wildly in scale across parameters (universal in deep nets — embeddings vs. deep layers), a single \(\eta\) is always wrong for most of them. AdaGrad (Duchi et al., 2011) introduced the fix: divide each parameter’s step by the accumulated magnitude of its own past gradients. Its flaw is that the accumulator only grows, so the effective rate decays to zero. RMSProp replaced the sum with an exponential moving average — and Adam combined that with momentum.

33.2 Adam: momentum on both moments

Adam tracks two exponential moving averages of the gradient \(g_t\): the first moment \(m_t\) (mean, i.e. momentum) and the second moment \(v_t\) (uncentered variance):

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

The \(\sqrt{\hat v_t}\) divisor is the adaptive part: parameters with large, noisy gradients take smaller steps; quiet parameters take larger ones. The bias correction (\(1-\beta^t\) terms) undoes the cold-start bias from initializing \(m_0=v_0=0\). Defaults \(\beta_1{=}0.9\), \(\beta_2{=}0.999\) are near-universal.

33.3 AdamW: decoupled weight decay

The original Adam folds L2 regularization into the gradient \(g_t\), which then gets divided by \(\sqrt{\hat v_t}\) — so parameters with large gradients are decayed less, coupling regularization to gradient scale. That is not what weight decay is supposed to do. AdamW (Loshchilov & Hutter, 2019) decouples it, applying decay directly to the weights:

\[ \theta_t = \theta_{t-1} - \eta\Big(\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon} + \lambda\,\theta_{t-1}\Big). \]

This single correction measurably improves generalization and is why AdamW — not Adam — is the standard in every modern LLM training recipe.

import numpy as np
def adamw_step(theta, g, m, v, t, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8, wd=0.01):
    m = b1*m + (1-b1)*g
    v = b2*v + (1-b2)*g*g
    mhat = m / (1 - b1**t)
    vhat = v / (1 - b2**t)
    theta = theta - lr*(mhat/(np.sqrt(vhat)+eps) + wd*theta)
    return theta, m, v

theta, m, v = np.array([1.0]), np.array([0.0]), np.array([0.0])
for t in range(1, 4):
    g = 2*theta              # gradient of theta^2
    theta, m, v = adamw_step(theta, g, m, v, t)
    print(f"step {t}: theta={theta[0]:.5f}")
step 1: theta=0.99899
step 2: theta=0.99798
step 3: theta=0.99697
WarningPitfall: Adam vs AdamW are not interchangeable

Calling Adam(weight_decay=...) in most frameworks applies coupled decay, not AdamW’s decoupled form. For Transformers always use the dedicated AdamW optimizer — the difference shows up as worse generalization, not as an error.

NoteKey takeaway

Adaptive optimizers extend SGD+momentum with a per-parameter learning rate derived from each gradient’s own variance. Adam combined momentum with RMSProp-style scaling and bias correction; AdamW fixed weight decay. Together they are the training engine of the Transformer era — the “AdamW” assumed in every scaling and language-model chapter.

← Back to: Optimization — GD, SGD, Momentum · Used in: Scaling Laws, Language Models