Adafactor Factored Second Moment

Problem

Perform one Adafactor update with factored second moments on a 2D parameter.

row_i = beta2 * row_i + (1 - beta2) * mean_j(g_ij ** 2) and col_j = beta2 * col_j + (1 - beta2) * mean_i(g_ij ** 2); the per-entry variance is row_i * col_j. Return [w_new, row_new, col_new] with w_ij -= lr * g_ij / (sqrt(row_i * col_j) + eps).

Hint

The outer product of row and column moments approximates the full second-moment matrix.

Test cases

case 1

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

expected: [[[0.945945946530314,1.9459459465303142],[2.945945946530314,3.945945946530314]],[0.925,0.925],[0.925,0.925]]

case 2

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

expected: [[[0.9858578644162691,-0.01999999996],[-0.01897366593701028,0.9821114561960017]],[2.5,12.5],[5,10]]

case 3

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

expected: [[[0.929835359440053,0.08432740284893459],[-0.1281025214030441,1]],[0.625,0.75],[0.8125,0.5625]]

Starter codePython

def adafactor_update(w, g, row, col, lr, beta2, eps):
    # Returns [w_new, row_new, col_new]
    # Your code here
    pass

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