Gradient Noise Scale and Critical Batch Size

Doubling a batch doubles the cost of a step and reduces the variance of the gradient estimate by half. The trade looks obviously worthwhile until the step count stops falling, and the point at which it stops is not a matter of implementation quality. It is determined by a ratio of two quantities that the training run can measure: the variance of the per-example gradient, and the squared norm of the full-batch gradient. That ratio has units of examples, and it is the batch size at which the noise in the gradient estimate is comparable to the gradient itself.

Treating it as a measurable quantity rather than a rule of thumb turns batch size selection into arithmetic.

The estimator, its bias and its noise

A mini-batch gradient $g_B = \frac{1}{B}\sum_{i \in B} \nabla \ell_i$ is an unbiased estimate of the full-batch gradient $G = \mathbb{E}[\nabla\ell]$, with covariance $\Sigma/B$ where $\Sigma$ is the covariance of the per-example gradient. The estimator’s error is therefore $|g_B - G| \approx \sqrt{\operatorname{tr}\Sigma / B}$ in expectation, since the squared error’s expectation is the trace of the covariance.

Define the gradient noise scale as the ratio

$$\mathcal{B} \equiv \frac{\operatorname{tr}\Sigma}{|G|^2}$$

then the relative error of the estimate is $\sqrt{\mathcal{B}/B}$: at $B = \mathcal{B}$ the noise has the same magnitude as the signal, and at $B \gg \mathcal{B}$ the estimate is close to $G$. The quantity is measured in examples, but it is an example count weighted by how much each example’s gradient contributes to the disagreement.

The definition uses the trace rather than the largest eigenvalue, which makes it an average over directions rather than a worst case. A gradient whose noise is concentrated in a few directions and negligible in others has a small trace even if the noisy directions have large variance. Optimisation cares about the worst direction, so the trace-based noise scale is an optimistic summary; it is nonetheless the one with a simple estimator and it is what large-batch analyses use.

Estimating it from two batch sizes

Both terms in the ratio are estimable from two independent estimates of the gradient at the same parameters, of sizes $B_{\text{small}}$ and $B_{\text{big}}$:

$$\hat{\mathcal{B}} = \frac{|g_{B_{\text{small}}} - g_{B_{\text{big}}}|^2}{B_{\text{small}}|g_{B_{\text{big}}}|^2}\cdot\frac{1}{1 - B_{\text{small}}/B_{\text{big}}}$$

The derivation: the difference of two unbiased estimates has covariance $\Sigma(1/B_{\text{small}} - 1/B_{\text{big}})$, so its squared norm estimates $\operatorname{tr}\Sigma (1/B_{\text{small}} - 1/B_{\text{big}})$, and dividing by the estimate of $|G|^2$ gives the ratio. The correction factor removes the bias introduced by using the larger batch’s norm in the denominator.

import math

def noise_scale(g_small, g_big, b_small, b_big):
    diff = math.sqrt(sum((a - b) ** 2 for a, b in zip(g_small, g_big)))
    norm = math.sqrt(sum(b * b for b in g_big))
    if norm == 0.0:
        return float("inf")
    ratio = (diff ** 2) / (b_small * norm ** 2)
    return ratio / (1.0 - b_small / b_big)

Two practical requirements follow from the arithmetic. The two batches must be drawn independently, or the difference is biased downward by shared examples. And $B_{\text{big}}$ must substantially exceed $B_{\text{small}}$, since the correction’s denominator approaches zero as they converge and amplifies the estimator’s error.

What the noise scale does to the step count

For a fixed total number of examples processed, the number of steps is $N/B$. If the per-step progress toward the optimum were proportional to the signal-to-noise ratio of the step direction, doubling $B$ would double the progress per step and halve the steps, making batch size free. It is not free, because the relationship between the estimate’s error and the distance the step moves is not linear near a minimum: the iterates perform a random walk whose radius is set by the noise, and the excess loss due to noise scales as the squared distance, i.e. as $1/B$. Doubling the batch halves the noise-induced excess loss but halves the number of steps, so the total examples required to reach a given excess loss is unchanged.

Combining the two regimes: while $B \ll \mathcal{B}$ the steps are dominated by noise and increasing $B$ reduces the number of steps roughly proportionally; while $B \gg \mathcal{B}$ the steps are dominated by the signal and increasing $B$ barely reduces the count. The transition is at $B \approx \mathcal{B}$, with the reduction in steps at the transition on the order of a small constant rather than an unbounded gain. This makes $\mathcal{B}$ the batch size above which the compute efficiency of increasing batch size is nearly gone, and below which there is something to gain.

The second consequence concerns the learning rate. If the rate is scaled up proportionally to $B$, the per-step displacement grows and the walk’s radius grows along with it; the useful rule is that the rate can rise only as far as the noise floor’s radius is permitted to rise, which ties the maximum useful batch size to the peak rate and to the schedule’s decay. A run at $B \gg \mathcal{B}$ with a rate scaled linearly in $B$ has a noise floor that is not reduced by the larger batch at all, since $\eta\sigma/\sqrt{B}$ grows as $\sqrt{B}$.

Why the noise scale grows over training

The denominator $|G|^2$ shrinks as the parameters approach a minimum: the full-batch gradient at a stationary point is zero. The numerator $\operatorname{tr}\Sigma$ measures per-example disagreement, which does not vanish — different examples have different gradients even at the optimum, since the minimum is not a point where every example’s gradient is zero.

The consequence is that $\mathcal{B}$ increases monotonically along a training trajectory, often by orders of magnitude. Early in training the gradient of the average loss is large and all examples agree on the coarse direction, so a small batch is a good estimate; late in training the average gradient is a residual of disagreeing examples and the estimate is dominated by sampling noise for any batch size that was adequate at the start.

Three practical statements follow. The optimal batch size is not a constant; a schedule that is optimal at step zero is wasteful at the end. A large batch chosen for early-training efficiency is above the late-training critical size, so the late phase’s steps are wasted on noise reduction that no longer helps. And a fixed batch size with a decaying rate spends the late phase at a noise radius that is small only because the rate is small, which is the mechanism discussed under schedules as step-size integrals.

Practical measurement points

Because $\mathcal{B}$ is a property of the trajectory, it should be measured where it will be used. The estimator above requires two gradients at the same parameters, which means a forward-backward at the sampled step with a larger batch than the training batch; a measurement every few thousand steps is cheap relative to the run. What to watch for:

• A measured $\mathcal{B}$ smaller than the current batch size at the end of training is expected and not an error.
• A $\mathcal{B}$ that decreases during training is a sign that the trajectory is leaving a minimum rather than approaching one, or that the loss is being evaluated on a distribution that is changing.
• The trace-based estimate is dominated by large-magnitude coordinates; parameters whose gradients are systematically larger (embeddings with rare tokens, the output head) contribute most of it. Per-group estimates are more informative than one scalar when the groups differ by orders of magnitude, and clipping as described in gradient clipping and norm scaling interacts with exactly those coordinates.

Limits of the model

The critical-batch picture assumes the loss surface is locally quadratic and that the noise’s covariance is fixed in time. Neither holds across a training run. The quadratic assumption fails wherever the curvature changes on the scale of the step, which is much of early training and the regions around saddle points. The fixed-covariance assumption fails because the batch’s examples are drawn from a finite dataset, so the per-example gradients at two points in training are correlated through the shared examples, and the estimator’s independence assumption across steps is violated.

There is also a boundary effect: the noise scale is defined in terms of directions, but the optimisation only cares about directions in which the loss changes. A gradient component in a flat direction can have large variance and contribute to the trace without affecting the loss at all, which makes $\mathcal{B}$ larger than the batch size the optimisation actually needs. Regularisation changes this: a gradient that the regulariser suppresses may be noise in the trace and signal in the loss, and the decoupled decay in decoupled weight decay and adaptive scaling applies a direction-free force that the noise analysis does not model.

Finally, the analysis computes the batch size at which a step’s direction is reliable, not the batch size at which the hardware is efficient. Those two are independent constraints, and the binding one is often the hardware, which is why large-batch recipes in practice pair a batch size near the measured critical size with a rate scaled to the schedule rather than to the batch.

The two constraints have different currencies, which is the practical resolution. Compute efficiency is measured in examples processed per unit of progress; hardware efficiency is measured in examples processed per second, and the latter improves with batch size in steps, as the batch crosses the size at which a matrix multiply saturates the device’s execution units. A batch below the hardware efficiency threshold leaves the device idle, and one above the critical batch size wastes examples on noise reduction. When the two do not overlap, the batch is usually chosen at the hardware threshold and the number of parameter updates is reduced accordingly, with the schedule’s total displacement adjusted to match.