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: The paper investigates the connection between loss-based flatness and generalization. The authors relate flatness to interpolation from representative data in the feature space, and use two notions to do that: 1) representativness and 2) feature robustness. They derive a new flatness metric -- a relative flatness -- and on a small set of empirical observations show that it correlates well with generalization and that it has some other desirable theoretical properties such as reparametrization invariance.