Title: Normalization Layers Are All That Sharpness-Aware Minimization Needs
Venue: NeurIPS 2023
Accepted: yes
Keywords: sharpness-aware minimization

flatness

generalization

normalization layers
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

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.

1 Introduction

Numerous works have been dedicated to studying the potential connection between flatness of minima and generalization performance of deep neural networks [34, 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-aware minimization (SAM) has become a popular approach to actively try to find minima with low sharpness using a min-max type algorithm [23]. SAM was found to be remarkably effective in enhancing generalization performance for various settings [23, 14, 7, 1, 33].

Several variants of SAM have been proposed, focusing both on enhanced performance [37, 35] and reduced computational cost [11, 20]. In particular, with the original aim of making SAM more efficient Mi et al. [42] propose a 'sparse' SAM approach, which applies SAM only to a 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 or even 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 [41, 46]. Frankle et al. [24] found for ResNets that the trainable affine parameters of the normalization layers

10

3

10

2

10

1

10 0

10 1

81

82

83

84

85

Test Accuracy (%)

all

OnlyNorm

no-norm

ASAM, p=2.0, elem.

ASAM, p=2.0, layer

ASAM, p=inf, elem.

SAM, p=2.0

SGD

all

OnlyNorm

no-norm

have remarkable representation capacity in their own right, whereas disabling them can 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 a network (often less than 0.1% of the total parameters) outperforms perturbing all parameters. We find that this can not be solely attributed to 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 a network ( SAM-ON , short for SAM-OnlyNorm ) 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 4)

Alternative sparse perturbation approaches do not result in similar performance as SAM-ON , especially not at the extreme sparsity levels of our method. (Section 5.1)

Similar to SAM-all , SAM-ON yields non-trivial adversarial robustness (Section 4.2). 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 (LayerNorm) [6] form an essential component of most convolutional [25, 29] and Transformer [51, 19] architectures, respectively. Across various works these normalization layers were shown to accelerate and stabilize training, reducing sensitivity to initialization and learning rate [13, 5, 60, 36]. But despite their widespread adoption and illustrated effectiveness, a conclusive explanation for their success is 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 BatchNorm, 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 of 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 to 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 a 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 [49].

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 [23, 14, 7]. Hence there have been several approaches to try to reduce the computational cost of SAM, such as ESAM which utilizes sharpness-sensitive 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' [11] and that only applying SAM to part of the parameters (SSAM) using e.g. a Fisher-information mask can lead to enhanced performance [42]. A common variant of SAM utilizes m -sharpness [23], which uses subbatches of size m and benefits performance [20, 2, 43] though nuances in its implementation vary [8]. Andriushchenko and Flammarion [2] argue that its success is not unique to settings with BatchNorm and hence cannot be attributed to 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 focus on the interplay between two popular aspects of neural network training, both of which we recapitulate here: In Sec. 3.1 we provide an overview of 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) [6], which are an essential building block of most convolutional [25, 29] and transformer [51, 19] architectures, respectively. Normalization layers transform an input x according to

where µ and σ 2 are the mean and variance, which are computed over the batch dimension in the case of BatchNorm, or over the embedding dimension, in the case of LayerNorm. BatchNorm is therefore sensitive to the choice of batch size [38, 49]. For BatchNorm, µ and σ are computed from the current batch-statistics during training, and running estimates are used at test time. In our experiments, we focus on the trainable parameters γ and β , 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 a neural network f w : R d -→ R k which is parameterized by a vector w as our model. The training dataset S train = { ( x 1 , y 1 ) , ... ( x n , y n ) } consists of input-output pairs which are drawn from the data distribution D and we write the loss function as l : R k × R k -→ R + . The goal is to learn a model f w with good generalization performance, i.e. low expected loss L D ( w ) = E ( x , y ) ∼ D [ l ( y , f w ( x ))] on the distribution D . The training loss can be written as L ( w ) = 1 n ∑ n i =1 l ( y i , f w ( x i )) . 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 ℓ p ball around w , i.e. max || ϵ || p ≤ ρ L ( w + ϵ ) -L ( w ) . This leads to the overall objective

In practice, SAM uses p = 2 and approximates the inner maximization by a single gradient step, yielding ϵ = ρ ∇ L ( w ) / ||∇ L ( w ) || 2 and requiring an additional forward-backward pass compared to SGD. The gradient is then re-evaluated at the perturbed point w + ϵ , giving the actual weight update

with learning rate α . Computing ϵ 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

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 in Eq. (2) is not invariant with respect to a rescaling of the weights that leaves f w invariant [18], ASAM [37], a partly scale-invariant version of SAM, was proposed, with the objective

where T w is a normalization operator, making the perturbation adaptive to the scale of the network parameters. Kwon et al. [37] choose T w to be diagonal with entries T i w = | w i | + η for weight parameters and T i w = 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:

We note that for T w equal to 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 a Euclidean distance measure between parameters. The approach can also be framed as a variant of ASAM, with T w being diagonal with entries T i w = 1 / √ 1 + ηf i and f i approximating the i th diagonal entry of the Fisher-matrix by the squared average batch-gradient, f i = ( ∂ w i L Batch ( w )) 2 . For our experiments, we additionally employ layerwise normalization. This is, we set the diagonal entries of T i w = || W layer [ i ] || 2 , which corresponds to a normalization with respect to the ℓ 2 -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 a 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 (basic+AutoAugment, denoted as +AA). We consider a range of SAM-variants which differ either in the perturbation model ( ℓ 2 or ℓ ∞ ) 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 to the BatchNorm parameters, for a WideResNet-28-10 (WRN-28) on CIFAR-100 in Figure 1. We observe that SAM-ON obtains higher accuracy than conventional SAM ( SAM-all ) for all SAM variants considered (more SAMvariants 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 elementwiseℓ 2 variants) or maintains similar performance as SAM-all (for all other variants). For variants which do not experience a 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-100. For each SAM-variant and dataset, we probe a set of pre-defined ρ -values (shown in Table 7 in the Appendix) with a ResNet-56 (RN-56) and fix the best-performing ρ for the other models to compare SAM-ON to SAM-all . We report mean accuracy and standard deviation over 3 seeds for CIFAR-100 in Table 1. On average, SAM-ON outperforms SAM-all for all considered

SAM-variants. Of these, layerwiseℓ 2 achieves the highest performance for most settings. We obtain similar results for CIFAR-10 (App. Table 10) and for more network architectures (App. Table 11).

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 to an overall batch-size of 512. Apart from ρ , all hyperparameters are shared for all SAM-variants and can be found in the Appendix in Table 8. 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 2. We observe that for layerwise ℓ 2 , the all variant achieves higher accuracy, whereas the SAM-ON models outperform their all counterparts for elementwise ℓ 2 and elementwise ℓ ∞ . 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.

Varying the batchsize. In Figure 4 (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, 2], 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 small-scale datasets like CIFAR, our aim here is not to outperform state-of-the-art, 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-all by a clear margin. For the elementwiseℓ 2 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 a ViT-T and CIFAR-10 as well (Table 3).

10

3

10

2

10

1

10 0

10 1

64

65

66

67

68

69

70

71

Test Accuracy (%)

all

ON

no-norm

ASAM, p=2.0, elem.

ASAM, p=2.0, layer

ASAM, p=inf, elem.

AdamW

SAM, p=2.0

all

ON

no-norm

ImageNet. For ImageNet, we train a ViT-S/32 from scratch with batchsize 128 on a single GPU for 300 epochs. In Table 4 we evaluate SAM-all , SAM-ON and the vanilla variant for both AdamW and Lion [15] as base optimizers on 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 AdamW, SAM-ON either outperforms SAM-all (e.g. w.r.t. adversarial robustness) or performs on par (e.g. for ID accuracy the numbers are within standard deviations). For Lion, SAM-ON always outperforms SAM-all , underlining that the diverse benefits of SAM can be achieved or even surpassed by only perturbing the normalization layers. We 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.

WRN-28

ResNeXt

85

83.0

SAM-ONB2+B3

SAM-ON

82.5

SAM-ONB2+B3

84

SAM-ON

SAM-all

82.0

SAM-all

83

2

81.5

ccura

82

81.0

SAM-ON B2

SAM-ONB2

80.5

81

SGD

SAM-ONB1

SGD

SAM-ONB1

80.0

1.0

1.2

1.4

1.6

1.8

2.0

1.0

1.2

1.4

1.6

1.8

2.0

Runtime(normalized bySGD time)

Runtime (normalized by SGD time)

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 128 on a single A100 with PyTorch. Since the normalization parameters at the earlier layers of the network require a gradient, potentially a full backpropagation pass has to be computed for the ascent-step of SAM-ON , even though only a tiny fraction of all parameters is perturbed. However, as discussed in [12] for a related setting, the gradients of the intermediate (no-norm) layers do not need to be stored or used for updating. This leads to computational gains of SAM-ON over SAM-all as reported in Figure 3.

In contrast, although future development of hardware for sparse operation may allow for acceleration, Mi et al. [42] report that their sparse perturbation approach does not at present lead to 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 test accuracy. When perturbing the normalization layers from Block 2 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 an interesting research direction for future work.

5 Towards Understanding SAM-ON

To gain a better understanding of SAM-ON , we study different hypotheses for the method's success. First, we investigate the role of sparsity by comparing SAM-ON to different sparsified perturbation approaches (Section 5.1), concluding that sparsity alone is not enough to explain its success. Then, we highlight that SAM-ON might in fact find sharper minima, while generalizing better than SAM-all (Section 5.2). We also show that SAM-ON - similar to 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 in 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 a WRN-28 with BatchNorm and the setting described in Section 4 for the ablation studies. Further ablation studies are presented in Appendix B.

10

2

10

1

10 0

10 1

95.00

95.25

95.50

95.75

96.00

96.25

96.50

Test Accuracy (%)

SGD

SAM

SSAM 50%

SSAM 99.93%

SAM-ON

SAM-rand

10 1

10 2

m

81

82

83

84

Test Accuracy (%)

all

ON

bs=128

bs=256

bs=512

all

ON

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 is simply due to perturbing fewer parameters. We therefore compare SAM-ON to other sparse perturbation approaches.

In order to determine the parameters which should be perturbed, one solution Mi et al. [42] propose is to compute a binary mask via an approximation of the Fisher matrix. We call this approach SSAM-F to 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 'relatively more stable but a 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 4.3).

We provide a comparison between SAM-ON , SSAM-F , and a random sparsity mask of the same sparsity level as SAM-ON for a) ResNet-18 on CIFAR-10 (main setting considered in [42]) in Figure 4 (left) and b) WRN-28 on CIFAR-100 in Table 5. We find that although SSAM-F can indeed perform on par or even outperform SAM-all at the medium to high sparsity levels recommended in [42] it 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 is not the sole cause for SAM-ON 's success.

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, 34] - indeed this was the initial motivation behind SAM in [23]. Andriushchenko and Flammarion [2] 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 a model's generalization performance [4]. In this section we therefore compare the sharpness of SAM-all to SAM-ON models (following the setup in [4], more details provided in Appendix B.10).

In Table 6, we report logit-normalized worst-case adaptive ℓ ∞ m -sharpness, the sharpness definition which achieved the best correlation with generalization error for CIFAR data in the large-scale study in [4]. We investigate a 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 hyperparameter-free method designed for accurate estimation of adversarial robustness. Instead of the input space, which is optimized in adversarial robustness, the method is adapted to the weight space for sharpness computation. In addition to the 20-step AutoPGD-sharpness, 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 logit-normalization, the sharpness definition reported in Table 6 corresponds exactly to the perturbation model that ASAM elementwise ℓ ∞ uses, and hence the 1-step sharpness reported should be fairly close the the objective that ASAM elementwise ℓ ∞ 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 ρ 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 SGD-trained 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 B.10 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 et al. [3] further showed that SAM-all leads to features of lower rank compared to vanilla optimizers. Following their setup we measure the feature rank for a WRN-28 and report our findings in Table 6 (final row). We find that SAM-ON also leads to features of lower rank compared to SGD, but observe no significant change compared to SAM-all .

5.3 SAM-ON can change the affine parameter values

In Figure 5 we show the distribution of the scaling parameter γ and the shift parameter β (as defined in Eq. 1) for a WRN-28 trained with SAM-ON and SAM-all . While there are only minor changes in the distribution of β , there is a clear shift in the distribution of γ towards larger values for SAM-ON . In Figure 9 in Appendix B.6 we study more SAM-variants and rediscover a pattern from Section 4.1: The distribution shift for γ only seems to appear for those variants, where the optimal ρ of SAM-ON shifts towards larger values compared to SAM-all (i.e. not for ASAM elementwiseℓ 2 ). We thus hypothesize that the perturbation model of ASAM elementwiseℓ 2 implicitly focuses mostly on perturbing the normalization layers already, which is why excluding them ( no-norm ) leads to a drastic performance decrease, whereas SAM-ON minimally changes the γ -distribution. In contrast, the other SAM-variants (for which the optimal ρ -value changes) may not perturb the normalization layers enough in SAM-all for them to be effective, which is why no-norm has little effect, while SAM-ON leads to a distinctive shift of the γ -distribution.

0.2

0.1

0.0

0.1

0.2

0.3

Value

0

250

500

Distribution of

0.2

0.1

0.0

0.1

0.2

0.3

Value

0

250

500

Distribution of

SAM  p=2 ON

SAM  p=2 all

6 Discussion and Conclusion

In recent years the method of sharpness-aware minimization (SAM) [23] has risen to prominence due to its demonstrated effectiveness and the community's long-standing interest in flat minima. In this work we show that only applying SAM to the normalization layers ( SAM-ON ) - typically 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 BatchNorm 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 to use varies. We do not see a consistent benefit of using reparameterization-invariant perturbation models compared to variants with fewer invariances. In particular, we find that layerwise ℓ 2 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 to 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, 24, 42] we focus on vision data. We provide a simple experiment in the language domain in Appendix B.5 indicating that the efficacy of SAM-ON might be preserved for language tasks, but more extensive benchmarking, as done by [7] is 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.

References

M. Abbas, Q. Xiao, L. Chen, P.-Y. Chen, and T. Chen. Sharp-MAML: Sharpness-aware model-agnostic meta learning. ICML , 2022.

M. Andriushchenko and N. Flammarion. Towards understanding sharpness-aware minimization. ICML , 2022.

M. Andriushchenko, D. Bahri, H. Mobahi, and N. Flammarion. Sharpness-aware minimization leads to low-rank features. NeurIPS , 2023.

M. Andriushchenko, F. Croce, M. Müller, M. Hein, and N. Flammarion. A modern look at the relationship between sharpness and generalization. preprint arXiv:2302.07011 , 2023.

S. Arora, Z. Li, and K. Lyu. Theoretical analysis of auto rate-tuning by Batch Normalization. ICLR , 2019.

J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer Normalization. Advances in NeurIPS Deep Learning Symposium , 2016.

D. Bahri, H. Mobahi, and Y. Tay. Sharpness-aware minimization improves language model generalization. ACL , 2022.

K. Behdin, Q. Song, A. Gupta, A. Acharya, D. Durfee, B. Ocejo, S. Keerthi, and R. Mazumder. mSAM: Micro-batch-averaged sharpness-aware minimization. preprint arXiv:2302.09693 , 2023.

P. Benz, C. Zhang, and I. S. Kweon. Batch normalization increases adversarial vulnerability and decreases adversarial transferability: A non-robust feature perspective. ICCV , 2021.

J. Bjorck, C. Gomes, B. Selman, and K. Q. Weinberger. Understanding batch normalization. NeurIPS , 2018.

A. Brock, S. De, S. L. Smith, and K. Simonyan. High-performance large-scale image recognition without normalization. ICML , 2021.

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.

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.

X. Chen, C.-J. Hsieh, and B. Gong. When Vision Transformers outperform ResNets without pre-training or strong data augmentations. ICLR , 2022.

X. Chen, C. Liang, D. Huang, E. 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. arXiv:2302.06675 , 2023.

F. Croce and M. Hein. Reliable evaluation of adversarial robustness with an ensemble of diverse parameter-free attacks. ICML , 2020.

N. Dimitriou and O. Arandjelovic. A new look at ghost normalization. preprint arXiv: 2007.08554 , 2020.

L. Dinh, R. Pascanu, S. Bengio, and Y. Bengio. Sharp minima can generalize for deep nets. ICML , 2017.

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. ICLR , 2021.

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.

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.

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.

P. Foret, A. Kleiner, H. Mobahi, and B. Neyshabur. Sharpness-aware minimization for efficiently improving generalization. ICLR , 2021.

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.

K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. CVPR , 2016.

S. Hochreiter and J. Schmidhuber. Simplifying neural nets by discovering flat minima. Advances in Neural Information Processing Systems , 1994.

S. Hochreiter and J. Schmidhuber. Flat minima. Neural Computation , 9(1):1-42, 1997.

E. Hoffer, I. Hubara, and D. Soudry. Train longer, generalize better: closing the generalization gap in large batch training of neural networks. NeurIPS , 2017.

G. Huang, Z. Liu, L. van der Maaten, and K. Q. Weinberger. Densely connected convolutional networks. CVPR , 2017.

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.

S. Ioffe and C. Szegedy. Batch Normalization: Accelerating deep network training by reducing internal covariate shift. ICML , 2015.

Y. Jiang*, B. Neyshabur*, H. Mobahi, D. Krishnan, and S. Bengio. Fantastic generalization measures and where to find them. ICLR , 2020.

J. Kaddour, L. Liu, R. Silva, and M. J. Kusner. When do flat minima optimizers work? NeurIPS , 2022.

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.

M. Kim, D. Li, S. X. Hu, and T. Hospedales. Fisher SAM: Information geometry and sharpness aware minimisation. ICML , 2022.

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.

J. Kwon, J. Kim, H. Park, and I. K. Choi. ASAM: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. ICML , 2021.

X. Lian and J. Liu. Revisit Batch Normalization: New understanding and refinement via composition optimization. AISTATS , 2019.

Y. Liu, S. Mai, X. Chen, C.-J. Hsieh, and Y. You. Towards efficient and scalable sharpness-aware minimization. CVPR , 2022.

I. Loshchilov and F. Hutter. Decoupled weight decay regularization, 2019.

K. Lyu, Z. Li, and S. Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. NeurIPS , 2022.

P. 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.

T. Möllenhoff and M. E. Khan. SAM as an optimal relaxation of Bayes. ICLR , 2023.

C. Na, S. V. Mehta, and E. Strubell. Train flat, then compress: Sharpness-aware minimization learns more compressible models. EMNLP , 2022.

H. Petzka, M. Kamp, L. Adilova, C. Sminchisescu, and M. Boley. Relative flatness and generalization. NeurIPS , 2021.

S. Santurkar, D. Tsipras, A. Ilyas, and A. Madry. How does batch normalization help optimization? NeurIPS , 2018.

K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image recognition, 2015.

A. Steiner, A. Kolesnikov, X. Zhai, R. Wightman, J. Uszkoreit, and L. Beyer. How to train your ViT? Data, augmentation, and regularization in Vision Transformers. TMLR , 2022.

C. Summers and M. J. Dinneen. Four things everyone should know to improve batch normalization. ICLR , 2020.

C. Szegedy, W. Zaremba, I. Sutskever, J. Bruna, D. Erhan, I. Goodfellow, and R. Fergus. Intriguing properties of neural networks. ICLR , 2014.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. NeurIPS , 2017.

Z. Wei, J. Zhu, and Y. Zhang. Sharpness-aware minimization alone can improve adversarial robustness. ICML AdvML-Frontiers Workshop, arXiv:2305.05392 , 2023.

R. Wightman. Pytorch image models. https://github.com/rwightman/ pytorch-image-models , 2019.

C. Xie and A. Yuille. Intriguing properties of adversarial training at scale. ICLR , 2020.

S. Xie, R. Girshick, P. Dollár, Z. Tu, and K. He. Aggregated residual transformations for deep neural networks. CVPR , 2017.

J. Xu, X. Sun, Z. Zhang, G. Zhao, and J. Lin. Understanding and improving layer normalization. NeurIPS , 2019.

Z. Yao, A. Gholami, K. Keutzer, and M. Mahoney. PyHessian: Neural networks through the lens of the Hessian. IEEE BigData, arXiv:1912.07145 , 2020.

M. Yazdanpanah, A. A. Rahman, M. Chaudhary, C. Desrosiers, M. Havaei, E. Belilovsky, and S. E. Kahou. Revisiting learnable affines for Batch Norm in few-shot transfer learning. CVPR , 2022.

S. Zagoruyko and N. Komodakis. Wide residual networks. In BMVC , 2016.

H. Zhang, Y. N. Dauphin, and T. Ma. Fixup initialization: Residual learning without normalization. ICLR , 2019.

J. Zhuang, B. Gong, L. Yuan, Y. Cui, H. Adam, N. Dvornek, S. Tatikonda, J. Duncan, and T. Liu. Surrogate gap minimization improves sharpness-aware training. ICLR , 2022.

Appendix

This appendix is structured as follows:

In Appendix A we provide more training details. In particular, we report the hyperparameters used for the CIFAR experiments in A.1 and for the ImageNet experiments in A.2. In A.3 we provide more details and a formal definition of the SAM-variants used throughout this paper.

In Appendix B we show additional experimental results for: CIFAR in B.1, ImageNet in B.3, and a machine translation task in B.5. In B.2 we provide additional ablation studies for sparse perturbation SSAM approaches and in B.4 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.7), SAM'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 on the sharpness evaluation and more ablations in B.10.

In Appendix C we provide a 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 , ∞} ) or in the definition of the normalization operator. We use SGD, the original SAM with no normalization and p = 2 , Fisher-SAM and the following ASAM-variants: elementwiseℓ ∞ , layerwiseℓ 2 , and elementwiseℓ 2 . For the ViT-experiments, we use AdamW instead of SGD. For each of the ASAM-variants, we normalize both bias and weight parameters and set η = 0 . Additionally, we employ the original ASAM-algorithm, where the bias parameters are not normalized and η = 0 . 01 . We train all models on a single GPU for 200 epochs, and m-sharpness is not employed (unless indicated otherwise). For ResNets, we follow [37] and adopt a learning rate of 0.1, momentum of 0.9, weight decay of 0 . 0005 and use label smoothing with a factor of 0.1. We use both basic augmentations (random cropping and flipping) and strong augmentations (basic+AutoAugment). For ViTs we use 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 ρ we considered for each method can be found in Table 7. The ResNet-networks we considered for the CIFAR-experiments in the main paper are ResNet56 (RN56) [25], ResNeXt-29-32x4d (RNxT) [55], and WideResNet-28-10 (WRN) [59]. We adopted the ViTs to CIFAR by setting the image-size to 32 and patch-size to 4.

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 elementwiseℓ 2 we used the hyperparameters from [23] and [37]. For the layerwise ℓ 2 and elementwiseℓ ∞ we tried two ρ -values per configuration and report the results of the better one (named ρ (reported) in the table). ρ (discarded) refers to the ρ 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 SAM-configuration and reported the best one.

A.3 SAM variants

Here, we provide a more comprehensive overview of the SAM-variants used throughout the experiments. To this end, we first recall the definition of the (A)SAM-perturbation (Eq. (5) in the main paper):

with the normalization operator T i w , which is diagonal for all variants. We note that SAM-ON can be formally defined as using the conventional (A)SAM-algorithm but setting all entries T i w = 0 if w i is not a normalization parameter. This leads to a change of the perturbation ϵ according to Eq. (5). Importantly, the magnitude of ϵ is still ρ , since both the nominator and the denominator of Eq. (5) change. We provide an overview over all (A)SAM-variants and their respective perturbation models in Table 9.

B Further Experimental Results

B.1 SAM-ON on CIFAR

We omitted the results for ResNet-like models on CIFAR-10 in the main paper. Those are thus reported in Table 10. Due to 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 ρ in Figure 6 for a WRN-28 and in Figure 7 for a ViT-S on CIFAR-100. We show results for various VGG-models [47] and DenseNet-100 [30] for CIFAR-10/100 in Table 11 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 for 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 to Figure 4 for ResNet-18 on CIFAR-10 data in the main paper we provide an accompanying Figure 8 which includes error bars

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 ρ = 0 . 2 on CIFAR100 (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 .

10

3

10

2

10

1

10 0

10 1

80.5

81.0

81.5

82.0

82.5

83.0

83.5

84.0

84.5

85.0

Test Accuracy (%)

all

ON

no-norm

ASAM, p=2.0, elem. eta=0.01

ASAM, p=2.0, elem.

ASAM, p=2.0, layer

ASAM, p=inf, elem.

Fisher, p=2.0 eta=1.0

SAM, p=2.0

SGD

all

ON

no-norm

10

3

10

2

10

1

10 0

10 1

64

65

66

67

68

69

70

71

Test Accuracy (%)

all

ON

no-norm

AdamW

ASAM, p=2.0, elem.

ASAM, p=2.0, elem. eta=0.01

ASAM, p=2.0, layer

ASAM, p=inf, elem.

Fisher, p=2.0 eta=1.0

SAM, p=2.0

all

ON

no-norm

10

2

10

1

10 0

10 1

95.25

95.50

95.75

96.00

96.25

96.50

Test Accuracy (%)

SGD

SAM

SSAM-F 50%

SSAM-D 50%

SSAM-F 99.93%

SSAM-D 99.93%

SAM-ON

SAM-rand

B.3 Finetuning from ImageNet-21k

Since ViTs are commonly trained on large-scale datasets and then fine-tuned, we investigate this scenario for SAM-ON . In particular, we consider a ViT-S pretrained on ImageNet-21k from [48]. We fine-tune for 9 epochs with SGD for a range of SAM-variants with their respective SAM-ON counterpart. For each setup, we probe three values of ρ and report the best result in Table 13. We find in this setting that SAM-ON performs on par with SAM-all although there are small differences across SAM variants: for layerwiseℓ 2 SAM-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.

B.4 Adversarial robustness

Here, we provide additional results and extend the discussion on adversarial robustness from Section 4.2. In a 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 [9, 54], it is interesting to investigate whether the robustness properties of SAM can be preserved when training with SAM-ON instead of SAM-all. In Table 15 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 a small margin. For a WRN-28-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 over SAM-all, but for the ASAM-elementwiseℓ ∞ 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 in SAM, illustrating again their special role.

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 16: 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 a 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.

B.6 Weight distribution after training

In order to get a better understanding of the impact of SAM-ON on γ and β (as defined in Eq. 1), we train a WideResNet-28-10 with different SAM-variants and both SAM-ON and all . We show the distribution of | w i | , i.e. the parameter magnitudes, at the end of training for different layer types in Figure 9. Different to the discussion in Section 5.3, we show the y -axis on log-scale, in order to inspect more nuanced differences. For elementwise ℓ 2 there is no strong change in the distribution of the BatchNorm parameters between all and SAM-ON . For elementwise ℓ ∞ , layerwise ℓ 2 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 we observed when comparing the optimal ρ -value for all and SAM-ON in Table 7: The optimal ρ of elementwise ℓ 2 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 ℓ 2 variant showed a strong performance decrease in no-norm (Figure 1), 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.

B.7 Removing the affine parameters

Frankle et 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, we set γ = 1 and β = 0 and train the remaining parameters with SAM. The results are shown for a WRN-28 in Figure 10, where it can be seen that fixing the normalization parameters ( fix-norm ) does not lead to a 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 a ResNet-101 and a 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.

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 are 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 11, 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 regularization and stochasticity to the normalization parameters. To this end, we apply dropout solely on the normalization layers (Figure 11, left) and find that this is not the case.

0.0

0.5

1.0

10 1

10 3

SGD

Distribution of

0.0

0.5

1.0

10 1

10 3

Distribution of

0.0

0.5

1.0

10 2

10 4

Distribution of FC weights

0.0

0.5

1.0

10 1

10 3

SAM

0.0

0.5

1.0

10 1

10 3

0.0

0.5

1.0

10 1

10 3

OnlyNorm

all

0.0

0.5

1.0

10 1

10 3

elementwise p=2

0.0

0.5

1.0

10 1

10 3

0.0

0.5

1.0

10 2

10 4

all

OnlyNorm

0.0

0.5

1.0

10 1

10 3

elementwise p=

0.0

0.5

1.0

10 1

10 3

0.0

0.5

1.0

10 1

10 3

all

OnlyNorm

0.0

0.5

1.0

i

10 1

10 3

layerwise p=2

0.0

0.5

1.0

i

10 1

10 3

0.0

0.5

1.0

wi

10 1

10 3

OnlyNorm

all

10

2

10

1

10 0

81.5

82.0

82.5

83.0

83.5

84.0

Test Accuracy (%)

fix-norm, SAM, p=2.0

SAM, p=2.0 all

SAM, p=2.0 OnlyNorm

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

/ dropout rate

80

81

82

83

84

Test Accuracy (%)

WRN28-10 CIFAR100 - Dropout on Normalization layers

SAM all

SAM ON

SGD dropout-ON via params

SGD dropout-ON via activations

10

1

10 0

10 1

77

78

79

80

81

82

83

84

85

Test Accuracy (%)

WRN28-10 CIFAR100 - Weight decay

SAM all

SAM ON

SGD

no wd norm, SAM all

no wd norm, SGD

no wd, SAM all

no wd, SAM ON

no wd, SGD

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 error, while in other cases people use it as a 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 generalization of deep neural networks and the flatness of the loss-surface [26, 32, 22, 34, 18]. Recently, Andriushchenko et al. [4] conducted a large-scale study for a range of models, datasets, and sharpnessdefinitions, finding that 'while there definitely exist restricted settings where correlation between sharpness and generalization is significantly positive (e.g., for ResNets on CIFAR-10 with a specific combination of augmentations and mixup) it is not true anymore when we compare all models jointly' and concluding 'that one should avoid blanket statements like flat minima generalize better '. In order to evaluate sharpness, we therefore adopt their setup and choose the best-performing sharpness measure for CIFAR from their study, which is logit-normalized elementwise-adaptive worst-caseℓ ∞ -m -sharpness. This is, m -sharpness s m w is defined as the largest possible change in loss within the adaptive perturbation model defined in 4,

where T i w = | w i | , p = ∞ and D m returns data batches of size m . ρ here denotes the size of the ball over which sharpness is evaluated and is not to be confused with the ρ 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 ℓ 2 -norm in order to mitigate the scale-sensitivity of classification losses. In practice, Andriushchenko et al. [4] compute s m w over a subset of the train set of size 1024 and use m = 128 , i.e. average 8 batches. We use a subset of size 2048 in order to obtain more reliable sharpness estimates, and adopt m = 128 . The maximization in (6) is performed with AutoPGD [16], a hyperparameter-free method designed for accurate estimation of adversarial robustness. It is to note that except for the logit-normalization, the sharpness definition reported in Table 6 corresponds exactly to the perturbation model that ASAM elementwise ℓ ∞ uses, and hence the 1-step sharpness reported should be fairly close the the objective that ASAM elementwise ℓ ∞ actually minimizes during training. While ASAM elementwise ℓ ∞ yields slightly smaller sharpness values than the conventional SAM algorithm, the differences are rather small when compared to the significantly sharper SAM-ON models. For the results in Table 6 in the main paper we tuned the sharpness radius ρ such that we obtain sharpness values similar to those reported to yield the highest correlation in Andriushchenko et al. [4]. In Table 18 we report sharpness values for a ResNeXt-model, in addition to the WRN-28 from the main paper. In all cases the 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 a WRN-28. SAM-ON is sharper than SAM-all with respect to most metrics, 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 et al. [4].

C Convergence Analysis

We provide in 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 in [2, 42].

Our assumptions:

Assumption C.1. We assume function f : R n → R to be L -smooth: there exists L > 0 such that

Assumption C.2. There exists M > 0 for any sample x i such that

Remark C.3 . If Assumption C.1 holds ( L -smoothness), then ∀ v, w ∈ R n :

This well-known result can be derived using the fundamental theorem of calculus and CauchySchwartz.

Remark C.4 . Assumption C.2 guarantees that the variance of the stochastic gradient is less than M .

SAM-ON. In the following we shall denote the true gradient as ∇ f ( w ) and the noisy observation gradient as g ( w ) . The gradient of the loss of the i th training example is denoted as g x i ( w ) . We partition the neural network parameters layer-wise as w = { w N , w A } , with w N ∈ R n N , w A ∈ R n A , n = n N + n A , where w N represent the normalization layer parameters and w A all other layers. The

iteration for w N is:

and for w A is:

For term B we obtain:

Theorem C.5. Assuming C.1 and C.2, h ≤ 1 /L , we obtain:

with w ∗ the optimal solution to f ( w ) .

Proof. From Assumption C.1 and thus Remark C.3 it follows that:

Taking the double expectation gives (because unbiased gradient and Assumption C.2 and Remark C.4):

For term A we obtain using Assumption C.1:

Using xy ≤ 1 2 ∥ x ∥ 2 2 + 1 2 ∥ y ∥ 2 2 and Assumption C.1 we get for C :

Plugging this into (17) gives:

In T iterations we obtain using a telescoping sum:

Peer reviews:
comments: 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.

STRENGTHS:
The SAM-ON proposed method yields several benefits:
1. improves upon SAM-all and carious other SAM-variants
2. achieves this by perturbing only the normalization layers, which correspond to a small percentage of the total parameters

WEAKNESSES:
Weaknesses of this paper include:
1. gains in Table 1 are marginal
2. originality of approach is lacking given prior literature on SAM variants
3. 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?
4. it is unclear whether SAM-ON is truly better than other SAM variants in out of distribution robustness, e.g. WILDS benchmark
5. results are only shown for CIFAR and ImageNet datasets; more datasets and benchmarks need to be considered

QUESTIONS:
see weaknesses above

One 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-27

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAs 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nSee Weaknesses.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 9, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nYes. This paper discusses the limitation of this paper, where they only focus on the vision data.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: 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.

STRENGTHS:
Strengths:
1. This paper focus on an important problem, SAM is very important for improving generalization.
2. The experimental finding is very interesting and can efficiently reduce the computation.
3. The authors try to evaluate the performance of proposed method on various datasets and models.

WEAKNESSES:
Weakness:
1. 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.
2. 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.
3. Although SAM-ON can reduce the computation, the training efficiency cannot be significantly improved.

QUESTIONS:
Question:
1. Do you try to run the experiments for multiple times with different seeds.
2. 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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-08