The Limiting Dynamics of SGD: Modified Loss, Phase Space Oscillations, and Anomalous Diffusion

Abstract

In this work we explore the limiting dynamics of deep neural networks trained with stochastic gradient descent (SGD). As observed previously, long after performance has converged, networks continue to move through parameter space by a process of anomalous diffusion in which distance travelled grows as a power law in the number of gradient updates with a nontrivial exponent. We reveal an intricate interaction between the hyperparameters of optimization, the structure in the gradient noise, and the Hessian matrix at the end of training that explains this anomalous diffusion. To build this understanding, we first derive a continuous-time model for SGD with finite learning rates and batch sizes as an underdamped Langevin equation. We study this equation in the setting of linear regression, where we can derive exact, analytic expressions for the phase space dynamics of the parameters and their instantaneous velocities from initialization to stationarity. Using the Fokker-Planck equation, we show that the key ingredient driving these dynamics is not the original training loss, but rather the combination of a modified loss, which implicitly regularizes the velocity, and probability currents, which cause oscillations in phase space. We identify qualitative and quantitative predictions of this theory in the dynamics of a ResNet-18 model trained on ImageNet. Through the lens of statistical physics, we uncover a mechanistic origin for the anomalous limiting dynamics of deep neural networks trained with SGD.

arXiv comment78 pages, 9 figures, Neural Computation 2024

References in corpus (28)

• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• Neural Tangent Kernel: Convergence and Generalization in Neural Networks
• One weird trick for parallelizing convolutional neural networks
• Don't Decay the Learning Rate, Increase the Batch Size
• Deep Neural Networks as Gaussian Processes
• A Mean Field View of the Landscape of Two-Layers Neural Networks
• Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent
• Stochastic Dynamical Structure (SDS) of Nonequilibrium Processes in the Absence of Detailed Balance. II: construction of SDS with nonlinear force and multiplicative noise
• Three Factors Influencing Minima in SGD
• Structure of Stochastic Dynamics near Fixed Points
• Empirical Analysis of the Hessian of Over-Parametrized Neural Networks
• Gradient Descent Happens in a Tiny Subspace
• Eigenvalues of the Hessian in Deep Learning: Singularity and Beyond
• An Investigation into Neural Net Optimization via Hessian Eigenvalue Density
• On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport
• On the Relation Between the Sharpest Directions of DNN Loss and the SGD Step Length
• Fluctuation-dissipation relations for stochastic gradient descent
• On the Origin of Implicit Regularization in Stochastic Gradient Descent
• Stochasticity helps to navigate rough landscapes: comparing gradient-descent-based algorithms in the phase retrieval problem
• The Implicit Regularization of Stochastic Gradient Flow for Least Squares
• The Full Spectrum of Deepnet Hessians at Scale: Dynamics with SGD Training and Sample Size
• A Tail-Index Analysis of Stochastic Gradient Noise in Deep Neural Networks
• Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability
• On the Validity of Modeling SGD with Stochastic Differential Equations (SDEs)
• Optimization and Generalization of Shallow Neural Networks with Quadratic Activation Functions
• Neural Mechanics: Symmetry and Broken Conservation Laws in Deep Learning Dynamics
• Reconciling Modern Deep Learning with Traditional Optimization Analyses: The Intrinsic Learning Rate
• Anomalous diffusion dynamics of learning in deep neural networks

Cited by in corpus (3)

• Weight fluctuations in (deep) linear neural networks and a derivation of the inverse-variance flatness relation
• Noise Balance and Stationary Distribution of Stochastic Gradient Descent
• Anti-Correlated Noise in Epoch-Based Stochastic Gradient Descent: Implications for Weight Variances in Flat Directions