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

Penalty and decay under adaptive scaling

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}}

  • A large gradient history produces weak effective decay; a small history produces strong effective decay.
  • The effective shrinkage from a single \lambda varies by coordinate and step (a 100× disparity is demonstrated in 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

Coupled and decoupled updates

The comparison holds the model, tuned \eta, and \lambda = 0.1 fixed while changing coupled to decoupled decay:

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 increasing \lambda has little effect.

Learning-rate and decay sweeps

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:

  • For AdamW, the best displayed \lambda is the same at every \eta, so the two hyperparameters can be searched more independently.
  • Coupled: the optimum wanders, and lives 2–3 orders of magnitude lower.
  • Fully tuned, both reach similar loss; decoupling makes the two hyperparameters easier to tune independently.

What weight decay is actually doing at scale

One-epoch LLM training has little classical overfitting, yet many reported configurations use \lambda = 0.1. Several scale-control effects help explain this choice.

  • Normalized layers are scale-invariant → gradient noise pushes norms up, decay pulls down → equilibrium: constant rotation per step (Kosson et al., 2024).
  • Through this equilibrium, \lambda influences 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 receive decay (96% of parameters).
  • Normalization parameters and biases are exempt because they directly set normalization scales.
  • Embeddings are exempt because their gradients are sparse while decay applies 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

Section 9.9 uses the same matrix/non-matrix split: Muon assigns matrices a different optimizer rather than only a different decay setting.

Optimizer-state memory

AdamW stores two additional numbers per parameter:

                                           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 parameters still require an fp32 master copy in this configuration.
  • 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 changes \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.