Taxonomizing local versus global structure in neural network loss landscapes

1 Introduction . Among the many approaches to understanding the behavior of neural network ( NN ) models , the study of their loss landscapes [ 1 , 2 ] has proven to be particularly fruitful . Indeed , analyzing loss landscapes has helped shed light on the workings of many popular techniques , including large-batch training [ 3 , 4 ] , adversarial training [ 5 ] , residual connections [ 6 ] , and BatchNorm [ 7 ] . One particular concept of recent interest is the so-called sharpness of local minima [ 3 , 5 , 8–10 ] . While sharpness can be measured by first-order sensitivity measures , such as the Jacobian or Lipschitz constant , it is more appropriately measured by second-order sensitivity measures , typically via the Hessian spectrum [ 11 ] . It has been observed that in some cases NNs generalize well when they converge to a relatively flat , i.e. , non-sharp , local minimum [ 3 ] . While such local sharpness measures can provide insight , their focus on the local geometry of the loss landscape neglects the global structure of the loss landscape ( namely , precisely the sort of structure that statistical mechanics approaches to learning aim to quantify [ 12 , 13 ] ) . Indeed , it is well-known that existing sharpness-based metrics can be altered ( trivially ) by reparameterization tricks or ( more interestingly ) by taking algorithmic steps which have the effect of changing the local structures on the loss landscape [ 5 , 8 , 14 ] . For example , [ 5 ] shows that adversarial training can decrease the magnitude 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . of Hessian eigenvalues and bias the model towards a locally smooth area , even though adversarial training can reduce clean test accuracy [ 15 ] . Similarly , [ 14 ] shows that Hessian eigenvalues become smaller with reduced ℓ2 regularization , even though increased ℓ2 regularization is known to reduce overfitting and improve training , if used properly . More general considerations would suggest ( and indeed our own empirical results , e.g. , as reported in Figure 7 , demonstrate ) that by training to data with noisy labels , one can find models that generalize poorly and yet simultaneously lie in very “ flat ” regions of the loss landscape , with small Hessian eigenvalues , and vice versa . These observations ( and other observations we describe below ) indicate that the previously-observed empirical correlation between very local metrics like sharpness and more global properties like generalization performance may be correlative and not causative , i.e. , they may be due to the confounding factor that results in the published literature are on reasonably-good models trained to reasonably-good data , rather than due to some fundamental properties of deep NNs . They also raise the question of how to capture more global properties of the loss landscape . Motivated by these considerations , we are interested in understanding local properties/structure versus global properties/structure of the loss landscape of realistic NN models . While similar ideas underlie work that adopts a statistical mechanics perspective [ 2 , 12 , 13 , 16 ] , here we are interested in adopting an operational machine learning ( ML ) perspective , where we employ metrics that have been used within ML as “ experimental probes ” to gain insight into local versus global properties . To do so , we employ the following metrics . • First , we consider Hessian-based metrics , including the largest eigenvalue and the trace of the Hessian . These metrics try to capture local curvature properties of the loss landscape [ 11 ] . • Second , we use mode connectivity [ 17 , 18 ] —in particular , the connectivity between trained models . This metric tries to capture how well-connected different local minima are to each other . • Third , we use CKA similarity [ 19 ] to capture a correlation-like similarity between the outputs of different trained models . Averaging the CKA over several pairs of models can be thought of as approximating so-called overlap integrals frequently appearing in statistical mechanics [ 12 , 13 , 20 ] . We have considered many other metrics , but these three seem to be particularly useful for identifying global structure versus local structure in loss landscapes . Informally , mode connectivity , as its name suggests , captures connectivity , where well-connected models exhibit a single “ rugged basin ” with low-energy / low-loss , potentially non-linear , paths through the loss landscape , ( i.e. , continuous chains of models ) all achieving a small loss value . We expect this property to be important since the connectivity of local minima indicates efficiency of the training dynamics to explore the loss landscape , without becoming stuck at saddle points or in a “ bad ” local minimum . Similarly , CKA captures similarity , where an ensemble of good models will produce roughly similar outputs . These two types of metrics are different and complementary ; and both of them are very different than Hessian-based metrics , which clearly capture much more local information . Here we briefly summarize our main contributions . • We design an experimental setup based on two control parameters , a temperature-like parameter that correlates with the magnitude of SGD noise during training , e.g. , batch size ( in most figures ) , learning rate , or weight decay , and a load-like parameter that measures the relationship between model size and data quantity and/or quality , e.g. , the amount of data , size of intermediate layers , amount of exogenously-introduced label noise , etc . By training thousands of models , under a variety of settings , and by measuring local and global metrics of the loss landscape , we identify four distinct phases in temperature-load space , with relatively sharp transitions between them . • Using global connectivity ( measured by mode connectivity ) and local flatness ( measured by the Hessian ) , we taxonomize loss landscapes into four categories , which are pictorially represented in Figure 1 , labelled Phase I through Phase IV . For reasons observed in our empirical results in Section 3 , it is often convenient to further divide Phase IV into two subcategories , depending on whether the trained models produce similar representations ( as measured by CKA similarity ) . If the loss landscape satisfies the first property , we say it is globally well-connected ; and if the loss landscape also satisfies the second property , we say it is globally nice . Depending on whether the Hessian eigenvalues are large or small , we say the loss landscape is locally sharp or locally flat . • Based on these results , as well as measured model quality , e.g. , test accuracy , we empirically demonstrate that the global ( but not necessarily local ) structure of a loss landscape is well-correlated with good generalization performance , and that the best generalization occurs in the phase associated with a locally flat , globally nice loss landscape . We demonstrate these results on a range of computer vision and natural language processing benchmarks ( CIFAR-10 , CIFAR-100 , SVHN , and IWSLT 2016 De-En ) and various models ( ResNet , VGG , and Transformers ) . We also vary the amount of data , the number of noisy labels , etc. , to study both the effect of the quantity of data and the quality of data on changing the loss landscape . • We observe the well-known double descent phenomenon [ 21 , 22 ] in our experiments , which exhibits itself as a “ bad fluctuation ” between the different phases ( e.g. , see the transition that separates Phase I and II from Phase III and IV in Figure 4a ) . Our empirical observations on double descent corroborates recent theoretical analysis [ 23 , 24 ] , which views the phenomenon as a consequence of a transition between qualitatively different phases of learning [ 13 ] . Computing connectivity and similarity requires comparing multiple distinct models . This could be computationally expensive , especially if model training is expensive . For many reasonably-sized models , however , the metrics we consider are sufficiently tractable so as to be useful , e.g. , during model training . Moreover , the phase transitions and the metrics that we use to determine the phases lead to practical tools that can diagnose typical failure modes in training , which we will discuss towards the end of the paper . In this short conference version , we focus on the main message , and we provide a much more thorough discussion on prior work in the full version online [ 25 ] , in which we also talk about related papers in the study of loss landscapes and statistical mechanics of learning . In order that our results can be reproduced and extended , we have open-sourced our code.1 2 Setup . In the sequel , we consider training a NN fθ : Rdin → Rdout , with trainable parameters θ , to a dataset consisting of n datapoint/label pairs Strain = { ( x1 , y1 ) , . . . , ( xn , yn ) } . Our nominal training objective is to minimize a loss function of the form L ( θ ) = 1 n ∑ ( x , y ) ∈Strain ℓ ( fθ ( x ) , y ) + λ∥θ∥22 . ( 1 ) Here ℓ is a loss function , typically chosen to be the cross entropy loss . The parameter λ is the weight decay parameter , which controls the level of ℓ2 regularization . We consider optimizing NN models using standard minibatch SGD , with iterates of the form θ ← θ − ηg̃ ( θ ) , g̃ ( θ ) = 1 B B∑ j=1 ∇θℓ ( fθ ( xij ) , yij ) , ( 2 ) where η is the learning rate , 0 < B ≤ n is the batch size , and the indices i1 , . . . , iB of each minibatch are sampled without replacement from { 1 , . . . , n } . For classification tasks , we consider also the training/testing accuracy , which is simply the fraction of correctly classified points , acctrain ( θ ) = 1 n ∑ ( x , y ) ∈Strain 1 ( fθ ( x ) = y ) , and similarly for acctest ( θ ) on a given test set Stest . 1https : //github.com/nsfzyzz/loss_landscape_taxonomy We now briefly introduce the main metrics and control parameters which we will consider . Temperature and load . In the sequel , a load-like parameter of a loss landscape refers to some quantity related to the amount and/or quality of data , relative to the size of the model . Specifically , we vary either i ) model size ( e.g. , width , which captures the size of an internal representation of the data ) , for fixed training set size n , ii ) training set size n , for fixed model size , or iii ) the “ quality ” of training data , which is varied by randomizing a fraction α of the training labels . Each of these control parameters directly induces a different loss landscape by changing the data Strain and/or architecture fθ for which the loss L ( θ ) is being computed . For example , we expect that increasing width will result in a smoother loss landscape [ 26 ] ; we shall see this effect with CKA similarity in the transition from Phase IV-A to IV-B . The second control parameter we vary in our experiments is a temperature-like parameter , representing the amount of noise introduced in the SGD iterates ( 2 ) . Most commonly , we take this to be the batch size B , although we will also use the learning rate η and the weight decay parameter λ . Increasing temperature corresponds to smaller batch size , and large learning rate or weight decay . Varying the temperature does not directly define a different loss function L ( θ ) , but rather it indirectly induces a different effective loss function . This is because , at different temperatures , the iterates of SGD concentrate on different regions of the loss landscape . Due to the noise in the stochastic optimization , the training dynamics may not be able to “ see ” certain features of the loss landscape . CKA similarity . To measure the similarity of two NN representations , we use the centered kernel alignment ( CKA ) metric , proposed in [ 19 ] . For a NN fθ , let Fθ = [ fθ ( x1 ) · · · fθ ( xm ) ] ⊤ ∈ Rm×dout denote the concatenation of the outputs2 of the network over a set of m randomly sampled datapoints . Then the ( linear ) CKA similarity between two parameter configurations θ , θ′ is given by s ( θ , θ′ ) = Cov ( Fθ , Fθ′ ) √ Cov ( Fθ , Fθ ) Cov ( Fθ′ , Fθ′ ) , ( 3 ) where for X , Y ∈ Rm×d , we define Cov ( X , Y ) = ( m − 1 ) −2tr ( XX⊤HmY Y ⊤Hm ) , and Hm = Im − m−111⊤ is the centering matrix . The CKA similarity is known to be an effective way to compare the overall representations learned by two different trained NNs [ 19 ] . Rather than computing the similarity directly on the original training points , we measure CKA on a perturbed training set comprised of Mixup samples [ 27 ] ; this can reduce trivial similarity that occurs when the models are trained to exactly or approximately zero training error . See also Appendix A.4.1 in the full paper for the ablation study on different perturbed training sets . Mode connectivity . For two parameter configurations θ , θ′ , computing mode connectivity involves finding a low-energy curve γ ( t ) , t ∈ [ 0 , 1 ] , for which γ ( 0 ) = θ , γ ( 1 ) = θ′ , such that ∫ L ( γ ( t ) ) dt is minimized [ 17 , 18 ] . A number of techniques have been proposed to find such curves γ . In this work , we use the technique proposed in [ 17 ] , which parameterizes the Bezier curve with k + 1 bends , given by γϕ ( t ) = ∑k j=0 ( k j ) ( 1− t ) k−jtjθj for t ∈ [ 0 , 1 ] , where θ0 = θ , θk = θ′ , and ϕ = { θ1 , . . . , θk−1 } are trainable parameters of additional models , defining “ bends ” on the curve γϕ ( t ) . We use Bezier curves with three bends ( k = 2 ) . Given the curve γϕ ( t ) , we define the mode connectivity of the models θ , θ′ to be mc ( θ , θ′ ) = 1 2 ( L ( θ ) + L ( θ′ ) ) − L ( γϕ ( t∗ ) ) , ( 4 ) where t∗ maximizes the deviation t 7→ | 12 ( L ( θ ) + L ( θ ′ ) ) − L ( γϕ ( t ) ) | . There are three possibilities for mode connectivity . If mc ( θ , θ′ ) < 0 , then 12 ( L ( θ ) +L ( θ ′ ) ) < L ( γϕ ( t∗ ) ) , which means there is a “ barrier ” of high loss between θ , θ′ ; in this case , we will say that the loss landscape is poorly-connected or simply say that mode connectivity is poor . If mc ( θ , θ′ ) > 0 , then this implies a curve of low loss connecting θ , θ′ , but it also implies that the training failed to locate a reasonable optimum , i.e. , L ( θ ) and L ( θ′ ) are large . If mc ( θ , θ′ ) ≈ 0 , then we will say that the loss landscape is well-connected or simply say that the mode connectivity is good . Note that for all the experiments except neural machine translation , we use the training error ( 0-1 loss ) when computing mode connectivity , so that mode connectivity is always normalized to the range of [ −100 , 100 ] . We provide additional details 2For this work , we focus on the similarity of representations at the output layer , i.e. , after the softmax is applied , although the CKA similarity can be used to compare the representations at any layer . on this procedure , as well as an ablation study on different mode connectivity hyperparameters , in Appendix A.4.2 of the full paper [ 25 ] . Hessian . The Hessian at a given point θ0 in parameter space is represented by the matrix∇2θL ( θ0 ) . To summarize the Hessian in a single scalar value , we report the dominant eigenvalue λmax ( ∇2θL ( θ0 ) ) and/or the trace tr ( ∇2θL ( θ0 ) ) , calculated using the PyHessian software [ 11 ] . ℓ2 distance . We will also occasionally report the ℓ2 distance between two parameter configurations ∥θ − θ′∥2 as a measure of similarity between models , although we typically find that the CKA similarity is a more informative measure .

Summary: This paper takes a large number of neural networks that span range of different parameters (e.g. width), train them in a large number of ways (e.g. dataset, batch size), and studies several metrics that characterize the local (e.g. Hessian spectral norm) and global (weight vector pair-wise connectivity on potential non-linear paths). By doing a large scale study of this kind, the authors are able to identify several distinct phases of loss landscapes, dividing them phenomenologically into (globally {well,poorly}-connected) x (locally {flat, sharp}). The authors also develop a simple mathematical toy model that can model this taxonomy.