Noise Scale From Gradient Norms

Problem

Estimate the gradient noise scale from aggregate gradient norms.

Given the sum of per-sample squared gradient norms, the sample count n, and the squared norm of the full-batch gradient, estimate the per-sample variance as (sum_sq_norms - n * mean_norm_sq) / (n - 1), clamp it at 0, and return variance / mean_norm_sq. Return 0.0 when n < 2 or mean_norm_sq <= 0.

Hint

The n - 1 divisor makes the variance estimate unbiased.

Test cases

case 1

in: [12,6,1]

expected: 1.2

case 2

in: [40,5,4]

expected: 1.25

case 3

in: [1,1,1]

expected: 0

case 4

in: [4,4,2]

expected: 0

Starter codePython

def noise_scale_from_norms(sum_sq_norms, n, mean_norm_sq):
    # Your code here
    pass

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