Relative Flatness and Generalization

1 Introduction . Flatness of the loss curve has been identified as a potential predictor for the generalization abilities of machine learning models [ 6 , 10 , 11 ] . In particular for neural networks , it has been repeatedly observed that generalization performance correlates with measures of flatness , i.e. , measures that quantify the change in loss under perturbations of the model parameters [ 4 , 8 , 16 , 21 , 34 , 39 , 41 , 44 ] . In fact , Jiang et al . [ 14 ] perform a large-scale empirical study and find that flatness-based measures have a higher correlation with generalization than alternatives like weight norms , margin- , and optimization-based measures . It is an open problem why and under which circumstances this correlation holds , in particular in the light of negative results on reparametrizations of ReLU neural networks [ 5 ] : these reparameterizations change traditional measures of flatness , yet leave the model function and its generalization unchanged , making these measures unreliable . We present a novel and rigorous approach to understanding the connection between flatness and generalization by relating it to the interpolation from representative samples . Using this theory we , for the first time , identify conditions under which flatness explains generalization . At the same time , we derive a measure of relative flatness that simplifies to ridge/Tikhonov regularization for ordinary least squares [ 36 ] , and resolves ∗equal contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the reparametrization issue for ReLU networks [ 5 ] by appropriately taking the norm of parameters into account as suggested by Neyshabur et al . [ 25 ] . Formally , we connect flatness of the loss surface to the generalization gap Egen ( f , S ) = E ( f ) − Eemp ( f , S ) of a model f : X → Y from a model classH with respect to a twice differentiable loss function ` : Y × Y → R+ and a finite sample set S ⊆ X × Y , where E ( f ) = E ( x , y ) ∼D [ ` ( f ( x ) , y ) ] and Eemp ( f , S ) = 1 |S| ∑ ( x , y ) ∈S ` ( f ( x ) , y ) . That is , Egen ( f , S ) is the difference between the risk E ( f ) and the empirical risk Eemp ( f , S ) of f on a finite sample set S drawn iid . according to a data distribution D on X × Y . To connect flatness to generalization , we start by decomposing the generalisation gap into two terms , a representativeness term that quantifies how well a distribution D can be approximated using distributions with local support around sample points and a feature robustness term describing how small changes of feature values affect the model ’ s loss . Here , feature value refers to the implicitly represented features by the model , i.e. , we consider models that can be expressed as f ( x ) = ψ ( w , φ ( x ) ) = g ( wφ ( x ) ) with a feature extractor φ and a model ψ ( which includes linear and kernel models , as well as most neural networks , see Fig . 1 ) . With this decomposition , we measure the generalization ability of a particular model by how well its interpolation between samples in feature space fits the underlying data distribution . We then connect feature robustness ( a property of the feature space ) to flatness ( a property of the parameter space ) using the following key identity : Multiplicative perturbations in feature space by arbitrary matrices A ∈ Rm×m correspond to perturbations in parameter space , i.e. , ψ ( w , φ ( x ) +Aφ ( x ) ) = g ( w ( φ ( x ) +Aφ ( x ) ) ) = g ( ( w + wA ) φ ( x ) ) = ψ ( w + wA , φ ( x ) ) . ( 1 ) Using this key equation , we show that feature robustness is approximated by a novel , but natural , loss Hessian-based relative flatness measure under the assumption that the distribution can be approximated by locally constant labels . Under this assumption and if the data is representative , then flatness is the main predictor of generalization ( see Fig . 2 for an illustration ) . This offers an explanation for the correlation of flatness with generalization on many real-world data distributions for image classification [ 14 , 21 , 26 ] , where the assumption of locally constant labels is reasonable ( the definition of adversarial examples [ 35 ] even hinges on this assumption ) . This dependence on locally constant labels has not been uncovered by previous theoretical analysis [ 26 , 37 ] . Moreover , we show that the resulting relative flatness measure is invariant to linear reparameterization and has a stronger correlation with generalization than other flatness measures [ 14 , 21 , 26 ] . Other measures have been proposed that similarly achieve invariance under reparameterizations [ 21 , 37 ] , but the Fisher-Rao norm [ 21 ] is lacking a strong theoretical connection to generalization , our measure sustains a more natural form than normalized sharpness [ 37 ] and for neural networks , it considers only a single layer , given by the decomposition of f ( including the possibility of choosing the input layer when φ = idX ) . An extended comparison to related work is provided in Appdx . A . The limitations of our analysis are as follows . We assume a noise-free setting where for each x ∈ X there is a unique y = y ( x ) ∈ Y such that Px , y∼D ( y|x ) = 1 , and this assumption is also extended to the feature space of the given model , i.e. , we assume that φ ( x ) = φ ( x′ ) implies y ( x ) = y ( x′ ) for all x , x′ ∈ X and write y ( x ) = y ( φ ( x ) ) . Moreover , we assume that the marginal distribution DX is described by a density function pD ( x ) , that f ( x ) = ψ ( w , φ ( x ) ) = g ( wφ ( x ) ) is a local minimizer of the empirical risk on S , and that g , ψ , φ are twice differential . Quantifying the representativeness of a dataset precisely is challenging since the data distribution is unknown . Using results from density estimation , we derive a worst-case bound on representativeness for all data distributions that fulfill mild regularity assumptions in feature space φ ( X ) , i.e. , a smooth density function pφ ( D ) such that ∫ z∈φ ( X ) ∣∣∇2 ( pφ ( D ) ( z ) ||z||2 ) ∣∣ dz and ∫z∈φ ( X ) pφ ( D ) ( z ) /||z||m dz are well-defined and finite . This yields a generalization bound incorporating flatness . In contrast to the common bounds of statistical learning theory , the bound depends on the feature dimension . The dimension-dependence is a result of the interpolation approach ( applying density estimation uniformly over all distributions that satisfy the mild regularity assumptions ) . The bound is consistent with the no-free-lunch theorem and the convergence rate derived by Belkin et al . [ 2 ] for a model based on interpolations . In practical settings , representativeness can be expected to be much smaller than the worst-case bound , which we demonstrate by a synthetic example in Sec . 6 . Generally , it is a bound that remains meaningful in the interpolation regime [ 1 , 3 , 24 ] , where traditional measures of generalization based on the empirical risk and model class complexity are uninformative [ 22 , 42 ] . Contribution . In summary , this paper rigorously connects flatness of the loss surface to generalization and shows that this connection requires feature representations such that labels are ( approximately ) locally constant , which is also validated in a synthetic experiment ( Sec . 6 ) . The empirical evaluation shows that this flatness and an approximation to representativeness can tightly bound the generalization gap . Our contributions are : ( i ) the rigorous connection of flatness and generalization ; ( ii ) novel notions of representativeness and feature robustness that capture the extent to which a model ’ s interpolation between samples fits the data distribution ; and ( iii ) a novel flatness measure that is layerand neuron-wise reparameterization invariant , reduces to ridge regression for ordinary least squares , and outperforms state-of-the-art flatness measures on CIFAR10 . 2 Representativeness . In this section , we formalize when a sample set S is representative for a data distribution D. Partitioning the input space . We choose a partition { Vi |i = 1 , . . . , |S| } of X such that each element of this partition Vi contains exactly one of the samples xi from S. The distribution can then be described by a set of densities pi ( x ) = 1αi · pD ( x ) · 1Vi ( x ) with support contained in Vi ( where 1Vi ( x ) = 1 if x ∈ Vi and 0 otherwise ) and with normalizing factor αi = ∫ Vi pD ( x ) dx . Then the risk decomposes as E ( f ) = ∑|S| i=1 αi · Ex∼pi [ ` ( f ( x ) , y ( x ) ) ] . Since xi ∈ Vi for each i , we can change variables and consider density functions λ∗i ( ξ ) = pi ( xi + ξ ) with support in a neighborhood around the origin of X . The risk then decomposes as E ( f ) = |S|∑ i=1 αi · Eξ∼λ∗i [ ` ( f ( xi + ξ ) , y ( xi + ξ ) ) ] . ( 2 ) Starting from this identity , we formalize an approximation to the risk : In a practical setting , the distribution pD is unknown and hence , in the decomposition ( 2 ) , we have unknown densities λ∗i and unknown normalization factors αi . We assume that each neighborhood contributes equally to the loss , i.e. , we approximate each αi with 1|S| . Then , given a sample set S and an |S|-tuple Λ = ( λi ) 1≤i≤|S| of “ local ” probability density functions on X with support supp ( λi ) in a neighborhood around the origin 0X , we call the pair ( S , Λ ) -representative for D with respect to a model f and loss ` if |ERep ( f , S , Λ ) | ≤ , where ERep ( f , S , Λ ) = E ( f ) − |S|∑ i=1 1 |S| · Eξ∼λi [ ` ( f ( xi + ξ ) , y ( xi + ξ ) ) ] . ( 3 ) If the partitions Vi and the distributions λi are all chosen optimal so that the approximation αi = 1|S| is exact and λi = λ∗i , then ERep ( f , S , Λ ) = 0 by ( 2 ) . If the support of each λi is decreased to the origin so that λi = δ0 is a Dirac delta function , then ERep ( f , S , Λ ) = Egen ( f , S ) equals the generalization gap . For density functions with an intermediate support , the generalization gap can be decomposed into representativeness and the expected deviation of the loss around the sample points : Egen ( f , S ) = ERep ( f , S , Λ ) + |S|∑ i=1 1 |S| · Eξ∼λi [ ` ( f ( xi + ξ ) , y ( xi + ξ ) ) − ` ( f ( xi ) , yi ) ] The main idea of our approach to understand generalization is to use this equality and to control both representativeness and expected loss deviations for a suitable |S|-tuple of distributions Λ . From input to feature space . An interesting aspect of -representativeness is that it can be considered in a feature space instead of the input space . For a model f = ( ψ ◦ φ ) : X → Y , we can apply our notion to the feature space φ ( X ) ( see Fig . 1 for an illustration ) . This leads to the notion of - representativeness in feature space defined for an |S|-tuple Λφ = ( λφi ) 1≤i≤|S| of densities on φ ( X ) by replacing xi with φ ( xi ) in ( 3 ) , which we denote by EφRep ( f , S , Λφ ) . By measuring representativeness in a feature space , this becomes a notion of both data and feature representation . In particular , it assumes that a target output function y ( φ ( x ) ) also exists for the feature space . We can then decompose the generalization gap Egen ( f ) of f = ( ψ ◦ φ ) into EφRep ( f , S , Λ φ ) + 1 |S| |S|∑ i=1 Eξ∼λφi [ ` ( ψ ( φ ( xi ) + ξ ) , y ( φ ( xi ) + ξ ) ) − ` ( f ( xi ) , yi ) ] The second term is determined by how the loss changes under small perturbations in the feature space for the samples in S. As before , for λi = δ0 the term in the bracket vanishes and EφRep ( f , S , Λφ ) = Egen . But the decomposition becomes more interesting for distributions with support of nonzero measure around the origin . If the true distribution can be interpolated efficiently in feature space from the samples in S with suitable λφi so that E φ Rep ( f , S , Λ φ ) ≈ 0 , then the term in the bracket approximately equals the generalization gap and the generalization gap can be estimated from local properties in feature space around sample points .

Summary: This paper develops a theoretical approach to tackling the relationship between flatness of the loss landscape around trained models and generalization. To accomplish this, the authors need to make certain assumptions about the data such as locally constant labels as well as its representativeness. The authors show that flatness with respect to the last layer parameters is related to feature robustness, and ultimately show that generalization performance can be bounded by their flatness quantity. The paper then presents a few experiments on synthetic datasets or small models to test the performance of the theory. They find that their measure performs similarly to others such as PacBayes and Trace.