Large Batch Training › Batch Size Limits

The linear scaling rule describes how to adjust the learning rate when you scale up batch size. But there is a separate question: is there a limit to how large a batch you should use, regardless of learning rate scaling?

The answer is yes, and this limit is sometimes called the critical batch size.

The critical batch size is the point at which further increasing the batch size no longer improves the speed of convergence (measured in number of examples processed). Below this size, larger batches speed up convergence. Above it, you get no benefit from more parallelism.

Intuitively, this happens because beyond a certain batch size, the gradient estimate is already accurate enough. The bottleneck for learning is no longer gradient noise but the optimization geometry itself: the curvature of the loss landscape, saddle points, and movement across a complex high-dimensional surface. Adding more samples to a batch that is already providing a low-variance gradient estimate is simply redundant.

Think of it this way: if the true gradient points northeast, a batch that gives you an estimate pointing northeast-by-north is good enough. Doubling the batch size to get an estimate pointing northeast-by-northeast gives you a more accurate estimate but does not let you take a bigger step. You are still limited by the curvature of the loss surface, which determines the maximum safe step size.

Formally, the critical batch size is related to the ratio of gradient noise to gradient signal. McCandlish et al. (2018) formalized this as:

$$B_{\text{crit}} = \frac{\text{tr}(H \Sigma)}{g^T H g}$$

where:

• $B_{\text{crit}}$: the critical batch size
• $H$: the Hessian of the loss (capturing curvature)
• $\Sigma$: the per-sample gradient covariance matrix
• $g$: the true gradient vector
• $\text{tr}(\cdot)$: the matrix trace operator

This expression measures the ratio of gradient noise (numerator) to gradient signal (denominator) in a curvature-aware way. The critical batch size is larger when the gradient noise is high relative to the signal, and smaller when the signal is strong.

The numerator, $\text{tr}(H \Sigma)$, captures the total noise in the gradient weighted by curvature. Noise in directions of high curvature matters more because those directions are the ones where accurate gradient estimation most strongly affects convergence. The denominator, $g^T H g$, is the curvature-weighted gradient signal, measuring how much the gradient points in useful (high-curvature) directions. The ratio gives a curvature-sensitive signal-to-noise measure, and the critical batch size is the point at which increasing batch size no longer improves this ratio relative to the cost of the extra computation.

In practice, you cannot easily compute this quantity directly because you cannot compute the full Hessian for large models. But you can estimate the critical batch size empirically: train with progressively larger batches (adjusting the learning rate proportionally) and measure how many examples you need to process to reach a target validation loss. The batch size at which this number stops decreasing is approximately the critical batch size.

Research has found that the critical batch size varies widely across tasks and architectures:

• For image classification on CIFAR-10 with small models, critical batch sizes can be as low as a few hundred.
• For ImageNet with ResNet-50, the critical batch size is roughly 1,000 to 2,000 (much lower than the 8,192 that works with careful tuning, suggesting the extra scaling is not free from a convergence-per-example perspective).
• For language model pre-training (GPT-style), the critical batch size grows as training progresses. Early in training, when gradients are large and noisy, large batches help. Later in training, when the model is near convergence, smaller batches would suffice.
• For transformer-based language models, Kaplan et al. (2020) found that the critical batch size scales with the number of tokens processed, not just model size, suggesting a dynamic rather than fixed optimal batch size.

The implication for practitioners is important: there is an efficiency frontier. At batch sizes below the critical value, you are leaving GPU parallelism on the table. At batch sizes above it, you are wasting compute on redundant gradient averaging. Optimal large-scale training aims to hit this frontier.

However, hardware constraints often force you above the critical batch size anyway. When you have 10,000 GPUs, you need an enormous per-step batch just to utilize all the hardware. In these cases, the extra batch scaling beyond the critical size is a tax paid for hardware utilization. Engineers accept some inefficiency (in examples processed per useful gradient step) in exchange for faster wall-clock time.

Even at or below the critical batch size, large batch training can produce models that generalize slightly worse than small batch training with the same total compute budget. This is the "generalization gap" of large batch training, documented by Keskar et al. (2017).

The mechanism is the sharpness of minima. Large batch training, with its accurate gradient estimates, tends to converge to sharp minima where the loss is low but curvature is high. Small batch training, with its inherent noise, naturally explores the loss landscape more and finds flatter minima. Flat minima generalize better because small perturbations to the weights (representing test distribution shift) cause only small increases in the loss.

Several approaches address this gap:

• Ghost batch normalization: Run batch normalization with a smaller "ghost" batch size even when processing a larger batch, preserving the regularization effect of BN statistics computed on smaller groups. This decouples batch normalization from the training batch size.
• Linear warmup: Helps the optimizer find a good basin before scaling to the full learning rate, giving it time to identify flat regions.
• Sharpness-aware minimization (SAM): Explicitly penalizes sharp minima during training, regardless of batch size, by seeking parameters whose neighborhood has uniformly low loss.
• Data augmentation: Increases effective noise through input perturbations, partially compensating for the reduced gradient noise of large batches. More aggressive augmentation with large batches can close much of the generalization gap.
• Gradient noise injection: Deliberately adding noise to gradients to simulate the effect of smaller batches, while keeping the large batch for hardware efficiency.

In practice, for language model pre-training at the scales used in production (batch sizes of millions of tokens), the generalization gap is usually small and manageable with standard techniques. The model capacity and data volume dominate, and the batch size effect becomes a secondary concern.