Normalization Layers Are All That Sharpness-Aware Minimization Needs
Authors: Maximilian Mueller
Tiffany Joyce Vlaar
David Rolnick
Matthias Hein
Venue: NeurIPS 2023 poster (2023)
Keywords: sharpness-aware minimization
flatness
generalization
normalization layers
TL;DR: We show that only applying SAM to the normalization layers of the network enhances performance for ResNet and Vision Transformer architectures compared to applying SAM to all parameters.

Abstract: Sharpness-aware minimization (SAM) was proposed to reduce sharpness of minima and has been shown to enhance generalization performance in various settings. In this work we show that perturbing only the affine normalization parameters (typically comprising 0.1% of the total parameters) in the adversarial step of SAM can outperform perturbing all of the parameters. This finding generalizes
to different SAM variants and both ResNet (Batch Normalization) and Vision Transformer (Layer Normalization) architectures. We consider alternative sparse perturbation approaches and find that these do not achieve similar performance enhancement at such extreme sparsity levels, showing that this behaviour is unique to the normalization layers. Although our findings reaffirm the effectiveness
of SAM in improving generalization performance, they cast doubt on whether this is solely caused by reduced sharpness.

Normalization Layers Are All That Sharpness-Aware Minimization Needs
Maximilian Müller University of Tübingen and Tübingen AI Center
Tiffany Vlaar
McGill University and Mila Quebec AI Institute tiffany. vlaar@mila quebec
maximilian wwmlmmmlleeeee muellerdusii. .uniiiiiiiiiiie uni-tuebingen.de
David Rolnick McGill University and Mila Quebec AI Institute drolnickocs. mmgill.cc
Matthias Hein
University Tübingen and Tübingen AI Center matthias hwwnhuniniin... cactouunittuebinnn -tuebingendee
Abstract
Sharpness-aware minimization (SAM) was proposed to reduce sharpness of minima and has been shown to enhance generalization performance various settings, In this work we show that perturbing only the affine normalization parameters (typically comprising 0.1% of the total parameters) in the adversarial step of SAM can outperform perturbing all of the parameters. This finding generalizes different SAM variants and both ResNet (Batch Normalization) and Vision Transformer (Layer Normalization) architectures. We consider alternative sparse perturbation approaches and find that these do achieve similar performance enhancement at such extreme sparsity levels, showing that this behaviour is unique the normalization layers. Although our findings reaffirm the effectiveness of SAM in improving generalization performance, they cast doubt whether this is solely caused by reduced sharpness.
Introduction
Numerous works have been dedicated to studying the potential connection between flatness of minima
and generalization performance of deep neural networks [3 18 21 45 4 Several aspects of training
are thought to affect sharpness, but how these interact with each other remains an ongoing area of research. Recently, sharpness-awe -wwree aware minimization (SAM) has become a popular approach actively try to find minima with low sharpness using a min-max type algorithm remarkably effective in enhaneing generalization performance for various settings 23
[23]. SAM was found to be
14
33
Several variants of SAM have been proposed, focusing both on enhanced performance 37 35 and reduced computational cost [1] 20) In particular, with the original aim of making SAM more efficient Mi et 42 propose 'sparse" SAM approach, which applies SAM only to select number of parameters. Although they do not actually succeed in reducing wall-clock time, they make the surprising observation that using 50% (in some settings even up to 95%) sparse perturbations can maintain
enhance performance compared to applying SAM to all parameters. They thus hypothesize that "complete perturbation on all parameters will result in suboptimal minima".
Similar to the effect of SAM 23
14 normalization layers are thought to reduce sharpness Frankle et 24 found for ResNets that the trainable affine parameters of the normalization layers
41 46
Code is provided at http:///gthh..commmmllleme//hhhp
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

ASAM, p=2.0, elem.
Test Accuracy (%6
84
ASAM, p=2.0, layer
ASAM, p=inf, elem.
83
SAM, p=2.0
SGD
OnlyNorm
81
no-norm
10
10a
10
109
10
P
Figure 1: The interplay of normalization layers with SAM: Perturbing only normalization layers (OnlyNorm, dashed) improves generalization performance, while omitting them in the perturbation (no-norm, dotted) can harm training. WideResNeet28-11 trained with different MAriiniiants on CIFAR-100. Best seen in color.
have remarkable representation capacity in their own right, whereas disabling them reduce performance. Inspired by this we focus on the interplay between SAM and normalization layers and show that for various settings perturbing exclusively the normalization layers of network (often less than 0. 1% of the total parameters) outperforms perturbing all parameters. We find that this can not be solely attributed possible benefits of sparse perturbation approaches and highlight the unique role played by the normalization layer affine parameters, As our main contributions we show that:
Applying SAM only to the normalization layers of network (SAM-ON short for SAM-OnlyNormi enhances performance on CIFAR data compared to applying SAM to the full network (SAM-all) and also performs competitively on ImageNet. We corroborate the remarkable generalization performance of SAM-ON for ResNet and Vision Transformer architectures, across different SAM variants, and for different batch sizes. (Section
Alternative sparse perturbation approaches do not result in similar performance as SAM-ON, especially not at the extreme sparsity levels of method. (Section
Similar SAM-all, SAM-ON yields non-trivial adversarial robustness (Section 4.2p. It also reduces the feature-rank, but the sharpness-reducing qualities of SAM-all are not fully preserved (Section
5.2
2 Related Work
Normalization layers. Batch Normalization (BatchNorm) 31 and Layer Normalization (Layer-
Norm) (6] form an essential component of most convolutional [25] 29 and Transformer [5] architectures, respectively. Across various works these normalization layers were shown to accelerate and stabilize training, reducing sensitivity initialization and learning rate [13 5 60 36). But despite their widespread adoption and illustrated effeectiveness, a conclusive explanation for their success still elusive. The original motivation for BatchNorm as reducing internal covariance shift (31] has been disputed [46]. The hypothesis that normalization layers enhance smoothness is supported through both empirical and theoretical analyses 46 10 41, though also not completely undisputed [57]. Unlike LayerNorm, BatchNorm is sensitive to the choice of batch size [38, 49). Ghost Narch where BatchNorm statistics from disjoint subsets of the batch are used, is found to regularize and generally enhance generalization 28 49] even though it reduces smoothness [17
Affine parameters. There are relatively few papers that study the role the trainable affine parameters of the normalization layers. Frankle et al. (24) were able to obtain surprisingly high performance on vision data by only training the BatchNorm layers, illustrating the expressive power of the affine parameters, which potentially achieve this by sparsifying activations. For BatchNorm in ResNets, disabling the affine parameters was shown empirically reduce generalization performance 24]. but for LayerNorm in Transformers to not affect or even improve performance [56]. For few-shot transfer tasks disabling the BatchNorm affine parameters during pretraining was found to enhance performance (58) Further, many other aspects of training will have non-trivial effect, e.g. applying weight decay to the BatchNorm affine parameters was found to increase performance for ResNets but
harm performance in other settings (4)].

---

Sharpness-aware minimization. SAM was developed to try to actively seek out minima with low sharpness [23] Training with SAM may lead to increased sparsity of active neurons [14] and models which are more compressible 44 SAM has been shown to be effective in enhancing generalization performance in various settings, but also increases the computational overhead compared to base optimizers [3 14 7 Hence there have been several approaches to try to reduce the computational cost of SAM, such as ESAM which utilizes sharpness-ennsttie data selection and perturbs only a randomly selected fraction of parameters [20]. Related work shows that "only employing 20% of the batch to compute the gradients for the ascent step, [can] result in equivalent performance" [ and that only applying SAM part of the parameters (SSAM) using e.g. a Fisher-information mask can lead to enhanced performance 42 common variant of SAM utilizes m-sharpness [23), which uses subbatches of size m and benefits performance [20 2 4
though nuances in its implementation vary []. Andriushchenko and Flammarion [2] argue that its success is not unique to settings with BatchNorm and hence cannot be attributed the Ghost BatchNorm effect. Andriushchenko et al. 3 further show that SAM leads to low-rank features. We discuss different SAM variants in Section 3.2
3 Background: SAM and Normalization Layers
In this paper, we foeus the interplay between two popular aspects of neural network training, both of which we recapitulate here: In Sec. 3.1 we provide an overview normalization layers, in particular BatchNorm and LayerNorm, and in Sec. 3.2 of Sharpness- Aware Minimization variants.
3.1 BatchNorm and LayerNorm
Modern neural network architectures typically incorporate normalization layers. In this work we will focus on Batch Normalization (BatchNorm) 31 and Layer Normalization (LayerNorm) [ which
are an essential building block of most convolutional 25 29 and transformer [5I] 19 architectures, respectively. Normalization layers transform an input x according
x 1 + 3
N (x) = 7
where ja and o² are the mean and variance, which are computed over the batch dimension in the case of BatchNorm, over the embedding dimension, in the case of LayerNorm. BatchNorm is therefore sensitive to the choice of batch size [38, [4). J For BatchNorm, f and a are computed from the current batch-statistics during training, and running estimates are used test time. In our experiments, we focus on the trainable parameters and 3, which perform an affine transformation of the normalized input.
3.2 SAM and its variants
We recapitulate SAM (23), ASAM [37 and Fisher-SAM (35) with their respective perturbation models. To this end, we consider neural network fw RRd H which is parameterized (1),
yn)} consists of drawn from the data distribution D and we write the loss function as R x R -> R4. The goal is to learn a model fw with good generalization performance, i.e. low expected loss (w( = E(xxy)= [y,, fw (x)) on the distribution D. The training loss can be written as L(w) 1 Σn ((y., fx,,)) Conventional SGD-like optimization methods minimize (a regularized version of) L by stochastic gradient descent. SAM aims at additionally minimizing the worst-case sharpness of the training loss in a neighborhood defined by an Ep ball around w, i.e. |a L(w +c) L(w). This leads to the overall objective
by a vector w as our model. The training dataset Strain input-output pairs which
n
min max L(w + c). Se
(2)
In practice, SAM uses p- 2 and approximates the inner maximization by single gradient step, L(w)ll2 and requiring an additional Yorward-backward pass compared to SGD. The gradient is then re-evaluated at the perturbed point + giving the actual weight update
yielding e = PV L(w)/ / ////
V + +e)
with learning rate a. Computing e separately for the batch of each GPU in multi-GPU settings and then averaging the resulting perturbed gradients for the update step in Eq. 3 has been shown to
3

---

increase SAM's performance [23]. This method is called m-sharpness, with m being the number of samples on each GPU. Since the perturbation model Eq. 2 not invariant with respect to a partly scale-invariant version of
rescaling of the weights that leaves fw invariant (18). ASAM 37
SAM, was proposed, with the objective
min 1 S
max
L(w +))
where Tw is a normalization operator, making the perturbation adaptive to the scale of the network parameters. Kwon et al. [37] choose T, to be diagonal with entries Ty |,o,] + 11 for weight parameters and Th = 1 for bias parameters, called elementwise normalization. is typically set to 0.01. As with SAM, the inner maximization is solved by a single gradient step:
L(w) for =
Co.: = wi.giig V (L(()) for 0o.
(5)
1IT (
We note that for The equal the identity matrix and p 2, this is equivalent to the original SAM formulation. Recently, Kim et al. (35) proposed to use a distance metric induced by the Fisher information instead of Euclidean distance measure between parameters. The approach can also be framed as a variant of ASAM, with Tw being diagonal with entries Tü - 1/v1- and fa approximating the jth diagonal entry of the Fisher-matrix by the squared average batch-gradient, fi = (ww, Batch (w))?. For our experiments, we additionally employ layerwise normalization. This is, set the diagonal entries of Tü [ which corresponds normalization with respect to the Ez-norm of a layer, similar to (39
4 SAM-ON: Perturbing Only the Normalization Layers
We study the effect of applying SAM (and its variants) solely to the normalization layers of considered model. We will refer to this approach as SAM-ON (SAM-OnlyNorm) throughout this paper and provide a convergence analysis for SAM-ON in Appendix C We find that SAM-ON obtains enhanced generalization performance compared to conventional SAM (denoted as SAM-all) for ResNet architectures with BatchNorm (Section 4.1 and Vision Transformers with LayerNorm (Section 4.2) on CIFAR data and performs competitively on ImageNet. For comparison, we also study the reverse of SAM-ON, i.e. we exclude the affine normalization parameters from the adversarial SAM-step, which we shall refer to as no-norm.
Training set-up. We use SGD with momentum, weight decay, and cosine learning rate decay as our base optimizer for ResNet architectures and employ label smoothing to adopt similar settings as in the literature 37 For Vision Transformers we employ AdamW 40) as our base optimizer on CIFAR and for ImageNet we additionally use Lion [15]. We use both basic augmentations (random cropping and flipping) and strong augmentations AbusieeAutoAugmmnnt denoted as +AA). We consider a range of SAM-aaaiants which differ either in the perturbation model (f2 or too) or in the definition of the normalization operator. We train models for 200 epochs, and do not employ m-sharpness unless indicated otherwise. Complete training details are described in Appendix A
4.1 BatchNorm and ResNet
CIFAR. We showcase the effect of SAM-ON, i.e. only applying SAM the BatchNorm parameters, for a WideRessNet-2-110 (WRN-28) on CIFAR- in Figure ] We observe that SAM-ON obtains higher accuracy than conventional SAM (SAM-all) for all SAM variants considered (more SAM- variants are shown in Figure 6 in the Appendix). In contrast, excluding the affine BatchNorm parameters from the adversarial step (no-norm) either significantly decreases performance (for clementwise-E2 variants) or maintains similar performance as SAM-all (for all other variants). For variants which do not experience performance drop for no-norm, the ideal SAM perturbation radius shifts towards larger values, indicating that the perturbation model cannot perturb the BatchNorm parameters enough when all parameters are used. To study if the benefits of only perturbing the normalization layers extends to other settings, we train more ResNet-like models on CIFAR- 10 and CIFAR- For each SAM-variant and dataset, we probe a set of pre-defined p-values (shown in Table 7 in the Appendix) with ResNet-56 (RN-56) and fix the best-performing p for the other models to compare SAM-ON to SAM-all. We report mean accuracy and standard deviation over 3 seeds for 100 in Table [ On average, SAM-ON outperforms FAl-all for all considered

---

Table 1: SAM-ON improves over SAM-all for BatchNorm and ResNets:
Test accuracy for ResNet-
like models on CIFAR-100. Bold values mark the better performance between SAM-ON and SAM-all within a SAM-variant, and underline highlights the overall best method per model and augmentation.
RN-56
RNxT 55
WRN-28
[59
variant
all
ON
all
ON
all
ON
SGD
72.82
80.16
80.71
SAM
75.07=00.
75.58
81.79#0.4
82.22-0022 #0.2
83. +0.3
84.19
basic aug.
el. 82
75.05-60.1
76.25
81.26-10.2
82.30 t0.3
82.38
83.67-20.3
cl. l2, orig.
75.54=0.7
76.07
82.15 #0.3
81.90=0.4
83.67#0.1
83.53+00.2
el. 80
75.36+0.1
76.10
81.02+0.6
82.38 L0.3
83.25#0.2
84.14
Fisher
75.01
75.65
81.55
82.21 +0
83.37
84.01 #0.1
layer. l
74.63+0.1
76.03#0..3
81.66+002
82.52 +0.2
83.23+0.2
84.05
basic aug. + AA
SGD
75.26
80.31 EU:3
83.62=V.
SAM
76.33
76.02=t0.3
82.33
83.19 t0.2
85.30+00.1
=5.4
el. lg
76.51
76.0440.3
82.00
83.20 L0.1
84.80-0..3
85.43 +00.3
el. f2, orig.
76.49=0.2
76.58=00.4
82.78=0.1
82.87=0.3
85.25+004
85.41 #0.1
el. Loo
74.89-10.4
76.19-10.4
82.33
83.11 L0.2
85.28
85.46 10.1
Fisher
76.67=0011
76.25=+0.2
82.56
83.28
85.09+0.3
85.35
layer. l
76.23
76.93
82.61 #0.3
83.32 10.2
#0.2
85.32-00.3
85.95
SAM-variants. Of these, layerwise-lz achieves the highest performance for most settings. We obtain similar results for CIFAR-10 (App. Table and for more network architectures (App. Table IT].
ImageNet. For ImageNet, we adopt the timm training script [53]. We train a ResNet-50 for 100 epochs on eight 2080-Ti GPUs with m 64, leading an overall batch-size of 512. Apart from p. all hyperparameters are shared for all SAM-variants and can be found in the Appendix in Table We select the most promising SAM-ON variants and compare them against the established methods (SGD, SAM, ASAM elementwise (2). The results are shown in Table We observe that for layerwise ég. the all variant achieves higher accuracy, whereas the SAM-ON models outperform their all counterparts for elementwise & and clementwise foe. All SAM-ON variants outperform the previously established methods (SGD, SAM, ASAM). For reference, we also show the values reported for ESAM 20 and GSAM [61, two SAM-variants we did not include in our study.
Table 2: ImageNet top-1 accuracy for ResNet-50. ESAM 20 and GSAM 61 values taken
from the respective papers.
SGD
SAM
ESAMGSAM
elem. &
elem.
foo
layer &
all
all
all
all
ON
all
ON
all
ON
77.03
77.65
11
77.05
77.20
77.
00.05
77.82
77.45 0.04 77.82
00001
78.14
77.87
Varying the batchsize. In Figure (right) we report the performance of a WRN-28 on CIFAR- 100 with SAM-ON and SAM-all for a range of batch-sizes and values of m, where m is the batch-size per accelerator, as discussed in Section 3.2. Similar to the findings in 23 22 we confirm that lower values of m lead to better performance within each batch-size. Importantly, SAM-ON outperforms
SAM-all for all combinations of batch-size and m, illustrating that Ghost BatchNorm 28 49 (see discussion in Section 2) does not play a role in the success of SAM-ON.
4.2 LayerNorm and Vision Transformer
CIFAR. To study the effectiveness of SAM-ON beyond ResNet architectures and BatchNorm, we train ViTs from scratch on CIFAR data with AdamW as the base optimizer (Figure 2 Although ResNet architectures are known to outperform Vision Transformers when trained from scratch on scale datasets like CIFAR, our aim here is not to outperform vate-of-the-as but rather to study if the benefits of SAM-ON extend to substantially different training settings. Remarkably, we find that the same phenomena occur: The SAM-ON variants outperform their conventional counterparts SAM-al/ by a clear margin. For the elementwise-fy variants there is a strong drop in accuracy for no-norm, whereas for the other SAM variants the optimal perturbation radius / shifts towards larger values. We show that this extends to ViT-T and CIFAR-10 10 as well (Table 3)
5

---

71
ASAM, p=2.0, elem.
Test Accuracy (%)
ASAM, p=2.0, layer
69
ASAM, p-inf, elem.
68
66
all
65
ON
no-norm
64
10-3
10
10a
10
Figure 2: The interplay of normalization layers with SAM for VIT training VIT-S trained with different SAM-variants on CIFAR-100. Like for ResNets, perturbing only normalization layers (OnlyNorm) improves generalization performance, while omitting them in the perturbation (no-norm) can harm training.
Table 3: SAM-ON improves over SAM-all for LayerNorm and ViTs: Shown are ViT models on CIFAR-I0 and CIFAR-100. Bold values mark the better performance between SAM-ON and SAM-all within a SAM-variant, and underline highlights the overall best method per model.
CIFAR-I0
CIFAR-100
VII-T
VTT-S
VII-1
VIT-S
variant
all
ON
all
ON
all
ON
all
ON
Adam W
89.1920.2
90.34 90.81 92.71
63.79=0.0
64.372U2
SAM
88.92=0.0
92.22 +0.1
=0.
=0.0
64.70=#04
69.84
66.58+0.3
69.13
el. E2
90.02 92.52 10.3
92.40
93.97
10.2
Lo.1
66.98
70.31 10.2
basic aug.
el. E2, orig.
90.22+0.2 92.68 10.2
0.. 94.01 10.5
92.05
65.81=0.6
71.22
71.25
10.3
67.23
70.42 10.5
el. fxc
89.82=0.3 92.66 92.666 +.2
92.66-0-200.6--00
93.68=0.2
65.11=0.3
68.11
+0.7
66.55
69.19-0011
Fisher
89.08 92.03
91.13
92.49=0.1
64.70=0.4
69.55
66.59
69.30
layer. 8
89.51 93.08 93.00
40.
21.21
91.21
94.02
=0.1
65.30=0.6
69.66
+0.1
67.39
70.0
ImageNet. For ImageNet, train VIT-S/32 from scratch with batchsize 128 on single GPU for 300 epochs. In Table 4 we evaluate SAM-all, SAM-ON and the vanilla variant for both Adam and Lion [15] base optimizers both ID and OOD datasets. Since Wei et al. [52] showed that SAM-trained models show non-trivial robustness to small adversarial perturbations, we also investigate the robustness of SAM-ON (last rows of Table 4). Both SAM-all and SAM-ON improve strongly over the base optimizer in all setups. For Adam W. SAM-ON either outperforms SAM-all (e.g. adversarial robustness) performs on par (e.g. for ID accuracy the numbers within standard deviations). For Lion, SAM-ON always outperforms SAM-all, underlining that the diverse benefits of SAM be achieved provide all experimental details in in Appendix A.2 and a more thorough evaluation and discussion of the adversarial robustness of SAM-all and SAM-ON in Appendix B.4
surpassed by only perturbing the normalization layers. We
Table 4: SAM-ON performs well on ImageNet: Training ViT-S/32 from scratch.
AdamW
Lion SAM-all
vanilla
SAM-all
SAM-ON
vanilla
SAM-ON
ImageNet
66.89=00.04
71.47=0011
71.37=0.026
68.20
71.90+0.19
70.64
ImageNetV2
48.43#0.48
53.61 #0111
53.67#0.29
50.20
54.20+0.27
55.38
#0.09
ImageNetR
20.04
31.56=0.48
32.98
20.61
32.17=0.11
33.87
=0.47
adv. rob, OOD
ImageNetA
0.72=0.15
5.21
5.19=0.18
5.45=0.19
5.01
5.77
10.21
ImageNetSketch
13.68:0.24
18.5040.44
19.35
14.47
18.22
20.48
10.12
ObjectNet
11.32+0..99
13.75#0.12
13.55=+0.25
12.06
13.93=t0.40
15.35 10.13
E,, = 0.25
19.6740.47
37.53
41.16
10124
22.01
38.5240.66
43.12
20.97
f, é 0.50
5.47=0.11
17.71 #0.61
22.72
6.63=0.16
19.03+0.92
24.27
loc. 025
33.45
18.0810.14
49.34
10.08
35.31 10.08
49.57
51.37 10.99
loc. E 0.5/255
14.98 10.18
29 29.68 10.09
32.46
+0.15
10.1
31.68
34.23 1.73

---

ResNext
85
83.0
SAM-ON B2+82
SAM-ON
84
SAM-ON
Accuracy (%)
SAM-ON
SAM=Ill
82.0
SAM-all
Accuracy
81.5
SAM-ON
82
81.0
SAI/-ON
SAM-ON
80.5
81
SAM-ON B1
SAM-ON B1
80.0
1.0
1.2
1.4
1.6
1.8
2.0
1.0
1.2
1.4
1.6
1.6
2.0
Runtime (normalized by SGD time)
Runtime (normalized time)
Test accuracy vs. normalized wall-clock runtime for SAM and different variations of SAM-ON for WRN-28 (left) and ResNeXt (right) on CIFAR-100 100. Only perturbing selected normalization parameters (e.g. those from block 3, those from block 2 and 3) can lead to further computational gains. Reported values are averaged over three random seeds.
Figure 3: Computational gains of SAM-ON over SAM-all:
4.3 Computational savings
In Figure 3 we report the wall-clock time of training a WRN-28 (left) and a ResNeXt (right) with batchsize I28 on a single A100 with PyTorch. Since the normalization parameters at the earlier layers of the network require gradient, potentially full backpropagation pass has be computed for the ascent-step of SAM-ON, even though only tiny fraction of all parameters is perturbed. However, as discussed [22 for a related setting, the gradients of the intermediate (no-norm) layers do need to be stored or used for updating. This leads to computational gains of SAM-ON over SAM-all as reported in Figure
In contrast, although future development of hardware for sparse operation may allow for acceleration, Mi al. 42 report that their sparse perturbation approach does not at present lead reduced wall-clock time. Their approach also suffers from additional computational cost associated with selecting the mask, and hence is outshined by SAM-ON both in computational cost and generalization performance (as discussed in Section 5.1
We also show results for perturbing only the normalization layers of selected blocks (Block 1-3) of the network, and interestingly the main benefits of SAM-ON seem to arise from the later normalization layers, allowing for further computational savings. When perturbing only the normalization layers from the last block (B3), the biggest computational gains can be achieved (reducing the additional cost of SAM by more than 50%), without much loss in accuracy. When perturbing the normalization layers from Block and 3 (B2+B3), the test accuracy even slightly improves over SAM-ON for both models, while the runtime is still significantly lower. We have not investigated this variant of SAM-ON thoroughly, i.e. in combination with other ASAM perturbation models and more network architectures, but think that this is interesting research direction for future work.
5 Towards Understanding SAM-ON
To gain better understanding of SAM-ON, we study different hypotheses for the method's 's success. First, we investigate the role of sparsity by comparing SAM-ON to different sparsified perturbation approaches (Section (.I)). concluding that sparsity alone is not enough explain its success. Then, we highlight that SAM-ON might in fact find sharper minima, while generalizing better than SAM-all (Section (.2)). We also show that SAM-ON similar SAM-all reduces the feature-rank compared to vanilla optimizers (Section 5.2). Further, we showcase that depending on the perturbation method, SAM-ON can induce a significant shift the distribution of the normalization parameters (Section 5.3) and relate this to the no-norm results from Section 4 Unless stated otherwise, we use WRN-28 with BatchNorm and the setting described in Section for the ablation studies. Further ablation
studies are presented in Appendix B
7

---

all
96.25-
Test Test Accuracy (%)
84
ON
Test Accuracy
96,00
83
95.75
95.50
82
bs 128
95.25
SGD
www.............................
SSAM 50%
SAM-ON SAM-rand
bs 256
95.00
SAM
SSAM 99.93%
b5 512
10
10
10
16a
10a
10
(with different sparsity levels) and random mask SAM-rand (same sparsity level 99.93% as SAM-ON) sparse perturbation approaches on CIFAR- 10 for 18.. Right: SAM-ON improves over SAM for various batch-sizes (bs) and values of m, where smaller values m tend to improve performance. WRN-28, CIFAR-100.
Figure 4: Left: SAM-ON outperforms SSAM-F 42
5.1 The effect of sparsified perturbations
Mi et al. 42 propose a sparsified SAM (SSAM) approach which only considers part of the parameters for the SAM perturbation step. They find that using 50% perturbation sparsity SSAM can outperform SAM-all, and still perform competitively with up to 99% sparsity in certain settings. This raises the question whether the enhanced performance of SAM-ON over SAM-all simply due perturbing fewer parameters. We therefore compare SAM-ON to other sparse perturbation approaches.
In order to determine the parameters which should be perturbed, is to compute a binary mask via an approximation of the Fisher matrix. We call this approach
solution Mi et al. [42 propose
SSAM-F avoid confusion with Fisher-SAM introduced in Section 3.2 Mi et al. [42] also propose a dynamic mask sparse perturbation approach (SSAM-D), but do not clearly favour either method. Since they consider SSAM-F 'reltvvrrll more stable but bit time-consuming, while [SSAM-D] is more efficient", we will for fair comparison focus on SSAM-F here and provide results for SSAM-D in Appendix B.2 According to Mi et al. [42) neither approach improves wall-clock time in practice. while we find that SAM-ON does (see Section
We provide comparison between SAM-ON, SSAM-F, and random sparsity mask of the same sparsity level as SAM-ON for a) ResNet- 18 on CIFAR-10 (main setting considered in 422 in Figure 4 (left) and b) WRN-28 CIFAR-100 in Table 5 on par or even outperform SAM-all at the medium to high sparsity levels recommended in [42 is less successful than SAM-ON. Moreover, for the sparsity levels of SAM-ON, which are above 99.9% in both settings, a random mask performs poorly. These results suggest that sparsity not the sole cause for SAM-ON's success.
We find that although SSAM-F can indeed perform
Table SAM-ON outperforms other sparse perturbation approaches: Although SSAM-F with different sparsity levels can outperform SAM-all CIFAR 100 with WRN-28, it is less effective than SAM-ON, especially when probed very high sparsity levels.
SAM 0%
SAM-ON
Random Mask
SSAM-F
Sparsity
99.95%
99.95% 80.97-10.2
50%
99.95%
Test Accuracy (%)
83.11
84.19.102
83.94
83.14
5.2
Sharpness and feature-rank of SAM-ON
The improved generalization performance of SAM-trained models is often attributed to finding flatter minima [27 indeed this was the initial motivation behind SAM in (23. Andriushchenko and Flammarion however cast doubt on this explanation, and argue that the benefits of SAM might stem from a favorable implicit bias induced by the method. A recent study furthermore found that sharpness often does not correlate well with model's generalization performance (4) In this section we therefore compare the sharpness of SAM-all SAM-ON models (following the setup 4 more details provided in Appendix B.10).

---

In Table 6 we report logit- gegiitnormalzzed worst-case adaptive foc m-sharpness, the sharpness definition which achieved the best correlation with generalization error for CIFAR data in the large-scale study in A. We investigate WRN-28 on CIFAR-100 trained with SGD, SAM, and ASAM- both with their respective all and ON variants.
Worst-case sharpness is measured with 20 steps of AutoPGD [16]. a hyperparameeter-free method designed for accurate estimation of adversarial robustness. Instead of the input space, which is opti- mized in adversarial robustness, the method is adapted to the weight space for sharpness computation. In addition to the 20-step SeooPGD-sharpnesss we also report 1-step sharpness, i.e. the change in loss obtained when performing only a single gradient step, like in the perturbation step of SAM. It is to note that except for the nogit-normalization, the sharpness definition reported in Table 6 corresponds exactly to the perturbation model that ASAM elementwise &o uses, and hence the 1 1-step sharpness reported should be fairly close the the objective that ASAM elementwise for actually minimizes during training. For both sharpness definitions, we pick / such that we get sharpness values similar to those that lead to good correlation with generalization in 4 This is, for 1-step sharpness we have to pick p larger than for the 20-step sharpness to get to similar loss changes.
We observe that the SAM-ON models obtain the best generalization performance, while having significantly higher sharpness than their SAM-all counterparts, and sometimes even higher than that of the baseline teDinrained model. This supports the claims of [2] that although SAM-all might in fact find flatter minima, we should be careful in attributing this as the sole reason for its improved generalization performance. In Appendix [010 we report more sharpness metrics and find that SAM-ON is sharper than SAM-all for most metrics, although there exist exceptions.
A recent study by Andriushchenko al. further showed that SAM-all leads features of lower rank compared to vanilla optimizers. Following their setup we measure the feature rank for WRN-28 and report our findings in Table (final row). We find that SAM-ON also leads features of lower rank compared to SGD, but observe no significant change compared to SAM-all.
Table 6: SAM-ON models are sharper, yet generalize better. Shown is oggitnnormalieed loc. m-ssaappeess from [4] averaged over three models per method for a WRN-28 on CIFAR-100, and the feature rank like investigated in 3 (last row).
SGD
SAM
ASAM-el.-1
all
ON
all
ON
Test Accuracy (%)
30.71=10.2
83.11 10.3
84.19 #0.2
83.25
84.14
20 steps, p = 0.003
0.071 =0.00
0.048
10.001
00000
0.048
10.001
0.078=00000
Shappn pnepe
20 steps, p = 0.005
0.201 +0.001
0.139+0.004
+0.004
0.124
0..002 0.002
20 steps, p 0.007
0.4433+0.000
0.309
0.005
0.255
0.255=0.000
0.580
I step. P 0.007
0.117 1U.02
0.098
0.008
1.170
0.095
0.142
8 1 step. = 0.01
0.204 0.204-10.005 100000
0.183 0000
0.315 10.000
0.170=0.00
.003
0.271 1:00019
1 step. = 0.03
0.809=00.00
0.769 +0.017
0.843=#0.007
0.724 0.000
4..012
Feature Rank
4004=12
17
3728 =11
3613 =17
37254111
5.3 SAM-ON can change the affine parameter values
In Figure 5 show the distribution of the scaling parameter 7 and the shift parameter 3 (as defined in Eq. [ for a WRN-28 trained with SAM-ON and SAM-all. While there are only minor changes in the distribution of B. there is a clear shift in the distribution of y towards larger values for SAM-ON. In Figure in Appendix [B.6] we study more SAM-variants and rediscover a pattern from Section 4.1 The distribution shift for y only seems to appear for those variants, where the optimal p of SAM-ON shifts towards larger values compared to SAM-all (i.e. not for ASAM elementwise-f.) We thus hypothesize that the perturbation model of ASAM elementwise- E2 implicitly focuses mostly on perturbing the normalization layers already, which is why excluding them (no-norm) leads to drastic performance decrease, whereas SAM-ON minimally changes the 7-distribution. the other SAM-variants (for which the optimal p-value changes) may not perturb the normalization layers enough in SAM-all for them to be effective, which why no-norm has little effect, while SAM-ON leads to distinctive shift of the ddiitribbtiion.
contrast,
9

---

Distribution of
Distribution B
SAM p=2 SAM
500
500
all
250
250
-02
-0.1
0.0
0.1
0.2
0.3
-0.2
-0.1
0.1
0.3
Value
Value
Figure 5: SAM-ON changes the weight distribution of the normalization layer weights of a WRN-28 towards larger values. More SAM-variants are shown in Figure iin the appendix.
Discussion and Conclusion
In recent years the method of sharpness-aware minimization (SAM) has risen to prominence due to its demonstrated effectiveness and the community's 's Long-standing interest in flat minima. In this work we show that only applying SAM the normalization layers (SAMONN -ypiccalyy less than 0.1% of the total parameters can significantly enhance performance compared to regular SAM (SAM-all). We show results on CIFAR and ImageNet for ResNet and Vision Transformers with BarchNorm and Layernorm, respectively. Although the use of sparsified perturbations was recently shown to benefit generalization [42], we show that the success of SAM-ON cannot be attributed to sparsity alone: targeting the normalization layers clearly improves over other masked sparsity approaches.
We find that while SAM-ON outperforms SAM-all in almost all settings, the optimal SAM variant use varies. We do not see consistent benefit of using reparameterizzaiooniivvaa invariinnn perturbation models compared to variants with fewer invariances. In particular, we find that layerwise E in combination with SAM-ON reaches the highest accuracy for many settings.
While SAM-ON improves generalization, it loses some of SAM's sharpness-reducing qualities. Although perhaps surprising given SAM's original motivation, this finding relates naturally the literature. [4) find that sharpness does not always correlate well with generalization performance. Further, 2 question if sharpness is the sole driving factor behind SAM's success in enhancing generalization. We lend support to this question, showing that SAM-like methods can generalize better, without significant sharpness reduction.
In summary, we demonstrate benefits of SAM beyond reducing sharpness and highlight the special role played by the normalization layers. More investigation into the interplay of SAM and other aspects of training are needed to fully understand where the methods gains come from.
Limitations. Similar to the main inspirations for this work 23 provide a simple experiment in the language domain in Appendix B.5 of SAM-ON might be preserved for language tasks, but more extensive benchmarking, by 7 required. Further, more work is required to try to leverage the perturbation sparsity for reduced computational cost beyond the gains obtained in this work, and to fully explore the benefits of perturbing subsets of the normalization layers building on the findings in Section 4.3
24
we focus on vision data. We indicating that the efficacy
done
Acknowledgements
We acknowledge support from the Deutsche Forschungsgemeinschattt (DFG, German Research Foundation) under Germany's Excellence Strategy (EXC number 2064/1, Project number 390727645), as well as from the Carl Zeiss Foundation in the project "Certification and Foundations of Safe Machine Learning Systems in Healthcare". We also thank the European Laboratory for Learning and Intelligent Systems (ELLIS) for supporting Maximilian Müller. We are grateful for support from Canada CIFAR Chairs Program and US National Science Foundation award tel: 19860 In addition, we acknowledge material support from NVIDIA and Intel in the form of computational resources and are grateful for technical support from the Mila IDT team maintaining the Mila Compute Cluster.
10

---

References
[1] M. Abbas, Q. Xiao, L. Chen, P.-Y. Chen, and T. Chen. model-agnostic meta learning. ICML, 2022.
Sharp-MAML: Sharpness-aware
[2] M. Andriushchenko and N. Flammarion. Towards understanding sharpness-aware minimization. ICML, 2022.
[3] M. Andriushchenko, D. Bahri, H. Mobahi, and N. Flammarion. Sharpness-aware minimization leads to low-rank features. NeurIPS, 2023.
[4] M. Andriushchenko, F. Croce, M. Müller, M. Hein, and N. Flammarion. A modern look the relationship between sharpness and generalization, preprint (rXiv:2332.07011 2023.
[5] S. Arora, Z. Li, and K. Lyu. Theoretical analysis of auto rate-tuning by Batch Normalization. ICLR, 2019.
[6] L. Ba, R. Kiros, and G. Hinton. Layer Normalization. Advances in NeurlPS Deep Learning Symposium, 2016.
[7] D. Bahri, H. Mobahi, and Y. Tay. Sharpness-awart minimization improves language model generalization. ACI,, 2022.
[8] K. Behdin, Q. Song, A. Gupta, A. Acharya, D. Durfee, Ocejo, S. Keerthi, and R. Mazumder. mSAM: Micro-hatch-averaged sharpness-aware minimization. preprint arXiv:2302.0906 2023.
[9] Benz, C. Zhang, and Kweon. Batch normalization increases adversarial vulnerability and decreases adversarial transferability: A non-robust feature perspective. ICCV, 2021.
[10] J. Bjorck, C. Gomes, B. Selman, and K. Q. Weinberger. Understanding batch normalization. NeurlPS, 2018.
[11] A. Brock, De, S. L. Smith, and K. Simonyan. High-performance large-scale image recognition without normalization. ICML, 2021.
[12] B. Chen, P. Li, B. Li, C. Lin, C. Li, M. Sun, J. Yan, and W. Ouyang. BN-NAS: Neural architecture search with Batch Normalization. ICCV, 2021.
[13] M. X. Chen, O. Firat, A. Bapna, M. Johnson, W. Macherey, G. Foster, L. Jones, M. Schuster, N. Shazeer, N. Parmar, A. Vaswani, J. Uszkoreit, L. Kaiser, Z. Chen, Y. Wu, and M. Hughes. The best of both worlds: Combining recent advances in neural machine translation. ACL, 2018.
[14] X. Chen, C.-J. Hsich, and B. Gong. When Vision Transformers outperform ResNets without pre-training strong data augmentations. ICLR, 2022.
[15] X. Chen, C. Liang, D. Huang, F. Real, K. Wang, Y. Liu, H. Pham, X. Dong, T. Luong, C.-J. Hsieh, Y. Lu, and Q. V. Le. Symbolic discovery of optimization algorithms. arXivv2302.06677, 2023.
[16] F. Croce and M. Hein. Reliable evaluation of adversarial robustness with ensemble of diverse parameter-free attacks. ICML, 2020.
[17] N. Dimitriou and O. Arandjelovic. A new look at ghost normalization. preprint arXiv: 2007.08554, 2020.
[18] L. Dinh, R. Pascanu, S. Bengio, and Y. Bengio. Sharp minima can generalize for deep nets. ICML, 2017.
[19] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, Gelly, Uszkoreit, and N. Houlsby. An image worth 16x16 words: Transformers for image recognition at scale. ICLR, 2021.
[20] J. Du, H. Yan, J. Feng, J. T. Zhou, L. Zhen, R. S. M. Goh, and V. Tan. Efficient sharpness aware minimization for improved training of neural networks. ICLR, 2022.
[21] G. K. Dziugaite and D. M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. Uncertainty in AI, 2017.
[22] G. K. Dziugaite, A. Drouin, B. Neal, N. Rajkumar, E. Caballero, L. Wang, I. Mitliagkas, and
D. M. Roy. In search of robust measures of generalization. NeurIPS, 2020.

---

[23] P. Foret, A. Kleiner, H. Mobahi, and B. Neyshabur. Sharpness improving generalization. ICLR, 2021.
Swaree aware minimization for efficiently
[24] J. Frankle, D. J. Schwab, and A. S. Morcos. Training BatchNorm and only BatchNorm: On the expressive power of random features in CNNs. ICLR, 2021.
[25] K. He, X. Zhang, Ren, and Sun. Deep residual learning for image recognition. CVPR, 2016.
[26] S. Hochreiter and J. Schmidhuber. Simplifying neural nets by discovering flat minima. Advances in Neural Information Processing Systems, 1994.
[27] S. Hochreiter and J. Schmidhuber, Flat minima. Neural Computation, 9(1): 42, 1997.
[28] E. Hoffer, I. Hubara, and D. Soudry. Train longer, generalize better: closing the generalization gap in large batch training of neural networks. NeurlPS, 2017.
[29] G. Huang, Z. Liu, L. van der Maaten, and K. Q. Weinberger. Densely connected convolutional networks. CVPR, 2017.
[30] G. Huang, Z. Liu, L. van der Maaten, and K. Q. Weinberger. Densely connected convolutional networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2017.
[31] S. loffe and C. Szegedy. Batch Normalization: Accelerating deep network training by reducing internal covariate shift. ICML. 2015.
[32] Y. Jiang B. Neyshabur*, H. Mobahi, D. Krishnan, and S. Bengio. Fantastic generalization measures and where to find them. ICLR, 2020.
[33] Kaddour, L. Liu, Silva, and M. J. Kusner. When do flat minima optimizers work? NeurIPS, 2022.
[34] N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang. On large-batch training for deep learning: Generalization gap and sharp minima. ICLR, 2017.
[35] M. Kim, D. Li, S. X. Hu, and T. Hospedales. Fisher SAM: Information geometry and sharpness aware minimisation. ICML, 2022.
[36] J. Kohler, H. Daneshmand, A. Lucchi, M. Zhou, K. Neymeyr, and T. Hofmann. Exponential convergence rates for Batch Normalization: The power of length-direction decoupling in non-convex optimization. AISTATS, 2019.
[37] J. Kwon, J. Kim, H. Park, and I. K. Choi. ASAM: Adaptive sharpness-awarr minimization for scale-invariant learning of deep neural networks. ICML, 2021.
[38] X. Lian and J. Liu. Revisit Batch Normalization: New understanding and refinement via composition optimization. AISTATS, 2019.
[39] Y. Liu, S. Mai, X. Chen, C.-J. Hsich, and You. Towards efficient and scalable sharpness- aware minimization. CVPR, 2022.
[40] I. Loshchilov and F. Hutter. Decoupled weight decay regularization, 2019.
[41] K. Lyu, Z. Li, and S. Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. NeurlPS, 2022.
[42] Mi, L. Shen, T. Ren, Y. Zhou, X. Sun, R. Ji, and D. Tao. Make sharpness-aware minimization stronger: A sparsified perturbation approach. NeurIPS, 2022.
[43] T. Möllenhoff and M. E. Khan. SAM as an optimal relaxation of Bayes. ICLR, 2023.
[44] C. Na, S. V. Mchta, and E. Strubell. Train flat, then compress: Sharpness-aware minimization learns more compressible models. EMNLP, 2022.
[45] H. Petzka, M. Kamp. L. Adilova, C. Sminchisescu, and M. Boley. Relative flatness and generalization. NeurIPS, 2021.
[46] S. Santurkar, D. Tsipras, A. Ilyas, and A. Madry. How does batch normalization help optimiza- tion? NeurIPS, 2018.
[47] K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image recognition, 2015.
[48] A. Steiner, A. Kolesnikov, X. Zhai, R. Wightman, J. Uszkoreit, and Beyer. How to train your VIT? Data, augmentation, and regularization in Vision Transformers. TMLR, 2022.
12

---

[49] C. Summers and M. J. Dinneen. Four things everyone should know to improve batch normaliza- tion. ICLR, 2020.
[50] C. Szegedy, W. Zaremba, I. Sutskever, J. Bruna, D. Erhan, Goodfellow, and R. Fergus. Intriguing properties of neural networks. ICLR, 2014.
[51] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Kaiser, and
I. Polosukhin. Attention is all you need. NeurIPS, 2017.
[52] Z. Wei, J. Zhu, and Y. Zhang. Sharpness-awe aware minimization alone can improve adversarial robustness. ICML Fromtirro Workshop, arXiv:2305.0553 2023.
[53] R.
Wightman. pytorch- -aage-modde.. 2019.
Pytorch image models.
httt::///gitubb..o//rrrgggg/
[54] C. Xie and A. Yuille. Intriguing properties of adversarial training at scale. ICLR, 2020.
[55] Xie, Girshick, Dollár, Tu, and K. He. Aggregated residual transformations for deep
neural networks. CVPR, 2017.
[56] J. Xu, X. Sun, Z. Zhang, G. Zhao, and J. Lin. Understanding and improving layer normalization. NeurIPS, 2019.
[57] Z. Yao, A. Gholami, K. Keutzer, and M. Mahoney, PyHessian: Neural networks through the lens of the Hessian. IEEE BigData. arXiv:1912.07145, 2020.
[58] M. Yazdanpanah, A. A. Rahman, M. Chaudhary, Desrosiers, M. Havaei, E. Belilovsky, and
S. E. Kahou. Revisiting learnable affines for Batch Norm in few-shot transfer learning. CVPR,
2022.
[59] Zagoruyko and N. Komodakis. Wide residual networks. In BMVC, 2016.
[60] H. Zhang, Y. N. Dauphin, and T. Ma. Fixup initialization: Residual learning without normaliza- tion. ICLR, 2019.
[61] Zhuang, B. Gong, L. Yuan, Cui, H. Adam, N. Dvorekk S. Tatikonda, J. Duncan, and Liu. Surrogate gap minimization improves sharpness-aware training. ICLR, 2022.
13

---

Appendix
This appendix structured as follows:
In Appendix A we provide more training details. In particular, we report the hyperparameters used for the CIFAR experiments in[A.] JA.JJam and for the ImageNet experiments in A.2 A..3
we provide more details and a formal definition of the SAM- MAM-variants used throughout this paper.
In Appendix we show additional experimental results for: CIFAR in B.1 ImageNet B33 and a machine translation task in B.5 In B.2 we provide additional ablation studies for sparse perturbation SSAM approaches and in B. we extend the discussion on adversarial robustness. To gain a better understanding of SAM-ON, we further investigate: the weight distribution shift induced by SAM-ON (B.6). the effect of SAM when fixing the normalization parameters during training [B.7p. p. SAM's 's performance when only training the normalization layers (B.8). and ablations on weight decay and dropout (B.9). Finally, we provide an extended discussion
the sharpness evaluation and more ablations in B.10
In Appendix C we provide convergence analysis for SAM-ON.
A Training Details
A.1 CIFAR training details
For our CIFAR experiments, we consider a range of SAM-variants which differ either in the norm (p (2, xo)) or in the definition of the normalization operator. We use SGD, the original SAM with no normalization and 2, Fisher-SAM and the following ASAM-variants: elementwise- lx. layerwise-22, and clementwise-fz. For the ViT-experiments, we use AdamW instead SGD. For each of the ASAM-variants, we normalize both bias and weight parameters and set 7 0. Additionally, we employ the original ASAM-algorithe and 17 0.01. We train all models (unless indicated otherwise). For ResNets, we follow [37 and adopt learning rate of 0.1, momentum of 0.9, weight decay of 0.0005 and use label smoothing with factor of 0.1. We use both basic augmentations (random cropping and flipping) and strong augmentations ViTss AdamW with learning rate 0.0001, batchsize 64 and only strong augmentations, the other settings remain unchanged. The ResNet results were computed on 2080ti-GPUs and the ViT results on A100s. The values of p we considered for each method can be found in Table 7 The ResNet-networks we considered for the CIFAR-experiimnts in the main paper are ResNet56 (RN56) 25], ResNeeX--29-5x4d (RNxT) (55), and WideResNet-28- (WRN) (59 We adopted the ViTs CIFAR by setting the image-size to 32 and patch-size to 4.
where the bias parameters are not normalized single GPU for 200 epochs, and m- m-sharpness is not employed
Auto-Amentt For
Table 7: Search-space for The values used for the the experiments in Tables 13 and 10 are marked in bold.
CIFAR-IO
CIFAR-100 RN
CIFAR-10/100
SAM
all
0.05, 0.1. 0.25
0.05. 0.1. 0.5. 1.
0.025, 0.05, 0.1, 0.25, 0.5
SAM
ON
0.1, 0.5.
0.1, 1., 0.5, 2,5,5 10.
1,, 5., 10., 25
el. f2
all
0.5, 1, 2,
1,,
el. la
ON
0.5. 1, 2.
0.5. 2.5, 5., 10.
2.5,
el. l2. orig.
all
0.1, 0.5, 1, 5,
0.5, 2,55,
0.5, 1., 2.5.
el. 12, orig.
ON
0.1, 0.5, 1, 10
0.5. 2.5,, 2.5.
2.,,
el. lx
all
0.001, 0.005, 0.01, 0.05
0.001, 0.005, 0.01, 0.05
0.0005, 0.001, 0.0025, 0.005, 0.01
el. bx
ON
0.01, 0.025, 0.05, 0.1
0.01, 0.05, 0.1, 0.5
0.025, 0.05, 0.1, 0.25, 0.5
layer lg
all
0.01, 0.025, 0.05, 0.1 0.005,
0.001, 0.01. 0.05, 0.1
0.001, 0.0025, 0.005, 0.01, 0.025
layer la
ON
0.05, 0.1, 0.25, 0.5,
0.1, 0.2, 0.5, 1.
0.05, 0.1, 0.25, 0.5,
Fisher
all
0.05, 0.1. 0.5, 1.5
0.05, 0.1.
0.05,5 0.1. 0.5. 1.5 0.1, 0.5.1 .5.
Fisher
ON
0.5.1 10
0.1. 0.5.1 10
14

---

A.2 ImageNet training details
Table 8 shows the hyperparameters for all variants used for ImageNet training. For the ResNet-50 with SGD, SAM and clementwise-f2 we used the hyperparameters from [23] and [37]. For the layerwise &g and clementwise-(>> better one (named p (reported) the table). p (discarded) refers to the p value we probed, but found to perform worse than the other one. For the ViT-S (additional fine-tuning experiments in Appendix B.3 we tried at least three values of per SAn-coniiurrttion and reported the best one.
tried two p-values per configuration and report the results of the
Table 8: Hyperparameters for training on ImageNet. Top: ResNet-50 from scratch, center: VIT-S from scratch, bottom: finetuning the ViT-S.
param
SGD SAM elem. 62
ResNet-50
elem. foc
layer
all
all
all onlyNorm
all onlyNorm
all
onlyNorm
train epochs warm-up epochs
90
3
cool-down epochs batch-size augmentation
10
512
inception-style 0.2 Cosine 0.0001
Ir Ir decay
weight decay p (reported)
0.05
0.001
0.005 0.05
0.005 0.05
0.05
p (discarded)
0.01
0.5
Input Resolution m GPU Type
224 x 224
64 8 2080-ti
param
AdamW
AdamW4SAM
VIT-S scratch
Lion all
Lion+SAM
all
all
onlyNorm
all
onlyNorm
train epochs
300
warm-up epochs
10 0 128
cool-down epochs
batch-size augmentation
inception-style
Ir decay weight decay p (reported)
0.001 Cosine 0.1
15 10,20
1
10
p (discarded)
0.05,0.1.0..2
0.5,2
5,20
Input Resolution m GPU Type
224 x 224
128
1xA100
param
SGD
SAM onlyNorm
elem. & onlyNorm
VIT-S FT
elem. too all onlyNorm
layer & onlyNorm
all
all
all
all
train epochs
warm-up epochs
cool-down epochs batch-size augmentation
0
896 inception-style
Ir Ir decay weight decay
0.017
Cosine
0.0001
P (reported)
0.01
0.1
0.1
44
10
10
10-3
(discarded) p (discarded)
0.1
100
0.01
0.1
10
10-3
10-2
10-2
0.001
10
10
10-1
10-4
10-1
Input Resolution m GPU Type
224 x 224
128
7 A100
15

---

A.3 SAM variants
Here, we provide more comprehensive overview of the SAM-variants used throughout the exper- iments. To this end, we first recall the definition of the -ertr--pettrrationn (Eq. (5) in the main paper):
T (w) 62 VIT V (w)ll2
for 2,
Eoc = pTl (VL())))))))))))))))))))) for
with the normalization operator T which is diagonal for all variants. We note that SAM-ON can be formally defined as using the conventional (AloAM-aloorithm but setting all entries Tü = if w not normalization parameter. This leads to change of the perturbation according to Eq. (5) Importantly, the magnitude e is still p. since both the nominator and the denominator Eq. change. We provide overview over all (A)sAM-variants and their respective perturbation models
in Table
Table 9: The definition of T for the considered SAM-variants.
variant
T
p
all
2
SAM
if uy is normalization parameter
ON
2
clse
all
|waa [uss
2
0
el. l
if uy normalization parameter
ON
2
+ r
uy a weight parameter
0.01
el. E2, orig.
++1 r
if wy a bias parameter
|uvl
if uu is normalization weight
ON
1 0
if uy is normalization bias clse
2
0.01
cl. 600
00
[wy
if wi is a normalization parameter
ON
00
else
all
bayerfil lla
2
layer lz
W mngatullz
if w; is normalization parameter else
ON
2
all
+ v Lnason + (0w, I.Boooh (w))²)
2
Fisher
if my a normalization parameter else
ON
B
Further Experimental Results
B.1 SAM-ON CIFAR
omitted the results for ResNet-like models on CIFAR-10 in the main paper. Those are thus reported in Table 10 Due the already very high accuracies, the differences between SAM-ON and SAM-all are smaller, yet on average SAM-ON is still clearly the better method. We further plot all considered SAM-variants for different values of p in Figure for WRN-28 and in Figure 7 for ViT-S on CIFAR- 100. We show results for various VGG-models [47] and DenseNet- [0 for CIFAR-10/100 in Table TT and observe that SAM-ON consistently improves over SAM-all.
B.2
Additional ablation studies for sparse SAM
In this section we provide additional ablation studies sparsified perturbation approaches as discussed in Section 5.1. Mi et al. 42 proposed two sparsified SAM (SSAM) approaches: Fisher SSAM (SSAM-F) and Dynamic SSAM (SSAM-D). As an extension Figure for ResNet- on CIFAR- 10 data in the main paper provide accompanying Figure g which includes error bars
16

---

Table 10: SAM-ON improves over SAM-all for BatchNorm and ResNets on CIFAR-10: Test
accuracy for ResNet-like models on CIFAR-10. Bold values mark the better performance between SAM-ON and SAM-all within SAM-variant, and underline highlights the overall best method per model and augmentation
SAM variant
RN-56
RNxT
WRN-28
all
onlyNorm
all
onlyNorm
all
onlyNorm
SGD SAM
94.28
95.37
40.1
96.20
94.94
95.18 10.1
96.35 40.2
96.48 L0.1
97. 97.08-10.1
97.10 10.0
basic aug.
elem. E
94.96 1:0.1
94.94+0022
96.41 #0.1
96.53
96.98+002
97.06 $0.0
elem. E2, orig.
95.14
95.21 #0.1
=..2
96.40=0.1
96.41
97.10=0.1
97.07=0.1
elem. lo
94.93-0.11
94.96
96.06
96.22
96.95 40.2
97.00 10.1
Fisher
95.01 10.1
95.03 10.1
96.31 20.0
96.55 10.00
96.95 1..0
97.13 L0.1
layer. &
94.95+002
95.07=
96.07=0.3
96.46 t0.1
97.02
96.96-90.1
SGD
94.70=001
96. 19
97.01 20.0
basic aug. + AA
SAM
95.25+0011
95,40 =0.1
96.98=0.1
97.22 +0.3
97.57
97.58*0.0
elem. &
95.12 1:0.0
94.82 10.2
97.01 10.0
97.21 1.0.1
97. 97.61 4:0.0 4:0.0
97.69 L0.0
elem. E2, orig.
95.39
95.60 10011
97.24=0.0
97.33 1:0.1
97.60 0.00
97. 97.56=*0.0
elem. lo0
12
95.48-
96.70=0.2
96.91 +00.
97.52=00.1
97.62#0011
Fisher
19
95.38
96.7740.0
97.24
97.53-40.0
97.65 20.11
layer. l
95.43 10..3
95.28
96.80
96.88 L.0.1
97.60 0.0
97.48
Table 11: SAM-ON improves over SAM-all for BatchNorm and more ResNet models: Bold values mark the better performance between SAM-ON and SAM-all within SAM-variant, and underline highlights the overall best method per model and augmentation
VGG-13
VGG-16
VGG-19
DenseNet-100
SAM variant
all
onlyNorm
all
oalyNorm
all
onlyNorm
all
only
SGD
21.02
74.4
73.40=02
77.00
SAM
76.74=0.2
77.57
75.81
76.86
74.08=066
75.60
79. 79.42
79.90
CIFAR- 100
elem.
76.65
22.4
75.95
76.45
74.7220.
75.12
78.90:9.2
79.83 $3.3
elem. 1 0.01
77. 77.27+002
77.37 t.0.2
76.65
76.66
75.00=0.5
75.44
79.94
80.14
elem. &x
76. 82
77.62 =002
75.43
76.68 +0.1
72,71=02
74.50*004
79.17
79.64
Fisher, 1.
76.76
77.68 10.1
75.85
76.99
74.03
74.96 1.033
79.68
80.38 10.5
layer. &
76.76+002
77.91
75.99+002
77.12 t0.2
74.65=0.s
75.28 6.22
78. 78.25
79.86 50.3
SGD
94.29
93.85
93.82
94.51
SAM
91.88=0011
95.19=0022
91.9
95.02 h0.1
91.58=00
94.81
95.8g
95.89
CIFAR-10
elem.
$4.9770.11
95.08 10.0
05.01
95.02 1:0.1
94.6820.0
94.9922011
76.76/
1:02
elem. 1 0.01
94.95
95.13 #0011
94.8710011
95.12 +0.1
94.6620.1
94.87
95,92
20.1
elem. &x
91.96+0.1
95.06
91.74
94,91
91.68=0.
94.73
95.56
95.91#001
Fisher, 1.
95.07=0.0
95.17
94.77
95.10
94.55=00
94.91
95.65
96.00
layer.
94.78
95.09-0011
94.54
95.08
66.21 66.21
94.96
95.48
95.82
and comparisons with the dynamic sparse perturbation approach (SSAM-D) 42 We also provide
additional results for SSAM-D for a WideResNet-28 on CIFAR-100 data in Table 12 We found optimal performance for SSAM for 50% sparsity and م = 0.1 on CIFAR- 10 and 100 (as also observed in 42 for slightly different training settings). We find that although both SSAM approaches can perform on par or even outperform regular SAM, they are less effective than our SAM- ON approach. The generalization gap increases even further when considering the same high sparsity levels as for SAM-ON.
0.2 on CIFAR-
Table 12: Although SSAM-F and SSAM-D [42 with different sparsity levels can outperform SAM-all on 100 with WRN-28, they are less effective than SAM-ON.
SAM
SAM-ON
SAM-rand
SSAM-F
SSAM-D
Sparsity
0%
99.95%
99.95%
50%
99.95%
50%
99.95%
Accuracy
83.11
84.1940a
80.97
83.94
83.14
83.53
81.01
17

---

85.0
p=2.0,
84.5
ASAM, p=2.0, elem.
84.0
ASAM, p-2.0, layer
Accuracy
p-inf,
Fisher, p=2.0 etam1.0
SAM, p-2.0
all no-norm
ON
80.5
10
10
10
10
Figure 6: All considered SAM-variants and their SAM-ON counterpart for WRN-28 on CIFAR-100.
71
70
ASAM, elem.
ASAM, p=2.0, elem. eta=0.01
69
ASAM, p-2.0, layer
68
p=inf, Fisher, e.a=1.0
67
SAM, p=2.0
66
all ON
65
64
10-2
Figure 7: All considered SAM-variants and their SAM-ON counterpart for a VIT-S on CIFAR-100.
96.50
Test Accuracy (86
96.25
SGD SAM
96.00
SSAM-F 50%
SSAM-D 50%
95.75
SSAM-F 99.93%
95.50
SSAM 99.93%
SAM:
95.25
SAM-rand
10-2
10
10a
10a
Figure 8: SAM-ON outperforms SSAM-F and SSAM-D random mask SAM-rand (same sparsity level 99.93% SAM-ON) sparse perturbation approaches CIFAR-10 10 for ResNet-18.
42
(with different sparsity levels) and
B.3 Finetuning from ImageNet-21k
Since ViT' are commonly trained on large-scale datasets and then fine-tuned, we investigate this scenario for SAM-ON. In particular, we consider ViT-S pretrained on ImageNet-21k from (48). We fine-tune for 9 epochs with SGD for range of SAM-variants with their respective SAM-ON counterpart. For each setup, probe three values of p and report the best result in Table We find in this setting that SAM-ON performs on par with SAM-all although there are small differences across SAM variants: for layerwise-f, SAM-ON ON performs slightly worse, whereas for all other variants SAM-ON performs equally well or slightly better than SAM-all. In all cases SAM-ON outperforms plain SGD.
Table 13: Results for ImageNet-I1 1k fine-tuning of a ViT-S-224 from ImageNet-214 model.
SGD
SAM
ASAM elem. &
ASAM layer. L
ASAM elem. L
all
ON
all
ON
all 81.79
ON 81.75
all
ON
81.62
81.75
81.75
81.73
81.75
81.84
81.84
18

---

B.4 Adversarial robustness
Here, we provide additional results and extend the discussion on adversarial robustness from Sec- tion 4.2 In study by Wei et al. [52] SAM-trained models showed non-trivial robustness to small adversarial perturbations 50 Since there are several works highlighting the role of normalization layers for adversarial robustness p [54) it is interesting to investigate whether the robustness proper- ties of SAM can be preserved when training with SAM-ON instead of SAM-all. In Table we report the adversarial robustness of the ViT-S trained from scratch on ImageNet (as reported in Section 4.2 evaluated with the two white-box attacks from APGD, but for more radii). The SAM-ON models are not only better than the base optimizer, but consistently outperform the SAM-all models by small margin. For a WRN-28-14 10 on CIFAR-100 the differences are less pronounced and often within the standard deviation (reported over 3 seeds in Table 14). SAM-ON also improves for the elAA-eeemmentwisse.... the all-variant is slightly better than the ON-variant. Overall, we find that in order to get SAM-like improvements for adversarial robustness (as shown in (52)) it is enough to only perturb the normalization layers SAM, illustrating again their special role.
SAM-all, but
Table 14: Adversarial robustmess CIFAR-100: Reported is robust accuracy (in %) for a WRN-28 trained from scratch on CIFAR- 100. Adversarial robustness is evaluated with the two whitebox APGD attacks from autoattack [6].
SGD
SAM
ASAM-el.-boo
threat model
6
all
ON
all
ON
l 0.10
18. 14.1410111
28.14
1.09
31.28
10.50
30.33 30.16.122
f
0.20
2..11
+.39
6.62
=0.07
6.63 +0.12
6.10
Loo 1/255
10.29
17.96 1. 1.08
19.56
#0.33
20.69 #0.81
18.63 L0.30
too 2/255
0.67+
1.96+0.17
2.16
#0.07
2.62 +0.01
2.05+0.17
Clean acc.
80.7+0.2
$3.1
84.2 +0.2
+3.2
84. #0.2
Table 15: Adversarial robustness ImageNet: Reported is robust accuracy (in %) for a ViT-S trained from scratch on ImageNet, as reported in Table Adversarial robustness evaluated with the two whitebox APGD attacks from autoattack 16.
AdamW
Lion
c
vanilla
SAM-all
SAM-ON
vanilla
SAM-all
SAM-ON
b 0.25
19.67=0.47
37.53
41.16 =0.24
22.01 #0.78
50.66
43.12 =0.97
& 0.50
5.4740.18
17.71 20.61
22.72
#0.25
6.63-40046
19.03=40.92
24.27 + 1..34
& 1.00
0.43
3.34
5.58
10.19
0.57400.07
3.98
6.64 10.69
too 0.25/255
+00.0
48, 48.08 00.14
49.34
#0.08
35.31 10008
49.52 40.60
51.37 #0.99
loo 0.5/255
14.18
29.68+0.00
32,46
10.15
15.10
40.62
34.23 1.77
foo 1/255
2.61 10.16
8.64
10.82
10.56
2.93
10.02 10.56
12.03 1.30
Clean acc.
60.00
71.47=0.1
71:37
68.20=0.02
71.90
72.64 22.6647 =0.14
B.5
Machine translation task
To probe the effectiveness of SAM-ON outside the vision domain, we apply it to the IWSLT' 14 DE-EN machine translation task, following the setup of Kwon et al. [37]. We report the resulting Bleu scores in Table [6] SAM-all and SAM-ON perform similar (within standard deviations reported over 3 random seeds), both improving over the vanilla optimizer. While being very limited in its scope, this experiment is first hint that SAM-ON might also be effective outside the vision domain. Proper evaluations, as for instance done in 7. are required to confirm this for large-scale settings.
Table 16: IWSLT-DE-EN Bleu scores. Reported over random seeds.
vanilla
SAM-all
SAM-ON
34.56 10.11
34.8310.110
34.95 10.15
19

---

B.6 Weight distribution after training
In order to get a better understanding of the impact of SAM-ON on 7 and 3 (as defined Eq. D. 10 with different SAM-variants and both SAM-ON and all. We show the distribution of kw; i.e. the parameter magnitudes, at the end of training for different layer types in Figure 9 Different to the discussion in Section 55.3 we show the y-axis on log-scale, in order to inspect more nuanced differences. For clementwise lz there is no strong change in the distribution of the BatchNorm parameters between all and SAM-ON. For elementwise loos layerwise lg and SAM, however, the magnitude of the BatchNorm parameters shifts clearly towards larger values, especially for the weight parameters. We note that this resembles a pattern observed when comparing the optimal p value for all and SAM-ON in Table 7 The optimal p of elementwise lz did not change much for ResNet architectures, whereas for the other considered methods, it shifted towards larger values for SAM-ON. Additionally and in contrast to the other methods, the elementwise l2 variant showed a strong performance decrease in no-norm (Figure D. indicating that it implicitly focuses on perturbing the BatchNorm layers already. We note that larger BatchNorm parameters do not necessarily indicate a functionally different network, since there are many reparameterization invariances in ReLU networks, some of which ASAM tries to leverage in its perturbation definition Eq. 4 Nevertheless, the scale of the network still has an impact on the training dynamics, since other methods like e.g. weight decay depend on it. We discuss the impact of weight decay further in Appendix B.9
we train a WedeResNet-28-
B.7 Removing the affine parameters
Frankle al. [24] found for SGD that fixing the normalization parameters typically decreases the generalization performance of networks. As an ablation, we therefore study the effect of SAM when the normalization weights are non-trainable. This is, set 7 1 and 3 O and train the remaining parameters with SAM. The results are shown for a WRN-28 in Figure I0, where it can be seen that fixing the normalization parameters (fix-norm) does not lead to decrease in the performance of SAM. We thus hypothesize that in certain settings, SAM might not leverage the expressive power of the normalization layers, which might contribute to the improved performance of SAM-ON.
B.8 Training BatchNorm and only BatchNorm
The affine parameters of the normalization layers are relatively understudied in the literature. Recently, Frankle et al. (24) were able to obtain surprisingly high performance for ResNet architectures by only training the BatchNorm layers (freezing all other parameters), illustrating their expressive power. We study the effect of SAM in this setting (i.e. when all parameters except for the BatchNorm layers are frozen) for ResNet-101 and WRN-28 on CIFAR-10 and find that SAM still aids generalization in this setting (Table 17).
Table 17: Effect of SAM when training only BatchNorm layers, for networks trained on CIFAR-10.
Model ResNet-101 WRN-28
SGD
SAM p 0.01
SAM & 0.05
78.75
78.63
79.27
63.49
64.48
62.70
B.9
Weight decay and dropout
Here, we explore potential connections of SAM-ON with weight decay and dropout. Since weight decay is sometimes applied to all network parameters, and sometimes normalization layers omitted, it is worth investigating if the benefits of SAM-ON can be attributed to its interaction with weight decay. To this end, we train a WRN-28 with SGD, SAM-all and SAM-ON, and apply weight decay to either all parameters, all except the normalization layers, or not at all (Figure [I] right). For each setting SAM-ON outperforms SAM, outlining that its success should not be attributed to the interaction with weight decay.
We further test if SAM-ON-like performance can be achieved by simply applying stronger regulariza- tion and stochasticity to the normalization parameters. To this end, we apply dropout solely on the normalization layers (Figure []] left) and find that this is not the case.
20

---

Distribution y
Distribution of
Distribution weights
10a
10a
10a
SGD
10a
104
10a
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
10a
OnlyNorn
10a
SAM
101
10a
101
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
elementwise D=2
10a
10a
10a
all
OnlyNorm
102
101
0.0
0.5
0.5
1.0
0.0
0.5
1.0
10a
10
clementwise p=a
OnlyNorm
10a
10a
10
101
0.0
0.5
0.0
0.5
1.0
0.0
0.5
1.0
10a
10a
OnlyNorm
10a
layerwise
10
10a
101
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
B:
Wi
Figure 9: SAM-ON leads shift the distribution of y.
84.0
83.5
Test Accuracy
83.0
82.5
fix-nor,, p=2.0
82.0
SAM, p=2.0 all
SAM, p-2.0 OnlyNorm
81.5
10
10F
P
Figure 10: When training with SAM, fixing = 1, 3 = (fix-norm) barely changes the performance of the network. WRN-28, 100AR-100.
21

---

WRN28-10 CIPAR100 Dropout on Normalization layers
WAN28-10 CIFARTOD Weight decay
SAM
Test (90)
SaM ON
5GD
wil norm, SAM al
norm, SAM
norm, SCD
wil SAM
no wil SAM ON
no wil SCD
dropout
Figure 11: Left: Applying dropout only to the normalization layers (blue/red) performs worse than SAM-ON. Right: SAM-ON improves over SAM-all irrespective of whether weight decay is applied to all parameters (green), all except the normalization layers (blue) or not at all (yellow).
B.10 Details on sharpness evaluation
For the following discussion, we note that the term generalization is sometimes used as the difference between train and test eπror, while in other cases people use it as synonym for test error. Since in CIFAR settings the models achieve train error close to zero, the two definitions become equivalent.
Many studies have attempted to better understand the possible connection between the general- ization of deep neural networks and the flatness of the loss-surface [26 32 22
[8] Recently, Andriushchenko et al. [4] conducted large-scale study for a range of models, datasets, and sharpness- definitions, finding that "while there definitely exist restricted settings where correlation between sharpness and generalization is significantly positive (e.g.. for ResNets on CIFAR-I0 with a specific combination of augmentations and mixup) is not true anymore when compare all models jointly" and concluding "that one should avoid blanket statements like flat minima generalize better". In order evaluate sharpness, we therefore adopt their setup and choose the best-performing sharpness measure for CIFAR from their study, which is logit-normalized clementwise-adape worst-case-loe- m-sharpness. This is, m-sharpness s is defined as the largest possible change in loss within the adaptive perturbation model defined in
m aD
Ex
max 1|T Sl
L(w + +e) -((w)
(6)
where T |w oo and Dm returns data batches of size m. p here denotes the size of the ball over which sharpness is evaluated and is not to be confused with the p from the SAM-algorithm. Like for ASAM (37], the motivation behind adaptive sharpness measures is to make them invariant to reparameterizations of the network. Further, the logit-outputs of the network are normalized with respect to their f2-norm in order to mitigate the scale-sensitiiity of classification losses. In practice, Andriushchenko et al. 44] compute st over subset of the train set of size 1024 and use m = 128, i.c. average 8 batches. We use a subset of size 2048 in order to obtain more reliable sharpness estimates, and adopt a hyperparameter-ffre method designed for accurate estimation of adversarial robustness. It is to note that except for the fogit-normalization, the sharpness definition reported in Table corresponds exactly to the perturbation model that ASAM elementwise for uses, and hence the 1 1-step sharpness reported should be fairly close the the objective that ASAM elementwise f.x actually minimizes during training. While ASAM elementwise &o yields slightly smaller sharpness values than the conventional SAM algorithm, the differences are rather small when compared the significantly sharper SAM-ON models. For the results in Table in the main paper we tuned the sharpness radius P such that we obtain sharpness values similar to those reported yield the highest correlation in Andriushchenko et al. [4]. In Table 18 we report sharpness values for ResNeXt-modee in addition to the wRN-28 from the main paper. all cases the SAM- SAM-ON models are sharper than the SAM-all models yet generalize better. In Table 19 we further report other sharpness measures without logit-normalization for WRN-28. SAM-ON is sharper than SAM-all with respect to most metries, although there exist some exceptions. It should however be stressed that many of those metrics did not show good correlation with generalization in the study by Andriushchenko al.
128. The maximization in (6) is performed with AutoPGD [16),
22

---

Table 18: Sharpness evaluation of both a WRN-28 and a ResNeXt. SAM-ON is sharper than SAM-all all cases. Shown is 20-step logit-normalized &o sharpness from averaged over three models method. Dataset considered is CIFAR-100.
SGD
SAM
ASAM-el.-fx
all
ON
all
ON
ResNeXt WRN-28
Test Accuracy (%)
80.71
83.
84.19
83.25
84.14
20 sharpness, p 0.003
0.071
0.048 00:01
0.090
0.048 0.048=0:001 =0:00
0.07g
sharpness, = 0.005
0.201 0.201=0.000 #0.001
0.139 000004
0.296=0.011
0.124#0.002 000002
0.283
sharpness, p 0.007
0.43344.000
0.309
100011
0.585
0. 0.255 #00005
0.580
Test Accuracy (%)
80.16 #0.3
81.79 #0.4
82 22 20.2
81.02 =0.6
82.38 20.3
L sharpness, p 0.001
0.036
0.029
0.034
0.026 0.026
0.034
sharpness, p = 0.003
0.164 005
0.117
1.10040.00 40.004
2.140-0.002 002
0.099 400010
0.147#00.001
sharpness, 0.005
0.383+0.011
0. 0.252 20.008
0.291 #0.005
0. 203 #0.021
0.312
Table 19: Additional sharpness measures. WRN-28 (no logitnorm).
SGD
SAM
ASAM-el.-
adaptive
all
ON
all
ON
Test Accuracy (%)
80.71+0.2
83.11 #0.3
84.19-0.2
83.25 +02
84.14
E avg, / 0.005
False
1.358=#0.040
0.515
+0.020
2.372#t0.071
0.569
20.012
2.141=00045
82 avg. = 0.1
True
0.042##0.001
0.019
200001
0.022
0.040
0.019 200001
too avg, P 0.01
False
2.643==0.007
1.264
4:0.088
3.455=#0000
1.304 100007
3.259=000031
too avg, 0.2
True
0.008-t0.001
0.035#00000
0.034+0.000
0.068=t0.003
0.031 1:0001
worst, P 0..05
False
0.501 40.008
0.655
0.701 0.701 #0.057
0. 0.768440.11
0.313 +0.044
worst, P = 0.25
True
0.065=t0.008
0.033
0.037=#.017
0.056=00.006
0.062=00001
worst, P 1e 05
False
0.149=00.00
0.055
4:0.002
0.040
0.050
4:0 002
0.123=0.007
f00 worst, = 0.004
True
0.537
0,262
1:0.009
0.600
0.255
4:0001
00.057
C Convergence Analysis
We provide this section a convergence analysis for SAM-ON in the non-convex setting. Using standard assumptions we obtain a theorem which resembles findings for closely related methods such as found 2
Our assumptions:
Assumption C.1. We assume function ʃ Rd - to be L-smooth: there exists L >> 0 such that
[Vf(e) V s L|le w|12, Ve, E R".
Assumption C.2. There exists M > 0 for any sample such that
Va (w)l < M, Vw E R..
(8)
Remark C.3. If Assumption C.I holds (I-smoothness)) then Vv, w E IR:
((() (f(w) Vf(w) w))l < |le w|
This well-known result can be derived using the fundamental theorem of calculus and Cauchy- Schwartz.
Remark C.4. Assumption C.2 guarantees that the variance of the stochastic gradient less than M.
SAM-ON. the following we shall denote the true gradient as V/(w) and the noisy observation
gradient as g(w). The gradient of the loss of the ith training example is denoted as Joz (w). We partition the neural network parameters layer-wise as w (wN; w. with wi E RON WA E RAA,, nN + nA where wy represent the normalization layer parameters and WA all other layers. The
23

---

iteration for wy is:
.N. (w²)
12212 12 w'y + p 11 .. (wt)
WNy N
1+1 WN
wNN h gN, (2/)
(10)
and for W. is:
1/2 w'x WAA
++1 w A
w'x h 9A.x.. (w(+1/2)
(11)
Theorem C.5. Assuming
C.1 and C.2 h S 1/L, we obtain:
T-1
2(f(w0)) - f(ww))
-
+ 2LhM + L2p² + Lh),
(12)
T
hT
with w the optimal solution f(w).
Proof. From Assumption and thus Remark
C.3 follows that:
f(w 1 f(ww) + Vf(w") (w w() +
tut - w'112
(13)
2
S f((w)) hv f(w') 9s. (w0t1/2)
Ir. ( 1|
+
(14)
2
=((w²) - AV /V 9s. (w(+1/2)
112
(I(/++((((w) (IVS((')) Ix: (w (1))/ [[ +2 (VS(w") In. (w*+1/2))
2
= =((w²) Lh²
|Vf(("))| + Lh
2 |VVf(w') - gz, (wx
2
- (1 - Lh)h (Vi(w') . 9x. (w*+1/3)
(15)
L
S <((w") 2
| (V=("")) )1( + TA7 (W("))) 9r: (1²)
+ Lh 119z; (w²) Ir. (w 1/ (1 Lhyh (Vf(w") Ir, (wx+1/2)
(16)
Taking the double expectation gives (because unbiased gradient and Assumption C.2 and Remark
LA [IV LA
Els(w* [ES(w0))
[h?llg(w") - g ((w (1 LADAE VS(w") g(w
(17)
For term A we obtain using Assumption
A < 1. || 2 =
(18)
For term we obtain:
B = E (V Vfa(w()) (9N(w(+2/2) (ox(uutt))).
(19)
(A(")( (2A(w)+1/2) =A(w") + [a(w)))]
+ V fN(W") (2N(w++1/2) IN (w') + eN (N()))
(20)
= E |I] [[V
+ (V/(( (g1(w++1/2) - (00ggg(www))) B EVVS((w') (gN(w(+1/3) - -xxx))/)))(((()))))
(21)
24

---

Using zy 2 || + 11ull2 and Assumption C.1 we get for C:
[C] - - 2 f(
+
L
I
/ w²1l -
2 [1V + I2p²
1
(22)
2
Plugging this into
gives:
Lh
[[s(w(+)]] s [E[[("))]
B [[ T.P 1 (1 )|2
2
BIV || 1222²
+ (1 Lh)h
(23)
S [[[[(">)] - ~ [ M hT. p T.)-
(24)
In T iterations obtain using telescoping sum:
f(w") f(w") s B[ f(w")
T-1
$ 2
[[V. + MT + p1 Lh)T.
(25)
This gives Theorem C.5
25


Reviews:
[{"review_id": "0SZO7OYAW3", "reviewer": "Reviewer_Ftd6", "summary": "This paper tries to analyze the effects of different layers for sharpness-aware minimization (SAM). The authors find that normalization layer plays an important role in the improvements of SAM and only perturbing the normalization layer can obtain a comparable result. Therefore, this paper proposes SAM-ON and the experimental results illustrate that SAM-ON can obtain a great performance on many different tasks with different models. ", "questions": "Question:\n1. Do you try to run the experiments for multiple times with different seeds.\n2. You say only perturbing normalization layers can obtain similar results with perturbing all parameters. My question is do you try to only perturb the parameters outside the normalization layer?", "limitations": "N/A", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Strengths:\n1. This paper focus on an important problem, SAM is very important for improving generalization.\n2. The experimental finding is very interesting and can efficiently reduce the computation.\n3. The authors try to evaluate the performance of proposed method on various datasets and models.", "weaknesses": "Weakness:\n1. Although the proposed method can improve the performance on CIFAR. But CIFAR is too simple for SAM and that is weak to illustrate the improvement of SAM-ON. However, the results about resnet on imagenet is too close for me and the improvement is minimal. I'm not sure whether you run the experiments multiple times with different seeds.\n2. The experiments on ViT is also a little weak for me. For ViT, cifar is too simple and you should train ViT form scratch on imagenet. Since SAM can achieve significant improvement on ViT+ImageNet and maybe you need to analyze it to obtain more convincing results.\n3. Although SAM-ON can reduce the computation, the training efficiency cannot be significantly improved. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and helpful suggestions. \nWe address your concerns below.\n\n- __ImageNet results__:  \nWe agree with your comments on the ImageNet results. As you suggested, we have run a ViT from scratch on ImageNet (with additional evaluations on ImageNet-R and ImageNet-sketch) with both the Lion and AdamW optimizers and find that SAM-ON consistently outperforms SAM-all in this setting (see Table 2 in attached rebuttal pdf), we will include this in the revised paper. Our ImageNet results at present are indeed for a single seed (all other experiments are with multiple seeds), as we prioritized showing results across different SAM variants. Following your comment, we have started runs with different seeds. Preliminary results show very little variability (standard deviations of 0.047% and 0.007% for SAM-all and SAM-ON with RN50 and elementwise $\\ell_\\infty$). We will ensure to include multiple seeds for all results in the revised paper. Further, we agree that the results for ImageNet and ResNet are somewhat less impressive than for ViT, but would like to highlight that other SAM variants such as SSAM, ESAM, GSAM, and FisherSAM also obtain only marginal (or no) improvements in this setting, so this is within expectations. \n\n- __“CIFAR is too simple for SAM”__:  \n  We would like to draw the reviewer’s attention to e.g. Table 1 of our paper, where it can be seen that for CIFAR-100 regular SAM (denoted SAM-all) can improve by a large margin over the vanilla optimizer. For several model types, CIFAR-100 is far from being solved. Therefore, the consistent improvement of SAM-ON over SAM-all across different models (see Table 1 in rebuttal pdf for additional VGG-Nets and DenseNet) and SAM variants on CIFAR data is remarkable and we believe worth reporting. We hope this in combination with the novel ImageNet and ViT results addresses the reviewer’s concerns, but please do not hesitate to let us know if you believe further experiments are missing.\n\n- __“Although SAM-ON can reduce the computation, the training efficiency cannot be significantly improved”__:  \n  In Table 5 we observe that SAM-ON reduces wall-clock time compared to SAM by about 17%. We agree this is not a large improvement, although it does improve upon the sparse SAM approach by Mi et al., 2022. Further computational gains ($\\approx$ 30%) may be achieved for our approach by applying SAM to only _part of_ the normalization layers. Although this leads to some reduction in test accuracy, the approach still strongly outperforms the base optimizer (see Figure 1, bottom right in rebuttal pdf). We are happy to include this discussion in this paper, but should highlight here that the main aim of our work was to shed light on the remarkable success of SAM and to illustrate the important role of the normalization layers in obtaining its enhanced generalization performance and the connection with sharpness. Although the reduced computation obtained using SAM-ON is a positive side-effect, it was not the main aim of our work and hence was not listed as one of our main contributions on p2. \n\n- __“My question is do you try to only perturb the parameters outside the normalization layer?”__:  \n  Yes, this is denoted as no-norm (dotted lines) in Figure 1, 2, 6, and 7. no-norm typically performs similar to SAM-all, except for some perturbation variants like elementwise-$\\ell_2$, where we observed a drastic drop in test accuracy. We briefly discuss this in Section 4 and 5.3, but will extend this discussion in the revised paper.\n"}, {"review_id": "pPcAMgMyrd", "reviewer": "Reviewer_wExQ", "summary": "This paper relates the normalization layer with the Sharpness-Awareness Minimization (SAM). Surprisingly, this paper finds that in the perturbation stage, only perturbing the affine parameters of normalization layers in the networks leads to a better generalization performance. Later, the authors investigate the reason behind such surprising phenomenon, and find that SAM-ON(only perturbing the normalization layer) even increase the sharpness, which doubts the intuition that the effectiveness of SAM comes from sharpness minimization. Extensive experiments demonstrate the effectiveness of the methods this paper proposed.", "questions": "1. In fact, the authors don't explain why SAM with only perturbing normalization layer works better than original SAM. Therefore, I would like the authors to include some discussion or theoretical analysis on that.", "limitations": "Yes. This paper discusses the limitation of this paper, where they only focus on the vision data.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper identifies an important yet interesting problem: Is Sharpness-Awareness minimization effective because of minimizing sharpness? This paper found that only perturbing the normalization layer in the perturbation stage of SAM leads to both superior generalization performance and higher sharpness. The interact relationship between sharpness and generalization are quite doubtful.\n2. This paper investigates the SAM from the normalization layer and found the normalization might be the reason for the success of SAM, which is quite surprising.\n3. SAM with only perturbing the normalization layer largely saves the computational cost than the original SAM which needs perturbing all the parameters.", "weaknesses": "1. These experiments should be conducted on more network architectures, including VGG-Net and etc.\n2. The observation that SAM with only perturbing the normalization layer holds the higher sharpness than original SAM might be trivial. SAM are designed to reduce the sharpness, while a degraded version of SAM should hold a higher sharpness in principle. \n3. In this paper, it is clear that there is deep relationship between normalization layer and generalization, and there are indeed some papers on the relationship between normalization layer and generalization. Therefore, a related discussion should be included.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and helpful suggestions.\nWe address your concerns below.\n\n- __“more network architectures, including VGG-Net”__:  \n  Following your suggestion, we have provided results for more neural network architectures: DenseNet and various VGG-Nets in Table 1 in the rebuttal pdf. Across the different architectures and SAM variants we consistently observe that SAM-ON outperforms regular SAM.\n\n\n- __“observation that … [SAM-ON] .. holds the higher sharpness than original SAM might be trivial. SAM are designed to reduce the sharpness, while a degraded version of SAM should hold a higher sharpness in principle.”__:  \n  We agree this intuition is sensible, however prior work by Mi et al., 2022 hypothesizes that sparse SAM approaches can actually lead to a flatter landscape than regular SAM. It was thus worth investigating if SAM-ON (with its targeted, high sparsity level approach) would still be able to reduce sharpness. We will extend the discussion in the paper.\n\n\n- __“there is deep relationship between normalization layer and generalization .. Therefore, a related discussion should be included.”__  \n  We agree. Some discussion on the connection between normalization layers and their affine parameters with generalization can already be found in the related work section, but we are keen to extend this further. If you have any specific papers in mind please let us know and we would be happy to discuss them.\n\n\n- __“the authors don't explain why SAM with only perturbing normalization layer works better than original SAM.”__:  \n  We believe it is an important question for future work to understand the mechanism behind the success of not only SAM-ON but also regular SAM as our work lends support to the idea that sharpness-reduction may not be the sole source of SAM’s success. The search for alternative explanations for SAM’s success is an active area of research, for example recent work (that appeared on the arXiv after the NeurIPS submission deadline and hence not included in our related work section) suggests that the use of SAM prunes a significant number of activations and leads to low-rank features (arXiv:2305.16292). Our work contributes to this search by illustrating the success of only applying SAM to the normalization layers. We provide additional experiments in Figure 1 in the rebuttal pdf that explore potential connections of SAM-ON with weight decay and dropout. We find that SAM-ON’s improvements are not due to its interaction with weight decay (Figure 1, bottom left), nor can its success be mimicked by applying dropout solely to the normalization layers (Figure 1, top). We will include this discussion in the revised paper.\n"}, {"review_id": "uHqVjw2EgJ", "reviewer": "Reviewer_dJUQ", "summary": "As Sharpness-Aware Minimization (SAM) aims to regularize the flatness of the loss landscape for better generalization, this paper shows that only perturbing the normalization layers is sufficient to achieve this. To prove this, the authors first propose a method called SAM-ON (SAM-OnlyNorm). Then, they conduct experiments on CIFAR datasets across different models and variants of SAM, comparing the test accuracy with and without applying SAM-ON. The experiments (in most cases) show that SAM-ON outperforms vanilla SAM, as well as SAM variants with ON outperforming SAM variants. Furthermore, since SAM-ON only needs to perturb a few parameters, it can save computational resources. Finally, this paper conducts numerous experiments to understand SAM-ON and concludes that sharpness may not be the key to SAM.\n\n## post-rebuttal\nI've updated my score since my concerns are adequately addressed.", "questions": "See Weaknesses.", "limitations": "N/A", "rating": 9, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. This paper demonstrates a new understanding of the underlying mechanism of SAM, which is both novel and informative.\n2. The proposed SAM-ON can improve efficiency and generalization in most cases.\n3. The experiments are adequate to support the claims.\n4. The code is provided, which guarantees reproducibility.", "weaknesses": "1. SAM-ON cannot outperform SAM (or its variants) in some cases, showing that SAM-ON may not always be effective and decreasing the soundness of the claims that support SAM-ON.\n2. The improvement of SAM-ON on ImageNet seems to be consistently lower than on CIFAR datasets for different architectures. Therefore, this reviewer hypothesizes that SAM-ON only works for small datasets and is not very effective for large datasets.\n3. As the author acknowledges, this paper only investigates SAM and SAM on vision data. Thus, the analysis cannot support the proposed understanding that generalizes to other tasks, such as language tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and helpful suggestions. \nWe address your concerns below.\n\n1. __SAM-ON may not always be effective__  \n   We considered a wide range of different models (see Table 1 in rebuttal pdf for even more models), data augmentations, and datasets. For almost all of these a SAM-ON variant obtained the best test accuracy. There are indeed some exceptions, but even in those settings we find it remarkable that by only applying SAM to the normalization layers, this approach can so strongly outperform base optimizers. We believe this is a valuable contribution to the literature in the search for an enhanced understanding on the mechanisms behind SAM’s success.  \n\n3. __The improvement of SAM-ON on ImageNet seems to be lower than on CIFAR datasets__  \n   Additional results for ViTs trained from scratch on ImageNet can be found in Table 2 in the rebuttal pdf. We observe that SAM-ON outperforms SAM in this setting for both the AdamW and Lion optimizers, and both with respect to the ImageNet test set and OOD test sets. Further, we agree that the results for ImageNet and ResNet are less impressive than for ViT, but would like to highlight that other SAM variants such as SSAM, ESAM, GSAM, and FisherSAM also obtain only marginal (or no) improvements in this setting, so this is somewhat within expectations. We hope this addresses your concern. \n\n5. __language tasks__  \n    Thank you for the suggestion. We focused on vision data following the trend in the literature on SAM and its variants. However, a few papers do indeed study language tasks and we would be happy to include an experiment on e.g. the IWSLT’14 DE-EN machine translation task that was considered in the ASAM paper, in the camera-ready version. Please do not hesitate to let us know if there were specific experiments you had in mind. \n"}, {"review_id": "f9bVQ8iFCd", "reviewer": "Reviewer_sERj", "summary": "The paper proposes an adversarial perturbation method linked to SAM for the affine normalization parameters, in contrast to perturbing the full set of parameters. The results show this approach improves upon standard SAM and prior sparse SAM approaches.", "questions": "see weaknesses above\n\nOne of the questions this paper fails to answer is why sharpness is higher for SAM-ON, even though its formulation is motivated to minimize sharpness. This is not a complete enough story and opens up lots of questions that need to be addressed.", "limitations": "Somewhat discussed but limited. Would like to see a more thorough limitations discussion.", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The SAM-ON proposed method yields several benefits:\n1. improves upon SAM-all and carious other SAM-variants\n2. achieves this by perturbing only the normalization layers, which correspond to a small percentage of the total parameters\n", "weaknesses": "Weaknesses of this paper include:\n1. gains in Table 1 are marginal\n2. originality of approach is lacking given prior literature on SAM variants\n3. justification of SAM-ON is poor; the experiments in 5.2 and Table 7 try to poke at this but fall short as the authors do not provide a solid explanation as to what this can be attributed to; they only show that SAM-ON increases sharpness in L-inf sense; but what about other metrics of flatness? is there a more universal metric that can explain the effectiveness of SAM-ON?\n4. it is unclear whether SAM-ON is truly better than other SAM variants in out of distribution robustness, e.g. WILDS benchmark\n5. results are only shown for CIFAR and ImageNet datasets; more datasets and benchmarks need to be considered", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. Following your suggestion, we are happy to extend the limitations section of the paper. In case there is specific work you think would need to be included, we would much appreciate a pointer. Below we address your other comments.\n\n1. __“gains in Table 1 are marginal”__:   \nSAM-ON consistently strongly outperforms base optimizers and its gains compared to SAM are often >1% for ResNet architectures on CIFAR data. Similar or no gains were also observed in other SAM variant papers, such as ASAM and FisherSAM, so the amount of gain observed is within expectations. Further, we provide additional experiments on VGG-Nets and DenseNet for CIFAR data in Table 1 in the rebuttal pdf and observe that SAM-ON clearly outperforms SAM-all. Finally, our gains are even more impressive for ViTs (Table 3, main paper), which we believe to be very important as SAM has shown particular success in this domain (Chen et al., 2022).\n\n2. __“originality of approach is lacking”__:    \nTo illustrate the originality of our approach (as also corroborated by reviewer dJUQ), we briefly recap here the most closely related works and how they differ from our approach. The most related works are by: 1) Frankle et al., 2021, which specifically studies the BatchNorm parameters by only training BatchNorm, but in a very different setting: without SAM and all other parameters are frozen. 2) Mi et al., 2022 which propose a sparse SAM approach, but again the approach and aim is very different: they do not highlight the importance of specific parameter groups, whereas we stress the unique role played by the normalization layers allowing for an extremely sparse yet effective perturbation, and their aim is to reduce the computational cost, not to enhance generalization (we outperform them in both). 3) Much of the literature on SAM proposes new variants based on designing better perturbation models by taking reparametrization invariances (like ASAM) or loss geometry perspectives (like Fisher-SAM) into account. We think that the fact that our simple heuristic “perturb only the normalization layers” aids all of those methods is quite remarkable. We hope this addresses the reviewer’s concerns. If the reviewer is aware of further literature relating SAM to the normalization layers of a network we would much appreciate a specific pointer.\n\n3. __“other metrics of flatness”__:   \nAndriushchenko et al., 2023 performed a large-scale study on the relation between a range of sharpness measures and generalization, finding little correlation. For CIFAR data, \"_the best correlation [...] is achieved by_ $\\ell_\\infty$ _adaptive worst-case sharpness with logit normalization for a small_ $\\rho$\". Since this metric is additionally very close to the perturbation model of ASAM elementwise $\\ell_\\infty$ we chose to focus on this particular metric in our paper. Following your suggestion, we investigated more sharpness measures (Table 3 in rebuttal pdf) and find in alignment with our previous findings that SAM-ON is sharper than SAM-all with respect to most metrics (especially worst-case sharpness and the metrics that are optimized during SAM), although there exist some exceptions. Results for more models and $\\rho$ values are similar (omitted due to space limitations) and will be included in the revised paper.\n\n4. __“a universal metric that can explain the effectiveness of SAM-ON”__   \nThere is a long history of trying to find metrics that correlate well with generalization performance. One of these is sharpness, which we focused on in this paper. Our work lends evidence to recent findings (Andriushchenko et al., 2023) that sharpness may not always correlate well with generalization performance. The search for alternative explanations for SAM’s success is an active area of research, for example recent work (that appeared on arXiv after the NeurIPS submission deadline and hence not included in our related work section) suggests that the use of SAM prunes a significant number of activations and leads to low-rank features (arXiv:2305.16292). Our work contributes to this search for understanding the fundamental reasons for the success of SAM by illustrating the important role played by the normalization layers.\n\n5. __“OOD robustness”__:   \nIn Table 2 in the rebuttal pdf we have provided experiments for training Vision Transformers from scratch on ImageNet using SAM-all/SAM-ON with base optimizers AdamW and Lion, and evaluated those on OOD test sets (ImageNet-Sketch, ImageNet-R). We find consistent improvements of SAM-ON over SAM-all on ImageNet, ImageNet-Sketch, and ImageNet-R.\n\n6. __“more datasets and benchmarks”__:    \nFor the selection of our datasets and benchmarks we followed the literature on SAM and its variants. Additionally, we include in the rebuttal pdf (Table 1 and 2) results for OOD robustness as suggested by the reviewer, ViTs trained from scratch on ImageNet as suggested by reviewer Ftd6, and more models on CIFAR as suggested by reviewer wExQ. For all of these we observe consistent gains of SAM-ON over SAM-all. We hope this addresses your concern, but do let us know if there is anything that is not adequately supported.\n\n7. __“why is sharpness is higher for SAM-ON, even though its formulation is motivated to minimize sharpness.”__  \nOpinions differ on whether our findings that SAM-ON may increase sharpness are trivially true or unexpected, e.g. as mentioned by Reviewer wExQ: “The observation that .. [SAM-ON] .. holds the higher sharpness than original SAM might be trivial. SAM are designed to reduce the sharpness, while a degraded version of SAM should hold a higher sharpness in principle.” This is exactly why we believe it is important to study sharpness in our work because it is not obvious how we would expect it to behave with respect to SAM-ON. Our findings lend support to the idea that sharpness may not be the sole reason for SAM’s success and aid the search for alternative explanations (Andriushchenko et al., 2023)."}]

Author rebuttal:
We would like to thank you all for your time and useful comments. We are encouraged that you found that our paper addresses an interesting and important problem by demonstrating _"a new understanding of the underlying mechanism of SAM, which is both novel and informative”_ with _"extensive experiments [that] demonstrate the effectiveness of the methods"_.

Following your suggestions, in the pdf below we attach results for _“ViT [trained] from scratch on imagenet”_ with our SAM-ON method and also evaluate _“out of distribution robustness”_ (Table 2). For all datasets (ImageNet, ImageNet-R, ImageNet-sketch) and both Lion and AdamW optimizers, we observe that SAM-ON outperforms SAM-all. We also include results _“on more network architectures, including VGG-Net”_ (Table 1) and DenseNet on CIFAR data and again observe SAM-ON’s superiority across different SAM variants. 

We address your other comments below and look forward to engaging with you in the discussion period.


Meta-review: This paper discover an interesting phenomenon related to SAM optimizer which can highly improve its computational efficiency with comparable generalization performance or sometimes even better generalization than the original SAM algorithm. SAM is known to have a strong effect on generalization performance of deep neural networks, but it is computationally more expensive than SGD because in each iteration, SAM needs to compute the gradient twice. The key observation of this paper is that, instead of computing the gradient (that corresponds to the adversarial perturbation of the weights), one can perturb the affine parameters of batch normalization only, which can be around 0.1% of the number of weights in the models the paper studies. This is a huge reduction. Interestingly, the authors notice that other ideas that rely on a smaller number of parameters, such as sparse SAM, often lead to lower generalization performance compared to the proposed idea.

The ratings for the submission had a wide range, form 9 (Very Strong Accept) to 3 (Reject). I read all the reviews, specially the criticism raised by the reviewer sERj that gave the lowest rating. I also discussed these reviews with the Senior AC handling this submission. We eventually decided to accept the paper. The criticisms from  reviewer sERj were around lack of theory for the observed phenomena. However, this paper is an empirical paper, which discovers an interesting phenomenon and then backs it up with evidence spanning multiple datasets and different learning scenarios. As such, it is not necessary to have theory to get the paper accepted. This point was also mentioned as a reply to reviewer sERj by reviewer dJUQ. Based on this, and the fact that the submission has also received rating of 9 and 7, we recommend accept.
Decision: Accept (poster)