LAMB Layerwise Scaling

Problem

Perform one LAMB update on a scalar parameter.

Build the Adam-style update m_hat / (sqrt(v_hat) + eps) + weight_decay * w, then multiply it by the trust ratio ||w|| / ||update|| (1.0 when either norm is zero). Return [w - lr * ratio * update, ratio].

Hint

The trust ratio equalizes the update magnitude across layers.

Test cases

case 1

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

expected: [0.9,0.9900990295069111]

case 2

in: [0.5,0.1,0.01,0.5,5,0.01,0.9,0.999,1e-8,0]

expected: [0.495,2.0951048089323248]

case 3

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

expected: [-0.09999999800000003,1]

Starter codePython

def lamb_update(w, m, v, g, t, lr, beta1, beta2, eps, weight_decay):
    # Returns [w_new, ratio]
    # Your code here
    pass

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