Large Batch Training › The Linear Scaling Rule

The most influential contribution to large batch training was the systematic analysis of how learning rate should scale with batch size. The rule itself is simple to state, but its derivation illuminates exactly why it works and where it breaks down.

Consider SGD with batch size $B$ and learning rate $\eta$. The derivation below assumes the loss is locally linear over the region swept by $k$ consecutive steps, which is the key approximation. After one step, parameters change by:

$$\theta_{t+1} = \theta_t - \eta \cdot g_B$$

Now consider scaling the batch size by a factor $k$ to batch size $kB$. The new gradient estimate:

$$g_{kB} = \frac{1}{kB} \sum_{i \in \mathcal{B}'} \nabla_\theta \ell(x_i, y_i; \theta)$$

has variance $\frac{\sigma^2}{kB}$, which is $k$ times smaller than the variance of $g_B$.

The key intuition is this: with batch size $kB$, you need $k$ fewer steps to process the same amount of data. If you want the model after processing $n$ samples to be in the same state regardless of batch size, you need the cumulative weight update to be the same. One large-batch step with batch size $kB$ should replicate what $k$ small-batch steps with batch size $B$ would have achieved.

More formally, $k$ steps of SGD with learning rate $\eta$ and batch size $B$ produce:

$$\Delta\theta = -\eta \sum_{j=1}^{k} g_{B,j}$$

One step with batch size $kB$ and learning rate $\eta' = k\eta$ produces:

$$\Delta\theta' = -k\eta \cdot g_{kB}$$

When the loss landscape is locally linear (a first-order approximation), these two updates are equal in expectation, since $\mathbb{E}[g_{kB}] = \mathbb{E}[g_B] = \nabla L(\theta)$:

$$\mathbb{E}[\Delta\theta'] = -k\eta \cdot \nabla L(\theta) = \mathbb{E}[\Delta\theta]$$

This is the formal justification for the linear scaling rule.

The derivation is mathematically clean, but the intuition is worth dwelling on. The argument assumes that $k$ consecutive small-batch updates do not change the gradient significantly, so the gradient computed on the first mini-batch is approximately the same as the gradient on the second, third, and so on. This is the local linearity assumption: the loss function behaves like a flat plane over the small region swept by $k$ SGD steps.

When this assumption holds, the $k$ small-batch steps are equivalent to a single step in the average gradient direction, scaled by $k$. A single large-batch step computes this average gradient more accurately (lower variance) and takes a proportionally larger step (with $k\eta$), arriving at the same expected location.

This framing also shows what can go wrong. If the model parameters change substantially over $k$ steps (high curvature, large gradients, or large learning rate), then the gradient estimated at step 1 points in a different direction from the true gradient at step $k$. The equivalence breaks down, and simply scaling the learning rate linearly overestimates how large a step you can safely take.

The linear scaling rule is exact only in the limit where the loss is locally linear (equivalently, where the gradient does not change significantly within a single step). In practice, the loss landscape has curvature, and the approximation breaks down at two extremes.

At small batch sizes, the gradient noise is so high that the local linearity assumption is violated. The gradient estimated on one mini-batch points in a very different direction than the gradient estimated on the next, so the simple equivalence argument does not hold.

At very large batch sizes, a different problem emerges. As you keep scaling the learning rate up proportionally to batch size, eventually the learning rate becomes so large that a single step overshoots the local minimum. The update magnitude exceeds the scale over which the linear approximation is valid. This is why there is an upper limit to how far the linear scaling rule can be pushed.

The practical range where the rule works reliably varies by architecture and dataset. For ImageNet training with ResNets, it holds up to batch sizes of around 8,192. For large language models trained on massive corpora, the effective range depends on model size, architecture, and stage of training. The rule is a powerful starting point, not an infallible prescription.

Before the linear scaling rule was fully established, researchers proposed an alternative: the square root scaling rule, which suggests using learning rate $\eta \cdot \sqrt{k}$ when scaling batch size by $k$. This rule comes from a different analysis of noise dynamics.

The argument runs as follows. The variance of the gradient estimate scales as $\sigma^2 / B$. The standard deviation (noise amplitude) scales as $\sigma / \sqrt{B}$. If you want to keep the signal-to-noise ratio of your update constant as you increase batch size, you should scale the learning rate by $\sqrt{k}$, because the noise has decreased by $\sqrt{k}$.

In practice, the square root rule tends to be too conservative for moderate batch size increases. The linear rule matches empirical results better for scaling up to 10x or 100x the original batch size. For very large scaling factors (beyond 1,000x), neither rule alone is sufficient, and more sophisticated techniques are needed. The linear rule has become the dominant standard, but knowing the square root rule exists helps you understand why some older papers use different scaling conventions.

At the start of training, the network parameters are far from their final values. The loss landscape in this early phase is highly non-linear, with large curvature in many directions. The linear scaling rule's underlying assumption of local linearity is worst at the beginning.

This is why large batch training almost always uses a warmup phase. As we discussed in the Learning Rate Warmup chapter, you start with a small learning rate and gradually increase it to the target value over some number of steps. For large batch training, warmup is often essential for stable convergence.

The Facebook AI Research paper found that jumping directly to the linearly scaled learning rate caused instability in the first epoch. Using a warmup of 5 epochs resolved this. The intuition is that during warmup, the network moves from a random initialization to a region of parameter space where the local linearity assumption begins to hold, making the scaled learning rate safe. Put differently, warmup gives the optimizer time to find a basin before it starts taking large steps inside that basin.

The optimal warmup duration scales roughly with the batch scaling factor $k$. If you are scaling batch size by 8x, a warmup of 5 epochs may suffice. If you are scaling by 1,000x (as in large LLM training), you may need to warm up for thousands of steps. The general guideline is to warm up for at least 1-3% of total training steps, but this varies significantly across tasks.