Stochastic Gradient Flow in High Dimensions

• Stochastic Gradient Flow is a continuous-time stochastic process that approximates stochastic gradient descent through differential equations modeling both drift and intrinsic noise.
• SGF leverages stochastic differential equations to reveal algorithmic phenomena such as implicit regularization, double descent, and noise-induced phase transitions in high dimensions.
• Applications of SGF span machine learning, statistical physics, and control systems, offering insights into convergence, uncertainty quantification, and nonparametric estimation.

Stochastic Gradient Flow (SGF) is a class of continuous-time stochastic processes used to approximate the dynamics of stochastic gradient descent (SGD) and related iterative optimization procedures, particularly in machine learning, high-dimensional statistics, statistical physics, and stochastic control. SGF captures not only deterministic gradient flow but also stochastic fluctuations arising from data subsampling or intrinsic noise, and provides a mathematically principled framework for analyzing algorithmic behavior, uncertainty quantification, and emergent properties in high dimensions.

1. Mathematical Formulation and Key SDE Representations

SGF is typically formalized by a stochastic differential equation (SDE) whose drift term corresponds to the gradient of a target functional (such as risk or free energy), while the diffusion term encodes the stochasticity inherent in the underlying procedure. A prototypical SGF for parameter vector $\btheta^t \in \mathbb{R}^d$ is given by

$\de \btheta^t = -\,\alpha\Bigl( h_t(\btheta^t) +\frac{1}{\delta}\bX^\top \ell_t(\br^t;z) \Bigr)\de t + \sqrt{\frac{\tau}{\delta} \sum_{i=1}^n \bx_i\,\ell_t(r^t_i;z_i)^\top \de B^t_i} ,\quad \br^t=\bX\btheta^t,$

where $\delta = n/d$ is the sample-to-dimension ratio, $\tau = \eta/B$ defines an effective noise scale (learning rate over batch size), $h_t$ and $\ell_t$ encode regularizer and loss gradients, and $\{B^t_i\}$ are independent standard Brownian motions (Nishiyama et al., 6 Feb 2026).

For least squares, a canonical SGF is

$$d\beta(t) = \frac{1}{n}X^\top(y-X\beta(t))\,dt + Q_\epsilon(\beta(t))^{1/2} dW(t),$$

where $Q_\epsilon(\beta)$ depends on data-subsampling statistics and $W(t)$ is Brownian motion (Ali et al., 2020). In nonlinear or nonparametric contexts, the drift and diffusion may be highly state-dependent, and generalizations to measure-valued or infinite-dimensional settings are routine, e.g., in density-flow or field-theoretic models (Caluya et al., 2019, Kuehn et al., 2018, Carosso et al., 2019).

2. High-Dimensional Limit and Dynamical Mean-Field Theory (DMFT)

The analysis of SGF in the high-dimensional regime ($d, n \to \infty$, $n/d \to \delta$) reveals the emergence of closed low-dimensional (often scalar or matrix) integro-differential equations—a consequence of self-averaging and dynamical concentration. Dynamical mean-field theory (DMFT), originating in statistical physics, provides the principal analytic tool to rigorously derive these limit systems and to capture the time evolution of macroscopic order parameters such as covariance, response functions, and alignment statistics.

For generalized linear and two-layer mean-field network models, the DMFT reduction for SGF yields coupled equations for single-coordinate processes $(\theta^t, r^t)$: $$\begin{aligned} d\theta^t &= u^t\,dt - \Bigl( h_t(\theta^t) + \Gamma(t)\theta^t + \int_0^t R_\ell(t,s)\theta^s\,ds \Bigr)dt, \quad u \sim \mathrm{GP}(0, \delta^{-1}C_\ell), \ r^t &= w^t - \delta^{-1}\!\int_0^t R_\theta(t,s)\ell_s(r^s;z)\left(ds + \sqrt{\tau\delta}\,dB^s\right),\quad w\sim \mathrm{GP}(0, C_\theta), \end{aligned}$$ with additional order parameters $C_\theta$, $R_\theta$, $C_\ell$, $R_\ell$ self-consistently defined (Nishiyama et al., 6 Feb 2026). The macroscopic evolution notably depends on the data spectrum (e.g., through Marchenko–Pastur law or random feature kernels), noise intensity $\tau$, and nonlinear response kernels and curvatures.

3. Statistical and Algorithmic Properties

The SGF framework provides explicit analytical results on the risk, convergence, and generalization of stochastic optimization algorithms.

• Implicit Regularization and excess risk: In least squares, the SGF path closely matches ridge regression with time-dependent regularization $\lambda=1/t$. An explicit excess-risk bound,

$$\mathrm{Risk}(\beta^{sgf}(t)) - \mathrm{Risk}(\beta^{ridge}(1/t)) \le 0.6862 \, \mathrm{Var}(\beta^{ridge}(1/t)) + (\epsilon n/m) \sum_{i=1}^p \nu_i(t),$$

quantifies the effect of batch size $m$, step size $\epsilon$, and iteration time $t$ on the approximation of SGD by SGF (Ali et al., 2020). The residual fluctuations vanish as either $\epsilon \to 0$ or $m \to \infty$.

• Test risk and double descent: Path-integral analysis yields closed formulas for the test-risk difference between SGF and pure gradient flow, notably revealing an $O(\gamma)$ correction to generalization error attributable to stochasticity. In the weak-features regime, SGF accurately captures double-descent and its stochastic uplift below the interpolation threshold (Veiga et al., 2024).
• Generalization to nonlinear and nonconvex settings: For high-dimensional neural nets and Gaussian mixture models, DMFT shows that finite-batch SGF noise regularizes the landscape, suppresses overfitting spikes, and enables escape from poor local minima, effectively smoothing nonconvex objectives (Mignacco et al., 2020).

4. Stochastic Modified Flows and Mean-Field SDEs

To achieve precise correspondence with multi-step SGD statistics (beyond single-point marginals), "stochastic modified flows" (SMF) generalize classical diffusion approximations. The SMF SDEs are driven by cylindrical Brownian motions matched to the distribution of SGD stochasticity: $$dX_t^\eta = -\nabla( R(X_t^\eta) + \tfrac{\eta}{4}|\nabla R(X_t^\eta)|^2)dt + \sqrt{\eta} \int_\Theta G(X_t^\eta, \theta) W(d\theta,dt),$$ where $G(x,\theta)$ encodes micro-level fluctuations and $W$ is cylindrical noise (Gess et al., 2023). In the mean-field (infinite-width) regime, this approach extends to McKean–Vlasov SDEs governing the empirical distribution of network parameters.

These flows enjoy two vital properties: (i) maximal regularity (even with degenerate or low-rank empirical gradient covariances); (ii) exact matching of finite-dimensional statistics to SGD up to $O(\eta^2)$.

5. Relation to Stochastic Gradient Processes and Uncertainty Quantification

SGF is also employed as a continuous-time envelope for piecewise-deterministic Markov processes or stochastic gradient processes (SGP), especially in settings with discrete or continuously parameterized data indices. Here, parameter dynamics are coupled to a fast-mixing Markov process representing data or loss subsampling: $$\begin{cases} di(t) = \text{CTMC generator}, \ d\theta(t) = -\nabla \Phi_{i(t)}(\theta(t))dt, \end{cases}$$ with strong convexity/ergodicity results showing precise convergence to the full-gradient flow as mixing or learning rates tend to zero (Latz, 2020, Jin et al., 2021).

For model-free uncertainty quantification, SGF-driven systems allow construction of nonparametric bases (e.g., via diffusion maps) for linear and nonlinear tasks such as prediction, filtering, and response, circumventing the curse of dimensionality when appropriately truncated (Berry et al., 2014).

6. SGF in Infinite Dimensions, Physical Systems, and Control

SGF formalism generalizes to infinite-dimensional systems such as stochastic partial differential equations (SPDEs), field theory, and mean-field neural field models:

• Gradient flows in SPDEs and statistical physics: Heat equations, macroscopic energy flows, and phase-field models can be recast as SGFs in suitable Hilbert or Wasserstein metric spaces (Caluya et al., 2019, Peletier et al., 2014, Kuehn et al., 2018).
• Stochastic control and measure-valued flows: In entropy-regularized stochastic control, SGF governs the evolution of probability-measure–valued controls in the Wasserstein space, ensuring contractive convergence to optimal controls and revealing a precise connection to variational Bayesian inference and mean-field RL algorithms (Šiška et al., 2020).
• Stochastic RG and quantum field theory: SGF equations arise as Langevin/Fokker–Planck representations of functional RG evolution, clarifying connections between gradient flow, equilibrium distributions, and critical behavior (Carosso et al., 2019, Carosso, 2019).
• Emergence in fundamental physics: Recent work invokes SGF of field variables to interpolate between topological field theory and classical gravity, positing stochastic evolution in field space with phase transitions between ultraviolet (BF theory), pre-geometric, and IR (Einstein–Hilbert) regimes (Addazi et al., 22 May 2025).

7. Unified Perspective, Applications, and Open Directions

The SGF/DMFT formalism unifies gradient descent, online, and multi-pass SGD in the high-dimensional regime, capturing macroscopic evolution via low-dimensional kernels that encode data geometry, noise intensity, and model-specific nonlinearity (Nishiyama et al., 6 Feb 2026). This framework illuminates algorithmic phenomena such as:

• The interplay of drift and noise in implicit regularization and generalization performance.
• The emergence of lazy versus rich feature learning and the mechanisms underlying phase transitions ("double descent", early stopping windows).
• The systematic analysis of convergence rates, test/train error, and fluctuation-determined phase structure in both synthetic and real-world optimization landscapes.

Practical impacts include improved understanding of optimal learning rate and batch size schedules, design of nonparametric uncertainty quantification algorithms, architectural design principles for neural networks, and rigorous control of statistical and geometric properties in physical and control systems.

Open directions include rigorous infinite-time control, low-regularity extensions, common-noise mean-field limits for deep architectures, and further applications to complex systems where stochastic fluctuations play a fundamental dynamical or structural role.