Adafactor Relative Step Size

Problem

Perform one Adafactor update with factored second moments and relative step sizes.

Update 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). For each row take clip = max(eps2, RMS(w_i)) and set w_ij = w_ij - lr * clip * g_ij / (sqrt(row_i * col_j) + eps). Return [w_new, row_new, col_new].

Hint

The row RMS scales the step so small weights still move relative to their size.

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,0.001]

expected: [[[0.9145330371356326,1.9145330371356326],[2.808890061204824,3.808890061204824]],[0.925,0.925],[0.925,0.925]]

case 2

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

expected: [[[0.4950000000141421,-0.007071067797723341],[-0.006708203924014089,0.4936754446853201]],[2.5,12.5],[5,10]]

case 3

in: [[[0.01,0.02],[0.03,0.04]],[[0.1,0.1],[0.1,0.1]],[1,1],[1,1],0.1,0.5,1e-8,0.01]

expected: [[[0.009686903208163464,0.019686903208163464],[0.0292998942899164,0.039299894289916405]],[0.505,0.505],[0.505,0.505]]

Starter codePython

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

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