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 = d2l.zeros((feature_dim, 1)), d2l.zeros(1)
    v_w, v_b = d2l.zeros((feature_dim, 1)), d2l.zeros(1)
    return ((m_w, v_w), (m_b, v_b))

def adamw(params, states, hyperparams):
    beta1, beta2, eps = 0.9, 0.999, 1e-6
    for p, (m, v) in zip(params, states):
        with torch.no_grad():
            m[:] = beta1 * m + (1 - beta1) * p.grad
            v[:] = beta2 * v + (1 - beta2) * torch.square(p.grad)
            m_hat = m / (1 - beta1 ** hyperparams['t'])
            v_hat = v / (1 - beta2 ** hyperparams['t'])
            p[:] -= hyperparams['lr'] * (m_hat / (torch.sqrt(v_hat) + eps)
                                         + hyperparams['wd'] * p)
        p.grad.zero_()
    hyperparams['t'] += 1

One number, two meanings

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

curves = {}
for name, opt_cls in [('Adam + $\\ell_2$', torch.optim.Adam),
                      ('AdamW', torch.optim.AdamW)]:
    torch.manual_seed(0)
    model = d2l.TinyLM(len(data.vocab))
    optimizer = opt_cls(model.parameters(), lr=0.003, weight_decay=0.1)
    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).

torch.manual_seed(0)
model = d2l.TinyLM(len(data.vocab))
decay = [p for name, p in model.named_parameters()
         if p.ndim == 2 and 'emb' not in name]
no_decay = [p for name, p in model.named_parameters()
            if p.ndim != 2 or 'emb' in name]
optimizer = torch.optim.AdamW([
    {'params': decay, 'weight_decay': 0.1},
    {'params': no_decay, 'weight_decay': 0.0}], lr=0.003)
n_decay = sum(p.numel() for p in decay)
n = n_decay + sum(p.numel() for p in no_decay)
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.