Batch Size and Learning Dynamics

If SGD with learning rate $\eta$ and batch size $B$ produces a certain training trajectory (in the continuous-time SDE approximation), then SGD with learning rate $k\eta$ and batch size $kB$ produces approximately the same trajectory, for any scaling factor $k > 0$.

The effective noise temperature of SGD is:

$T_{\text{eff}} = \frac{\eta}{B} \cdot \text{tr}(C(\theta))$

Scaling both $\eta$ and $B$ by $k$ preserves $T_{\text{eff}}$.

What matters for the dynamics is the ratio $\eta/B$, not $\eta$ or $B$ individually. This ratio controls the magnitude of the noise injected per step. If you use $4\times$ larger batches, you must use $4\times$ larger learning rate to maintain the same effective noise.

This is the theoretical justification for the practice of scaling the learning rate linearly with batch size, used in large-scale distributed training (Goyal et al., 2017). Without this rule, increasing the batch size would reduce noise and change the implicit regularization of SGD.

The linear scaling rule breaks when: (1) the learning rate becomes so large that the continuous-time SDE approximation fails (discrete effects dominate), (2) the loss landscape is not smooth enough for the local diffusion approximation, or (3) during the initial transient before SGD reaches a stationary regime. In practice, a warm-up period is needed for large learning rates.

Let $B_{\text{noise}} = \text{tr}(C(\theta)) / \|\nabla L(\theta)\|^2$ be the gradient noise scale. The number of optimization steps $S$ to reach a target loss $\epsilon$ scales as:

$S(B) \approx S_{\min} \cdot \left(1 + \frac{B_{\text{noise}}}{B}\right)$

where $S_{\min}$ is the minimum steps achievable (at $B \to \infty$). The total compute (in units of samples processed) is $E(B) = B \cdot S(B)$, minimized near $B \approx B_{\text{noise}}$.

For $B \ll B_{\text{noise}}$: noise dominates. Doubling $B$ halves the number of steps with nearly the same total compute (linear speedup). For $B \gg B_{\text{noise}}$: signal dominates. Doubling $B$ barely reduces steps, so total compute roughly doubles (no speedup). The crossover is at $B \approx B_{\text{noise}}$.

$B_{\text{noise}}$ is measurable during training (estimate gradient variance from multiple mini-batches). It tells you the largest useful batch size: going beyond $B_{\text{noise}}$ wastes compute. McCandlish et al. (2018) measured $B_{\text{noise}}$ for language models and found it increases during training, explaining why larger batches become useful later.

The analysis assumes the gradient variance $C(\theta)$ is approximately constant, which changes as training progresses. The noise scale $B_{\text{noise}}$ itself varies over training, so the optimal batch size is not a single number but a trajectory.

In the SDE approximation of SGD, the stationary distribution concentrates on regions where the local loss is low and the Hessian eigenvalues are small. Specifically, the stationary density is approximately proportional to:

$p(\theta) \propto \exp\left(-\frac{2B}{\eta} L(\theta)\right) \cdot |\det H(\theta)|^{-1/2}$

where $H(\theta)$ is the Hessian. The second factor favors flat minima (small eigenvalues of $H$).

SGD noise acts like a temperature that helps the optimizer escape sharp minima (high curvature) but not flat minima (low curvature). The wider a minimum is, the harder it is for noise to push the iterate out. Higher noise temperature ($\eta/B$ large) means only the flattest minima are stable.

This provides a theoretical explanation for the empirical observation that small-batch SGD often generalizes better than large-batch SGD: flat minima are associated with better generalization because nearby points have similar loss (the function is stable under perturbation of parameters). This is an observation, not a theorem about generalization itself; the connection between flatness and generalization is debated.

The sharp/flat minima story has known weaknesses. Dinh et al. (2017) showed you can reparameterize a network to make any minimum arbitrarily sharp without changing the function it computes. The SDE approximation also breaks for large learning rates. Flatness measured by the Hessian trace may not correlate with generalization in all architectures.