Shampoo Preconditioner One Block

Problem

Perform one Shampoo update with diagonal preconditioners.

Accumulate L_i = L_i + sum_j(g_ij ** 2) and R_j = R_j + sum_i(g_ij ** 2), then update w_ij -= lr * g_ij / (L_i ** 0.25 * R_j ** 0.25 + eps). Return [w_new, L_new, R_new].

Hint

The fourth root of each accumulated row and column scale is the preconditioner.

Test cases

case 1

in: [[[1,2],[3,4]],[[0.5,0.5],[0.5,0.5]],[1,1],[1,1],0.1,1e-8]

expected: [[[0.959175171286947,1.959175171286947],[2.959175171286947,3.959175171286947]],[1.5,1.5],[1.5,1.5]]

case 2

in: [[[0,0],[0,0]],[[1,2],[3,4]],[0,0],[0,0],0.5,1e-8]

expected: [[[-0.1880301539472129,-0.31622776501683797],[-0.3772300279561973,-0.42294850448179844]],[5,25],[10,20]]

case 3

in: [[[1,1],[1,1]],[[1,0],[0,1]],[0,0],[0,0],0.1,1e-8]

expected: [[[0.900000001,1],[1,0.900000001]],[1,1],[1,1]]

Starter codePython

def shampoo_step(w, g, L, R, lr, eps):
    # Returns [w_new, L_new, R_new]
    # Your code here
    pass

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