Mini Batch Gradient Variance

Problem

The gradient noise in a mini-batch is captured by the population variance of the per-sample gradients, computed coordinatewise.

Given a list of gradient vectors, return a list with the variance of each coordinate.

Hint

Compute the mean first, then average squared deviations per coordinate.

case 1

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

expected: [1,1]

case 2

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

expected: [0,0]

case 3

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

expected: [0.6666666666666666,2]

Starter codePython

def mini_batch_gradient_variance(grads):
    # Your code here
    pass

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