Sharpness Minimization Algorithms Do Not Only Minimize Sharpness To Achieve Better Generalization
Authors: Kaiyue Wen
Zhiyuan Li
Tengyu Ma
Venue: NeurIPS 2023 oral (2023)
Keywords: Sharpness
Flatness
Generalization
Generalization Bound
SAM
TL;DR: For 2 layer ReLU networks, sharpness may not always imply generalization but sharpness minimization algorithms may still generalize even when non-generalizing flattest models exist.

Abstract: Despite extensive studies, the underlying reason as to why overparameterized
neural networks can generalize remains elusive. Existing theory shows that common stochastic optimizers prefer flatter minimizers of the training loss, and thus
a natural potential explanation is that flatness implies generalization. This work
critically examines this explanation. Through theoretical and empirical investigation, we identify the following three scenarios for two-layer ReLU networks: (1)
flatness provably implies generalization; (2) there exist non-generalizing flattest
models and sharpness minimization algorithms fail to generalize poorly, and (3)
perhaps most strikingly, there exist non-generalizing flattest models, but sharpness
minimization algorithms still generalize. Our results suggest that the relationship
between sharpness and generalization subtly depends on the data distributions
and the model architectures and sharpness minimization algorithms do not only
minimize sharpness to achieve better generalization. This calls for the search for
other explanations for the generalization of over-parameterized neural networks

Sharpness Minimization Algorithms Do Not Only Minimize Sharpness To Achieve Better Generalization
Kaiyue Wen Tsinghua University
Zhiyuan Li Stanford University
wenky200mails, edu. edu. en
Tengyu Ma Stanford University tengyuma@stanfordd
Abstract
Despite extensive studies, the underlying reason as to why overparameterized neural networks can generalize remains elusive. Existing theory shows that common stochastic optimizers prefer flatter minimizers of the training loss, and thus natural potential explanation is that flatness implies generalization. This work critically examines this explanation. Through theoretical and empirical investigation, we identify the following three scenarios for two-layer ReLU networks: (1) flatness provably implies generalization; (2) there exist non-generalizing flattest models and sharpness minimization algorithms fail to generalize poorly, and (3) perhaps most strikingly, there exist non-generalizing flattest models, but sharpness minimization algorithms still generalize. Our results suggest that the relationship between sharpness and generalization subtly depends on the data distributions and the model architectures and sharpness minimization algorithms do not only minimize sharpness to achieve better generalization. This calls for the search for other explanations for the generalization of over parameterized neural networks.
Introduction
remains mysterious why stochastic optimization methods such as stochastic gradient descent (SGD) can find generalizable models even when the architectures are overparameterized (Zhang et al., 2016; Gunasekar et al., 2017; Li et al., 2017: Soudry et al., 2018; Woodworth al., 2020). Many empirical and theoretical studies suggest that generalization is correlated with or guaranteed by the flatness of the loss landscape at the learned model (Hochreiter Schmidhuber, 1997; Keskar al., 2016; Dziugaite & Roy, 2017; Jastrzebski et al., 2017; Neyshabur et al., 2017; Wu et al., 2018; 81 Jiang et al., 2019; Blanc et al., 2019; Wei Ma, 2019a,b; HaoChen et al., 2020; Foret et al., 2021; Damian et al., 2021; Li et al., 2021; Ma & Ying, 2021; Ding al., 2022: Nacson et al., 2022; Wei et al., 2022; Lyu et al., 2022: Norton Royset, 2021; Wu & Su, 2023). Thus, a natural theoretical question
Question 0. Does the flatness of the minimizers always correlate with the generalization capability?
The answer to the question turns out to false. First, Dinh al. (2017) theoretically construct very sharp networks with good generalization. Second, recent empirical results (Andriushchenko et al., 2023b) find that sharpness may not have a strong correlation with test accuracy for collection of modern architectures and settings, partly due to the same there exist sharp models with good generalization. We note that, technically speaking, Question is ill-defined without specifying the collection of models which the correlation is evaluated. However, those sharp but generalizable models appear to be the main cause for the non-correlation.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

All Flattest Minimizers
Sharpness Minimization
Architecture
Generalize Well.
Algorithms Generalize.
2-layer w/o Bias
/ (Theorem 3.1)
2- 2-layer w/ Bias
X (Theorem 4.1)
2-layer w/ simplified BatchNorm
(Theorem 3.2)
2 2-1ayer w/ simplified LayerNorm
X (Theorem 5.1)
Table 1: Overview of Our Results. Each row in the table corresponds to one architecture. 'Tee second column indicates whether all flattest minimizers of training loss generalize well. / indicates that all (near) flattest minimizers of training loss provably generalize well and X indicates that there provably exists flattest minimizers that generalize poorly. The third column indicates whether the sharpness minimization algorithms generalize well in our experiments. Results in row 2 and 4 deny Question 1 and Question 2 respectively.
Observing the existing theoretical and empirical evidence, is natural to ask the one-side version of Question O, where we are only interested in whether sharpness implies generalization but not vice versa.
Question 1. Do all the flattest neural network minimizers generalize well?
Though there are some theoretical works that answer Question affirmatively for simplified linear models (Li et al., 2021; Ding et al., 2022; Nacson et al., 2022: Gatmiry et al., 2023), the answer to Question for standard neural networks remains unclear. Those theoretical results linking generalization sharpness for more general architectures typically also involve other terms in generalization bounds, such as parameter dimension or norm (Neyshabur et al., 2017; Foret et al., 2021; Wei & Ma, 2019a,b; Norton & Royset, 2021), thus do not answer Question 1 directly.
Our first contribution a theoretical analysis showing that the answer to Question can be false, even for simple architectures like 2-layer ReLU networks. Intriguingly, we also find that the answer to Question subtly depends on the architectures of neural networks. For example, simply removing the bias in the first layer turns the aforementioned negative result into a positive result, as also shown in the Theorem 4.3 of Wu & Su (2023) (that the authors only came to be aware of after putting this work online).
More concretely, we show that for the parity xor problem with mean square loss and with data sampled from hypercube {,1, 1] all flattest 2-layer ReLU neural networks without bias provably generalize. However, when bias added, for the same data distribution and loss function, there exists flattest minimizer that fails to generalize for every unseen data. Since adding bias in the first layer can be interpreted as appending a constant input feature, this result suggests that the generalization of the flattest minimizer is sensitive to both network architectures and data distributions.
Recent theoretical studies (Wu al., 2018; Blanc et al., 2019; Damian al., 2021; et al., 2021; Arora et al., 2022; Wen et al., 2022; Nacson et al., 2022; Lyu et al., 2022; Bartlett et al., 2022; Li al., 2022) also show that optimizers including SGD with large learning rates or label noise and Sharpness-Aware Minimization (SAM, Foret et al. (2021)) may implicitly regularize the sharpness of the training loss landscape. These optimizers are referred to sharpness minimization algorithmss in this paper. Because Question 1 is not always true, it is then natural to hypothesize that sharpness- minimization algorithms will fail for architectures and data distributions where Question is not true.
Question 2.
Will sharpness mininization algorithm fail to generalize when there exist non- generalizing flattest minimizers?
priori, the authors were expecting that the answer to Question 2 is affirmative, which means that a possible explanation is that the sharpness minimization algorithm works if and only if for certain architecture and data distribution, Question is true. However, surprisingly, we also answer this question negatively for some architectures and data distributions. In other words, we found that sharpness-minimizatee
algorithms can still generalize well even when the answer to Question is false. The result is consistent with our theoretical discovery that for many architectures, there exist both non- gonn-eneraaizing and generalizing flattest minimizers of the training loss. We show empirically that Sharpness-minimiizee architectures.
algorithms can find different types of minimizers for different
Our results summarized in Table 1. We show through theoretical and empirical analysis that the relationship between sharpness and generalization can fall into three different regimes depending on the architectures and distributions. The three regimes include:

---

Scenario 1. Flattest minimizers of training loss provably generalize and sharpness minimization
algorithms find generalizable models. This regime (Theorems 3.1 and 3.2) includes 2-layer ReLU MLP without bias and 2-layer ReL.U MLP with simplified BatchNorm (without mean subtraction and bias). We answer both the Question and Question affirmatively this scenario."
Scenario 2. There exists a flattest minimizer that has the worst generalization over all minimizers.
Also, sharpness minimization algorithms fail to find generalizable models. This regime includes 2 layer ReLU MLP with bias. We deny Question while affirm Question 2 in this scenario.
Scenario 3. There exist flattest minimizers that do not generalize but the sharpness minimization algorithm still finds the generalizable flattest model empirically. This regime includes 2-layer ReLU MLP with a simplified LayerNorm (without mean subtraction and bias). In this scenario, the sharpness minimization algorithm relies other unknown mechanisms beyond minimizing sharpness to find a generalizable model. We deny both Question and Question in this scenario.
Setup
Rademacher Complexity. Given data S = {2x) the empirical Rademacher complexity of function class F is defined as (F) 1 (e /e E (x;( Architectures. As summarized in Table I, we will consider multiple network architectures and discuss how architecture influences the relationship between sharpness and generalization. For each model fo parameterized by 8, we will use d to denote the input dimension and m denote the network width, We will now describe the architectures in detail.
2-MLP-No-Bier B fnobias (1) = W2relu (W1z) with 0 = (W1, W2).
2-MLP-Bias. fine MLP-Bias as fo
fo chins (x) W2relu (Wiæ + 61) with 0 = (W1, b1, W2). We additionally de- fhias, (x) Wprelu Warelu (W1z b1).
2-MLP-Sim-BN. (ab { be(w))
W2SBN, (relu (Wi + b1), frelu (Waa + b1))) where the simplified BatchNorm SBN is defined as Vm, n N, Vi x, R" Im),SBN, (z, {x. (xa)ke(m)) > (
1/2 and 0 (W1, Wi,bb,1 W2).
re)b(W 2-MLP-Sim-LNN f2) W2 max (Wi++h where is sufficiently small posi-
tive constant.
Surprisingly, our results show that the relationships between sharpness and generalization are strikingly different among these simple yet similar architectures.
Data Distribution. We will consider a simple data distribution as our testbed. Data distribution Pxor is a joint distribution over data point x and label y. The data point is sampled uniformly from the hypercube 1) and the label satisfies = x[1]x[2]. Many of our results, including our generalization bound in Section 3 and experimental observations can be generalized to broader family of distributions (Appendix B).
Loss. We will mean squared error muse for training and denote the training loss L. In Ap- pendix B, we will show that all our theoretical results and empirical observations hold for logistic loss with label smoothing probability p > O. We will also consider zero one loss Pr(yfe(x) > 0) for evaluating the model. We will use interpolating model to denote the model with parameter 0 that minimizes L.
Definition 2.1 (Interpolating Model), A model fo interpolates the dataset {( 9:) and only if Vi, (oxz) = ys-
Sharpness. Our theoretical analysis focuses on understanding the sharpness of the trained models. Precisely, for model fo parameterized by f, dataset { ) and loss function e, we will use the trace of Hessian of loss function, Tr L())) to measure how sharp the loss is at 6, which is a proxy for the sharpness along a random direction (Wen et al., 2022), or equivalently, the expected increment of loss under a random gaussian perturbation (Foret et al., 2021: Orvieto et al., 2022).
Tr(V²I(()) is not the only choice for defining sharpness, but theoretically many sharpness mini- mization algorithms have been shown to minimize this term over interpolating models. In particular,
The condition for Question is not satistied and thus the answer to Question 2 is affirmative.
3

---

175
Sharpness
1.0
250
Sharpness
1.0
Valid acc
Valid acc
Sharpness
150
Train acc
Sharpness
200
acc
0.9
0.8
0.8
125
Acc
150
0.7
100
100
0.6
0.6
50
75
0.5
10a
10a
10a
10a
10a
10a
Epoch
Epoch
(a) Baseline
(b) 1-SAM
Figure 1: Scenario I. We train a 2-1ayer MLP with ReLU activation without bias using gradient
descent with weight decay and I-SAM on Pxor with dimension d = 30 and training set size n = 100. In both cases, the model reaches perfect generalization. Notice that although weight decay doesn't 't explicitly regularize model sharpness, the flatness the model decreases through training, which consistent with our Lemma 3.1 relating sharpness to the norm of the weight.
under the assumptions that the minimizer of the training loss form smooth manifold Cooper (2018); Fehrman et al. (2020), Sharpness-Awan Minimization (SAM) (Foret et al., 2021) with batch size 1 and sufficiently small learning rate and perturbation radius (Wen et al., 2022; Bartlett et al., 2022), or Label Noise SGD with sufficiently small learning rate 77 (Blanc et al., 2019; Damian et al., 2021; Li et al., 2021), prefers interpolating models with small trace of Hessian the loss. Hence, we choose analyze trace of Hessian of the loss and will use SAM with batch size (we denote by 1-SAM) as our sharpness minimization algorithm in our experiments.
Notations. We use Tr to denote the trace of matrix and æf to denote the value of the i-th coordinate of vector z. We will use C to represent element-wise product. We use 1 as the (coordinate-wir wise) indicator function, for example, 1 [z > 0) is vector of the same length as whose j-th entry is 1 if x[i] > 0 and otherwise. We will use (zz) to hide logarithmic multiplicative factors.
3 Scenario I: All Flattest Models Generalize
3.1 Flattest models provably generalize
When the architecture 2-MLP-No-Bias, we will show that the flattest models can provably generalize, hence answering Question affirmatively for this architecture and data distribution Pxor-
Theorem 3.1. For any E (0,1) and input dimension d, for = S2 (dlog ($)), with probability at least 1- over the random draw of training set {(2,,y) H from Pro let L(0) A 1 E Cuse(fnohiie (x:), V1) be the training loss for 2-MLP-No-Biar holds that for all 6* E arg Li0= Tr 2 (0)), we have that
Exyy [euse (Spobins (z). < (d/n))
Theorem 3.1 shows that for Pxors flat models can generalize under almost linear sam- ple complexity with respect to the input dimension. We note that Theorem 3.1 implies that Prz.gmPoor [ (x)o)iss<))))) > 0] S (d/n). because if (nobias S 0, it holds that luse (fpobias .v) 2 1. This shows that the model can classify the input with high accuracy. The major proof step relating sharpness to the norm of the weight itself.
Lemma 3.1. Define Oc A {0 (W1, W2) Ex W1,ill (2, < C}. Under the setting of Theorem 3.1, there exists a absolute constant C' independent of and ö, such that with probability
at least - 5, arg min min Tr (.) C Oc and RR (.s(( pnobias E θc}) s 2 (Va/n))
We would like to note that similar results of Theorem 3.1 and lemma 3.1 have also been shown in prior work Wu Su (2023) (that the authors were not aware of before the first version of this work was online).
The almost linear complexity Theorem 3.1 not trivial. For example, Wei et al. (2019) shows that learning the distribution will require ((d2) samples for Neural Tangent Kernel (NTK) (Jacot et al.,

---

25
Sharpness
1.0
Valid aéc
0.9
Sharpness
0.8
W
2
11,13:
18,581
-18,582
0.02
0.7
-14.363
-14.363
0.03
0.6
13.768
13.771
0.03
0
0.5
-12.601
12.601
0.01
10a
10a
10a
Epoch
(a) 2-layer MLP with simplified BN
(b) Weights of the four neurons with the largest norm in the first Layer
Figure 2: Interpretable Flattest Solution We train 2-layer MIP with simplified BN using 1-SAM on Pxor with dimension = 30 and training set size n = 100. After training, we find that the model indeed interpretable. In Figure 2b, we inspect the weight of the four neurons of the four largest neurons in the first layer and we observe that the four neurons approximately extract features tx[ + 22]].
2018). In contrast, our result shows that learning the distribution only requires (dd samples as long as the flatness of the model is controlled.
Beyond reducing model complexity, flatness may also encourage the model to find a more interpretable solution. We prove that under stronger than i.i.d condition over the training set, the near flattest interpolating model with architecture 2-MLP-Sim- will provably generalize and the weight of the first layer will be centered the first two coordinates of the input, i.e., ||Wi,,[3 (3 dlla cliWil
Condition 1 (Complete Training Set Condition). There exists set S (-1,,,77
(- 1] such that the linear space spanned by S S = {8i 82 81, 82 E S} has rank d and the training set is {xx v) x E R4, x13 : d] E S, [1]]] x E [-1,13, =x[]] x x[2]}.
Condition for any width and any O, there exists constant k. >> O, such that for any width-m 2-MLP-Sim-BN. BN, fsbnn satisfying fg interpolates the training set and Tr (V²)(0)) S K inf ((r)=)))= Tr (V²I(0')). holds that Vx E {-1,11d, [<[]==[2] - [o(x)] < e and that Vi E [m]] [Wr, d]Il2 S Wi.ll2.
Theorem 3.2. Given any training set {a4,,
One may notice that in Theorem 3.2 we only consider the approximate minimizer of sharpness. This because the gradient of output with respect to W1, b1. despite never being zero, will converge to zero as the norm of Wi, bi converges cc.
Condition may seem stringent. In practice (Figure 2b), find it necessary for 1-SAM find generalizable solution. We hypothesize that this condition is mainly technical. Theorem 3.2 shows that sharpness minimization may guide the model to find an interpretable and low-rank representation. Similar implicit bias of SAM has also been discussed in Andriushchenko et al. (2023a) The proof is deferred to Appendix B.1
3.2
SAM empirically finds the flattest model that generalizes
We use I-SAM to train 2-MLP-No-BBI -Bis on data distribution Pxor to verify our Theorem 1 (Figure 1). As expected, the model interpolates the training set and reaches a flat minimum that generalizes perfectly to the test set.
We then verify our Theorem 3.2 by training 2-layer MLP with simplified BN on data distribution Pxor (Figure 2a). Here we do not enforce the strong theoretical Condition 1. However, we still observe that SAM finds a flat minimum that generalizes well. We then perform detailed analysis of the model and find that the model is indeed interpretable. For example, the four largest neurons in first layer approximately extract features (relu(cræ(1) + c2æ(2)) C1, 02 E (-1, }}} (Figure 2b). Also, the first columns of the weight matrix of the first layer, corresponding to the useful features (relu(c)< x [1] + (2)(2)) C1, 02 E (-1,1)}, have norms 42.47 and 42.48, while the largest column
norm of the rest of the weight matrix is only 5.65.
5

---

4 Scenario I:: Both Flattest Generalizing and Non-generalizing Models Exist, and SAM Finds the Former
4.1 Both generalizing and non-generalizing solutions can be flattest
In previous section, we show through Theorems 3.1 and 3.2 that sharpness benefits generalization under some assumptions. It is natural to ask whether it is possible to extend this bound to general architectures. However, in this section, we will show that the generalization benefit depends on model architectures. In fact, simply adding bias to the first layer of 2 2-MLP-No-Biar makes non- vacuous generalization bound impossible for distribution Pxor This then leads to a negative answer to Question
Definition 4.1 (Set extreme points). A finite set S C IRd is a set of extreme points if and only if for any S, is vertex of the convex hull of S.
Definition 4.2 (Memorizing Solutions). A D-layer network is a memorizing solution for training dataset if (1) the network interpolates the training dataset, and (2) for any depth k E [D - 11, there an injection from the imput data to the neurons on depth k. such that the activations layer k for each input data is a one-hot vector with the non-zero entry being the corresponding neuron.
Theorem 4.1. For any D 2 2, if the input data points {x+) of the training set form a set of extreme points (Definition 4. 4.1), then there exists a width n layer D MLP-Bias that is memorizing solution (Definition f.r the training dataset and has minimal sharpness over all the interpolating solutions.
As one may suspect, these memorizing solutions can have poor generalization performance.
Proposition 4.1. For data distribution Pxoro for any number of samples n, there exists width-n 2- MLP-Bias that memorizes the training set as in Theorem 4. I, reaches minimal sharpness over all the interpolating models and has generalization error max n//22 0} measured by zero error.
This corollary shows that flat model can generalize poorly. Comparing Theorems 3.1 and 4.1, one may observe the perhaps surprising difference caused by slightly modifying the architectures (adding bias or removing the BatchNorm). To further show the complex relationship between sharpness and generalization. the following theorem suggests, despite the existence of memorizing solutions, there also exists flattest model that can generalize well.
Proposition 4.2. For data distribution Pxar. for any number of samples n. there exists a width-n 2- MLP-Bias that interpolates the training dataset, reaches minimal sharpness over all the interpolating models, and has zero generalization ervor measured by zero one error.
The flat solution constructed is highly simple. It contains four activated neurons, each corresponding to one feature in =x|1] + x[2] (Equation (5)).
Proof sketch. For simplicity, we will consider 2-MLP-Bias here. The construction of the memorizing solution in Theorem 4.1 is follows (visualized in Figure 3). As the input data points form set of extreme points (Definition 4.1), for each input data point 2i there exists a vector ||w;]| = 1, w, E Rdd such that vj ش i, w x; w xz. We can then choose
Wi [v rv 1 w /E]T b IV VVrilll (-w, 2, + e) _ W2 [sign((: //l//Tak..
Here r, = (1 + 1)2/2 and is a sufficiently small positive number. Then it holds that relu(W1 b,) Vralysees where e; is the i-th coordinate vector. This shows there is one-to-one correspondence between the input data and the neurons. It is easy to verify that the model interpolates the training dataset. Furthermore. for Pxor and sufficiently small e, for any input & {xt hijj it holds that relu(W jz + b1) = O. Hence the model will output the same label 0 for all the data points outside the training set. This indicates Proposition 4.1.
To show the memorization solution has minimal sharpness, we need the following lemma that relates the sharpness and the Jacobian of the model.
Lemma 4.1. For mean squared error loss Imse. if model So is differentiable and interpolates dataset (Z,, v)) he(n then Tr ( ((2)()) 2 Ei=n vofo(x:))1

---

Figure 3: Visualization of Memorization Solu- tions. This is an illustration of the memorizing solutions constructed in Theorem 4.1. Here the imput data points come from unit circle and are marked as dots. The shady area with the corre- sponding color represents the region where the corresponding neuron is activated. One can see that the network can output the correct label for each input data point in the training set as long the weight vector on the corresponding neuron properly chosen. Further, the network will make the same prediction 0 for all the input data points outside the shady area and this volume can be made almost as large as the support of the training set by choosing sufficiently small. Hence the model can interpolate the training set while generalizing poorly.
Wax + bg > o
W3X + D2
X1
W2X + b2 >
Proof of Lemma 4.1. By standard calculus, it holds that,
1
Tr L())
Tr Tr(V[[Uo(t) ))
n
2
Tr (oll(( vi) (Vooo(z;)) (VoSo(a:)))
==1
2 Σ Tr (Vofo()))) Σ Vofo(z)))
(1)
n
The first equation in Equation (I) use Vi, fa(;;) = ys. The proof is then complete.
After establishing Lemma 4.1, one can then explicitly calculate the lower bound of (o:))())) |?
condition on fo(x;) y/- For simplicity of writing, will view the bias term part of the weight matrix by appending to the input data point. Precisely, we will use notation 2 E Rd+11 to denote transformed input satisfying Vjj E [d], x[[]] = z,bi,, ad + 1] = 1 and W = [W1 b1] E (d++))) to denote the transformed weight matrix.
By the chain rule, we have,
[oJo(m)) = VVw; [w,,
1122 > 0))=4 relu (Wi=4) l|
W O 1 Wiiz > 01 relu (Wfz²) ) 112
(2)
Then by Cauchy-Schwarc inequality, we have
[Vofo(=;))1 = I o 1 2 > 0) || || ||relu (Win+) ||
2 2 || (W2 3 W > 0)) T relu (Wi++) = 21|
In Equation (3), we use condition fo(;;) = Finally, notice that the lower bound is reached when
We 3 [W > 0j = relu (Wfz{)/l /||=/
(4)
Condition Equation (4) is clearly reached for the memorization construction we constructed, where both sides of the equation are equal to / |p.l/||=?[[e This completes the proof of Theorem 4.1.
However, the memorization network is not the only parameter that can reach the lower bound. For example, for distribution Pxur. if parameter 0 satisfies,
Vi,j E {0, 1). W W1,2++411 = rr(( 1)), (-1>"... 0j. b, + j + 1] = -r, W2 + yl = (-1)++3 /r.
(5)
vk > 4, W.A ........ 0,,b, [k] = 0. [k[ O.

---

180
Sharpness
1.0
200
Sharpness
1.0
Valid acc
Valid acc
160
0.9
Sharpness
Sharpness
150
140
0.8
0.8
120
0.7
100
0.7
100
0.6
0.6
50
80
0.5
0.5
10a
10a
10a
10a
10a
10a
Epoch
Epoch
(a) Baseline
(b) 1-SAM
Figure 4: Scenario II. We train 2-layer MLP with ReLU activation with Bias using gradient descent with weight decay and I-SAM on Pxor with dimension d = 30 and training set size n 100. One can clearly observe distinction between the two settings. The minimum reached by 1-SAM is flatter but the model fails to generalize and the generalization performance even starts degenerate after 4000 epochs. The difference between Figures lb and 4b indicates a small change in the architecture can lead to a large change in the generalization performance.
600
Sharpmess
1.0
Sharpness
1.0
Valid acc Train
600
Valid acc
Sharpness
580
Sharpness
0.8
400
0.8
Acc
560
200
0.6
540
0
101
103
10
10a
103
10a
Epoch
Epoch
(a) Baseline
(b) 1 1-SAM
Figure 5: Scenario II with Softplus Activation. We train a 2-layer MLP with Softplus activation
(SoftPlus(zz) log(1 + e)) with bias using gradient descent with weight decay and 1-SAM on P with dimension = 30 and training set size = 100. We observe similar phenomenon as Figure 4.
with (d + 1)2/4, then for any & E 111, it holds that relu(Wia: + b1) ==1/2==(1]==[[] and fo(x) z[1] x x/2]. Hence it is possible for Equation (5) to hold while the model has perfect generalization performance.
4.2 SAM empirically finds the non-generalizing solutions
In this section, we will show that in multiple settings, SAM can find solutions that have low sharpness but fail generalize compared to the baseline full batch gradient descent method with weight decay. This proves that flat minimization can hurt generalization performance. However, one should note that Question is not denied for the current architectures.
Converged models found by SAM fail to generalize. We perform experiments on data distribution Pxor in Figure 4. We apply small learning rate gradient descent with weight decay and observe that the converged model found by SAM has a much lower sharpness than the bascline. However, the generalization performance of SAM is much worse than the baseline. Moreover, the generalization performance even starts to degenerate after 4000 epochs. We conclude that this scenario, sharpness minimization empirically hurt generalization performance.
baseline
1-SAM may fail to generalize with other activation functions. natural question whether phenomenon that 1 1-SAM fails to generalize is limited to ReLU activation. In Figure 5, we show empirically that I-SAM fails to generalize for 2-layer networks with softplus activation trained on the same dataset, although there no known guarantee for the existence of memorizing solutions.

---

Sharpmess
1.0
Sharpness
1.0
Valid acc
Valid acc
0.9
Sharpness
Sharpness
5
0.8
0.8
Acc
4
0.7
0.6
4-
0.6
2
2
0.5
10a
10a
10a
10
10a
102
10a
Epoch
Epoch
Standard Training
Projected Training
Figure 6: Scenario III. We train two-layer ReLU networks with simplified LayerNorm on data distribution Pxor with dimension 30 and sample complexity n 100 using 1-SAM. In Figure 6a, we use standard training. In Figure 6b, we restricted the norm of the weight and the bias of the first layer as 10, to avoid minimizing the sharpness by simply increasing the norm. We can see that in both cases, the models almost perfectly generalize.
5 Scenario II: Both Flattest Generalizing and Non-generalizing Models Exist, and SAM Finds the Latter
5.1 Both generalizing and non-generalizing solutions can be flattest
Despite the surprising contrary between Theorems and 4.1, experiments show that Question consistently hold. However, we will provide counterexample in this section. Specifically, we will consider data distribution Pxor and 2-layer ReLU MIP with simplified LampliieddLayerieiee LayerNorm. One can first show both generalizing and non-generalizing solutions exist similar to Theorem 4.1 and propositions 4.1 and 4.2.
Theorem 5.1. If the input data points {x;} of the training set form set of extreme points (Defini- tion 4. 1), for sufficiently small c, then there exists a width-n 2-MLP-Sim-LN with hyperparameter e that is memorizing solution (Definition 4.2) for the training datasel and has minimal sharpness over all the interpolating solutions.
Proposition 5.1. For data distribution P Pxor. for sufficiently small for any number of samples n, there exists a width-n 2-MLP-Sim-LN with hyperparameter that memorizes the training set as in Theorem 4. 1, reaches minimal sharpness over all the interpolating models and has generalization ervor max n2/2d, 00 measured by zero one error
Proposition 5.2. For data distribution P Pxor. for sufficiently small e, for any number of samples n, there exists width-n 2-MLP-Sim-LN with hyperparameter that interpolates the training dataset, reaches minimal sharpness over all the interpolating models, and has zero generalization error measured by zero one error.
The construction and intuition behind Theorem 5.1 and propositions 5.1 and 5.2 similar to that of Theorem 4.1 and propositions 4.1 and 4.2. The proof is deferred to Appendix B.
5.2
SAM empirically finds generalizing models
Notice in Section 5.1 our theory makes the same prediction in Section 4. However, strikingly, the experimental observation is reversed (Figure 6). Now running SAM can greatly improve the generalization performance till the model perfectly generalizes. This directly denies Question 2 as now we have a scenario in which sharpness minimization algorithms can improve generalization till perfect generalization while there exists flattest minimizer that will generalize poorly.
6 Discussion and Conclusion
We present theoretical and empirical evidence for (1) whether sharpness minimization implies generalization subtly depends on the choice of architectures and data distributions, and (2) sharpness minimization algorithms including SAM may still improve generalization even when there exist flattest models that generalize poorly. Our results suggest that low sharpness may not be the only cause of the generalization benefit of sharpness minimization algorithms.

---

ACKNOWLEDGEMENTS
The authors would like to thank the support from NSF IIS 2045685.
References
Maksym Andriuscchenko, Dara Bahri, Hossein Mobahi, and Nicolas Flammarion. Sharpness- aware minimization leads to low-rank features. arXiv preprint 1XX2:2305..6692, 2023a.
Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flam- marion. A modern look at the relationship between sharpness and generalization. arXiv preprint arXiv:2302.07011, 2023b.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on edge of stability in deep learning. arXiv preprint vr.Xiv:2205097725 2022.
Peter L Bartlett, Philip M Long, and Olivier Bousquet. The dynamics of sharpness-aware minimiza- tion: Bouncing across ravines and drifting towards wide minima. arXiv preprint 1rXiv:22/0.01513, 2022.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process. arXiv preprint (XXxv:990400900) 2019.
Yaim Cooper. The loss landscape of overparameterized neural networks. arXiv:1804. 10200, 2018.
arXiv preprint
Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global minimizers, 2021.
Lijun Ding, Dmitriy Drusvyatskiy, and Maryam Fazel. Flat minima generalize for low-rank matrix recovery. arXiv preprint arXiv:2203.037556 2022.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 1019-1028. JMLR. org, 2017.
Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. arXiv preprint arXiv: 11008. 2017.
Benjamin Fehrman, Benjamin Gess, and Arnulf Jentzen. Convergence rates for the stochastic gradient descent method for non-convex objective functions. Journal of Machine Learning Research, 21: 136, 2020.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-awar aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021.
Khashayar Gatmiry, Zhiyuan Li, Ching-Yar Yao Chuang, Sashank Reddi, Tengyu Ma, and Stefanic Jegelka. The inductive bias of flatness regularization for deep matrix factorization. arXiv preprint arXiv:2306. 13239, 2023.
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. In Advances in Neural Information Processing Systems, pp. 6151-6159, 2017.
Jeff Z HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. arXiv preprint arXiv:2006.08688) 2020.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):144, 1997.
Jean Honorio and Tommi Jaakkola. Tight bounds for the expected risk of linear classifiers and pac-bayes finite-sample guarantees. In Artificial Intelligence and Statistics, pp. 384-392, PMLR, 2014.
10

---

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization neural networks. In Advances neural information processing systems, pp. 8571-8580, 2018.
Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv:1711.0462 2017.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv preprint arXiv: 1912.02/8, 2019.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.0450 2016.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations. arXiv preprint arXiv:
pp. 2-47, 2017.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss?-a mathematical framework. In International Conference on Learning Representations, 2021.
Zhiyuan Li, Tianhao Wang, and Dingli Yu. Fast mixing of stochastic gradient descent with normal-
ization and weight decay. Advances in Neural Information Processing Systems, 55:9233-9244 2022.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normaliza- tion layers: Sharpness reduction. arXiv preprint arXiv:2206. 07085, 2022.
Chao Ma and Lexing Ying. On linear stability of sgd and input-smoothness of neural networks. Advances in Neural Information Processing Systems, 34:
2021.
Jifí Matoušek. On variants of the johnson-lindenstrane lemma. Random Structures Algorithmss 33(2):142-156, 2008.
Mor Shpigel Nacson, Kavya Ravichandran, Nathan Srebro, and Daniel Soudry. Implicit bias of the step size linear diagonal neural networks. International Conference on Machine Learning, pp. (6270-16295.. PMLR, 2022.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring general- ization in deep learning. In Advances in Neural Information Processing Systems, pp. 5947-5956, 2017.
Matthew D Norton and Johannes O Royset. Diametrical risk minimization: Theory and computations. Machine Learning, pp. 1-19, 2021.
Antonio Orvieto, Anant Raj, Hans Kersting, and Francis Bach. Explicit regularization in over- parametrized models via noise injection. arXiv preprint 3rX1v:2206.04613, 2022.
Alessandro Rinaldo. 36-709: Advanced probability theory. https: hstps://www... stat. cmu, edu/ "arinaldoo/Teaching/S3 /tse//6779//111/ Bttp//Srriedd_LLetee 2019.
Shai Shalev-Shwartz and Shai Ben- BeviD.vid. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent separable data. The Journal of Machine Learning Research, 19(1): 2822-2878, 2018
Nathan Srebro, Karthik Sridharan, and Ambuj Tewari. Smoothness, low noise and fast rates. In J. Lafferty, C. Williams, J. Shawe-Taa Taylor, R. Zemel, and A. Culotta (eds.), Advances in Neural Information Processing Systems, volume 23. Curran Associates, Inc., 2010. URL https::/proccede/ 7CC199051ee22eaaa111177744443332 Paper .df.

---

Colin Wei and Tengyu Ma. Data-dependent sample complexity of deep neural networks via lipschitz augmentation. Advances in Neural Information Processing Systems, pp. 9722-9733, 2019a.
Colin Wei and Tengyu Ma. Improved sample complexities for deep networks and robust classification via an all-layer margin. arXiv preprint arXiv: arXiv:1910.0 042244 2019b.
Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. In Advances in Neural Information Processing Systems, pp- 9709-99 9709-9721, 2019.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: case study on approximating turing machines with transformers, Advances in Neural Information Processing Systems, 15:12071-- 2022.
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How does sharpness-aware minimization minimize sharpness? arXiv preprint urXiv:2211.05722 2022.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. arXiv preprint ar arXiv:2002.0901 09277, 2020.
Lei Wu and Weijie J Su. The implicit regularization of dynamical stability in stochastic gradient descent. arXiv preprint arXiv:2305. 17490, 2023.
Lei Wu, Chao Ma, et al. How sgd selects the global minima in peer-parameterized learning: A dynamical stability perspective. Advances in Neural Information Processing Systems, 31, 2018.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint arXiv: /6/ 1.03330, 2016.
12


Reviews:
[{"review_id": "o2AKMXwsfq", "reviewer": "Reviewer_fkTA", "summary": "This paper explores the relationship between sharpness and generalization in overparameterized neural networks. The authors challenge the popular belief that flatness of the loss function implies better generalization and show that sharpness minimization algorithms do not always lead to better generalization. Through theoretical and empirical investigation, the authors identify three different scenarios for two-layer ReLU networks, showing that the relationship between sharpness and generalization is subtle and depends on the data distributions and the model architectures. The paper's contributions include a theoretical analysis of the relationship between sharpness and generalization.", "questions": "1. In abstract, what does this sentence “there exist non-generalizing flattest models and sharpness minimization algorithms fail to generalize poorly” mean?", "limitations": "The authors have adequately addressed the limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "The authors identify three scenarios for two-layer ReLU networks and provide theoretical and empirical evidence to support their findings. I believe the subtle details in architecture and distribution located in this paper are remarkable. Overall, this paper makes significant contributions to the field of generalization theory and challenges the popular belief that the flatness of the loss function implies better generalization.  ", "weaknesses": "1.\t One potential weakness of the paper is that it only examines the relationship between sharpness and generalization for two-layer ReLU networks (personally, I think this is acceptable since it is very hard to analyze more complex scenarios). Additionally, the authors only examine the CIFAR-10 and CIFAR-100 datasets, and it would be valuable to examine their findings on other datasets to determine if their results generalize to other domains. \n2.\tThe expression of the article is not very friendly and a little difficult to understand.\n3.\tThere are some grammatical and formatting problems, some of which are shown as follows.\n1)\tPictures should be resized to make them look more coherent. In addition, the text in the image should be consistent with the font within the paper.\n2)\tThere are some spelling errors, such as “subtlely”.\n3)\tReferences are poorly formatted. Please use a uniform format for references. \n4)\tFormulas and variables should be consistent with the font within the paper.\n5)\tOperators should be carefully defined.\n6)\tThere are some formatting issues. For example, lowercase letters as the first letter of a sentence.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s positive review! Below are our responses to the reviewer’s concerns:\n\n* **W1.** We would like to clarify that we only perform experiments on synthetic tasks instead of CIFAR10 or CIFAR100 in the paper. However, we added an experiment on MNIST in the rebuttal phase. See the general response for a detailed discussion.\n* **W2.** Thank you for the helpful advice! We will fix and correct the formatting issues in future versions.\n* **Q1.** We made a typo in the quoted sentence — there should be no “poorly” at the end of the sentence. The corrected version is “there exist non-generalizing flattest models and sharpness minimization algorithms fail to generalize”.\n"}, {"review_id": "gObGKHPbD5", "reviewer": "Reviewer_2a5p", "summary": "This work aims to examine different cases for which flatness implies generalization. The conclusion heavily depends on the underlying model, while there are three different scenarios that contradict with each other. Additionally, the theoretical guarantees study the generalization error at the solution, however the empirical guarantees consider the generalization error at the last iteration of the SAM algorithm. Despite the convergence of the SAM algorithm, the training horizon may also affect generalization. With alternative sampling schedules of the data-set, it can be possible to minimize the training horizon (for instance see Repeated Random Sampling for Minimizing the Time-to-Accuracy of Learning by Okanovic et al.). I think it would be useful to examine if the SAM algorithm generalizes better when considering alternative sampling schedules. In such a case, the batch-selection policy could be more significant than the sharpness minimization. As the author also mention flattest solutions are not directly associated with efficient generalization in general.\n\nFurther, there is a deep literature for the edge of stability phenomenon (Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability by Cohen et al.). In this case, the sharpness increases at termination (Understanding Gradient Descent on the Edge of Stability in Deep Learning by Arora et al., Section 2. Related Works), however good generalization performance have been observed in some cases. These cases might be considerd by the authors to provide a more complete picture for the empirical part of the paper, and additional conclusion about sharpness minimization and generalization. \n\n  ", "questions": "In theorem 3.1 the variable $\\theta^*$ should be a random variable that depends on the training data-set. Is this true? In such a case the authors may consider to use the notation of the conditional expectation in the statement of the Theorem to make this fact clear.", "limitations": "The authors have addressed the limitations sufficiently.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The results are original and provide a partial evidence between sharpness and generalization, however the conclusion is not clear. Theorem 3.1 is an interesting result, however it is limited to a simple architecture but it is valuable.", "weaknesses": "Some parts of the paper are not clear enough. For instance, there exist statements like \"All flattest models generalize\" while the statement holds only for certain cases. It appears that the theoretical guarantees study generalization at the solution, however in practice the generalization of the algorithm might be different (see also the Summary). The main question of the paper has been considered in prior works as well. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the effort and time invested in the review. However, there are a lot of misunderstandings in the review and we clarify them below one by one.\n1. To our best knowledge, there is no reported correlation between progressive sharpening and better generalization. Arora et al. 2022, shows that normalized GD reduces sharpness after reaching edge of stability and simultaneously improve generalization.\n\n2. The goal of our work is to investigate sharpness and generalization. We agree with the reviewer that in practice many design choices including the batch sampling scheme can affect generalization but we choose to investigate the minimal setting where the data are i.i.d sampled and both the sharpness and the training loss are minimized. As shown in the paper, even this minimal setting would require much investigation.\n\n3. We believe that our work has made a solid step toward revealing the relationship between sharpness and generalization by answering the questions listed in the introductions which are not answered by any previous works that the authors know of. \n\n4. Regarding the title of the sections, the authors have already stated in the introduction that the relationship may subtly depend on architectures and data distributions. The authors have also clarified in both the introduction and section title that statements including `All flattest model generalize` can only be stated under the scenario defined by architectures and data distributions.\n\n5. It is true that the variable is a random variable depending on the training data set. However, the randomness is already explicitly considered in `1 − δ over the random draw of training set`  on line 143.\n"}, {"review_id": "aSgaXZtVlP", "reviewer": "Reviewer_ex27", "summary": "This paper studies the connections between sharpness of the training loss and generalization performance under a simplified MLP architecture. In particular, authors show scenarios of when flat solution do/do not generalize, as well as cases when Sharpness-Aware Minimization (SAM) does/does not generalize.", "questions": "- Maybe some visualization might help address some of my concerns in the previous part? For example visualizing the SAM/GD trajectories?\n\n- Additional experiments either on noisy data settings or more realistic deep networks might help fill some of the gaps in the paper.\n\n", "limitations": "n/a", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The theoretical insights of the paper are interesting and hopefully can help better understand the relationship between flatness and generalization.", "weaknesses": "I think the paper has several shortcomings, outlined below.\n\n- Model setup: I think the model setup is too simplistic. The input space is binary, which I think is fine. However, no noise is considered in the model and the output label $y=x[1]x[2]$ satisfies a very specific structure. Therefore, it is not clear to me why results discussed in the paper would generalize to broader model setups, and the observations discussed here are not just an artifact of the simplistic model. As shown in Lemma 3.1, the sharpness under the paper's setup is closely related to the weight norm, which may or may not be the case in a more realistic setup. \n\n- Experiments: I don't find the experiments on SAM convincing. How is SAM implemented? How are the algorithm hyper-parameters chosen? How is SAM initialized? Maybe changing the SAM parameters and initialization would change the outcome of the experiments.  Can the observation of SAM converging to non-generalizing flat solution be replicated in a more realistic setup? Is weight decay used for SAM, as specially, in this case sharpness seems to be closely related to the weight norm.\n\n- The authors miss [1]. As this paper discusses the connections between generalization and sharpness/SAM, the authors should discuss and compare their results to [1], specially as [1] discusses noisy setups.\n\n\n[1] Behdin, K., & Mazumder, R. (2023). Sharpness-aware minimization: An implicit regularization perspective. arXiv preprint arXiv:2302.11836.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the review and comment! Regarding the reviewer’s concern, there seem to be some misunderstandings that we wish to clarify.\n1. **W1: The generality of the results** \n\nWhile our positive result theorem 3.1 only holds for a 2-layer ReLU network without bias, our main results, which are negative for generalization for flattest models, hold for more general settings (see Theorem 4.1) including the noisy labels setting and multi-layer networks. \n\n2. **W1: Relationship between sharpness and weight norm**\n\nWe would also like to clarify that in all the setups excluding the 2-layer network without bias, sharpness is not necessarily related to weight norm. Lemma 4.1 only holds for data distribution with zero mean and a two-layer network without bias. For a two-layer network with bias, on the synthetic setup we consider, constraining the norm of the weight can actually induce generalization. Proposition 4.1 shows that constraining the sharpness can’t induce generalization, which shows that in this setup, sharpness is actually very different from the weight norm. \n\n3. **W1: Noisy setups and additional related works**\n\nExperimentally SAM does have a generalization benefit over SGD when labels are noisy. But analyzing the generalization benefit of SAM for noisy labels is more complicated because the best performance is usually achieved by early-stopped SAM and interpolating the noise in the labels usually degrades the generalization. We leave understanding the generalization benefit of SAM with noisy labels as future work. \n\nThe authors would like to thank the reviewer for pointing out the missing citation [1] and will add it in the following version. Compared with [1], our works are different in three ways:\n\n(a). [1] considers a two-layer relu net which is essentially linear in its trainable parameters, because the second layer is fixed to be all 1's, and the authors of [1] assume the activations of the first layer don't change throughout the training, while we consider a general two-layer relu net; \n\n(b). Because the model is linear as mentioned in (a) and the loss is l2 loss, all the solutions have the same sharpness (or more explicitly the same hessian) in [1], regardless of whether having noise. In contrast, in our setup, models can have different sharpness; \n\n(c). In the setting of [1], SAM and GD converge to the same solution and there is no generalization bound provided, regardless of having noise or not, while we show the flattest model in certain settings can always have good generalization (thm 3.1).\n\n4. **Experiments.**\n\na. **Weight decay.** We would like to stress that weight decay is not used for SAM because it is already proved that constraining the norm can lead to generalization in many setups (see Figure 4a for example) and our goal is to investigate the relationship between sharpness and generalization so adding weight decay may confound the relationship.\n\nb. **Hyperparameter.**  All the hyperparameters are listed in detail in Table 2 on page 30 of the appendix and we use standard initialization offered by Pytorch. SAM is implemented using an open-source library and the implementation can also be found in the appended material `code/sam.py`.\n\nRegarding the reviewer’s question, the authors perform a hyperparameter search on SAM for the experiments in Figure 4b. We search through learning rate from [3e-2, 1e-2, 3e-3, 1e-3] and perturbation radius from [0.01, 0.03, 0.05, 0.10]. We find out that using small initialization (scaling down the first layer by a factor of 0.01), the model can find a perfect generalizing solution using 1-SAM. We believe this is however caused by the implicit bias of small initialization, which might have a similar effect as norm constraint, as vanilla SGD also reaches good generalization performance using this initialization. However, using the standard initialization, the best model we found only reaches 0.84 validation accuracy and the hyperparameter is exactly the one we adopted in the paper. We believe this shows that what we have shown is not an artifact due to poor hyperparameter and we would also want to stress that we are only claiming that SAM can find a flattest but non-generalizing solution, which is inherently an existential claim.\n\nc. **Extension.** The authors have already done some extended experiments on other synthetic setups (see Figure 7 in the appendix) and also added an experiment on MNIST to further strengthen the paper according to the reviewer’s advice. Please see the general response for a detailed discussion.\n\nd. **Visualization.** We agree that visualization could help the readers better understand what happens during training. Unfortunately, we do not have a good idea of how to visualize a high dimensional trajectory except plotting certain statistics like train/valid acc/loss, sharpness, etc.\n"}, {"review_id": "vi4fp1dW0T", "reviewer": "Reviewer_ZroU", "summary": "This paper sheds light on the relationship between three concepts: (a) generalization (b) flatness and (c) explicitly inducing flatness via SAM. The paper identifies that the relationship is nuanced and is based on certain architectural properties of the model.\n\n Specifically, depending on the architecture, there are three possible regimes:\n1. A regime where _every flat minimizer generalizes well_ and SAM finds these minimizers.\n2. A regime where _there exists flat minimizers that generalize poorly_ but SAM does _not_ find these minimizers.\n3. A regime where _there exists flat minimizers that generalize poorly_ and SAM _does_ find these minimizers.\n\nThe paper shows that \n\n- Regime (1) happens for 2-layer RELU networks _without_ bias\n- Regime (2) happens for 2-layer RELU networks _with_ bias\n- Regime (3) happens for 2-layer RELU networks _with_ bias, _with_ a simplified Layer Norm.\n\nNote that in all cases, the existence of a flat solution that generalizes/memorizes is proven theoretically; while the effect of SAM is demonstrated empirically.", "questions": "1. Def 4.1 for the memorizing solutions appears to be incomplete and can be more formal. First, by _input data_, please clarify that you are referring to the training data. Furthermore, you also require the additional criterion that the activations of the model on _test_ data are all zero --- is this right? If not, the definition appears to be incomplete as even the raw inputs must more or less satisfy the orthogonality assumption with high probability? i.e., when we sample from a high-dimensional hypercube, two samples must more or less be orthogonal. So this alone shouldn't characterize a memorization solution? \n\n2. Could you provide intuition for what the \"flattest solution\" in Eq 4 satisfies? It seems to require that _the activated top-layer weights have correspondence to the activations themselves_. I don't think Fig 3 captures this specific structure?\n\n3. Do you empirically find that SAM learns a solution whose parameters resemble Eq 4 in any way?\n\n4. The motivation behind the choice of loss function can be made more transparent. In particular, the main loss function used is squared error loss. In the appendix, the authors also use logistic loss _but with label smoothing_. If I am not wrong these choices ensure that the flattest minimizer exists and is unique up to some permutations. However, with pure logistic loss, the flattest minimizer would become ill-defined since there would be many different flat minimizers \"at $\\infty$\"? I believe this is the reason behind a line of work in PAC-Bayes literature that tries to \"normalize\" flatness in a well-defined way. In other words, this work \"sidesteps\" this annoying issue; this is okay, as it still helps outline an elegant insight. But this needs to be clearly stated so that readers can put this result in context. \n\n### Minor suggestions\n\n1. I'd suggest making the specific notion of trace-of-hessian-based flatness clear in the abstract and in the beginning of the introduction. Since there are many notions of flatness, it would be good to be upfront about which notion is being discussed here.\n\n2. In the introduction, I'd suggest clarifying that the SAM-related results are purely empirical.\n\n3.  Section 4 title seems incorrect. In this regime, SAM finds the _latter_ non-generalizing solution not the former one.\n\n4. Line 167 makes an abrupt reference to a theorem that's not discussed until then. \n\n5. It's not clear until we get to the discussion of SAM as to whether the proof is empirical or not. \n\n6. In Sec 5.1, instead of simply referring to Thm 3.1 and 4.1, I would add a verbal description so that it's easier to follow.\n\n\n### Summary\n\nOverall, I found this paper to be insightful, rigorous and clear in its message. The community would find this valuable given the confusing nature of literature in flatness. \n\n### Post-rebuttal update\n\nI have increased my score since the authors addressed my response satisfactorily. \n**I'd also like to respectfully express my strong disagreement with concerns raised in the other reviews.**\n- Regarding \"Two-layer network is too simplistic\": When a certain truth has not been proven in any setting to begin with, and when someone makes headway by providing the first proof in a simple setting, it is unfair and counter-productive to hold it against them. For without that headway, there is little hope for more general analyses. In fact, IMO proofs in more general settings (like a multi-layered network) get so unwieldy as to become uninsightful. Furthermore, in this scenario, the setting is a 2-layer network which does not reduce to a linear function of its parameters. IMO, the statements in this paper are already highly non-trivial to identify, let alone prove. \n- Regarding considering other alternate versions of SAM, this is an interesting reformulation of the problem. Thanks to this work, we can be better posed to answer more important questions. Hence, again I don't want to fault this paper for not proving its elegant theory in various sophisticated versions of SAM. I would share this review's concern had the considered setting been far removed from practice, and promised no tangible way of extending to other situations --- this doesn't seem to be the case here, in my humble opinion.", "limitations": "The paper lists reasonable limitations in the conclusion.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The questions formalized in the paper are novel, and well-formulated. \n2. I gained a much clearer understanding of some generally unclear and confusing concepts: (a) flatness, (b) generalization and (c) explicitly inducing flatness. All these ideas are dealt with care and rigor. \n3. The introduction is well-written and the chain of arguments are explained clearly. Despite being a theoretical and nuanced paper, I found it extremely easy to follow the arguments in each section. I thoroughly enjoyed reading the paper!\n4. Rather than unnecessarily complicate the theory, the authors come up with a reasonably minimal settings of a 2-layer _non-linear_ model and support their arguments with either a formal result or with a rightly-designed empirical result. \n5. That there are situations where the flattest solution does not generalize well is surprising from a theoretical point of view.", "weaknesses": "I currently do not have any major concerns about this paper. Most of my comments are follow-up questions, which I enumerate in the next section. My most important question --- which may or may not hint at a weakness --- is below:\n\n### Major question\n\nQ1a: Could you be more explicit about any parameter count dependence in the Rademacher complexity term of Thm 3.1 and Lemma 3.1? I am unsure if this is orthogonal to the main point of the result, so I'd also appreciate some clarity on that.  \n\nQ1b: On that note, I'd also like to know why this bound would break down upon the introduction of bias into the architecture. Would one of the terms within the Rademacher complexity end up scaling with the number of datapoints?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the detailed review and comment! We would like to answer the reviewers’ questions here.\n\n1. **Q1a.** Our bound is independent of the width of the model. Under the setting of Thm 3.1 and Lemma 3.1, we show our result by using sharpness as an upper bound for the parameter norm square using a uniform bound approach (see Lemmas A.5, A.6, and A.10).\n2. **Q1b.** When the bias is introduced, both Lemma A.5 and A.6 will break. In this case, the sharpness can no longer be used as a valid upper bound for the weight norm and hence the flattest model will no longer fall into the function class defined in Lemma 3.1 with high probability.\n3. **Questions 1.**  You are correct that ‘input data’ should be replaced by ‘training data’. We don’t need the activations on testing data to be all zero because our definition of memorizing solution only involves training data and does not necessarily imply bad generalization (c.f. Proposition 4.1, it takes some extra effort to show there exists a memorizing solution with bad generalization). Also, we would like to note that the construction of memorizing solution doesn’t need orthogonality within the input data. See Figure 3 as an example. We will try to clarify the statement in future versions. Thank you!\n3. **Questions 2.** Here equation 4 characterizes the following correspondence: for each train data, among the activated neurons, the weights of the second layer are proportional to the activations (up to a factor of train data norm). The intuition behind this condition is that the 2-homogeneity of the network cause the gradient of the first and second layer weight ($W_1’$ and $W_2$) must have a constant product in the norm given an input and output pair. Equation 4 is then derived by balancing these two gradients. This proof technique could be useful for future theoretical analysis of homogeneous models.\n\\\nFigure 3 and the memorizing solution we constructed is a special case where both sizes of equation 4 are vectors with only 1 nonzero coordinate. This can be done by choosing the first layer in a similar fashion as in Figure 3 which ensures all the input training data will activate exactly 1 neuron. Other than showing the activation is one-hot for each sample, figure 3 doesn't directly correspond to this additional structure, and shows no information about the second layer because the activated region of each hidden neuron is solely decided by its incoming weight, which belongs to the first layer. \n\n4. **Questions 3.** Yes, the converged model in Figure 4b reaches almost minimal sharpness and we observe strong alignment in $\\mathrm{relu}(W_1x_1 + b_1)$ and $W_2 \\odot 1(W_1x_1 + b_1 > 0)$. The average cosine similarity between the two vectors is above 0.99.\n\n5. **Questions 4.** We agree. We sidestep the infinity issue by using the soft label in the logistic loss. We will make clear the discussion here in future versions. Thank you!\n\nWe would like to thank the reviewer again for the advice on writing and clarifying. We will adopt these suggestions in future versions.\n"}]

Author rebuttal:
We would like to thank the reviewers for the detailed reviews. Below we address a common question (by reviewers fkTA and ex27) on data distribution, whether the phenomena observed on synthetic datasets in this paper extend to vision datasets, by providing additional MNIST experiments. 

**MNIST experiments**. 
We observe that on the MNIST dataset when running SAM for a sufficiently long time, the validation accuracy will also drop while the sharpness decrease as in the synthetic setup, showing the generality of our result. We also observe that the converged accuracy of 1-SAM, in this case, is lower than the converged accuracy of SGD with weight decay regularization, which is also consistent with our findings in the synthetic setup.

Due to time constraints (we can only use batch size 1 for 1-SAM), we choose to downsample the MNIST image to 7$\times$7. We use a training sample size of 1000 and a 2-layer ReLU network with a width of 2000 here. The hyperparameters we choose and the resulting training curves are shown in the appended pdf. 

Finally, we want to stress that even without additional vision experiments, our theoretical analysis has already excluded the possibility of proving a generalization bound for all flattest interpolating deep relu networks and shows that the relationship between flatness/sharpness minimization algorithm and generalization subtly depends on the architecture. For simplicity of presentation, we choose to focus on the two-layer net in the main paper and leave more studies on more realistic settings for future work.


Meta-review: This paper provides a rigorous study of the relationship between flat minima, generalization, and the sharpness-aware optimizers such as SAM and presents intriguing and thought-provoking analyses/experiments/observations on the interaction of these three notions. Specifically, the authors show that whether flat minima are optimal or not (in terms of achieving the best generalization performance) depends on the choice of architecture and data distribution. In particular, they prove that for certain configuration of the above, the flattest minima provably correspond to the minima that generalize [and sharpness-aware optimizers can indeed find the flattest minima]. Perhaps more interestingly, the authors show that even when the flattest minima do not generalize well, SAM can still improve generalization. The latter observation is interpreted as there might be other regularization effects in SAM at work that are important for generalization and yet need to be understood.

All reviewers find the findings of this submission intriguing. and rate for accept. In concordance with them, I believe this paper is a timely research study into one of the biggest recent mysteries of generalization of deep neural networks, and hence recommend accept.
Decision: Accept (oral)