Adam Full Update Scalar

Problem

One full Adam update: m_new = b1 m + (1 - b1) g, v_new = b2 v + (1 - b2) g^2, m_hat = m_new / (1 - b1^t), v_hat = v_new / (1 - b2^t), x_new = x - lr * m_hat / (sqrt(v_hat) + eps).

Given x, m, v, g, b1, b2, t, lr, and eps, return [x_new, m_new, v_new].

Hint

Bias-correct both moments before forming the ratio.

Test cases

case 1

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

expected: [0.900000002,0.04999999999999999,0.0002500000000000002]

case 2

in: [0,1,1,-1,0.9,0.9,5,0.01,1e-8]

expected: [-0.012501373437368462,0.8,1]

case 3

in: [2,0.5,0.25,2,0.5,0.5,3,0.5,0]

expected: [1.5416507514858944,1.25,2.125]

Starter codePython

def adam_full_update_scalar(x, m, v, g, b1, b2, t, lr, eps):
    # Your code here
    pass

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