Error Feedback Step

Problem

One step of error feedback for compressed SGD. Add the carried error to the gradients, keep only the top-k entries by absolute value (ties by row then column) as the sparse update, and return [sparse, new_error] where new_error is the residual of the effective gradient that was not transmitted.

The signature is error_feedback_step(gradients, error, k). Both are 2D matrices of the same shape.

Hint

Carrying the compression error keeps stale gradient mass from being lost forever.

Test cases

case 1

in: [[[1,2],[3,4]],[[0,0],[0,0]],2]

expected: [[[0,0],[3,4]],[[1,2],[0,0]]]

case 2

in: [[[1,0],[0,1]],[[0.5,0],[0,0.5]],1]

expected: [[[1.5,0],[0,0]],[[0,0],[0,1.5]]]

case 3

in: [[[0,0],[0,0]],[[0,0],[0,0]],1]

expected: [[[0,0],[0,0]],[[0,0],[0,0]]]

case 4

in: [[[1,2]],[[0,1]],2]

expected: [[[1,3]],[[0,0]]]

case 5

in: [[[2]],[[1]],1]

expected: [[[3]],[[0]]]

Starter codePython

def error_feedback_step(gradients, error, k):
    # Your code here
    pass

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