AdamW (Single Update Step)

Adam with L2 regularisation folded into the gradient doesn't actually regularise the way you'd expect — the adaptive denominator scales the penalty differently for every parameter, so weights with large gradients get barely any decay. AdamW fixes it by decoupling the decay from the gradient path entirely.

1m = beta1 * m + (1 - beta1) * grad
2v = beta2 * v + (1 - beta2) * grad²
3
4m_hat = m / (1 - beta1^t)
5v_hat = v / (1 - beta2^t)
6
7param = param - learning_rate * (m_hat / (sqrt(v_hat) + epsilon) + weight_decay * param)

Task: write adamw_step(params, grads, m, v, learning_rate, beta1, beta2, epsilon, weight_decay, timestep) returning [new_params, new_m, new_v], rounded to 4 decimal places.

• timestep starts at 1 and is used only for the bias correction.
• Return the moments before bias correction — that's what carries to the next step.

The bias correction matters most early on. Both moments start at zero, so without dividing by 1 - beta^t the first few updates would be far too small — and with beta2 = 0.999 that bias takes hundreds of steps to fade on its own.