Large Batch Training › Why Batch Size Matters

Michael BrenndoerferJanuary 27, 20266 min read
Language AI HandbookAll chapters
Part 32 of 67Training Optimization
Chapter 234Large Batch Training › Why Batch Size Matters
Reading section 2 of 12 · Why Batch Size Matters · 6 minSection 2 of 12 · 6 min

Book contents

Section 2 of 12 · 6 min read

Large Batch Training › Why Batch Size MattersOne section · Sections in this chapter2 / 12 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Before examining the mechanics of scaling, it helps to understand what gradient descent computes and why batch size changes its character.

The Gradient and Its Estimator

The true objective you want to minimize is the loss averaged over your entire dataset:

L(θ)=1N∑i=1Nℓ(xi,yi;θ)L(\theta) = \frac{1}{N} \sum_{i=1}^{N} \ell(x_i, y_i; \theta)

where:

  • L(θ)L(\theta): the full-dataset loss as a function of model parameters θ\theta
  • NN: total number of training examples
  • ℓ(xi,yi;θ)\ell(x_i, y_i; \theta): the loss on the ii-th example (xi,yi)(x_i, y_i)

The true gradient of this loss with respect to the parameters is:

∇θL(θ)=1N∑i=1N∇θℓ(xi,yi;θ)\nabla_\theta L(\theta) = \frac{1}{N} \sum_{i=1}^{N} \nabla_\theta \ell(x_i, y_i; \theta)

Computing this exactly would require a pass over all NN examples for each parameter update. For datasets with millions or billions of samples, that is computationally intractable.

Stochastic gradient descent avoids this by computing a noisy estimate using a mini-batch B\mathcal{B} of size BB:

gB=1B∑i∈B∇θℓ(xi,yi;θ)g_B = \frac{1}{B} \sum_{i \in \mathcal{B}} \nabla_\theta \ell(x_i, y_i; \theta)

where:

  • gBg_B: the mini-batch gradient estimate
  • BB: the batch size
  • B\mathcal{B}: the set of BB randomly sampled example indices

This estimate is unbiased: its expected value equals the true gradient. But it has variance that depends inversely on BB:

Var[gB]=1B⋅σ2\text{Var}[g_B] = \frac{1}{B} \cdot \sigma^2

where σ2\sigma^2 is the per-sample gradient variance. Doubling the batch size halves the gradient variance. With a larger batch, your estimate of the true gradient direction becomes more accurate. This is simply the central limit theorem applied to gradient estimation: the sample mean of BB independent random vectors has variance 1/B1/B times the variance of a single sample.

Mini-batch SGD

Mini-batch stochastic gradient descent computes parameter updates from a small random subset of the training data rather than the full dataset. Each such subset is called a mini-batch. The update provides an unbiased but noisy estimate of the true gradient, trading computational efficiency for stochasticity.

Noise and the Quality of Updates

The noise in the gradient estimate plays a dual role, and understanding both sides is essential for making good batch size decisions.

On the positive side, gradient noise acts as a form of implicit regularization. Noisy updates prevent the optimizer from settling into sharp minima, those loss landscape features with high curvature in many directions. Sharp minima tend to correspond to solutions that overfit to training data. The parameters that minimize the training loss in a sharp minimum are very specific: small perturbations to the weights increase the training loss significantly. When you deploy the model, the test distribution differs slightly from training, and this difference looks like a perturbation to the weights. Sharp minima therefore generalize poorly.

Flat minima, where the loss is low and relatively insensitive to small perturbations, generalize better precisely because they tolerate the implicit shift between train and test. Gradient noise helps find these flat regions because the stochastic updates effectively prevent the optimizer from getting trapped in narrow, sharp basins. This connection between batch size, noise, and generalization was studied in depth by Keskar et al. (2017), who found that models trained with large batches consistently converged to sharper minima than those trained with small batches under the same budget.

On the negative side, noise means you sometimes step in the wrong direction. If the gradient noise is large relative to the true signal, you spend many updates wandering randomly rather than making consistent progress toward the minimum. With batch size 1, each update might point nearly opposite to the true gradient if you happen to sample an outlier example. The optimizer makes progress on average (the estimate is unbiased) but wastes many steps on unhelpful updates.

This tradeoff gives batch size its characteristic effect: small batches take many individually noisy steps, making slower but ultimately more stable progress, while large batches take fewer but more accurate steps, converging faster per epoch but risking poor generalization.

Curvature and the Role of the Learning Rate

The appropriate learning rate depends on the curvature of the loss landscape. In regions of high curvature (sharp narrow valleys), a large learning rate causes divergent oscillations as each update overshoots the local minimum and bounces to the other side. In regions of low curvature (flat plateaus), a small learning rate makes progress unnecessarily slow.

When you increase the batch size, you reduce gradient noise. This changes the effective signal-to-noise ratio of each update. If you keep the learning rate fixed while increasing the batch size, the optimizer takes fewer but more accurate steps per epoch. The effective step size relative to the gradient signal has changed, and the training dynamics shift. Specifically, with a fixed learning rate, each more-accurate large-batch step covers less distance in gradient-normalized space than the noisy small-batch step would have.

This is the core insight behind learning rate scaling rules: batch size and learning rate are not independent settings. They are coupled through the statistics of the gradient estimate. Increasing batch size without adjusting learning rate is like making more accurate measurements and then using fewer of them without increasing your confidence accordingly.

Historical Context and Motivation

Understanding where the linear scaling rule came from helps contextualize why it works when it does and fails when it does not. For most of the 2010s, practitioners treated batch size as primarily a hardware constraint. You used the largest batch that fit in GPU memory, typically 32 to 256, and accepted the associated learning rate.

The breakthrough came when distributed training made truly large batches possible. A cluster of 16 GPUs, each processing a batch of 32, produces an effective batch of 512 per step. A cluster of 512 GPUs produces an effective batch of 16,384. At these scales, the question of how to set the learning rate became urgent, because the naive answer (use the same learning rate as the single-GPU setting) clearly did not work.

The Facebook AI Research team's 2017 paper "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour" provided the systematic answer. They trained ResNet-50 on ImageNet using 256 GPUs in one hour, matching the accuracy of a week-long single-GPU run. The key contributions were the linear scaling rule, the warmup strategy, and careful attention to batch normalization behavior at large scale. This work catalyzed a wave of research and engineering that ultimately enabled the training of GPT-3, PaLM, and similar models at scales previously considered impossible.