AdamW

Dive into Deep Learning · §9.7

Decoupled weight decay and the modern default
why \ell_2 through Adam fails · AdamW · what decay really does · what it costs

A promise from §3.7

Motivation

Under SGD, penalty and decay are the same operation:

\mathbf{x}_{t+1} = \mathbf{x}_t - \eta\,(\mathbf{g}_t + \lambda \mathbf{x}_t) = (1 - \eta\lambda)\,\mathbf{x}_t - \eta\,\mathbf{g}_t.

Under Adam, the penalty goes through the preconditioner (schematically — the appendix has the exact recurrence):

\underbrace{\frac{\eta\lambda}{\sqrt{\hat{v}_{t,i}} + \epsilon}\, x_{t,i}}_{\textrm{$\ell_2$ through Adam}} \qquad\textrm{vs.}\qquad \underbrace{\eta\lambda\, x_{t,i}}_{\textrm{weight decay}}

• Large gradient history → barely decayed; quiet coordinate → crushed.
• One \lambda, re-priced per coordinate, per step — nobody chose those prices (100× disparity demo: appendix ch. 25).

AdamW: decay outside the preconditioner

\mathbf{x}_{t+1} = (1 - \eta\lambda)\,\mathbf{x}_t - \eta\, \frac{\hat{\mathbf{m}}_t}{\sqrt{\hat{\mathbf{v}}_t} + \epsilon}

Loss gradient → preconditioner. Decay → applied directly, scaled by the schedule (Loshchilov & Hutter, 2019).

One term added to Adam from §9.6:

def init_adamw_states(feature_dim):
    m_w, m_b = jnp.zeros((feature_dim, 1)), jnp.zeros(1)
    v_w, v_b = jnp.zeros((feature_dim, 1)), jnp.zeros(1)
    return [(m_w, v_w), (m_b, v_b)]

def adamw(params, grads, states, hyperparams):
    beta1, beta2, eps = 0.9, 0.999, 1e-6
    for i, (p, (m, v), g) in enumerate(zip(params, states, grads)):
        m = beta1 * m + (1 - beta1) * g
        v = beta2 * v + (1 - beta2) * jnp.square(g)
        m_hat = m / (1 - beta1 ** hyperparams['t'])
        v_hat = v / (1 - beta2 ** hyperparams['t'])
        params[i] = p - hyperparams['lr'] * (
            m_hat / (jnp.sqrt(v_hat) + eps) + hyperparams['wd'] * p)
        states[i] = (m, v)
    hyperparams['t'] += 1
    return params[0], params[1]

One number, two meanings

Same model, same tuned \eta, same \lambda = 0.1 — coupled vs. decoupled:

def coupled(lr, wd):
    return optax.chain(optax.add_decayed_weights(wd), optax.adam(lr))

curves = {}
for name, tx in [('Adam + $\\ell_2$', coupled(0.001, 0.1)),
                 ('AdamW', optax.adamw(0.001, weight_decay=0.1))]:
    model = d2l.TinyLM(len(data.vocab), rngs=nnx.Rngs(0))
    optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)
    curves[name] = d2l.train_lm(model, data, optimizer, num_steps=1000)
d2l.plot(list(range(0, 1000, 25)), [smooth(c) for c in curves.values()],
         'step', 'training loss', legend=list(curves))

• Decoupled: trains as if the decay were absent.
• Coupled: stuck a full nat higher. Progress shrinks \sqrt{\hat{\mathbf{v}}} → the penalty comes back amplified; once it dominates, Adam normalizes it and the dial stops responding.

A grid, twice

3\times 3 learning rate × weight decay, held-out loss, on a training slice small enough to overfit. Best \lambda per row in bold; each panel’s best cell boxed in red:

• AdamW: the bold column is the same at every \eta — one line search.
• Coupled: the optimum wanders, and lives 2–3 orders of magnitude lower.
• Fully tuned, both reach similar loss — decoupling buys tunability.

What weight decay is actually doing at scale

One-epoch LLM training has little classical overfitting — yet \lambda = 0.1 is universal. Why?

• Normalized layers are scale-invariant → gradient noise pushes norms up, decay pulls down → equilibrium: constant rotation per step (Kosson et al., 2024).
• Through it, \lambda sets the effective learning rate — decayed runs reach lower training loss, not a train/val trade (D’Angelo et al., 2024). Plus: keeps bf16 out of divergence.
• \eta\lambda = a timescale: \tau = B/(\eta\lambda D) epochs; scale \lambda with B, D instead of re-tuning (Bergsma et al., 2025).

What not to decay

The census populations of §9.6, treated differently:

• Matrices — decay (96% of parameters).
• Norms and biases — exempt: they set the normalization scales.
• Embeddings — exempt: sparse gradients, decay every step; OLMo 2 traced spikes to decay grinding embedding norms down (1/\|\mathbf{x}\| in LayerNorm’s gradient).

model = d2l.TinyLM(len(data.vocab), rngs=nnx.Rngs(0))

def decay_mask(params):
    return jax.tree_util.tree_map_with_path(
        lambda path, p: p.ndim == 2 and 'emb' not in str(path), params)

optimizer = nnx.Optimizer(
    model, optax.adamw(0.001, weight_decay=0.1, mask=decay_mask),
    wrt=nnx.Param)
params = nnx.state(model, nnx.Param)
n_decay = sum(int(p.size) for p, m in zip(
    jax.tree.leaves(params), jax.tree.leaves(decay_mask(params))) if m)
n = sum(int(p.size) for p in jax.tree.leaves(params))
print(f'decayed: {n_decay} of {n} parameters')

decayed: 396800 of 412188 parameters

The same matrices / non-matrices split returns in §9.9 — Muon gives the matrices a different optimizer, not just a different decay.

The memory bill

Two extra numbers per parameter — do the arithmetic once:

                                           B/param    TinyLM  7B model
fp32 weights, grads, m, v                       16     6.6MB     112GB
bf16 weights + grads, fp32 master, m, v         16     6.6MB     112GB
the same + fp32 gradient accumulator            20     8.2MB     140GB

• bf16 does not shrink the bill; it adds an fp32 master copy.
• 7B model ≈ 140 GB, 12 of 20 B/param is optimizer state.
• Shrink it: Adafactor (factored v), 8-bit states. Or shard it: ZeRO (→ ch. 11, §29.6).

Recap

• \ell_2 through Adam ≠ weight decay: the preconditioner re-prices \lambda per coordinate, backwards. AdamW decouples: one \lambda, one meaning.
• Decoupled knobs tune separately: best \lambda independent of \eta in our grid.
• At scale, decay = effective-LR control via the noise–decay equilibrium, on timescale 1/(\eta\lambda).
• Decay matrices only; exempt embeddings, norms, biases.
• AdamW parameter ≈ 20 bytes in mixed precision; optimizer state dominates.