Adam Optimizer Step

Problem

Perform one Adam optimizer step. At step t with gradients g and moment estimates m and v: m = beta1*m + (1-beta1)*g, v = beta2*v + (1-beta2)*g^2, then use bias-corrected m_hat = m/(1-beta1^t), v_hat = v/(1-beta2^t) and update p = p - lr*m_hat/(sqrt(v_hat) + eps).

Return [new_params, new_m, new_v]; empty vectors return [[], [], []].

Hint

Update moments, correct for initialization bias, then take the step.

Test cases

case 1

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

expected: [[0.9900000009999999],[0.009999999999999998],[0.000010000000000000011]]

case 2

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

expected: [[0.900000002,2.099999998],[0.04999999999999999,-0.04999999999999999],[0.0002500000000000002,0.0002500000000000002]]

case 3

in: [[1],[0.2],[0.1],[0.01],2,0.01,0.9,0.999,1e-8]

expected: [[0.9974153900445542],[0.11000000000000001],[0.01003]]

case 4

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

expected: [[],[],[]]

Starter codePython

import math
def adam_step(params, grads, m, v, t, lr, beta1, beta2, eps):
    # Returns [new_params, new_m, new_v]
    # Your code here
    pass

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