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 serial 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) = S_{\min} \cdot (B + B_{\text{noise}})$

These two quantities trace out a Pareto frontier. $S(B)$ is monotonically decreasing in $B$ (large batches minimize wall-clock time). $E(B)$ is monotonically increasing in $B$ (small batches minimize total compute). $B \approx B_{\text{noise}}$ is the knee of this frontier: below it, serial steps explode while compute barely falls; above it, compute grows roughly linearly in $B$ while steps barely decrease.

For $B \ll B_{\text{noise}}$: noise dominates. Doubling $B$ nearly halves the serial step count at roughly constant total compute (near-linear speedup in time, near-free). For $B \gg B_{\text{noise}}$: signal dominates. Doubling $B$ barely reduces steps, so total compute roughly doubles with almost no time reduction (pure waste). At the knee $B \approx B_{\text{noise}}$, you pay roughly double the compute to halve the serial time: a balanced tradeoff, not a compute minimum.

$B_{\text{noise}}$ is measurable during training (estimate gradient variance from multiple mini-batches). It marks the Pareto knee between compute efficiency and wall-clock speed: below it, time gains are cheap; above it, time gains are expensive in compute. McCandlish et al. (2018), Section 2.2, frame this as the tradeoff curve between serial steps and total examples, and measure $B_{\text{noise}}$ for language models, finding that it increases during training, which is 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.

Under a constant-noise approximation of the SGD diffusion (Mandt, Hoffman, Blei 2017), the stationary density takes the Gibbs-like form:

$p(\theta) \propto \exp\left(-\frac{2B}{\eta \cdot \text{tr}(C)} L(\theta)\right)$

A common heuristic extension adds a curvature-dependent prefactor $|\det H(\theta)|^{-1/2}$ to argue that flatter minima (smaller Hessian eigenvalues) have higher stationary mass. This is a useful baseline, not a rigorous result: real SGD has state-dependent gradient covariance $C(\theta)$, which breaks detailed balance and the Gibbs form (Chaudhari and Soatto, 2018, "SGD performs variational inference").

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 picture is clean only when the noise is isotropic and state-independent, which is an idealization.

The Gibbs heuristic gives a theoretical narrative 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. This is an observation, not a theorem about generalization itself, and the underlying stationary-density formula is an approximation, not a rigorous identity for real SGD.

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. The state-dependent gradient covariance in real SGD violates the constant-noise assumption that makes the Gibbs form exact. Flatness measured by the Hessian trace may not correlate with generalization in all architectures.