Rethinking the limiting dynamics of SGD: modified loss, phase space oscillations, and anomalous diffusion

1 INTRODUCTION . Deep neural networks have demonstrated remarkable generalization across a variety of datasets and tasks . Essential to their success has been a collection of good practices on how to train these models with stochastic gradient descent ( SGD ) . Yet , despite their importance , these practices are mainly based on heuristic arguments and trial and error search . Without a general theory connecting the hyperparameters of optimization , the architecture of the network , and the geometry of the dataset , theory-driven design of deep learning systems is impossible . Existing theoretical works studying this interaction have leveraged the random structure of neural networks at initialization [ 1 , 2 , 3 ] and in their infinite width limits in order to study their dynamics [ 4 , 5 , 6 , 7 , 8 ] . Here we take a different approach and study the training dynamics of pre-trained networks that are ready to be used for inference . By leveraging the mathematical structures found at the end of training , we uncover an intricate interaction between the hyperparameters of optimization , the structure in the gradient noise , and the Hessian matrix that corroborates previously identified empirical behavior such as anomalous limiting dynamics . Not only is understanding the limiting dynamics of SGD a critical stepping stone to building a complete theory for the learning dynamics of neural networks , but recently there have been a series of works demonstrating that the performance of pre-trained networks can be improved through averaging and ensembling [ 9 , 10 , 11 ] . Combining empirical exploration and theoretical tools from statistical physics , we identify and uncover a mechanistic explanation for the limiting dynamics of neural networks trained with SGD . 2 DIFFUSIVE BEHAVIOR IN THE LIMITING DYNAMICS OF SGD . A network that has converged in performance will continue to move through parameter space [ 12 , 13 , 14 , 15 ] . To demonstrate this behavior , we resume training of pre-trained convolutional networks while tracking the network trajectory through parameter space . Let θ∗ ∈ Rm be the parameter vector for a pre-trained network and θk ∈ Rm be the parameter vector after k steps of resumed training . We track two metrics of the training trajectory , namely the local parameter displacement δk between consecutive steps , and the global displacement ∆k after k steps from the pre-trained initialization : δk = θk − θk−1 , ∆k = θk − θ∗ . ( 1 ) As shown in Fig . 1 , neither of these differences converge to zero across a variety of architectures , indicating that despite performance convergence , the networks continue to move through parameter space , both locally and globally . The squared norm of the local displacement ‖δk‖22 remains near a constant value , indicating the network is essentially moving at a constant instantaneous speed . This observation is quite similar to the “ equilibrium '' phenomenon or “ constant angular update '' observed in Li et al . [ 17 ] and Wan et al . [ 13 ] respectively . However , these works only studied the displacement for parameters immediately preceding a normalization layer . The constant instantaneous speed behavior we observe is for all parameters in the model and is even present in models without normalization layers . While the squared norm of the local displacement is essentially constant , the squared norm of the global displacement ‖∆k‖22 is monotonically growing for all networks , implying even once trained , the network continues to diverge from where it has been . Indeed Fig . 1 indicates a power law relationship between global displacement and number of steps , given by ‖∆k‖22 ∝ kc . As we ’ ll see in section 8 , this relationship is indicative of anomalous diffusion where c corresponds to the anomalous diffusion exponent . Standard Brownian motion corresponds to c = 1 . Similar observation were made by Baity-Jesi et al . [ 14 ] who noticed distinct phases of the training trajectory evident in the dynamics of the global displacement and Chen et al . [ 15 ] who found that the exponent of diffusion changes through the course of training . A parallel observation is given by Hoffer et al . [ 18 ] for the beginning of training , where they measure the global dis- placement from the initialization of an untrained network and observe a rate ∝ log ( k ) , a form of ultra-slow diffusion . These empirical observations raise the natural questions , where is the network moving to and why ? To answer these questions we will build a diffusion based theory of SGD , study these dynamics in the setting of linear regression , and use lessons learned in this fundamental setting to understand the limiting dynamics of neural networks . 3 RELATED WORK . There is a long line of literature studying both theoretically and empirically the learning dynamics of deep neural networks trained with SGD . Our analysis and experiments build upon this literature . Continuous models for SGD . Many works consider how to improve the classic gradient flow model for SGD to more realistically reflect momentum [ 19 ] , discretization due to finite learning rates [ 20 , 21 ] , and stochasticity due to random batches [ 22 , 23 ] . One line of work has studied the dynamics of networks in their infinite width limits through dynamical mean field theory [ 24 , 25 , 26 , 27 ] , while a different approach has used stochastic differential equations ( SDEs ) to model SGD directly , the approach we take in this work . However , recently , the validity of this approach has been questioned . The main argument , as nicely explained in Yaida [ 28 ] , is that most SDE approximations simultaneously assume that ∆t → 0+ , while maintaining that the learning rate η = ∆t is finite . The works Simsekli et al . [ 29 ] and Li et al . [ 30 ] have questioned the correctness of the using the central limit theorem ( CLT ) to model the gradient noise as Gaussian , arguing respectively that the heavy-tailed structure in the gradient noise and the weak dependence between batches leads the CLT to break down . In our work , we maintain the CLT assumption holds , which we discuss fur- ther in appendix A , but importantly we avoid the pitfalls of many previous SDE approximations by simultaneously modeling the effect of finite learning rates and stochasticity . Limiting dynamics . A series of works have applied SDE models of SGD to study the limiting dynamics of neural networks . In the seminal work by Mandt et al . [ 31 ] , the limiting dynamics were modeled with a multivariate Ornstein-Uhlenbeck process by combining a first-order SDE model for SGD with assumptions on the geometry of the loss and covariance matrix for the gradient noise . This analysis was extended by Jastrzębski et al . [ 12 ] through additional assumptions on the covariance matrix to gain tractable insights and applied by Ali et al . [ 32 ] to the simpler setting of linear regression , which has a quadratic loss . A different approach was taken by Chaudhari and Soatto [ 33 ] , which did not formulate the dynamics as an OU process , nor assume directly a structure on the loss or gradient noise . Rather , this analysis studied the same first-order SDE via the Fokker-Planck equation to propose the existence of a modified loss and probability currents driving the limiting dynamics , but did not provide explicit expressions . Our analysis deepens and combines ideas from all these works , where our key insight is to lift the dynamics into phase space . By studying the dynamics of the parameters and their velocities , and by applying the analysis first in the setting of linear regression where assumptions are provably true , we are able to identify analytic expressions and explicit insights which lead to concrete predictions and testable hypothesis . Stationary dynamics . A different line of work avoids modeling the limiting dynamics of SGD with an SDE and instead chooses to leverage the property of stationarity . These works [ 28 , 34 , 35 , 36 ] assume that eventually the probability distribution governing the model parameters reaches stationarity such that the discrete SGD process is simply sampling from this distribution . Yaida [ 28 ] used this approach to derive fluctuation-dissipation relations that link measurable quantities of the parameters and hyperparameters of SGD . Liu et al . [ 35 ] used this approach to derive properties for the stationary distribution of SGD with a quadratic loss . Similar to our analysis , this work identifies that the stationary distribution for the parameters reflects a modified loss function dependent on the relationship between the covariance matrix of the gradient noise and the Hessian matrix for the original loss . Empirical exploration . Another set of works analyzing the limiting dynamics of SGD has taken a purely empirical approach . Building on the intuition that flat minima generalize better than sharp minima , Keskar et al . [ 37 ] demonstrated empirically that the hyperparameters of optimization influence the eigenvalue spectrum of the Hessian matrix at the end of training . Many subsequent works have studied the Hessian eigenspectrum during and at the end of training . Jastrzębski et al . [ 38 ] , Cohen et al . [ 39 ] studied the dynamics of the top eigenvalues during training . Sagun et al . [ 40 ] , Papyan [ 41 ] , Ghorbani et al . [ 42 ] demonstrated the spectrum has a bulk of values near zero plus a small number of larger outliers . Gur-Ari et al . [ 43 ] demonstrated that the learning dynamics are constrained to the subspace spanned by the top eigenvectors , but found no special properties of the dynamics within this subspace . In our work we also determine that the top eigensubspace of the Hessian plays a crucial role in the limiting dynamics and by projecting the dynamics into this subspace in phase space , we see that the motion is not random , but consists of incoherent oscillations leading to anomalous diffusion . 4 MODELING SGD AS AN UNDERDAMPED LANGEVIN EQUATION . Following the route of previous works [ 31 , 12 , 33 ] studying the limiting dynamics of neural networks , we first seek to model SGD as a continuous stochastic process . We consider a network parameterized by θ ∈ Rm , a training dataset { x1 , . . . , xN } of size N , and a training loss L ( θ ) = 1N ∑N i=1 ` ( θ , xi ) with corresponding gradient g ( θ ) = ∂L ∂θ . The state of the network at the kth step of training is defined by the position vector θk and velocity vector vk of the same dimension . The gradient descent update with learning rate η , momentum β , and weight decayλ is given by vk+1 = βvk − g ( θk ) − λθk , θk+1 = θk + ηvk+1 , ( 2 ) where we initialize the network such that v0 = 0 and θ0 is the parameter initialization . In order to understand the dynamics of the network through position and velocity space , which we will refer to as phase space , we express these discrete recursive equations as the discretization of some unknown ordinary differential equation ( ODE ) , sometimes referred to as a modified equation as in [ 44 , 20 ] . While this ODE models the gradient descent process even at finite learning rates , it fails to account for the stochasticity introduced by choosing a random batch B of size S drawn uniformly from the set of N training points . This sampling yields the stochastic gradient gB ( θ ) = 1S ∑ i∈B∇ ` ( θ , xi ) . To model this effect , we make the following assumption : Assumption 1 ( CLT ) . We assume the batch gradient is a noisy version of the true gradient such that gB ( θ ) − g ( θ ) is a Gaussian random variable with mean 0 and covariance 1SΣ ( θ ) . Incorporating this model of stochastic gradients into the previous finite difference equation and applying the stochastic counterparts to Euler discretizations , results in the standard drift-diffusion stochastic differential equation ( SDE ) , referred to as an underdamped Langevin equation , d [ θ v ] = [ v − 2η ( 1+β ) ( g ( θ ) + λθ + ( 1− β ) v ) ] dt+ [ 0 0 0 2√ ηS ( 1+β ) √ Σ ( θ ) ] dWt , ( 3 ) where Wt is a standard Wiener process . This is the continuous model we will study in this work : Assumption 2 ( SDE ) . We assume the underdamped Langevin equation ( 3 ) accurately models the trajectory of the network driven by SGD through phase space such that θ ( ηk ) ≈ θk and v ( ηk ) ≈ vk . See appendix A for further discussion on the nuances of modeling SGD with an SDE .

Summary: The paper uses a stochastic differential equation approximation of stochastic gradient descent momentum to model the neural networks dynamics. At the theoretical level the study focuses on linear regression. They showed that there is agreement between the results in the simpler model and deep neural networks.  The most relevant fact is that the networks can have a super or sub diffusive behavior at the end of training. Unfortunately, their theory cannot explain this finding nor the practical consequences on the test loss.