AdamW

Adam with weight decay applied directly to the weights instead of injected into the gradient.

Equation

$$m_t={\beta }_1{m}_{t-1}+(1-{\beta }_1)g_t,v_t={\beta }_2{v}_{t-1}+(1-{\beta }_2)g_t^2$$ $${\hat{m}}_t=\frac{m_t}{1-{\beta }_1^t},{\hat{v}}_t=\frac{v_t}{1-{\beta }_2^t},{\theta }_t={\theta }_{t-1}-\eta \left(\frac{{\hat{m}}_t}{\sqrt{{\hat{v}}_t}+ϵ}+\lambda {\theta }_{t-1}\right)$$

L2 (Adam) sets $g_t+=\lambda \theta$, so the penalty passes through $1/\sqrt{\hat{v}}$, parameters with large historical gradients get less regularization, which is backwards. Decoupled (AdamW) keeps the $\lambda \theta$ term outside the adaptive scaling, making decay independent of gradient history and of $\eta$‘s adaptive rescaling.

From scratch

class AdamW:
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.95), eps=1e-8, wd=0.1):
        self.p = list(params); self.lr, self.eps, self.wd = lr, eps, wd
        self.b1, self.b2 = betas
        self.m = [torch.zeros_like(x) for x in self.p]
        self.v = [torch.zeros_like(x) for x in self.p]
        self.t = 0
 
    @torch.no_grad()
    def step(self):
        self.t += 1
        bc1 = 1 - self.b1 ** self.t
        bc2 = 1 - self.b2 ** self.t
        for x, m, v in zip(self.p, self.m, self.v):
            if x.grad is None:
                continue
            m.mul_(self.b1).add_(x.grad, alpha=1 - self.b1)
            v.mul_(self.b2).addcmul_(x.grad, x.grad, value=1 - self.b2)
            x.mul_(1 - self.lr * self.wd)                  # decoupled decay
            x.addcdiv_(m / bc1, (v / bc2).sqrt() + self.eps, value=-self.lr)

Notes

• Exclude biases, LayerNorm/norm gains, and often embeddings from wd: decaying 1-D params costs accuracy.
• Transformer default: betas=(0.9, 0.95), wd=0.1; the CV default beta2=0.999 is too slow-moving for spiky LLM gradients.
• State is $2\times$ params in fp32: the dominant memory cost after activations.
• Pairs with Learning Rate Schedules (warmup is near-mandatory) and Gradient Clipping.