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: The paper proposed an experiment framework for evaluating loss landscape of neural networks during training. The experimental setup systematically vary two types of hyperparameters: 1) temperature type, where at high temperatures, larger learning rate or smaller batch size will induce more noise to training, and 2) load type, including data quantity/quality and model width. Neural networks are trained with different hyperparameter combinations according to the load-temperature pairs. Then 3 metrics including mode connectivity, Hessian, and output CKA similarity evaluated on perturbed training data. Jointly consider the 3 metrics (local and global) evaluated on all loss landscapes, the authors classify the loss landscapes into different phases. One of the phases (IV with high connectivity, low Hessian and high CKA) corresponds to best test accuracy.