Momentum Update

Problem

One step of SGD with momentum: v_new = momentum * velocity - lr * grad and w_new = w + v_new.

The signature is momentum_update(w, grad, velocity, lr, momentum). Return [w_new, velocity_new].

Hint

Decay the old velocity, subtract the scaled gradient, then add the new velocity to the weights.

case 1

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

expected: [[0.99,1.98],[-0.01,-0.02]]

case 2

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

expected: [[0.4],[0.4]]

case 3

in: [[-1,0,1],[0,1,2],[0.5,0.5,0.5],0.01,0.5]

expected: [[-0.75,0.24,1.23],[0.25,0.24,0.23]]

case 4

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

expected: [[6.8],[1.8]]

Starter codePython

def momentum_update(w, grad, velocity, lr, momentum):
    # Your code here
    pass

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