Critical Batch Size

September 2026 – Vladislav Kruglikov

An Empirical Model of Large-Batch Training introduces the gradient noise scale, a useful way to think about the largest batch size that still improves training efficiency.

Batch size and data parallel communication

In data parallelism, each worker runs forward and backward on its local batch, then workers communicate gradients. The compute work grows with batch size. Doubling the batch roughly doubles the forward and backward work for that step. The gradient payload does not grow with batch size because it has one value per model parameter, not one value per activation or training example.

Therefore, an all-reduce communicates roughly the same number of bytes for a batch of 32 and a batch of 256. A larger batch makes each step more compute-heavy while the communication per step stays constant. This can improve compute-communication overlap and reduce the total number of gradient synchronizations needed to reach a target loss.

If increasing batch size by a factor of XX also reduces the required number of optimizer steps by roughly XX, the same amount of example-level work is completed in fewer steps. Instead of paying for XX gradient all-reduces, the larger batch pays for one. It saves roughly X−1X-1 communication rounds. The backward computation per example is not eliminated. The saving comes from fewer optimizer steps and fewer fixed communication costs.

Why batch size cannot grow forever

The true gradient is the average gradient over all training examples. Computing it exactly is usually too expensive, so training uses a mini-batch estimate instead. A small batch gives a noisy estimate. A larger batch averages more examples and produces an estimate closer to the true gradient.

This trades computation for gradient accuracy. A small fraction of the data can already capture most of the useful gradient signal. Capturing the remaining small improvement may require a disproportionately large number of extra examples. This is only a loose Pareto-style intuition, not a universal 80/20 law. The exact relationship depends on the data distribution, model, and point in training.

The gradient noise scale roughly measures the randomness in gradient estimates. Below that scale, gradient noise dominates and larger batches can make useful progress faster. Above it, the noise is already sufficiently suppressed, so extra examples make the gradient only slightly more useful.

Think of the gradient noise scale as the effective number of examples needed in one step to make good progress. If the gradient noise scale is roughly 1,000, a batch near 1,000 can tame most of the noise:

This can also be visualized on a loss-surface contour plot, where the level curves connect points with equal loss. The true gradient points downhill. Mini-batch gradients are noisy arrows around that direction. Once a large batch already points close to the true gradient, doubling the batch may move the arrow only a few degrees closer while doubling the compute. The exact angle depends on the problem, but the noise falls only with the square root of batch size, which explains the diminishing return.

This also helps explain why reinforcement learning can use global batch sizes of millions while simple handwritten digit classification can use batches of tens or hundreds. Reinforcement-learning gradients are highly stochastic: they depend on sampled actions, environment transitions, rewards, and long trajectories. More samples are needed to suppress that noise. Digit classification has a simpler supervised signal and often lower gradient variance, so much smaller batches can already give a useful estimate. The exact batch size still depends on the model, dataset, and training stage.

Batch variance and learning-rate scaling

The optimizer wants the mean gradient over the whole data distribution. A batch gives only an estimate of that mean. Different batches produce different estimates, spread around the true mean. The wider this spread, the more likely the estimated gradient is to be far from the true one. We therefore want to reduce the variance of the batch gradient so that its estimate is more likely to be close to the mean.

«Gradient noise is the deviation of an individual or batch gradient from the true mean gradient.»

For one training example, consider one component of its gradient, gig_i. A batch of size BB produces the average gradient component:

gˉ=1B∑i=1Bgi\bar{g} = \frac{1}{B}\sum_{i=1}^{B}g_i

If examples are independent and identically distributed, scaling a random variable by 1/B1/B scales its variance by 1/B21/B^2, while the variance of an independent sum is the sum of the individual variances. Therefore:

Var⁡(gˉ)=Var⁡(1B∑i=1Bgi)=1B2∑i=1BVar⁡(gi)=1B2 Bσg2=σg2B\operatorname{Var}(\bar{g}) = \operatorname{Var}\left(\frac{1}{B}\sum_{i=1}^{B}g_i\right) = \frac{1}{B^2}\sum_{i=1}^{B}\operatorname{Var}(g_i) = \frac{1}{B^2}\,B\sigma_g^2 = \frac{\sigma_g^2}{B}

The data distribution and the current model parameters determine σg\sigma_g. Batch size controls how much averaging suppresses that variability. Taking the square root gives the standard deviation of the batch gradient:

Std⁡(gˉ)=σgB\operatorname{Std}(\bar{g}) = \frac{\sigma_g}{\sqrt{B}}

Increasing the batch size by a factor of four reduces the standard deviation of the batch gradient by a factor of two. This is the source of the diminishing return: every additional example helps, but its marginal reduction in noise becomes smaller.

The important quantity is the standard deviation of the batch gradient, σg/B\sigma_g / \sqrt{B}. Consider the same batch size for two tasks. The values below are illustrative, but they show the main point:

TaskPer-example gradient std. σg\sigma_gBatch size BBBatch gradient std.
MNIST110241/1024=0.031251 / \sqrt{1024} = 0.03125
RL10102410/1024=0.312510 / \sqrt{1024} = 0.3125

The batch size is the same, but the resulting gradient spread is ten times larger for RL because its data and gradients are more stochastic. To reach the same spread, RL needs a larger batch.

If your goal is to keep the STD of each SGD update unchanged, then for plain SGD the update is the batch gradient multiplied by the learning rate:

Δθ=−ηgˉ\Delta\theta = -\eta\bar{g}

Multiplying a random variable by a constant multiplies its standard deviation by the absolute value of that constant. Because Std⁡(gˉ)=σ/B\operatorname{Std}(\bar{g}) = \sigma / \sqrt{B}, the update's STD is:

Std⁡(Δθ)=ηStd⁡(gˉ)=ησB\operatorname{Std}(\Delta\theta) = \eta\operatorname{Std}(\bar{g}) = \frac{\eta\sigma}{\sqrt{B}}

Here, η\eta is the learning rate and σ\sigma is the standard deviation of this gradient component across individual examples. If you increase BB by a factor of nn and the learning rate by n\sqrt{n}, then:

(n η)σnB=ησB\frac{(\sqrt{n}\,\eta)\sigma}{\sqrt{nB}} = \frac{\eta\sigma}{\sqrt{B}}

So, if you increase the batch size by 4×4\times and the learning rate by 2×2\times, the update STD stays the same at the same weights with independent samples. But the expected update also becomes two times larger, so the optimizer can still overshoot. Preserving the STD does not guarantee stability or optimal training.

Critical batch size and diminishing returns

Informally, the critical batch size is the largest batch size that still gives nearly linear training speedup. If doubling the batch almost halves the optimizer steps needed to reach a target loss, scaling is still efficient. Past the turning point, doubling the batch no longer comes close to halving those steps, so the extra compute gives little speedup.

This refers to optimization efficiency, not steps per epoch. For a fixed dataset, doubling batch size always halves the number of steps in an epoch by definition. Hardware throughput can also saturate when a device is full, but that is a separate limit. The optimization critical batch size can be lower than the hardware-saturation batch size.

The critical batch size is the scale at which increasing batch size stops giving proportional improvement in training time to a target loss. Below it, a larger batch can reduce the number of steps enough to offset the extra computation. Above it, each larger batch adds forward and backward work without reducing the steps to target loss by the same factor. Training then becomes slower or less compute-efficient, even though it performs fewer gradient all-reduces.

To preserve similar optimization behavior while scaling batch size in the useful regime, the learning rate is often scaled roughly linearly with batch size. This rule is empirical and only applies over a limited range. It must be checked for the model and training schedule rather than extended to arbitrarily large batches.

Fit the architecture to the data

For a given experiment, treat the dataset as fixed. It defines the distribution the model must fit. In that sense, the data is the world of the problem: it is given, and changing it changes the problem itself.

Once the data distribution is fixed, compare architectures on that same distribution. The goal is to find an architecture whose capacity and inductive bias fit the data well. We tune the architecture for the data, rather than tuning the data to make one architecture look better. Keeping this order makes architecture comparisons meaningful.

References