WHY FLATNESS DOES AND DOES NOT CORRELATE WITH GENERALIZATION FOR DEEP NEURAL NETWORKS 

1 INTRODUCTION . Among the most important theoretical questions in the field of deep learning are : 1 ) What characterizes functions that exhibit good generalization ? , and 2 ) Why do overparameterized deep neural networks ( DNNs ) converge to this small subset of functions that do not overfit ? Perhaps the most popular hypothesis is that good generalization performance is linked to flat minima . In pioneering works ( Hinton & van Camp , 1993 ; Hochreiter & Schmidhuber , 1997 ) , the minimum description length ( MDL ) principle ( Rissanen , 1978 ) was invoked to argue that since flatter minima require less information to describe , they should generalize better than sharp minima . Most measures of flatness approximate the local curvature of the loss surface , typically defining flatter minima to be those with smaller values of the Hessian eigenvalues ( Keskar et al. , 2016 ; Wu et al. , 2017 ; Zhang et al. , 2018 ; Sagun et al. , 2016 ; Yao et al. , 2018 ) . Another commonly held belief is that stochastic gradient descent ( SGD ) is itself biased towards flatter minima , and that this inductive bias helps explain why DNNs generalize so well ( Keskar et al. , 2016 ; Jastrzebski et al. , 2018 ; Wu et al. , 2017 ; Zhang et al. , 2018 ; Yao et al. , 2018 ; Wei & Schwab , 2019 ; Maddox et al. , 2020 ) . For example Keskar et al . ( 2016 ) developed a measure of flatness that they found correlated with improved generalization performance when decreasing batch size , suggesting that SGD is itself biased towards flatter minima . We note that others ( Goyal et al. , 2017 ; Hoffer et al. , 2017 ; Smith et al. , 2017 ; Mingard et al. , 2021a ) have argued that the effect of batch size can be compensated by changes in learning rate , complicating some conclusions from Keskar et al . ( 2016 ) . Nevertheless , the argument that SGD is somehow itself biased towards flat minima remains widespread in the literature . In an important critique of local flatness measures , Dinh et al . ( 2017 ) pointed out that DNNs with ReLU activation can be re-parameterized through a simple parameter-rescaling transformation . Tα : ( w1 , w2 ) 7→ ( αw1 , α −1w2 ) ( 1 ) where w1 are the weights between an input layer and a single hidden layer , and w2 are the weights between this hidden layer and the outputs . This transformation can be extended to any architecture having at least one single rectified network layer . The function that the DNN represents , and thus how it generalizes , is invariant under parameter-rescaling transformations , but the derivatives w.r.t . parameters , and therefore many flatness measures used in the literature , can be changed arbitrarily . Ergo , the correlation between flatness and generalization can be arbitrarily changed . Several recent studies have attempted to find `` scale invariant '' flatness metrics ( Petzka et al. , 2019 ; Rangamani et al. , 2019 ; Tsuzuku et al. , 2019 ) . The main idea is to multiply layer-wise Hessian eigenvalues by a factor of ‖wi‖2 , which renders the metric immune to layer-wise re-parameterization . While these new metrics look promising experimentally , they are only scale-invariant when the scaling is layer-wise . Other methods of rescaling ( e.g . neuron-wise rescaling ) can still change the metrics , so this general problem remains unsolved . 1.1 MAIN CONTRIBUTIONS . 1 . For a series of classic image classification tasks ( MNIST and CIFAR-10 ) we show that flatness measures change substantially as a function of epochs . Parameter re-scaling can arbitrarily change flatness , but it quickly recovers to a more typical value under further training . We also demonstrate that some variants of SGD exhibit significantly worse correlation of flatness with generalization than found for vanilla SGD . In other words popular measures of flatness sometimes do and sometimes do not correlate with generalization . This mixed performance problematizes a widely held intuition that DNNs generalize well fundamentally because SGD or its variants are themselves biased towards flat minima . 2 . We next study the correlation of the Bayesian prior P ( f ) with the generalization performance of a DNN that converges to that function f . This prior is the weighted probability of obtaining function f upon random sampling of parameters . Motivated by a theoretical argument derived from a non-uniform convergence generalization bound , we show empirically that logP ( f ) correlates robustly with test error , even when local flatness measures miserably fail , for example upon parameter re-scaling . For discrete input/output problems ( such as classification ) , P ( f ) can also be interpreted as the weighted `` volume '' of parameters that map to f . Intuitively , we expect local flatness measures to typically be smaller ( flatter ) for systems with larger volumes . Nevertheless , there may also be regions of parameter space where local derivatives and flatness measures vary substantially , even if on average they correlate with the volume . Thus flatness measures can be viewed as ( imperfect ) local measures of a more robust predictor of generalization , the volume/prior P ( f ) . 2 DEFINITIONS AND NOTATION . 2.1 SUPERVISED LEARNING . For a typical supervised learning problem , the inputs live in an input domain X , and the outputs belong to an output space Y . For a data distribution D on the set of input-output pairs X × Y , the training set S is a sample of m input-output pairs sampled i.i.d . from D , S = { ( xi , yi ) } mi=1 ∼ Dm , where xi ∈ X and yi ∈ Y . The output of a DNN on an input xi is denoted as ŷi . Typically a DNN is parameterized by a vector w and trained by minimizing a loss function L : Y × Y → R , which measures differences between the output ŷ ∈ Y and the ground truth output y ∈ Y , by assigning a score L ( ŷ , y ) which is typically zero when they match , and positive when they don ’ t match . Minimizing the loss typically involves using an optimization algorithm such as SGD on a training set S. The generalization performance of the DNN , which is theoretically defined over the underlying ( typically unknown ) data distribution D but is practically measured on a test set E = { ( x′i , y′i ) } |E| i=1 ∼ D|E| . For classification problems , the generalization error is practically measured as ( E ) = 1|E| ∑ x′i∈E 1 [ ŷi 6= y′i ] , where 1 is the standard indicator function which is one when its input is true , and zero otherwise . 2.2 FLATNESS MEASURES . Perhaps the most natural way to measure the flatness of minima is to consider the eigenvalue distribution of the Hessian Hij = ∂2L ( w ) /∂wi∂wj once the learning process has converged ( typically to a zero training error solution ) . Here for simplicity we use L ( w ) instead of L ( ŷ , y ) as ŷ is parameterized by w. Sharp minima are characterized by a significant number of large positive eigenvalues λi in the Hessian , while flat minima are dominated by small eigenvalues . Some care must be used in this interpretation because it is widely thought that DNNs converge to stationary points that are not true minima , leading to negative eigenvalues and complicating their use in measures of flatness . Typically , only a subset of the positive eigenvalues are used ( Wu et al. , 2017 ; Zhang et al. , 2018 ) . Hessians are typically very expensive to calculate . For this reason , Keskar et al . ( 2016 ) introduced a computationally more tractable measure called `` sharpness '' : Definition 2.1 ( Sharpness ) . Given parameters w′ within a box in parameter space Cζ with sides of length ζ > 0 , centered around a minimum of interest at parameters w , the sharpness of the loss L ( w ) at w is defined as : sharpness : = maxw′∈Cζ ( L ( w ′ ) − L ( w ) ) 1 + L ( w ) × 100 . In the limit of small ζ , the sharpness relates to the spectral norm of the Hessian ( Dinh et al. , 2017 ) : sharpness ≈ ∥∥∣∣ ( ∇2L ( w ) ) ∣∣∥∥ 2 ζ2 2 ( 1 + L ( w ) ) × 100 . The general concept of flatness can be defined as 1/sharpness , and that is how we will interpret this measure in the rest of this paper . 2.3 FUNCTIONS AND THE BAYESIAN PRIOR . We first clarify how we represent functions in the rest of paper using the notion of restriction of functions . A more detailed explanation can be found in appendix C. Here we use binary classification as an example : Definition 2.2 ( Restriction of functions to C ) . ( Shalev-Shwartz & Ben-David , 2014 ) Consider a parameterized supervised model , and let the input space be X and the output space be Y , noting Y = { 0 , 1 } in binary classification setting . The space of functions the model can express is a ( potentially uncountably infinite ) set F ⊆ Y |X | . Let C = { c1 , . . . , cm } ⊂ X . The restriction of F to C is the set of functions from C to Y that can be derived from functions in F : FC = { ( f ( c1 ) , . . . , f ( cm ) ) : f ∈ F } where we represent each function from C to Y as a vector in Y |C| . For example , for binary classification , if we restrict the functions to S + E , then each function in FS+E is represented as a binary string of length |S|+ |E| . In the rest of paper , we simply refer to `` functions '' when we actually mean the restriction of functions to S + E , except for the Boolean system in section 5.1 where no restriction is needed . See appendix C for a thorough explanation . For discrete functions , we next define the prior probability P ( f ) as Definition 2.3 ( Prior of a function ) . Given a prior parameter distribution Pw ( w ) over the parameters , the prior of function f can be defined as : P ( f ) : = ∫ 1 [ M ( w ) = f ] Pw ( w ) dw . ( 2 ) where 1 is an indicator function:1 [ arg ] = 1 if its argument is true or 0 otherwise ; M is the parameterfunction map whose formal definition is in appendix B . Note that P ( f ) could also be interpreted as a weighted volume V ( f ) over parameter space . If Pw ( w ) is the distribution at initialization , the P ( f ) is the prior probability of obtaining the function at initialization . We normally use this parameter distribution when interpreting P ( f ) . Remark . Definition 2.3 works in the situation where the spaceX and Y are discrete , where P ( f ) has a prior probability mass interpretation . This is enough for most image classification tasks . Nevertheless , we can easily extend this definition to the continuous setting , where we can also define a prior density over functions upon random initialization , with the help of Gaussian Process ( GP ) ( Rasmussen , 2003 ) . For the GP prior see appendix D. However , in this work , we focus exclusively on the classification setting , with discrete inputs and outputs .

Summary: In binary image classification setting, the paper shows that when using an attack set (augmenting the training set with additional intentionally mislabeled datapoints) to vary the generalization performance of a fixed neural network (NN) architecture, the (approximation of) Bayesian prior $P(f)$ of the outputs of the trained NN on the test set correlates _very_ well with generalization. The effect persists when using different optimizers, and different number of training epochs.  In contrast, the popular measure of flatness (approximation of the spectral norm of the Hessian of the training loss) correlates worse with generalization in these settings when using SGD, and does not correlate with it at all when using variants like Adam or Entropy-SGD.  The authors provide some preliminary explanation of the effect and propose the Bayesian prior as a more reliable generalization proxy.