Stochastic Gradient Variance

Problem

Estimate the total variance of a stochastic gradient from per-sample gradients.

Compute the mean gradient over the samples and the trace of the unbiased sample covariance, dividing each coordinate's squared deviations by n - 1. Return the trace, or 0.0 for fewer than two samples.

Hint

The trace of the covariance is the sum of per-coordinate variances.

Test cases

case 1

in: [[[1,2],[1.2,1.8],[0.8,2.2]]]

expected: 0.08

case 2

in: [[[1],[5]]]

expected: 8

case 3

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

expected: 0

case 4

in: [[[1,0]]]

expected: 0

Starter codePython

def stochastic_gradient_variance(grads):
    # Your code here
    pass

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