AdamW step (decoupled weight decay)

Background

AdamW is Adam with decoupled weight decay — the optimiser behind GPT-2, GPT-3, Llama, and essentially every modern transformer training run. It computes the same adaptive moment update as Adam, then applies an L2-style shrinkage term directly to the parameters instead of folding it into the gradient. That one change (Loshchilov & Hutter, 2017) is what makes weight decay regularise as intended in adaptive optimisers.

Problem statement

Implement adamw_step(params, grads, m, v, t, lr, beta1, beta2, eps, weight_decay) that applies one AdamW update, in place, to every parameter tensor. For each parameter $p$ with gradient $g$, moment buffers $m, v$, and weight-decay coefficient $\lambda$:

$$\begin{aligned} m &\leftarrow \beta_1\, m + (1-\beta_1)\, g \\ v &\leftarrow \beta_2\, v + (1-\beta_2)\, g^2 \\ \hat{m} &= \frac{m}{1-\beta_1^{\,t}}, \qquad \hat{v} = \frac{v}{1-\beta_2^{\,t}} \\ p &\leftarrow p - \mathrm{lr}\left(\frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon} + \lambda\, p\right) \end{aligned}$$

The decay term $\lambda\, p$ sits outside the adaptive $\hat{m}/\sqrt{\hat{v}}$ scaling — it is not added to $g$ before the moments are computed.

Input

• params — list[np.ndarray], the model parameters. Mutated in place.
• grads — list[np.ndarray], same shapes as params; the gradient $g$ for each.
• m — list[np.ndarray], first-moment buffers (start at 0). Mutated in place.
• v — list[np.ndarray], second-moment buffers (start at 0). Mutated in place.
• t — int, $t \ge 1$, the current step (1-indexed). Drives bias correction.
• lr, beta1, beta2, eps — scalars: learning rate, first/second-moment decay rates, stabiliser $\epsilon$.
• weight_decay — scalar $\lambda \ge 0$, the decoupled decay coefficient. weight_decay=0 is plain Adam.

Output

Returns None. Updates params, m, and v in place, so the caller's references see the new values after the call.

Examples

Example 1 — decay shrinks a parameter even when the gradient is zero

Input:  params=[1.0], grads=[0.0], m=[0.0], v=[0.0],
        t=1, lr=0.1, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0.5
Output: params≈[0.95]

Explanation: with $g=0$ the Adam term is 0, so only the decay acts: $p \leftarrow p - \mathrm{lr}\,\lambda\, p = 1 - 0.1\cdot0.5\cdot1 = 0.95$. Plain Adam would leave $p$ unchanged — this shrinkage toward 0 is the weight decay.

Example 2 — the decay is decoupled (moments never see it)

Input:  params=[1.0], grads=[1.0], m=[0.0], v=[0.0],
        t=1, lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0.1
Output: m=[0.1]   (= (1-β₁)·g, NOT (1-β₁)·(g + λ·p))

Explanation: the first moment is $m = (1-\beta_1)g = 0.1\cdot1 = 0.1$. The decay term is added straight to the parameter step, never to $g$, so it never enters $m$ or $v$. The wrong "Adam-with-L2" form would instead give $m = (1-\beta_1)(g + \lambda p) = 0.1\cdot1.1 = 0.11$.

Constraints

• Within each tuple, params[i], grads[i], m[i], v[i] share the same shape.
• $t \ge 1$ and is the true step count — bias correction divides by $1-\beta^{\,t}$.
• Update in place (*=, +=, -=).
• $\epsilon$ is added after the square root: $\sqrt{\hat{v}} + \epsilon$.
• Weight decay is decoupled: apply $\lambda\, p$ to the parameter update directly; do not add $\lambda\, p$ to $g$ before computing $m, v$. weight_decay=0 recovers plain Adam exactly.
• Tests compare with tolerance atol ≈ 1e-9.

Notes

• Decoupled vs L2. Folding $\lambda p$ into $g$ routes the decay through the adaptive $1/\sqrt{\hat{v}}$ scaling, so parameters with large gradients (which need more regularisation) receive less decay — backwards. AdamW keeps decay outside that scaling, restoring the behaviour everyone expects from weight decay.
• State persistence. m and v are optimiser state reused across steps, so the in-place update is part of the contract — it builds directly on the Adam optimiser step.