Entropic gradient descent algorithms and wide flat minima

1 Introduction . The geometrical structure of the loss landscape of neural networks has been a key topic of study for several decades ( Hochreiter & Schmidhuber , 1997 ; Keskar et al. , 2016 ) . One area of ongoing research is the connection between the flatness of minima found by optimization algorithms like stochastic gradient descent ( SGD ) and the generalization performance of the network ( Baldassi et al. , 2020 ; Keskar et al. , 2016 ) . There are open conceptual problems in this context : On the one hand , there is accumulating evidence that flatness is a good predictor of generalization ( Jiang et al. , 2019 ) . On the other hand , modern deep networks using ReLU activations are invariant in their outputs with respect to rescaling of weights in different layers ( Dinh et al. , 2017 ) , which makes the mathematical picture complicated1 . General results are lacking . Some initial progress has been made in connecting PAC-Bayes bounds for the generalization gap with flatness ( Dziugaite & Roy , 2018 ) . The purpose of this work is to shed light on the connection between flatness and generalization by using methods and algorithms from the statistical physics of disordered systems , and to corroborate the results with a performance study on state-of-the-art deep architectures . Methods from statistical physics have led to several results in the last years . Firstly , wide flat minima have been shown to be a structural property of shallow networks . They exist even when training on random data and are accessible by relatively simple algorithms , even though coexisting with exponentially more numerous minima ( Baldassi et al. , 2015 ; 2016a ; 1We note , in passing , that an appropriate framework for theoretical studies would be to consider networks with binary weights , for which most ambiguities are absent . 2020 ) . We believe this to be an overlooked property of neural networks , which makes them particularly suited for learning . In analytically tractable settings , it has been shown that flatness depends on the choice of the loss and activation functions , and that it correlates with generalization ( Baldassi et al. , 2020 ; 2019 ) . In the above-mentioned works , the notion of flatness used was the so-called local entropy ( Baldassi et al. , 2015 ; 2016a ) . It measures the low-loss volume in the weight space around a minimizer , as a function of the distance ( i.e . roughly speaking it measures the amount of “ good ” configurations around a given one ) . This framework is not only useful for analytical calcuations , but it has also been used to introduce a variety of efficient learning algorithms that focus their search on flat regions ( Baldassi et al. , 2016a ; Chaudhari et al. , 2019 ; 2017 ) . In this paper we call them entropic algorithms . A different notion of flatness , that we refer to as local energy in this paper , measures the average profile of the training loss function around a minimizer , as a function of the distance ( i.e . it measures the typical increase in the training error when moving away from the minimizer ) . This quantity is intuitively appealing and rather easy to estimate via sampling , even in large systems . In Jiang et al . ( 2019 ) , several candidates for predicting generalization performance were tested using an extensive numerical approach on an array of different networks and tasks , and the local energy was found to be among the best and most consistent predictors . The two notions , local entropy and local energy , are distinct : in a given region of a complex landscape , the local entropy measures the size of the lowest valleys , whereas the local energy measures the average height . Therefore , in principle , the two quantities could vary independently . It seems reasonable , however , to conjecture that they would be highly correlated under mild assumptions on the roughness of the landscape ( which is another way to say that they are both reasonable measures to express the intuitive notion of `` flatness '' ) . In this paper , we first show that for simple systems in controlled conditions , where all relevant quantities can be estimated well by using the Belief Propagation ( BP ) algorithm ( Mezard & Montanari ( 2009 ) ) , the two notions of flatness are strongly correlated : regions of high local entropy have low local energy , and vice versa . We also confirm that they are both correlated with generalization . This justifies the expectation that , even for more complex architectures and datasets , those algorithms which are driven towards high-local-entropy regions would minimize the local energy too , and thus ( based on the findings in Jiang et al . ( 2019 ) ) would find minimizers that generalize well . Indeed , we systematically applied two entropic algorithms , Entropy-SGD ( eSGD ) and Replicated-SGD ( rSGD ) , to state-of-the-art deep architectures , and found that we could achieve an improved generalization performance , at the same computational cost , compared to the original papers where those architectures were introduced . We believe these results to be an important addition to the current state of knowledge , since in ( Baldassi et al . ( 2016b ) ) rSGD was applied only to shallow networks with binary weights trained on random patterns and the current work represents the first study of rSGD in a realistic deep neural network setting . Together with the first reported consistent improvement of eSGD over SGD on image classification , these results point to a very promising direction for further research . While we hope to foster the application of entropic algorithms by publishing code that can be used to adapt them easily to new architectures , we also believe that the numeric results are important for theoretical research , since they are rooted in a well-defined geometric interpretation of the loss landscape . We also confirmed numerically that the minimizers found in this way have a lower local energy profile , as expected . Remarkably , these results go beyond even those where the eSGD and rSGD algorithms were originally introduced , thanks to a general improvement in the choice for the learning protocol , that we also discuss ; apart from that , we used little to no hyper-parameter tuning . 2 Related work . The idea of using the flatness of a minimum of the loss function , also called the fatness of the posterior and the local area estimate of quality , for evaluating different minimizers is several decades old ( Hochreiter & Schmidhuber , 1997 ; Hinton & van Camp , 1993 ; Buntine & Weigend , 1991 ) . These works connect the flatness of a minimum to information theoretical concepts like the minimum description length of its minimizer : flatter minima correspond to minimizers that can be encoded using fewer bits . For neural networks , a recent empirical study ( Keskar et al. , 2016 ) shows that large-batch methods find sharp minima while small-batch ones find flatter ones , with a positive effect on generalization performance . PAC-Bayes bounds can be used for deriving generalization bounds for neural networks ( Zhou et al. , 2018 ) . In Dziugaite & Roy ( 2017 ) , a method for optimizing the PAC-Bayes bound directly is introduced and the authors note similarities between the resulting objective function and an objective function that searches for flat minima . This connection is further analyzed in Dziugaite & Roy ( 2018 ) . In Jiang et al . ( 2019 ) , the authors present a large-scale empirical study of the correlation between different complexity measures of neural networks and their generalization performance . The authors conclude that PAC-Bayes bounds and flatness measures ( in particular , what we call local energy in this paper ) are the most predictive measures of generalization . The concept of local entropy has been introduced in the context of a statistical mechanics approach to machine learning for discrete neural networks in Baldassi et al . ( 2015 ) , and subsequently extended to models with continuous weights . We provide a detailed definition in the next section , but mention here that it measures a volume in the space of configurations , which poses computational difficulties . On relatively tractable shallow networks , the local entropy of any given configuration can be computed efficiently using Belief Propagation , and it can be also used directly as a training objective . In this setting , detailed analytical studies accompanied by numerical experiments have shown that the local entropy correlates with the generalization error and the eigenvalues of the Hessian ( Baldassi et al. , 2015 ; 2020 ) . Another interesting finding is that the cross-entropy loss ( Baldassi et al. , 2020 ) and ReLU transfer functions ( Baldassi et al. , 2019 ) , which have become the de-facto standard for neural networks , tend to bias the models towards high local entropy regions ( computed based on the error loss ) . Extending such techniques for general architectures is an open problem . However , the local entropy objective can be approximated to derive general algorithmic schemes . Replicated stochastic gradient descent ( rSGD ) replaces the local entropy objective by an objective involving several replicas of the model , each one moving in the potential induced by the loss while also attracting each other . The method has been introduced in Baldassi et al . ( 2016a ) , but only demonstrated on shallow networks . The rSGD algorithm is closely related to Elastic Averaging SGD ( EASGD ) , presented in Zhang et al . ( 2014 ) , even though the latter was motivated purely by the idea of enabling massively parallel training and had no theoretical basis . The substantial distinguishing feature of rSGD compared to EASGD when applied to deep networks is the focusing procedure , discussed in more detail below . Another difference is that in rSGD there is no explicit master replica . Entropy-SGD ( eSGD ) , introduced in Chaudhari et al . ( 2019 ) , is a method that directly optimizes the local entropy using stochastic gradient Langevin dynamics ( SGLD ) ( Welling & Teh , 2011 ) . While the goal of this method is the same as rSGD , the optimization techniques involves a double loop instead of replicas . Parle ( Chaudhari et al. , 2017 ) , combines eSGD and EASGD ( with added focusing ) to obtain a distributed algorithm that shows also excellent generalization performance , consistently with the results obtained in this work . 3 Flatness measures : local entropy , local energy . The general definition of the local entropy loss LLE for a system in a given configuration w ( a vector of size N ) can be given in terms of any common ( usually , data-dependent ) loss L as : LLE ( w ) = − 1 β log ∫ dw′ e−βL ( w ′ ) −βγd ( w′ , w ) . ( 1 ) The function d measures a distance and is commonly taken to be the squared norm of the difference of the configurations w and w′ : d ( w′ , w ) = 1 2 N∑ i=1 ( w′i − wi ) 2 ( 2 ) The integral is performed over all possible configurations w′ ; for discrete systems , it can be substituted by a sum . The two parameters β and γ̃ = βγ are Legendre conjugates of the loss and the distance . For large systems , N 1 , the integral is dominated by configurations having a certain loss value L∗ ( w , β , γ ) and a certain distance d∗ ( w , β , γ ) from the reference configuration w. These functional dependencies can be obtained by a saddle point approximation . In general , increasing β reduces L∗ and increasing γ̃ reduces d∗ . While it is convenient to use Eq . ( 1 ) as an objective function in algorithms and for the theoretical analysis of shallow networks , it is more natural to use a normalized definition with explicit parameters when we want to measure the flatness of a minimum . We thus also introduce the normalized local entropy ΦLE ( w , d ) , which , for a given configuration w ∈ RN , measures the logarithm of the volume fraction of configurations whose training error is smaller or equal than that of the reference w in a ball of squared-radius 2d centered in w : ΦLE ( w , d ) = 1 N log ∫ dw′ Θ ( Etrain ( w ) − Etrain ( w′ ) ) Θ ( d− d ( w′ , w ) ) ∫ dw′ Θ ( d− d ( w′ , w ) ) . ( 3 ) Here , Etrain ( w ) is the error on the training set for a given configuration w and Θ ( x ) is the Heaviside step function , Θ ( x ) = 1 if x ≥ 0 and 0 otherwise . This quantity is upper-bounded by zero and tends to zero for d → 0 ( since for almost any w , except for a set with null measure , there is always a sufficiently small neighborhood in which Etrain is constant ) . For sharp minima , it is expected to drop rapidly with d , whereas for flat regions it is expected to stay close to zero within some range . A different notion of flatness is that used in Jiang et al . ( 2019 ) , which we call local energy . Given a weight configuration w ∈ RN , we define δEtrain ( w , σ ) as the average training error difference with respect to Etrain ( w ) when perturbing w by a ( multiplicative ) noise proportional to a parameter σ : δEtrain ( w , σ ) = Ez Etrain ( w + σz w ) − Etrain ( w ) , ( 4 ) where denotes the Hadamard ( element-wise ) product and the expectation is over normally distributed z ∼ N ( 0 , IN ) . In Jiang et al . ( 2019 ) , a single , arbitrarily chosen value of σ was used , whereas we compute entire profiles within some range [ 0 , σmax ] in all our tests .

Summary: The paper studies (1) the relationship between the flatness of minima and their generalization properties, and (2) the connection between two measures of flatness, known as local entropy and local energy. Through a series of experiments, the authors show that the two measures are highly correlated and correlate well with generalization. They also empirically show that Entropy-SGD and Replicated-SGD, when used to explicitly optimize the local entropy, are able to flatter and better minima (in terms of lower generalization errors).