Critical Batch Size & Gradient Noise

The problem, from first principles

Doubling the batch size does not always halve the number of steps. Small batches have noisy gradients, so averaging more examples helps a lot. Past a point the gradient is already accurate and extra examples only cost compute. The turning point is the critical batch size, and it can be estimated from gradients you already compute. The gradient noise scale is the ratio of the total variance of per-example gradients to the squared length of their mean: noise over signal. When the batch is much smaller than this ratio, steps are noise-dominated and a bigger batch helps nearly linearly. When it is much larger, steps are signal-dominated and a bigger batch barely helps.

From theory to code

Implement gradient_noise_scale and steps_to_target.

Constraints

• per_example_grads has shape (n, d), one gradient vector per example, n >= 2.
• gradient_noise_scale(per_example_grads) returns trace_sigma / ||g_mean||**2 where g_mean is the mean gradient and trace_sigma = sum over examples of ||g_i - g_mean||**2 / (n - 1) (the unbiased trace of the covariance).
• steps_to_target(batch_size, s_min, b_crit) is s_min * (1 + b_crit / batch_size), the number of steps needed to reach a fixed loss at that batch size.
• Return Python floats.

Hints