Critical Batch Size Estimate

Problem

Estimate the critical batch size from gradient norms measured at two batch sizes.

Using E||g_B||^2 = ||G||^2 + sigma^2 / B, solve sigma^2 = (gs_norm_sq - gb_norm_sq) * bs * bb / (bb - bs) and ||G||^2 = gb_norm_sq - sigma^2 / bb, then return sigma^2 / ||G||^2. Return 0.0 when bb <= bs, gs_norm_sq <= gb_norm_sq, or the signal is not positive.

Hint

The noise estimate is amplified by the harmonic difference of the two batch sizes.

Test cases

case 1

in: [3,1.25,32,256]

expected: 64

case 2

in: [5,2,16,128]

expected: 34.90909090909091

case 3

in: [1,1,32,64]

expected: 0

case 4

in: [2,3,64,32]

expected: 0

Starter codePython

def critical_batch_size(gs_norm_sq, gb_norm_sq, bs, bb):
    # Your code here
    pass

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