Sharpness-Aware Minimization for Efficiently Improving Generalization — part 2/4

Cleaned-text code points [15877, 31309). Bibliography omitted here; source.md includes it.

In order to assess SAM’s efficacy, we apply it to a range of different tasks, including image classification from scratch (including on CIFAR-10, CIFAR-100, and ImageNet), finetuning pretrained models, and learning with noisy labels. In all cases, we measure the benefit of using SAM by simply replacing the optimization procedure used to train existing models with SAM, and computing the resulting effect on model generalization. As seen below, SAM materially improves generalization performance in the vast majority of these cases.

### 3.1 Image Classification From Scratch

We first evaluate SAM’s impact on generalization for today’s state-of-the-art models on CIFAR-10 and CIFAR-100 (without pretraining): WideResNets with ShakeShake regularization ( Zagoruyko & Komodakis, 2016 ; Gastaldi, 2017 ) and PyramidNet with ShakeDrop regularization ( Han et al., 2016 ; Yamada et al., 2018 ) . Note that some of these models have already been heavily tuned in prior work and include carefully chosen regularization schemes to prevent overfitting; therefore, significantly improving their generalization is quite non-trivial. We have ensured that our implementations’ generalization performance in the absence of SAM matches or exceeds that reported in prior work ( Cubuk et al., 2018 ; Lim et al., 2019 )

All results use basic data augmentations (horizontal flip, padding by four pixels, and random crop). We also evaluate in the setting of more advanced data augmentation methods such as cutout regularization ( Devries & Taylor, 2017 ) and AutoAugment ( Cubuk et al., 2018 ) , which are utilized by prior work to achieve state-of-the-art results.

SAM has a single hyperparameter ρ \rho (the neighborhood size), which we tune via a grid search over { 0.01 , 0.02 , 0.05 , 0.1 , 0.2 , 0.5 } \{0.01,0.02,0.05,0.1,0.2,0.5\} using 10% of the training set as a validation set 3 3 3 We found ρ = 0.05 \rho=0.05 to be a solid default value, and we report in appendix C.3 the scores for all our experiments, obtained with ρ = 0.05 \rho=0.05 without further tuning. . Please see appendix C.1 for the values of all hyperparameters and additional training details. As each SAM weight update requires two backpropagation operations (one to compute ϵ ^ ​ ( 𝒘 ) \hat{\bm{\epsilon}}(\bm{w}) and another to compute the final gradient), we allow each non-SAM training run to execute twice as many epochs as each SAM training run, and we report the best score achieved by each non-SAM training run across either the standard epoch count or the doubled epoch count 4 4 4 Training for longer generally did not improve accuracy significantly, except for the models previously trained for only 200 epochs and for the largest, most regularized model (PyramidNet + ShakeDrop). . We run five independent replicas of each experimental condition for which we report results (each with independent weight initialization and data shuffling), reporting the resulting mean error (or accuracy) on the test set, and the associated 95% confidence interval. Our implementations utilize JAX ( Bradbury et al., 2018 ) , and we train all models on a single host having 8 NVidia V100 GPUs 5 5 5 Because SAM’s performance is amplified by not syncing the perturbations, data parallelism is highly recommended to leverage SAM’s full potential (see Section 4 for more details). . To compute the SAM update when parallelizing across multiple accelerators, we divide each data batch evenly among the accelerators, independently compute the SAM gradient on each accelerator, and average the resulting sub-batch SAM gradients to obtain the final SAM update.

As seen in Table 1 , SAM improves generalization across all settings evaluated for CIFAR-10 and CIFAR-100. For example, SAM enables a simple WideResNet to attain 1.6% test error, versus 2.2% error without SAM. Such gains have previously been attainable only by using more complex model architectures (e.g., PyramidNet) and regularization schemes (e.g., Shake-Shake, ShakeDrop); SAM provides an easily-implemented, model-independent alternative. Furthermore, SAM delivers improvements even when applied atop complex architectures that already use sophisticated regularization: for instance, applying SAM to a PyramidNet with ShakeDrop regularization yields 10.3% error on CIFAR-100, which is, to our knowledge, a new state-of-the-art on this dataset without the use of additional data.

Beyond CIFAR-{10, 100}, we have also evaluated SAM on the SVHN ( Netzer et al., 2011 ) and Fashion-MNIST datasets ( Xiao et al., 2017 ) . Once again, SAM enables a simple WideResNet to achieve accuracy at or above the state-of-the-art for these datasets: 0.99% error for SVHN, and 3.59% for Fashion-MNIST. Details are available in appendix B.1 .

To assess SAM’s performance at larger scale, we apply it to ResNets ( He et al., 2015 ) of different depths (50, 101, 152) trained on ImageNet ( Deng et al., 2009 ) . In this setting, following prior work ( He et al., 2015 ; Szegedy et al., 2015 ) , we resize and crop images to 224-pixel resolution, normalize them, and use batch size 4096, initial learning rate 1.0, cosine learning rate schedule, SGD optimizer with momentum 0.9, label smoothing of 0.1, and weight decay 0.0001. When applying SAM, we use ρ = 0.05 \rho=0.05 (determined via a grid search on ResNet-50 trained for 100 epochs). We train all models on ImageNet for up to 400 epochs using a Google Cloud TPUv3 and report top-1 and top-5 test error rates for each experimental condition (mean and 95% confidence interval across 5 independent runs).

As seen in Table 2 , SAM again consistently improves performance, for example improving the ImageNet top-1 error rate of ResNet-152 from 20.3% to 18.4%. Furthermore, note that SAM enables increasing the number of training epochs while continuing to improve accuracy without overfitting. In contrast, the standard training procedure (without SAM) generally significantly overfits as training extends from 200 to 400 epochs.

### 3.2 Finetuning

Transfer learning by pretraining a model on a large related dataset and then finetuning on a smaller target dataset of interest has emerged as a powerful and widely used technique for producing high-quality models for a variety of different tasks. We show here that SAM once again offers considerable benefits in this setting, even when finetuning extremely large, state-of-the-art, already high-performing models.

In particular, we apply SAM to finetuning EfficentNet-b7 (pretrained on ImageNet) and EfficientNet-L2 (pretrained on ImageNet plus unlabeled JFT; input resolution 475) ( Tan & Le, 2019 ; Kornblith et al., 2018 ; Huang et al., 2018 ) . We initialize these models to publicly available checkpoints 6 6 6 https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet trained with RandAugment (84.7% accuracy on ImageNet) and NoisyStudent (88.2% accuracy on ImageNet), respectively. We finetune these models on each of several target datasets by training each model starting from the aforementioned checkpoint; please see the appendix for details of the hyperparameters used. We report the mean and 95% confidence interval of top-1 test error over 5 independent runs for each dataset.

As seen in Table 3 , SAM uniformly improves performance relative to finetuning without SAM. Furthermore, in many cases, SAM yields novel state-of-the-art performance, including 0.30% error on CIFAR-10, 3.92% error on CIFAR-100, and 11.39% error on ImageNet.

### 3.3 Robustness to Label Noise

The fact that SAM seeks out model parameters that are robust to perturbations suggests SAM’s potential to provide robustness to noise in the training set (which would perturb the training loss landscape). Thus, we assess here the degree of robustness that SAM provides to label noise.

In particular, we measure the effect of applying SAM in the classical noisy-label setting for CIFAR-10, in which a fraction of the training set’s labels are randomly flipped; the test set remains unmodified (i.e., clean). To ensure valid comparison to prior work, which often utilizes architectures specialized to the noisy-label setting, we train a simple model of similar size (ResNet-32) for 200 epochs, following Jiang et al. (2019) . We evaluate five variants of model training: standard SGD, SGD with Mixup ( Zhang et al., 2017 ) , SAM, and ”bootstrapped” variants of SGD with Mixup and SAM (wherein the model is first trained as usual and then retrained from scratch on the labels predicted by the initially trained model). When applying SAM, we use ρ = 0.1 \rho=0.1 for all noise levels except 80%, for which we use ρ = 0.05 \rho=0.05 for more stable convergence. For the Mixup baselines, we tried all values of α ∈ { 1 , 8 , 16 , 32 } \alpha\in\{1,8,16,32\} and conservatively report the best score for each noise level.

As seen in Table 4 , SAM provides a high degree of robustness to label noise, on par with that provided by state-of-the art procedures that specifically target learning with noisy labels. Indeed, simply training a model with SAM outperforms all prior methods specifically targeting label noise robustness, with the exception of MentorMix ( Jiang et al., 2019 ) . However, simply bootstrapping SAM yields performance comparable to that of MentorMix (which is substantially more complex).

## 4 Sharpness and Generalization Through the Lens of SAM

### 4.1 m m -sharpness

Though our derivation of SAM defines the SAM objective over the entire training set, when utilizing SAM in practice, we compute the SAM update per-batch (as described in Algorithm 1 ) or even by averaging SAM updates computed independently per-accelerator (where each accelerator receives a subset of size m m of a batch, as described in Section 3). This latter setting is equivalent to modifying the SAM objective (equation 1 ) to sum over a set of independent ϵ \epsilon maximizations, each performed on a sum of per-data-point losses on a disjoint subset of m m data points, rather than performing the ϵ \epsilon maximization over a global sum over the training set (which would be equivalent to setting m m to the total training set size). We term the associated measure of sharpness of the loss landscape m m -sharpness .

To better understand the effect of m m on SAM, we train a small ResNet on CIFAR-10 using SAM with a range of values of m m . As seen in Figure 3 (middle), smaller values of m m tend to yield models having better generalization ability. This relationship fortuitously aligns with the need to parallelize across multiple accelerators in order to scale training for many of today’s models.

Intriguingly, the m m -sharpness measure described above furthermore exhibits better correlation with models’ actual generalization gaps as m m decreases, as demonstrated by Figure 3 (right) 7 7 7 We follow the rigorous framework of Jiang et al. (2019) , reporting the mutual information between the m m -sharpness measure and generalization on the two publicly available tasks from the Predicting generalization in deep learning NeurIPS2020 competition. https://competitions.codalab.org/competitions/25301 . In particular, this implies that m m -sharpness with m < n m<n yields a better predictor of generalization than the full-training-set measure suggested by Theorem 1 in Section 2 above, suggesting an interesting new avenue of future work for understanding generalization.

### 4.2 Hessian Spectra

Motivated by the connection between geometry of the loss landscape and generalization, we constructed SAM to seek out minima of the training loss landscape having both low loss value and low curvature (i.e., low sharpness). To further confirm that SAM does in fact find minima having low curvature, we compute the spectrum of the Hessian for a WideResNet40-10 trained on CIFAR-10 for 300 steps both with and without SAM (without batch norm, which tends to obscure interpretation of the Hessian), at different epochs during training. Due to the parameter space’s dimensionality, we approximate the Hessian spectrum using the Lanczos algorithm of Ghorbani et al. (2019) .

Figure 3 (left) reports the resulting Hessian spectra. As expected, the models trained with SAM converge to minima having lower curvature, as seen in the overall distribution of eigenvalues, the maximum eigenvalue ( λ max \lambda_{\max} ) at convergence (approximately 24 without SAM, 1.0 with SAM), and the bulk of the spectrum (the ratio λ max / λ 5 \lambda_{\max}/\lambda_{5} , commonly used as a proxy for sharpness ( Jastrzebski et al., 2020 ) ; up to 11.4 without SAM, and 2.6 with SAM).

## 5 Related Work

The idea of searching for “flat” minima can be traced back to Hochreiter & Schmidhuber (1995) , and its connection to generalization has seen significant study ( Shirish Keskar et al., 2016 ; Dziugaite & Roy, 2017 ; Neyshabur et al., 2017 ; Dinh et al., 2017 ) . In a recent large scale empirical study, Jiang et al. (2019) studied 40 complexity measures and showed that a sharpness-based measure has highest correlation with generalization, which motivates penalizing sharpness. Hochreiter & Schmidhuber (1997) was perhaps the first paper on penalizing the sharpness, regularizing a notion related to Minimum Description Length (MDL). Other ideas which also penalize sharp minima include operating on diffused loss landscape ( Mobahi, 2016 ) and regularizing local entropy ( Chaudhari et al., 2016 ) . Another direction is to not penalize the sharpness explicitly, but rather average weights during training; Izmailov et al. (2018) showed that doing so can yield flatter minima that can also generalize better. However, the measures of sharpness proposed previously are difficult to compute and differentiate through. In contrast, SAM is highly scalable as it only needs two gradient computations per iteration. The concurrent work of Sun et al. (2020) focuses on resilience to random and adversarial corruption to expose a model’s vulnerabilities; this work is perhaps closest to ours. Our work has a different basis: we develop SAM motivated by a principled starting point in generalization, clearly demonstrate SAM’s efficacy via rigorous large-scale empirical evaluation, and surface important practical and theoretical facets of the procedure (e.g., m m -sharpness). The notion of all-layer margin introduced by Wei & Ma (2020) is closely related to this work; one is adversarial perturbation over the activations of a network and the other over its weights, and there is some coupling between these two quantities.

## 6 Discussion and Future Work

In this work, we have introduced SAM, a novel algorithm that improves generalization by simultaneously minimizing loss value and loss sharpness; we have demonstrated SAM’s efficacy through a rigorous large-scale empirical evaluation. We have surfaced a number of interesting avenues for future work. On the theoretical side, the notion of per-data-point sharpness yielded by m m -sharpness (in contrast to global sharpness computed over the entire training set, as has typically been studied in the past) suggests an interesting new lens through which to study generalization. Methodologically, our results suggest that SAM could potentially be used in place of Mixup in robust or semi-supervised methods that currently rely on Mixup (giving, for instance, MentorSAM). We leave to future work a more in-depth investigation of these possibilities.

## 7 Acknowledgments