Adding an L2 penalty to the gradient doesn't work well with Adam: the penalty gets divided by Adam's adaptive denominator, so heavily-updated weights are barely regularised. AdamW (Loshchilov & Hutter, 2019) decays the weights directly instead.
Implement one step of each for NumPy arrays, matching PyTorch (bias-corrected, with step t starting at 1). Both return (param, m, v):
adam_l2_step(p, g, m, v, t, lr, wd, b1=0.9, b2=0.999, eps=1e-8): use g + wd * p as the gradient, then take a normal Adam step.adamw_step(p, g, m, v, t, lr, wd, b1=0.9, b2=0.999, eps=1e-8): first p = p * (1 - lr * wd), then take a normal Adam step using g alone.Then implement decay_after(steps, p0, lr, wd, method): starting from p0 (an array), apply steps steps with zero loss gradient and return the final parameter. With zero gradient, AdamW shrinks weights geometrically, while Adam+L2 moves each weight by about lr per step whatever its size.