SGD Parameter Update

Problem

One step of stochastic gradient descent on nested parameters.

params and grads are nested lists with identical shapes. Return a new params structure where every element is updated as p - lr * g.

Hint

SGD takes a step opposite the gradient, scaled by the learning rate.

case 1

in: [[[1,2],[3,4]],[[0.1,0.1],[0.2,0.2]],0.5]

expected: [[0.95,1.95],[2.9,3.9]]

case 2

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

expected: [[1]]

case 3

in: [[[0.5,-0.5]],[[0.25,-0.25]],2]

expected: [[0,0]]

case 4

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

expected: [[1.5,1.5]]

Starter codePython

def sgd_step(params, grads, lr):
    # Your code here
    pass

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