Gradient Noise Scale Estimate

Problem

Estimate the gradient noise scale from per-sample gradients.

Compute the mean gradient and the unbiased per-coordinate variance dividing by n - 1, sum the variances into a trace, and return trace / (||mean||^2 + eps). Return 0.0 when fewer than 2 samples are given or the mean is the zero vector.

Hint

Use the unbiased variance with n - 1 and normalize by the squared gradient norm.

Test cases

case 1

in: [[[1,2],[1.2,1.8],[0.8,2.2]],1e-12]

expected: 0.0159999999999968

case 2

in: [[[1],[5]],1e-12]

expected: 0.8888888888887901

case 3

in: [[[0,0],[0,0]],1e-12]

expected: 0

Starter codePython

def gradient_noise_scale(grads, eps):
    # Your code here
    pass

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