On Flat Minima, Large Margins and Generalizability

1 INTRODUCTION . Understanding under which conditions a neural network will generalize from seen to unseen data is crucial , as it motivates design choices and principles which can greatly improve performance . Complexity or generalization measures are used to quantify the properties of a neural network which lead to good generalization . Currently however , established complexity measures such as VC-Dimension ( Vapnik , 1998 ) or Rademacher Complexity ( Bartlett & Mendelson , 2002 ) do not correlate with the generalizability of neural networks ( e.g . see Zhang et al . ( 2016 ) ) . Hence many recommendations , such as reducing model complexity , early stopping , or adding explicit regularization are also not applicable or necessary anymore . Therefore , there is an ongoing effort to devise new complexity measures that may guide recommendations on how to obtain models that generalize well . A popular approach is to consider the flatness of the loss surface around a neural network . Hochreiter & Schmidhuber ( 1997 ) used the minimum description length ( MDL ) argument of Hinton & Van Camp ( 1993 ) to claim that the flatness of a minimum can also be used as a generalization measure . Motivated by this new measure Hochreiter & Schmidhuber ( 1997 ) , and more recently Chaudhari et al . ( 2019 ) , developed algorithms with explicit regularization intended to converge to flat solutions . Keskar et al . ( 2016 ) then presented empirical evidence that flatness relates to improved generalizability and used it to explain the behavior of stochastic gradient descent ( SGD ) with large and small-batch sizes . Other works since have empirically corroborated that flatter minima generalize better ( e.g . Jiang et al . ( 2019 ) ; Li et al . ( 2018 ) ; Bosman et al . ( 2020 ) ) . There are however various issues that are still unresolved , which makes using flatness for constructing practical deep learning recommendations difficult . For one , flatness is computationally expensive to compute . The most common way to compute the flatness is via the Hessian , which grows quadratically in the number of parameters ; this becomes too large when used with modern networks containing millions of parameters . It is also not clear to what extent flatness is a true measure of generalizability , capable of discerning which neural network will or will not generalize . Dinh et al . ( 2017 ) showed that reparametrizations affect flatness and a flat model can be made arbitrarily sharp without changing any of its generalization properties . In addition Probably Approximately Correct ( PAC-Bayes ) bounds that bound the generalizability in terms of the flatness are also either affected by rescaling , impossible to evaluate or loose ( Neyshabur et al. , 2017 ; Arora et al. , 2018 ; Petzka et al. , 2020 ) . While there have been solutions attempting to prevent issues around reparametrization ( Liang et al. , 2019 ; Tsuzuku et al. , 2019 ) , it remains to establish whether flatness is an epiphenomenon of stochastic gradient descent or other complexity measures as Achille et al . ( 2018 ) and Jastrzebski et al . ( 2018 ) are suggesting . This motivates investigating possible correlations to more well-understood measures of generalization that may help alleviate issues surrounding flat minima , while allowing flat minima to be used when appropriate . In this paper we will demonstrate a correlation to classification margins , which are a well-understood generalization measure . Margins represent the linearized distance to the decision boundaries of the classification region ( Elsayed et al. , 2018 ) . An immediate consequence of such a relationship is that to assess generalizability , we could now simply use a computationally cheap and more robust margin based complexity measure . Our contributions will demonstrate further practical implications of the relationship between margins and flatness which open doors to valuable future work such as a better understanding of why and when a model generalizes and more principled algorithm design . • We prove that under certain conditions flatness and margins are strongly correlated . We do so by deriving the Hessian trace for the affine classifier . Based on its form , we derive an expression in terms of classification margins which we show correlates well with the Hessian trace , with increasing training accuracy for various neural network architectures . By being able relate the two complexity measures , we are now able to provide various practical recommendations , and offer different perspectives on phenomena that may not be explainable without such a view . These are shown in the following contributions . • We use our insight to replace the misleading folklore that , unlike large-batch methods , small-batch methods are able to escape sharp minima ( Keskar et al. , 2016 ) . We instead employ a margin perspective and use our empirical results along with recent results by Banburski et al . ( 2019 ) and Hoffer et al . ( 2017 ) to argue that a large batch method was unable to train long enough to maximize the margins . With our explanation , we help reframe the small and large-batch discussion and build further intuition . • We show that once a neural network is able to correctly predict the label of every element in the training set it can be made arbitrarily flat by scaling the last layer . We are motivated by the relationship to margins which suffer from the same issue . We highlight this scaling issue because , in some instances , it may still be beneficial for algorithm design to be guided by convergence to flat regions . Hence , we need to account for scaling issues which make it difficult to use flatness to assess whether a network generalizes better than another . Other works have made connections between flatness and well-behaved classification margins via visualizations ( see Huang et al . ( 2019 ) ; Wang et al . ( 2018 ) ) , but they have not demonstrated a quantifiable relationship . Further work has used both the classification margins and flatness to construct PAC-Bayes bounds ( Neyshabur et al. , 2017 ; Arora et al. , 2018 ) , and have related flatness to increased robustness ( Petzka et al. , 2020 ; Borovykh et al. , 2019 ) however they did not show when and to what extent these quantities are related . We structure the paper as follows . In Section 2 , we discuss both our notation and our motivation choosing the cross-entropy loss and the Hessian trace as the flatness measure and provide further background on the classification margins . In Section 3 , we present our contribution showing a strong correlation between the margins and flatness by deriving . In Section 4 , we combine recent results based on classification margins to offer a different perspective on the misleading folklore on why larger-batch methods generalize worse . In Section 5 , we highlight that networks can be made arbitrarily flat . Lastly , we offer our thoughts and future work in the Section 6 . 2 PROBLEM SETTING . We first define the basic notation that we use for a classification task . We let X represent the input space and Y = { 1 , ... , C } the output space where C are the number of possible classes . The network architecture is given by φ : Θ × X → R|Y| where Θ is the corresponding parameter space . We measure the performance of a parameter vector by defining some loss function ` : RC × Y → R. If we have have a joint probability distribution D relating input and output space then we would like to minimize the expected loss LD ( θ ) = E ( x , y ) ∼D [ ` ( φ ( θ , x ) , y ) ] . Since we usually only have access to some finite dataset D , we denote the empirical loss by L̃D ( θ ) = 1|D| ∑|D| i=1 ` ( φ ( θ , xi ) , yi ) . If LD and L̃D are close , then we would say a model generalizes well , as we were able to train on a finite dataset and extrapolate to the true distribution . We will use the cross-entropy loss which is given by ` ( φ ( θ , x ) , y ) = − log ( Sy ( φ ( θ , x ) ) ) where the softmax function S : RC → RC is given by S ( a ) i = eai∑C j=1 e aj ( see Goodfellow et al . ( 2016 ) ) . The choice of the cross-entropy function as the loss function has a significant impact on how the flatness measure behaves . Unlike the multiclass mean squared error ( MMSE ) , exponential type losses such as the cross-entropy loss on neural networks have been shown to include implicit regularization which leads to margin maximizing solutions for neural networks ( Banburski et al. , 2019 ) . Also , various properties for flat minima which have been proven for the MMSE loss by Mulayoff & Michaeli are not applicable to the cross-entropy loss , further highlighting the fundamental differences between the loss functions . While the MMSE loss has shown some promise for many classification tasks ( Hui & Belkin , 2020 ) the cross-entropy loss is still the loss which is most used and was primarily used for the empirical evidence around flat minima ( Keskar et al. , 2016 ; Chaudhari et al. , 2019 ) , which motivates our choice . The qualitative description of a flat region was given by Hochreiter & Schmidhuber ( 1997 ) as “ a large connected region in parameter space where the error remains approximately constant '' . We measure the flatness by the trace of the Hessian of the loss with respect to the parameters ( in short the Hessian trace ) denoted by Tr ( Hθ ( L̃D ( θ ) ) ( Dinh et al. , 2017 ) . Since the Hessian is symmetric , the Hessian trace is equivalent to the sum of its eigenvalues which for a fixed parameter space is proportional to the expected increase of the second order approximation of the loss around a fixed minimum θ in a random direction θ′ with θ′ ∼ N ( θ , I ) . Since we apply flatness arguments only close to minima , we assume that all eigenvalues are positive and that the Hessian trace is a good measure of flatness Sagun et al . ( 2017 ) . Even though the Hessian is only an approximation of flatness , the Hessian is often preferred as it allows us to reason about various directions in parameter space via its eigenvectors and eigenvalues ( see Sagun et al . ( 2017 ) ; Chaudhari et al . ( 2019 ) ) and alleviates the issue of infinitely long but sharp ridges making a minimum infinitely flat ( Dinh et al. , 2017 ; Freeman & Bruna , 2016 ) . The Hessian has also been linked to feature robustness via its use in the second order approximation of the loss ( e.g . Petzka et al . ( 2020 ) ; Borovykh et al . ( 2019 ) ) and is a promising quantity to relate to the margins . As we are working with non-linear functions it is intractable to compute exact distances to the decision boundary , therefore we use a measure which is related to the linearized distance as described in Elsayed et al . ( 2018 ) . Under this view , larger margins are better because the data is further from the decision boundary . Specifically , we define the margins as in Neyshabur et al . ( 2017 ) : for some vector v ∈ RC and label y we let the margin of v be γ ( v , y ) = |vy − maxj 6=y vj | . Since we use the margin in different contexts we define the output margins γ ( φ ( θ , x ) , y ) and the margins of the model output after the softmax layer γ ( S ( φ ( θ , x ) ) , y ) . Due to the intuition of margins relating to the regularity of the classification regions , they have been proven and shown to be a good generalization measure for linear networks ( Langford & Shawe-Taylor , 2003 ) and later for neural networks ( see Bartlett et al . ( 2017 ) ; Jiang et al . ( 2018 ; 2019 ) ) when correctly adjusted . Due to results by Banburski et al . ( 2019 ) and Soudry et al . ( 2018 ) , Poggio et al . ( 2019 ) claimed that a large part of the mystery around generalizability has been solved , since standard optimization methods are maximizing the margin instead of memorizing data .

Summary: The paper presents empirical evidence that the output margin - as a measure of the confidence of a multiclass predictor - is strongly correlated to the Hessian trace when using cross-entropy loss with softmax. Moreover, the paper presents a method for estimating the Hessian trace using the input norm and softmax output. This estimation is inspired by linear classifiers and shows a strong correlation with the Hessian trace.