Generic Adam Update

Problem

Perform one Adam update on a scalar parameter at time step t (1-indexed).

m_new = beta1 * m + (1 - beta1) * g, v_new = beta2 * v + (1 - beta2) * g * g, then bias-correct m_hat = m_new / (1 - beta1^t) and v_hat = v_new / (1 - beta2^t). Finally x_new = x - lr * m_hat / (sqrt(v_hat) + eps). Return [x_new, m_new, v_new].

Hint

Bias correction divides by 1 minus beta raised to t; update x using the corrected moments.

Test cases

case 1

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

expected: [0.900000002,0.04999999999999999,0.0002500000000000002]

case 2

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

expected: [0.9976134845480366,0.14,0.01024]

case 3

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

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

case 4

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

expected: [0,0,0]

Starter codePython

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

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