Nesterov Momentum Single Step

Problem

Take one Nesterov momentum step given the gradient at the lookahead point.

Update v_new = mu * v - lr * g and x_new = x + v_new, returning [x_new, v_new]. The gradient g is assumed to already be evaluated at x + mu * v.

Hint

Velocity accumulates the momentum term and the negative gradient.

case 1

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

expected: [0.8,-0.2]

case 2

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

expected: [0.55,0.55]

case 3

in: [-2,-1,0.5,0.2,0.8]

expected: [-2.9,-0.9]

Starter codePython

def nesterov_single_step(x, v, g, lr, mu):
    # Returns [x_new, v_new]
    # Your code here
    pass

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