Relative Flatness and Generalization
Year: 2021

Abstract:
Flatness of the loss curve is conjectured to be connected to the generalization ability of machine learning models, in particular neural networks. While it has been empirically observed that flatness measures consistently correlate strongly with generalization, it is still an open theoretical problem why and under which circumstances flatness is connected to generalization, in particular in light of reparameterizations that change certain flatness measures but leave generalization unchanged. We investigate the connection between flatness and generalization by relating it to the interpolation from representative data, deriving notions of representativeness, and feature robustness. The notions allow us to rigorously connect flatness and generalization and to identify conditions under which the connection holds. Moreover, they give rise to a novel, but natural relative flatness measure that correlates strongly with generalization, simplifies to ridge regression for ordinary least squares, and solves the reparameterization issue.

heading: 1 Introduction
text: 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.

heading: 2 Representativeness
text: 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.

heading: 3 Feature Robustness
text: Having decomposed the generalisation gap into a representativness and a second term of loss deviation, we now develop a novel notion of feature robustness that is able to bound the second term for specific families of distributions Λ using key equation (1). Our definition of feature robustness for a model f = (ψ ◦ φ) : X → Y depends on a small number δ > 0, a sample set S and a feature selection defined by a matrix A ∈ Rm×m of operator norm ||A|| ≤ 1. With feature perturbations φA(x) = (I +A)φ(x) and
EφF (f, S,A) := 1
|S| |S|∑ i=1 [ `(ψ(φA(xi)), y[φA(xi)])− `(f(xi), yi) ] , (4)
the definition of feature robustness is given as follows. Definition 1. Let ` : Y × Y → R+ denote a loss function, and δ two positive (small) real numbers, S ⊆ X × Y a finite sample set, and A ∈ Rm×m a matrix. A model f(x) = (ψ ◦ φ)(x) with φ(X ) ⊆ Rm is called ((δ, S,A), )-feature robust, if
∣∣∣EφF (f, S, αA)∣∣∣ ≤ for all 0 ≤ α ≤ δ. More generally, for a probability distribution A on perturbation matrices in Rm, we define
EφF (f, S,A) = EA∼A [ EφF (f, S,A) ] ,
and call the model ((δ, S,A), )-feature robust on average over A, if ∣∣∣EφF (f, S, αA)∣∣∣ ≤ for 0 ≤ α ≤ δ.
Given a feature extractor φ, feature robustness measures the performance of ψ when feature values are perturbed (with constant feature extractor φ). This local robustness at sample points differs from the robustness of Xu and Mannor [40] that requires a data-independent partitioning of the input space. The matrix A in feature robustness determines which feature values shall be perturbed. For
each sample, the perturbation is linear in the expression of the feature. Thereby, we only perturb features that are relevant for the output for a given sample and leave feature values unchanged that are not expressed. For φ mapping into an intermediate layer of a neural network, traditionally, the activation values of a neuron are considered as feature values, which corresponds to a choice of A as a projection matrix. However, it was shown by Szegedy et al. [35] that, for any other direction v ∈ Rm, ||v|| = 1, the values 〈φ(x), v〉 obtained from the projection φ(x) onto v, can be likewise semantically interpreted as a feature. This motivates the consideration of general feature matrices A.
Distributions on feature matrices induce distributions on the feature space Feature robustness is defined in terms of feature matrices (suitable for an application of (1) to connect perturbations of features with perturbations of weights), while the approach exploiting representative data from Section 2 considers distributions on feature vectors, cf. (3). To connect feature robustness to the notion of -representativeness, we specify for any distributionA on matrices A ∈ Rm×m an |S|-tuple ΛA = (λi) of probability density functions λi on the feature space Rm with support containing the origin. Multiplication of a feature matrix with a feature vector φ(xi) defines a feature selection Aφ(xi), and for each z ∈ Rm there is some feature matrix A with φ(xi) + z = φ(xi) + Aφ(xi) (unless φ(xi) = 0). Our choice for distributions λi on Rm are therefore distributions that are induced via multiplication of feature vectors φ(xi) ∈ Rm with matrices A ∈ Rm×m sampled from a distribution on feature matrices A . Formally, we assume that a Borel measure µA is defined by a probability distribution A on matrices Rm×m. We then define Borel measures µi on Rm by µi(C) = µA({A | Aφ(xi) ∈ C}) for Borel sets C ⊆ Rm. Then λi is the probability density function defined by the Borel measure µi. As a result, we have for each i that
EA∼A [ `(ψ(φA(xi)), y(φA(xi))) ] = Ez∼λi [ `(ψ(φ(xi) + z), y(φ(xi) + z)) ] Feature robustness and generalization. With this construction and a distribution ΛA on the feature space induced by a distribution A on feature matrices, we have that
E(f) = Eemp(f, S) + EφRep(f, S,ΛA) + E φ F (f, S,A) (5)
Here,A can be any distribution on feature matrices, which can be chosen suitably to control how well the corresponding mixture of local distributions approximates the true distribution. The third term then measures how robust the model is in expectation over feature changes for A ∼ A. In particular, if EφRep(f, S,ΛA) ≈ 0, then Egen(f, S) ≈ E φ F (f, S,A) and the generalization gap is determined by feature robustness. We end this section by illustrating how distributions on feature matrices induce natural distributions on the feature space. The example will serve in Sec. 5 to deduce a bound on EφRep(f, S,ΛA) from kernel density estimation.
Example: Truncated isotropic normal distributions are induced by a suitable distribution on feature matrices. We consider probability distributions Kδ||φ(xi)|| on feature vectors z ∈ Rm in the feature space defined by densities kδ||φ(xi)||(0, z) with smooth rotation-invariant kernels, bounded support and bandwidth h:
kh(zi, z) = 1 hm · k ( ||zi − z|| h ) · 1||zi−z||<h (6)
with 1||zi−z||<h = 1 when ||z − zi|| < h and 0 otherwise, and such that ∫ z∈Rm kh(z0, z) dz = 1 for all z0. An example for such a kernel is a truncated isotropic normal distribution with variance h2σ2I , kh(zi, z) = N (zi, h2σ2)(z). The following result states that the densities in (6) can indeed be induced by distributions on feature matrices, which will enable us to connect feature robustness with -representativeness. Proposition 2. Let Sφ = {φ(xi) |xi ∈ S} be a set of feature vectors in Rm. With kh defined as in (6), let λi(z) = kδ||φ(xi)||(0, z) define an |S|-tuple Λδ of densities. Then there exists a distribution Aδ on matrices in Rm×m of norm less than δ such that for each i = 1, . . . , |S|,
EA∼Aδ [ `(ψ(φA(xi)), y(φA(xi))) ] = Eξ∼λi [ `(ψ(φ(xi) + ξ), y(φ(xi) + ξ)) ] The technical proof is deferred to the appendix, but we describe the distribution Aδ on matrices for later use: The desired distribution is defined on the set of matrices of the form rO for a real number r
and an orthogonal matrix O (i.e. OOT = OTO = I) as a product measure combining the (unique) Haar measure on the set of orthogonal matrices O(m) with a suitable distribution on R. The Haar measure on O(m) induces the uniform measure on a sphere of radius r via multiplication with a vector of length r [17], and we choose a measure on R to match the radial change of the kernel kh.

heading: 4 Relative Flatness of the Loss Surface
text: Flatness is a property of the parameter space quantifying the change in loss under small parameter perturbations, classically measured by the trace of the loss Hessian Tr(H), where H is the matrix containing the partial second derivatives of the empirical risk with respect to all parameters of the model. In order to connect feature robustness (a property of the feature space) to flatness, we present how key equation (1) translates to the empirical risk: For a model f(x,w) = ψ(w, φ(x)) = g(wφ(x)) with parameters w ∈ Rd×m and g : Rd → Y a function on a matrix product of parameters w and a feature representation φ : X → Rm and any feature matrix A ∈ Rm×m we have that
Eemp(w+wA, φ(S)) = 1
|S| |S|∑ i=1 `(ψ(w + wA, φ(xi)), yi)
= 1
|S| |S|∑ i=1 `(ψ(w, φ(xi) +Aφ(xi)), yi) = 1 |S| |S|∑ i=1 `(ψ(w, φA(xi)), yi)
(7)
Subtracting Eemp(w, φ(S)) = 1|S| ∑|S| i=1 `(ψ(w, φ(xi)), yi), we can recognize feature robustness (4) on the right side of this equality when labels are constant under perturbations of the features, i.e. y(φA(xi)) = yi. In other words, flatness Eemp(w + v, φ(S)) − Eemp(w, φ(S)) describes the performance of a model function on perturbed feature vectors while holding labels constant. We proceed to introduce a novel, but natural, loss Hessian-based flatness measure that approximates feature robustness, given that the underlying data distribution D satisfies the assumption of locally constant labels.
With ws = (ws,t)t ∈ R1×m denoting the s-th row of the parameter matrix w, we let Hs,s′(w, φ(S)) ∈ Rm×m denote the Hessian matrix containing all partial second derivatives of the empirical risk Eemp(w, φ(S)) with respect to weights in rows ws and ws′ , i.e.
Hs,s′(w, φ(S)) = [ ∂2Eemp(w, φ(S)) ∂ws,t∂ws′,t′ ] 1≤t,t′≤m . (8)
Definition 3. For a model f(x,w) = g(wφ(x)), w ∈ Rd×m, with a twice differentiable function g, a twice differentiable loss function ` and a sample set S, relative flatness is defined by
κφTr(w) := d∑ s,s′=1 〈ws,ws′〉 · Tr(Hs,s′(w, φ(S))), (9)
where Tr denote the trace and 〈ws,ws′〉 = wswTs′ the scalar product of two row vectors.
Properties of relative flatness (i) Relative flatness simplifies to ridge regression for linear models f(x,w) = wx ∈ R (X = Rd, g = id and φ = id) and squared loss: To see this, note that for any loss function `, the second derivatives with respect to the parameters w ∈ Rd computes to ∂2` ∂wi∂wj = ∂ 2` ∂(f(x,w))2xixj . For `(ŷ, y) = ||ŷ − y||
2 the squared loss function, ∂2`/∂ŷ2 = 2 and the Hessian is independent of the parameters w. In this case, κidTr = c · ||w||2 with a constant c = ∑ x∈S 2Tr(xx T ), which is the well-known Tikhonov (ridge) regression penalty.
(ii) Invariance under reparameterization: We consider neural network functions
f(x) = wLσ(. . . σ(w2σ(w1x+ b1) + b2) . . .) + bL (10)
of a neural network of L layers with nonlinear activation function σ. By letting φl(x) denote the composition of the first l − 1 layers, we obtain a decomposition f(x,wl) = gl(wlφl(x)) of the network. Using (9) we obtain a relative flatness measure κlTr(w) for the chosen layer.
For a well-defined Hessian of the loss function, we require the network function to be twice differentiable. With the usual adjustments (equations only hold almost everywhere in parameter space), we can also consider neural networks with ReLU activation functions. In this case, Dinh et al. [5] noted that the network function —and with it the generalization performance— remains unchanged under linear reparameterization, i.e., multiplying layer l with α > 0 and dividing layer k 6= l by α, but common measures of the loss Hessian change. Our measure fixes this issue in relating flatness to generalization since the change of the loss Hessian is compensated by multiplication with the scalar products of weight matrices and is therefore invariant under layer-wise reparameterizations [cf. 26]. It is also invariant to neuron-wise reparameterizations, i.e., multiplying all incoming weights into a neuron by a positive number α and dividing all outgoing weights by α [23], except for neuron-wise reparameterizations of the feature layer φl. Using a simple preprocessing step (a neuron-wise reparameterization with the variance over the sample), our proposed measure becomes independent of all neuron-wise reparameterizations. Theorem 4. Let σi denote the variance of the i-th coordinate of φl(x) over samples x ∈ S and V = diag ( σ1, . . . , σnl−1 ) . If the relative flatness measure κlTr is applied to the representation
f(x) = wLσ(. . . σ(wlV σ(V −1wl−1σ(. . . σ(w1x+ b1)) . . .) + V −1bl−1) + bl) . . .) + bL
then κlTr is invariant under all neuron-wise (and layer-wise) reparameterizations
We now connect flatness with feature robustness: Relative flatness approximates feature robustness for a model at a local minimum of the empirical risk, when labels are approximately constant in neighborhoods of the training samples (φ(x), y) ∈ φ(S) in feature space. Theorem 5. Consider a model f(x,w) = g(wφ(x)) as above, a loss function ` and a sample set S, and let Om ⊂ Rm×m denote the set of orthogonal matrices. Let δ be a positive (small) real number and w = ω ∈ Rd×m denote parameters at a local minimum of the empirical risk on a sample set S. If the labels satisfy that y(φδA(xi)) = y(φ(xi)) = yi for all (xi, yi) ∈ S and all ||A|| ≤ 1, then f(x, ω) is ((δ, S,Om), )-feature robust on average over Om for = δ 2 2mκ φ Tr(ω) +O(δ3).
Applying the theorem to Eq. 5 implies that if the data is representative, i.e., EφRep(f, S,ΛAδ) ≈ 0 for the distribution Aδ of Prop. 2, then Egen(f(·, ω), S) . δ 2 2mκ φ Tr(ω) +O(δ3). The assumption on locally constant labels in Thm. 5 can be relaxed to approximately locally constant labels without unraveling the theoretical connection between flatness and feature robustness. Appendix B investigates consequences from even dropping the assumption of approximately locally constant labels.

heading: 5 Flatness and Generalization
text: Combining the results from sections 2–4, we connect flatness to the generalization gap when the distribution can be represented by smooth probability densities on a feature space with approximately locally constant labels. By approximately locally constant labels we mean that, for small δ, the loss in δ||φ(xi)||-neighborhoods around the feature vector of a training sample xi is approximated (on average over all training samples) by the loss for constant label y(xi) on these neighborhoods. This and the following theorem connecting flatness and generalization are made precise in Appendix D.4. Theorem 6 (informal). Consider a model f(x,w) = g(wφ(x)) as above, a loss function ` and a sample set S, let m denote the dimension of the feature space defined by φ and let δ be a positive (small) real number. Let ω denote a local minimizer of the empirical risk on a sample set S. If the distribution D has a smooth density pφD on the feature space Rm with approximately locally constant labels around the points x ∈ S, then it holds with probability 1−∆ over sample sets S that
Egen(f(·, ω), S) . |S|− 2 4+m
( κφTr(ω)
2m + C1(p
φ D, L) + C2(p φ D, L)√ ∆ ) up to higher orders in |S|−1 for constants C1, C2 that depend only on the distribution in feature space pφD induced by φ, the chosen |S|-tuple Λδ and the maximal loss L.
To prove Theorem 6 we bound both -representativeness and feature robustness in Eq. 5. For that, the main idea is that the family of distributions considered in Proposition 2 has three key properties: (i) it
provides an explicit link between the distributions on feature matrices Aδ used in feature robustness and the family of distributions Λδ of -representativeness (Proposition 2) (ii) it allows us to bound feature robustness using Thm. 5; and (iii) it is simple enough that it allows us to use standard results of kernel density estimation (KDE) to bound representativeness.
Our bound suffers from the curse of dimensionality, but for the chosen feature space instead of the (usually much larger) input space. The dependence on the dimension is a result of using KDE uniformly over all distributions satisfying mild regularity assumptions. In practice, for a given distribution and sample set S, representativeness can be much smaller, which we showcase in a toy example in Sec. 6. In the so-called interpolation regime, where datasets with arbitrarily randomized labels can be fit by the model class, the obtained convergence rate is consistent with the no free lunch theorem and the convergence rate derived by Belkin et al. [2] for an interpolation technique using nearest neighbors.
A combination of our approach with prior assumptions on the hypotheses or the algorithm in accordance to statistical learning theory could potentially achieve faster convergence rate. Our herein presented theory is instead based solely on interpolation and aims to understand the role of flatness (a local property) in generalization: If the data is representative in feature layers and if the distribution can be approximated by locally constant labels in these layers, then flatness of the empirical risk surface approximates the generalization gap. Conversely, Equation 7 shows that flatness measures the performance under perturbed features only when labels are kept constant. As a result, we offer an explanation for the often observed correlation between flatness and generalization: Real-world data distributions for classification are benign in the sense that small perturbations in feature layers do not change the target class, i.e., they can be approximated by locally constant labels. (Note that the definition of adversarial examples hinges on this assumption of locally constant labels.) In that case, feature robustness is approximated by flatness of the loss surface. If the given data and its feature representation are further -representative for small ≈ 0, then flatness becomes the main contributor to the generalization gap leading to their noisy, but steady, correlation.

heading: 6 Empirical Validation
text: We empirically validate the assumptions and consequences of the theoretical results derived above 2. For that, we first show on a synthetic example that the empirical correlation between flatness and generalization decreases if labels are not locally constant, up to a point when they are not correlated anymore. We then show that the novel relative flatness measure correlates strongly with generalization, also in the presence of reparameterizations. Finally, we show in a synthetic experiment that while representativeness cannot be computed without knowing the true data distribution, it can in practice be approximated. This approximation—although technically not a bound anymore—tightly bounds the generalization gap. Synthetic data distributions for binary classification are generated by sampling 4 Gaussian distributions in feature space (two for each class) with a given distance between their means (class separation). We then sample a dataset in feature space Sφ, train a linear classifier
2 Code is available at https://github.com/kampmichael/relativeFlatnessGeneralization.
ψ on the sample, randomly draw the weights of a 4-layer MLP φ, and generate the input data as S = (φ−1(Sφx ), S φ y ). This yields a dataset S and a model f = φ ◦ ψ such that φ(S) has a given class separation. Details on the experiments are provided in Appdx. C.
Locally constant labels: To validate the necessity of locally constant labels, we measure the correlation between the proposed relative flatness measure and the generalization gap for varying degrees of locally constant labels, as measured by the class separation on the synthetic datasets. For each chosen class separation, we sample 100 random datasets of size 500 on which we measure relative flatness and the generalization gap. Fig. 3 shows the average correlation for different degrees of locally constant labels, showing that the higher the degree, the more correlated flatness is with generalization. If labels are not locally constant, flatness does not correlate with generalization.
Approximating representativeness: While representativeness cannot be calculated without knowing the data distribution, it can be approximated from the training sample S by the error of a density estimation on that sample. For that, we use multiple random splits of S into a training set Strain and a test set Stest, train a kernel density estimation on Strain and measure its error on Stest. Again, details can be found in Appx. C. The lower the class separation of the synthetic datasets, the harder the learning problem and the less representative a random sample will be. For each sample and its distribution, we compute the generalization gap and the approximation to the generalization bound. The
results in Fig. 4 show that the approximated generalization bound tightly bounds the generalization error (note that this approximation is technically not a bound anymore). Moreover, as expected, the bound decreases the easier the learning problems become.
Relative flatness correlates with generalization: We validate the correlation of relative flatness to the generalization gap in practice by measuring it for 110 different local minima—achieved via different learning setups, such as initialization, learning rate, batch size, and optimization algorithm— of LeNet5 [19] on CIFAR10 [18]. We compare this correlation to the classical Hessian-based flatness measures using the trace of the loss-Hessian, the Fisher-Rao norm [21], the PACBayes flatness measure that performed best in the extensive study of Jiang et al. [14] and the L2-norm of the weights. The results in Fig. 5 show that indeed relative flatness has higher correlation than all the competing measures. Of these measures, only the Fisher-Rao norm is reparameterization invariant but shows the weakest correlation in the experiment. In Appdx C we show how reparameterizations of the network significantly reduce the correlation for non-reparameterization invariant measures.

heading: 7 Discussion and Conclusion
text: Contributing to the trustworthiness of machine learning, this paper provides a rigorous connection between flatness and generalization. As to be expected for a local property, our association between flatness and generalization requires the samples and its representation in feature layers to be representative for the target distribution. But our derivation uncovers a second, usually overlooked condition. Flatness of the loss surface measures the performance of a model close to training points when labels are kept locally constant. If a data distribution violates this, then flatness cannot be a good indicator for generalization.
Whenever we consider feature representations other than the input features, the derivation of our results makes one strong assumption: the existence of a target output function y(φ(x)) on the feature space φ(X ). By moving assumptions on the distribution from the input space to the feature space, we achieve a bound based on interpolation that depends on the dimension of the feature layer instead of the input space. Hence, we assume that the feature representation is reasonable and does not lose information that is necessary for predicting the output. To achieve faster convergence rates independent of any involved dimensions, future work could aim to combine our approach of interpolation with a prior-based approach of statistical learning theory.
Our measure of relative flatness may still be improved in future work. Better estimates for the generalization gap are possible by improving the representativeness of local distributions in two ways: The support shape of the local distributions can be improved and their volume-parameter δ can be optimally chosen. Both improvements will affect the derivation of the measure of relative flatness as an estimation of feature robustness for the corresponding distributions on feature matrices. Whereas different support shapes change the trace to a weighted average of the Hessian eigenvalues, the volume parameter can provide a correcting scaling factor. Both approaches seem promising to us, as our relative measure from Definition 3 already outperforms the competing measures of flatness in our empirical validation.

heading: Acknowledgements
text: Cristian Sminchisescu was supported by the European Research Council Consolidator grant SEED, CNCS-UEFISCDI (PN-III-P4-ID-PCE-2016-0535, PN-III-P4-ID-PCCF-2016-0180), the EU Horizon 2020 grant DE-ENIGMA (688835), and SSF.
Mario Boley was supported by the Australian Research Council (under DP210100045).
We would like to thank Julia Rosenzweig, Dorina Weichert, Jilles Vreeken, Thomas Gärtner, Asja Fischer, Tatjana Turova and Alexandru Aleman for the great discussions.

Review summaries:
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.

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.

This paper describes a connection between flatness of minima and generalization in deep neural networks. The authors define a concept called feature robustness and show that it is related to flatness. This is derived through a straightforward observation that perturbations in feature space can be recast as perturbations of the model in parameter space. This allows the authors to define a flatness measure for minima in deep networks (this flatness measure is also invariant to rescalings of the layers in neural networks with positively homogenous activations). The authors combine their notion of feature robustness with epsilon representativeness of a function to connect flatness to generalization. They present empirical evaluations and comparison to other measures on CIFAR10.