Logarithmic landscape and power-law escape rate of SGD

Stochastic gradient descent ( SGD ) undergoes complicated multiplicative noise for the mean-square loss . We use this property of the SGD noise to derive a stochastic differential equation ( SDE ) with simpler additive noise by performing a random time change . In the SDE , the loss gradient is replaced by the logarithmized loss gradient . By using this formalism , we obtain the escape rate formula from a local minimum , which is determined not by the loss barrier height ∆L = L ( θs ) − L ( θ∗ ) between a minimum θ∗ and a saddle θs but by the logarithmized loss barrier height ∆ logL = log [ L ( θs ) /L ( θ∗ ) ] . Our escape-rate formula strongly depends on the typical magnitude h∗ and the number n of the outlier eigenvalues of the Hessian . This result explains an empirical fact that SGD prefers flat minima with low effective dimensions , which gives an insight into implicit biases of SGD . 1 INTRODUCTION . Deep learning has achieved breakthroughs in various applications in artificial intelligence such as image classification ( Krizhevsky et al. , 2012 ; LeCun et al. , 2015 ) , speech recognition ( Hinton et al. , 2012 ) , natural language processing ( Collobert & Weston , 2008 ) , and natural sciences ( Iten et al. , 2020 ; Bapst et al. , 2020 ; Seif et al. , 2021 ) . Such unparalleled success of deep learning hinges crucially on stochastic gradient descent ( SGD ) or its variants as an efficient training algorithm . Although the loss landscape is highly nonconvex , the SGD often succeeds in finding a global minimum . It has been argued that the SGD noise plays a key role in escaping from local minima ( Jastrzȩbski et al. , 2017 ; Wu et al. , 2018 ; 2020 ; Zhu et al. , 2019 ; Meng et al. , 2020 ; Xie et al. , 2021 ; Liu et al. , 2021 ) . It has also been suggested that SGD has an implicit bias that is beneficial for generalization . That is , SGD may help the network to find flat minima , which are considered to imply good generalization ( Keskar et al. , 2017 ; Hoffer et al. , 2017 ; Wu et al. , 2018 ) . How and why does SGD help the network escape from bad local minima and find flat minima ? These questions have been addressed in several works , and it is now recognized that the SGD noise strength and structure importantly affect the efficiency of escape from local minima . Our work follows this line of research , and add new theoretical perspectives . In physics and chemistry , escape from a local minimum of the ( free ) energy landscape due to thermal noise at temperature T has been thoroughly discussed ( Kramers , 1940 ; Langer , 1969 ) . When the ( free ) energy barrier is given by ∆E , the escape rate is proportional to e−∆E/T , which is known as the Arrhenius law . By analogy , in machine learning , escape from a local minimum of the loss function is considered to be determined by the loss barrier height ∆L = L ( θs ) −L ( θ∗ ) , where L ( θ ) denotes the loss function at the network parameters θ , θ∗ stands for a local minimum of L ( θ ) , and θs denotes a saddle point that separates θ∗ from other minima . If we assume that the SGD noise is uniform and isotropic , which is often assumed in machine-learning literature ( Jastrzȩbski et al. , 2017 ) , the escape rate is proportional to e−∆L/D , where D denotes the SGD noise strength . In this paper , we show that inhomogeneity of the SGD noise strength brings about drastic modification for the mean-square loss . It turns out that the escape rate is determined by the logarithmized loss barrier height ∆ logL = logL ( θs ) −logL ( θ∗ ) = log [ L ( θs ) /L ( θ∗ ) ] . In other words , the escape rate is determined by not the difference but the ratio of L ( θs ) and L ( θ∗ ) . This result means that even if the loss barrier height ∆L is the same , minima with smaller values of L ( θ∗ ) are more stable . Moreover , given the fact that the eigenvalue spectrum of the Hessian at a minimum consists of a bulk of almost zero eigenvalues and outliers ( Sagun et al. , 2017 ; Papyan , 2019 ) , our escape-rate formula implies that SGD prefers flat minima with a low effective dimension , where the effective dimension is defined as the number of outliers ( MacKay , 1992 ) and flatness is measured by a typical magnitude of outlier eigenvalues ( Keskar et al. , 2017 ) . The previous theories ( Jastrzȩbski et al. , 2017 ; Wu et al. , 2018 ; Zhu et al. , 2019 ; Meng et al. , 2020 ; Xie et al. , 2021 ; Liu et al. , 2021 ) have also successfully explained that SGD prefers flat minima , but not shown the preference of small effective dimensions . The logarithmized loss naturally explains the latter , and sheds light on implicit biases of SGD . Main contributions : We obtain the following main results : • We derive an equation for approximating the SGD noise in Eq . ( 6 ) . Remarkably , the SGD noise strength in the mean-square loss is shown to be proportional to the loss function , which is experimentally verified in Sec . 5.2 . A key ingredient in deriving Eq . ( 6 ) is the decoupling approximation given in Eq . ( 7 ) . This is a novel approximate method introduced in our analysis , and hence we experimentally verify it in Sec . 5.1 . • We derive a novel stochastic differential equation ( SDE ) in Eq . ( 14 ) via a random time change introduced in Eq . ( 13 ) . Although the original SDE ( 4 ) has a multiplicative noise , the transformed SDE ( 14 ) has a simple additive noise with the gradient of the logarithmic loss . This shows the convenience of the logarithmic loss landscape for understanding SGD . • We derive a novel form of SGD escape rate from a local minimum in Eq . ( 16 ) . Remarkably , the escape rate depends on the ratio between L ( θ∗ ) and L ( θs ) . In Sec . 5.3 , we experimentally test the validity of this result for linear regressions . • Our escape rate crucially depends on the flatness and the effective dimension , which shows that SGD has implicit biases towards flat minima with low effective dimension . We also show in Eq . ( 18 ) that a local minimum with an effective dimension n greater than a certain critical value nc becomes unstable . Related works : The role of the SGD noise structure has been discussed in some previous works ( Zhu et al. , 2019 ; Xie et al. , 2021 ; Liu et al. , 2021 ; Meng et al. , 2020 ; Wojtowytsch , 2021 ) . It was pointed out that the anisotropic nature of the SGD noise is important : the SGD noise covariance matrix is aligned with the Hessian of the loss function , which is beneficial for escape from sharp minima ( Zhu et al. , 2019 ; Xie et al. , 2021 ; Liu et al. , 2021 ) . These previous works , however , do not take the inhomogeneity of the SGD noise strength into account , and consequently , escape rates derived there depend exponentially on the loss barrier height , which differs from our formula . Compared with the anisotropy of the SGD noise , the inhomogeneity of the SGD noise strength has been less explored . In ( Meng et al. , 2020 ; Wojtowytsch , 2021 ) , the SGD dynamics under a statedependent noise is discussed . However , in these previous works , the connection between the noise strength and the loss function was not theoretically established , and the logarithmized loss landscape was not discussed . The instability due to large effective dimensions was also not shown . Another recent work ( Pesme et al. , 2021 ) observed that the noise is proportional to the loss for specific simple models . In our paper , such a result is derived for more generic models . Gürbüzbalaban et al . ( 2021 ) showed that SGD will converge to a heavy-tailed stationary distribution due to a multiplicative nature of the SGD noise in a simple linear regression problem . Our paper strengthens this result : we argue that such a heavy-tailed distribution generically appears for the mean-square loss . 2 BACKGROUND . 2.1 SETUP . We consider supervised learning . Let D = { ( x ( µ ) , y ( µ ) ) : µ = 1 , 2 , . . . , N } be the training dataset , where x ( µ ) ∈ Rd denotes a data vector and y ( µ ) ∈ R be its label . The network output for a given input x is denoted by f ( θ , x ) ∈ R , where θ ∈ RP stands for a set of trainable parameters with P being the number of trainable parameters ( Extension to the multi-dimensional label and output is straightforward ) . In this work , we focus on the mean-square loss L ( θ ) = 1 2N N∑ µ=1 [ f ( θ , x ( µ ) ) − y ( µ ) ] 2 = : 1 N N∑ µ=1 ` µ ( θ ) . ( 1 ) The training proceeds through optimization of L ( θ ) . In most machine-learning applications , the optimization is done via SGD or its variants . In SGD , the parameter θk+1 at the time step k + 1 is determined by θk+1 = θk − η∇LBk ( θk ) , LBk ( θ ) = 1 2B ∑ µ∈Bk ` µ ( θ ) , ( 2 ) where η > 0 is the learning rate , Bk ⊂ { 1 , 2 , . . . , N } with |Bk| = B is a mini-batch used at the kth time step , and LBk denotes the mini-batch loss . Since the training dataset D is randomly divided into mini-batches , the dynamics defined by Eq . ( 2 ) is stochastic . When B = N , the full training data samples are used for every iteration . In this case , the dynamics is deterministic and called gradient descent ( GD ) . SGD is interpreted as GD with stochastic noise . By introducing the SGD noise ξk = − [ ∇LBk ( θk ) −∇L ( θk ) ] , Eq . ( 2 ) is rewritten as θk+1 = θk − η∇L ( θk ) + ηξk . ( 3 ) Obviously , 〈ξk〉 = 0 , where the brackets denote the average over possible choices of mini-batches . The noise covariance matrix is defined as Σ ( θk ) : = 〈ξkξTk 〉 . The covariance structure of the SGD noise is important in analyzing the SGD dynamics , which will be discussed in Sec . 3.1 . 2.2 STOCHASTIC DIFFERENTIAL EQUATION FOR SGD . When the parameter update for each iteration is small , which is typically the case when the learning rate η is small enough , we can consider the continuous-time approximation ( Li et al. , 2017 ; Smith & Le , 2018 ) . By introducing a continuous time variable t ∈ R and regarding η as an infinitesimal time step dt , we have a SDE dθt = −∇L ( θt ) dt+ √ ηΣ ( θt ) · dWt , ( 4 ) where dWt ∼ N ( 0 , IP dt ) with In being the n-by-n identity matrix , and the multiplicative noise√ ηΣ ( θt ) · dWt is interpreted as Itô since the noise ξk in Eq . ( 3 ) depends on θk but not on θk+1 . Throughout this work , we consider the continuous-time approximation ( 4 ) with Gaussian noise . In machine learning , the gradient Langevin dynamics ( GLD ) is also considered , in which the isotropic and uniform Gaussian noise is injected into the GD as dθt = −∇L ( θt ) dt+ √ 2DdWt , ( 5 ) where D > 0 corresponds to the noise strength ( it is also called the diffusion coefficient ) ( Sato & Nakagawa , 2014 ; Zhang et al. , 2017b ; Zhu et al. , 2019 ) . The stationary probability distribution PGLD ( θ ) of θ for GLD is given by the Gibbs distribution PGLD ( θ ) ∝ e−L ( θ ) /D . We will see in Sec . 4 that the SGD noise structure , which is characterized by Σ ( θ ) , drastically alters the stationary distribution and the escape rate from a local minimum .

Summary: This paper studies the behavior of SGD around the minimum. Unlike many other works that simply treat SGD noise as a fixed noise, the authors characterizes the location-dependence of SGD noise, which gives drastically different escaping behavior. By some simplification of the noise covariance matrix, the authors are able to formulate an SGD with isotropic noise with a time change. The new SDE is a dynamics on the log-loss landscape with a simple additive noise. With this result, the escape rate and stationary distributions are derived, depending polynomially on the loss, instead of exponentially. Numerical experiments are conducted to justify the assumptions made in the analysis, and verify the theoretical conclusions in the case of linear regression.