SGD as a Stochastic Differential Equation

Why does doubling the batch size and doubling the learning rate leave training roughly unchanged? Why does a small batch generalise better than a large one at the same number of epochs? Why is there a "temperature" that people talk about as if SGD were a physical system?

All three have the same answer, and it comes from reading the discrete update as a discretisation of a continuous stochastic process.

The construction

The SGD update with learning rate \(\eta\) and batch size \(B\) is

where \(\Sigma\) is the per-example gradient covariance. The noise has zero mean, since the mini-batch gradient is unbiased, and covariance inversely proportional to batch size, since averaging \(B\) samples divides variance by \(B\).

Treating \(\eta\) as a time step and matching moments gives the stochastic modified equation

with \(W_t\) a Wiener process. This is the continuous-time model, made rigorous as a weak approximation with error \(O(\eta^2)\) per unit time by Li, Tai, and E, and used as the basis for the Bayesian reading of constant-rate SGD by Mandt, Hoffman, and Blei (Mandt et al., 2017, Stochastic Gradient Descent as Approximate Bayesian Inference, JMLR 18(134)).

What the equation tells you immediately

Only the ratio matters. The drift term has no \(\eta\) and no \(B\); the diffusion term depends on them only through \(\eta / B\). Two configurations with the same ratio follow the same continuous-time process and should train the same way. That is the linear scaling rule in one line: double the batch, double the learning rate, hold the trajectory. It is the reason the rule works, and it also predicts exactly where it stops working, namely when \(\eta\) is large enough that the \(O(\eta^2)\) discretisation error is no longer negligible.

Noise scale is a temperature. Under a quadratic approximation and constant \(\eta\), the process has a stationary distribution resembling a Gibbs measure at temperature proportional to \(\eta / B\), so SGD does not converge to a point; it settles into a distribution around a minimum whose width is set by that ratio. This is the formal content of the "SGD as approximate Bayesian inference" claim, and it explains why constant-learning-rate training plateaus at a noise floor rather than reaching the exact minimum.

Small batches generalise better for a mechanical reason. Higher temperature means the process cannot remain in narrow minima, because escape time from a basin grows exponentially in basin depth divided by temperature, a Kramers-type argument. Whether flatness causes generalisation is a separate and contested question, but the selection mechanism itself falls straight out of the SDE.

Learning-rate decay is annealing. Decaying \(\eta\) lowers the temperature, shrinking the stationary distribution around whatever basin the trajectory currently occupies. The decay phase of a schedule is not just "smaller steps"; it is a commitment step.

When it breaks

The approximation has four documented failure modes, and all four matter in practice.

The noise is not Gaussian. The derivation invokes a central limit argument over the mini-batch, but per-example gradients in deep networks are heavy-tailed, and empirical work argues the appropriate driving process is an \(\alpha\)-stable Lévy process rather than Brownian motion. That changes the escape-time analysis qualitatively, from exponential to polynomial in barrier height.

The covariance is state-dependent and anisotropic. \(\Sigma(\theta)\) is not a multiple of the identity; it is aligned with the local Hessian in ways that are the whole story for which minima are found. Treating it as isotropic, which most clean analyses do, discards the interesting part.

Large \(\eta\) breaks the discretisation. Real training runs at the edge of stability, where the step size is deliberately large enough that the discrete dynamics and the continuous flow diverge. The SDE describes the small-step limit of a process that is not run with small steps.

Adam is not this equation. Adaptive preconditioning changes both drift and diffusion, and the linear scaling rule empirically holds worse for Adam than for plain SGD, with square-root scaling often fitting better.

Use it as the right first-order mental model and the source of the scaling rules, not as a quantitative predictor. See critical batch size and learning rate schedules.