LARS Layerwise Update

Problem

Perform one LARS layer update with a local learning rate.

Compute eta = trust_coef * lr * ||w|| / (||g|| + wd * ||w|| + eps), then return w - eta * (g + wd * w) elementwise. The trust coefficient keeps the update norm proportional to the weight norm.

Hint

The denominator adds the decay-scaled weight norm to the gradient norm.

Test cases

case 1

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

expected: [0.857697507304884,2.1423024926951157]

case 2

in: [[1,1],[0.2,0.2],0.05,1,0.01,1e-8]

expected: [0.9500000016835876,0.9500000016835876]

case 3

in: [[3,4],[0,0],0.2,0.8,0,1e-8]

expected: [3,4]

Starter codePython

def lars_layer_update(w, g, lr, trust_coef, wd, eps):
    # Your code here
    pass

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