AdamW

Adam keeps two exponential moving averages per parameter — of the gradient and of its square — and steps along the first divided by the root of the second.

Both are initialised at zero, so early estimates are biased toward it. Dividing by $1 - \beta^t$ removes the bias exactly; with $\beta_2 = 0.999$ the correction is still a factor of two at step 700, which is why omitting it produces the famous need for a longer warmup.

The decoupling

Classic L2 adds $\lambda\theta$ to the gradient, so it enters $m$ and $v$ and is then divided by $\sqrt{\hat v}$ along with everything else. The effective decay becomes per-parameter and inversely proportional to gradient history. AdamW applies $\lambda\theta$ outside the adaptive step, as in eq. 2, so every parameter decays at the same rate.

The consequence for tuning is that $\eta$ and $\lambda$ decouple too: under L2 they interact, and a learning-rate sweep silently sweeps the regularisation strength with it.

import torch
from torch import Tensor


@torch.no_grad()
def adamw_step(
    p: Tensor, grad: Tensor, m: Tensor, v: Tensor, step: int,
    lr: float = 3e-4, betas: tuple[float, float] = (0.9, 0.95),
    eps: float = 1e-8, weight_decay: float = 0.1,
) -> None:
    b1, b2 = betas
    m.mul_(b1).add_(grad, alpha=1 - b1)
    v.mul_(b2).addcmul_(grad, grad, value=1 - b2)

    m_hat = m / (1 - b1**step)
    v_hat = v / (1 - b2**step)

    # decay the weight itself, untouched by the adaptive scale
    p.mul_(1 - lr * weight_decay)
    p.addcdiv_(m_hat, v_hat.sqrt().add_(eps), value=-lr)

$\beta_2 = 0.95$ rather than Adam’s default $0.999$ is near-universal for language models: the shorter window reacts faster to the variance spikes that precede a loss divergence. Keep the moments in fp32 whatever the compute dtype — $v$ in bf16 underflows to zero for small gradients, and the step then divides by $\varepsilon$.