Stochastic learning flow: moment equations for stochastic gradient descent at finite step size
Abstract
Analytical theories of learning in deep networks often assume gradient flow, discarding the finite step size and batch noise of stochastic gradient descent (SGD). Starting from the discrete update, we derive continuous-time equations for the mean and covariance of the parameter distribution under a Gaussian closure for a general stochastic learning, to first order in the learning rate η. Applying these stochastic learning flow (SLF) equations to SGD, we show that the mean descends an effective loss regularized by the covariance and loss curvature, while the covariance, unlike in stochastic differential equation approximations, becomes unstable at curvature 2/η. In SGD ensembles of linear and nonlinear networks, SLF predicts learning curves and mean trajectories more accurately than gradient flow and captures the covariance dynamics, failing where the parameter distribution becomes strongly non-Gaussian. In the full-batch limit, the edge of stability corresponds to a stationary point of the coupled sharpness-covariance dynamics, recovering self-stabilization and central flow, and batch-to-batch curvature fluctuations lower this threshold under SGD.
Related
- Stochastic gradient descent on the epigenetic landscape: a unified framework for cellular plasticity, tumor heterogeneity, and the asymptotic irrelevance of fitness
- Convex optimization on moment polytopes: Hadamard mirror descent and efficient algorithms for quantum functionals and other tensor parameters
- Camera-in-the-loop gradient descent hologram generation for Fraunhofer diffraction
- Hybrid Analytical Numerical and Machine Learning Frameworks for Solving Deterministic and Stochastic Differential Equations with Stability, Convergence, and Uncertainty Quantification
- Accelerating natural gradient descent for PINNs with randomized numerical linear algebra
BibTeX
@misc{schmid2026stochastic,
title = {{Stochastic learning flow: moment equations for stochastic gradient descent at finite step size}},
author = {Christian Schmid and James M Murray},
howpublished = {arXiv},
year = {2026},
month = nov,
doi = {10.48550/arxiv.2611.00473},
eprint = {2611.00473},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2611.00473},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Stochastic learning flow: moment equations for stochastic gradient descent at finite step size
AU - Christian Schmid
AU - James M Murray
PY - 2026
DA - 2026/11/01
PB - arXiv
DO - 10.48550/arxiv.2611.00473
UR - https://arxiv.org/abs/2611.00473
AB - Analytical theories of learning in deep networks often assume gradient flow, discarding the finite step size and batch noise of stochastic gradient descent (SGD). Starting from the discrete update, we derive continuous-time equations for the mean and covariance of the parameter distribution under a Gaussian closure for a general stochastic learning, to first order in the learning rate η. Applying these stochastic learning flow (SLF) equations to SGD, we show that the mean descends an effective loss regularized by the covariance and loss curvature, while the covariance, unlike in stochastic differential equation approximations, becomes unstable at curvature 2/η. In SGD ensembles of linear and nonlinear networks, SLF predicts learning curves and mean trajectories more accurately than gradient flow and captures the covariance dynamics, failing where the parameter distribution becomes strongly non-Gaussian. In the full-batch limit, the edge of stability corresponds to a stationary point of the coupled sharpness-covariance dynamics, recovering self-stabilization and central flow, and batch-to-batch curvature fluctuations lower this threshold under SGD.
ER -