AdamW Decoupled Weight Decay

Problem

Perform one AdamW update with decoupled weight decay.

Compute the Adam moments and bias-corrected estimates exactly as Adam, then update x_new = x - lr * (m_hat / (sqrt(v_hat) + eps) + weight_decay * x). Return [x_new, m_new, v_new]. Unlike Adam with L2 regularization, the decay is applied to x directly, not to the gradient.

Hint

Shrink x toward zero by lr * weight_decay * x in addition to the adaptive step.

Test cases

case 1

in: [1,0,0,0.5,1,0.1,0.9,0.999,1e-8,0.01]

expected: [0.899000002,0.04999999999999999,0.0002500000000000002]

case 2

in: [1,0.1,0.01,0.5,5,0.01,0.9,0.999,1e-8,0.1]

expected: [0.9966134845480366,0.14,0.01024]

case 3

in: [-2,-0.3,0.4,1.2,10,0.05,0.8,0.99,1e-8,0.02]

expected: [-1.998,-5.551115123125783e-17,0.41040000000000004]

Starter codePython

def adamw_update(x, m, v, g, t, lr, beta1, beta2, eps, weight_decay):
    # Returns [x_new, m_new, v_new]
    # Your code here
    pass

3 test cases · run your code against them in the editor