Momentum Update Step

Problem

Perform one momentum update step element-wise: v_new = momentum * v + grad and param_new = param - lr * v_new.

Return [new_params, new_velocities]; empty vectors return [[], []].

Hint

Update the velocity first, then apply it to the parameters.

case 1

in: [[1,2],[0.1,0.2],[0,0],0.1,0.9]

expected: [[0.99,1.98],[0.1,0.2]]

case 2

in: [[0],[1],[0.5],0.1,0.9]

expected: [[-0.145],[1.45]]

case 3

in: [[1],[-1],[1],1,0]

expected: [[2],[-1]]

case 4

in: [[],[],[],0.1,0.9]

expected: [[],[]]

Starter codePython

def momentum_update(params, grads, velocities, lr, momentum):
    # Returns [new_params, new_velocities]
    # Your code here
    pass

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