Technical Summary: Why SGD is not Brownian Motion

Problem Statement

Stochastic Gradient Descent (SGD) is the foundational optimization algorithm for deep learning, yet its theoretical understanding remains incomplete. A dominant line of research models SGD as a discretization of Langevin dynamics, treating minibatch noise as Brownian motion acting on a fixed potential (the average loss). This framework assumes a continuous-time limit and a specific $\sqrt{\eta}$ noise scaling to derive a stationary Gibbs distribution.

However, this paper argues that the standard Langevin approximation fails to faithfully reproduce discrete SGD dynamics at finite learning rates ($\eta$). The core discrepancy lies in the scaling of stochasticity: in the true discrete update, noise enters at order $\eta$ (via the minibatch loss), whereas the Langevin formulation imposes a $\sqrt{\eta}$ scaling to preserve diffusion in continuous time. Consequently, the Langevin approximation omits terms of order $\eta^2$ that are present in the exact discrete dynamics, leading to qualitatively incorrect predictions regarding the existence of stationary distributions and the behavior of parameter variance near critical points.

Methodology

The authors propose a new framework that treats SGD not as Brownian motion in a fixed potential, but as deterministic dynamics in a fluctuating loss landscape. The randomness arises intrinsically from the sampling of minibatches, which changes the objective function at every step.

The methodology proceeds as follows:

1. Exact Master Equation: Starting from the exact discrete SGD update rule $w_{n+1} = w_n - \eta \nabla L_n(w_n)$, the authors derive an exact master equation for the evolution of the parameter distribution $p_n(w)$.
2. Discrete Fokker–Planck Derivation: By expanding the master equation in powers of the learning rate $\eta$ and truncating at the second order, they derive a Discrete Fokker–Planck equation. This equation captures the evolution of the parameter density without invoking a continuous-time approximation.
3. Local Analysis: The framework is applied to the vicinity of a critical point of the average loss. Using a local quadratic approximation of the loss landscape, the authors analyze the dynamics of the mean and covariance of the parameter distribution in the eigenbasis of the mean Hessian.
4. Assumptions: To obtain tractable closed-form expressions, the analysis assumes:
   • Independent minibatch sampling with replacement.
   • Approximate diagonality of the Hessian fluctuations in the mean-Hessian eigenbasis.
   • An empirical relation where gradient noise covariance is proportional to the Hessian structure ($E[G_i G_j] \approx \gamma E[H_{ij}]$).

Key Contributions

1. Discrete Fokker–Planck Framework

The paper derives a second-order expansion of the parameter density evolution (Equation 4) that differs from the standard Langevin-based Fokker–Planck equation. Specifically, the discrete expansion includes an additional term proportional to $(\nabla \bar{L})^2$ at order $\eta^2$. The standard Langevin formulation omits this term because it assumes a $\sqrt{\eta}$ noise scaling, which is inconsistent with the discrete update at finite $\eta$.

2. Mismatch with Langevin Dynamics

The authors demonstrate that this omission is not merely a quantitative correction but can lead to qualitative differences. In a one-dimensional toy example, the standard Langevin approximation predicts a stationary state with finite variance under conditions where the exact discrete SGD dynamics predict exponential growth of variance (instability). This discrepancy arises because the Langevin truncation discards terms of the same formal order as the retained diffusion term.

3. Unified Picture of SGD Dynamics

The analysis reveals that SGD dynamics decompose into distinct regimes based on the eigenvalues ($\lambda_i$) of the mean Hessian:

• Confined (Sharp) Directions: For directions with sufficiently large positive curvature, the parameter variance converges to a finite stationary limit. The limiting variance is approximately proportional to the learning rate $\eta$ and independent of the curvature magnitude (consistent with the inverse variance–flatness relation).
• Diffusive (Flat) Directions: For nearly-flat directions ($\lambda_i \approx 0$), the variance does not converge to a stationary distribution within the observed time horizon. Instead, it grows linearly with time ($\propto n$), corresponding to effective diffusion along the valleys of the loss landscape.

4. Empirical Validation

The theoretical predictions are validated on neural network models in computer vision (MLP on MNIST) and natural language processing (NanoGPT on Shakespeare and WikiText-2).

• Small-scale experiments: Using exact dense Hessians, the authors confirm that the parameter covariance is approximately diagonal in the mean-Hessian eigenbasis and that the variance profiles match the theoretical predictions, showing a clear separation between saturating and non-saturating directions.
• Large-scale experiments: On a 6.6M-parameter NanoGPT model, using stochastic Lanczos to estimate Hessian spectra, the empirical variance saturation levels align with the discrete prediction (Equation 19) and deviate significantly from the standard Langevin prediction as the learning rate increases.

Results and Significance

The paper claims that SGD is more accurately understood as deterministic dynamics in a fluctuating loss landscape rather than Brownian motion in a fixed potential. This perspective resolves discrepancies between theory and practice observed at finite learning rates.

The significance of this work lies in:

• Correcting the Theoretical Basis: It highlights that the standard Langevin approximation is a poor approximation for discrete SGD at practical learning rates, missing critical $\eta^2$ terms that determine stability and stationarity.
• Explaining Dynamical Regimes: It provides a unified explanation for why SGD exhibits both confined behavior in sharp directions and persistent diffusion in flat directions, a phenomenon often observed empirically but not fully captured by equilibrium statistical physics models.
• Finite-Step Effects: It emphasizes the importance of discrete-time effects in optimization dynamics, suggesting that the "noise" in SGD is an intrinsic consequence of the algorithm's structure and the loss landscape geometry, rather than an external stochastic forcing term.

The authors note limitations, including the reliance on local quadratic approximations, assumptions about the independence of Hessian fluctuations, and the restriction to sampling with replacement. However, the empirical evidence across different architectures and tasks suggests that the derived dynamical regimes are a general property of high-dimensional optimization landscapes.