PaperlayerView exact revision on arXiv

Source-linked AI summary

Sharpness-Aware Minimization for Efficiently Improving Generalization

Pierre Foret, Ariel Kleiner, Hossein Mobahi, Behnam Neyshabur

arXiv:2010.01412v3cs.LGstat.ML

TL;DR

Modern overparameterized models can achieve low training loss while generalizing differently, motivating methods beyond training-loss minimization. The paper introduces Sharpness-Aware Minimization, which jointly minimizes loss and sharpness by seeking uniformly low-loss neighborhoods. SAM improves generalization across diverse tasks and models, provides label-noise robustness on par with specialized procedures, and motivates further study of per-data-point sharpness.

  • Problem

    Overparameterized models can have similar training loss but substantially different generalization, limiting training loss as a sufficient selection criterion.

  • Method

    SAM simultaneously minimizes training loss and loss sharpness by seeking parameters whose neighborhoods have uniformly low loss.

  • Results

    SAM improves generalization across diverse vision tasks and models, achieves novel state-of-the-art performance on several tasks, and matches specialized procedures’ robustness to label noise.

  • Takeaways & Limitations

    SAM offers an efficient, scalable, and effective approach that is complementary to existing techniques for improving model generalization.

  • Takeaways & Limitations

    The second-order terms omitted by SAM’s approximation degraded performance in an initial experiment, requiring further investigation.

Abstract

from arXiv · show

In today's heavily overparameterized models, the value of the training loss provides few guarantees on model generalization ability. Indeed, optimizing only the training loss value, as is commonly done, can easily lead to suboptimal model quality. Motivated by prior work connecting the geometry of the loss landscape and generalization, we introduce a novel, effective procedure for instead simultaneously minimizing loss value and loss sharpness. In particular, our procedure, Sharpness-Aware Minimization (SAM), seeks parameters that lie in neighborhoods having uniformly low loss; this formulation results in a min-max optimization problem on which gradient descent can be performed efficiently. We present empirical results showing that SAM improves model generalization across a variety of benchmark datasets (e.g., CIFAR-10, CIFAR-100, ImageNet, finetuning tasks) and models, yielding novel state-of-the-art performance for several. Additionally, we find that SAM natively provides robustness to label noise on par with that provided by state-of-the-art procedures that specifically target learning with noisy labels. We open source our code at \url{https://github.com/google-research/sam}.

1 INTRODUCTION

Overparameterized models can fit training data yet generalize differently, making training-loss minimization alone insufficient. The paper introduces SAM, which targets uniformly low-loss neighborhoods and improves generalization across tasks and models, including robustness to label noise.

  • Overparameterized neural networks can memorize training data, so selected parameters must generalize beyond the training set.
  • Training-loss minimization can select local or global minima with similar loss but substantially different generalization performance.
  • Efficient algorithms that directly seek flatter minima and improve generalization across state-of-the-art models have remained elusive.
  • SAM simultaneously minimizes loss value and loss sharpness by seeking parameters whose neighborhoods have uniformly low loss.This approach is presented as an efficient and easily implemented procedure.
  • SAM improves generalization across widely studied vision tasks and models, including CIFAR-10, CIFAR-100, ImageNet, and finetuning tasks.The paper reports novel state-of-the-art performance on several studied tasks.
  • SAM provides robustness to label noise on par with state-of-the-art procedures designed specifically for noisy-label learning.

2 SHARPNESS-AWARE MINIMIZATION (SAM)

SAM replaces minimizing training loss at a single parameter value with minimizing loss across a neighborhood, targeting parameters with uniformly low loss and curvature. The method approximates this inner maximization efficiently and applies stochastic gradient descent to the resulting objective.

  • Overparameterized models can have multiple minima with similar training loss but substantially different generalization, making direct training-loss minimization insufficient.
  • SAM selects parameters whose neighborhoods have uniformly low training loss, combining low loss with low sharpness rather than optimizing loss only at the selected parameters.The sharpness term measures how quickly training loss increases under nearby parameter perturbations.
  • The SAM objective adds a neighborhood-wise loss maximization to training loss and uses an L2 regularization term with hyperparameter λ.The paper motivates the formulation using a generalization bound and substitutes the bound-dependent function with λ||w||2.
  • A first-order Taylor expansion approximates the inner maximization, producing a perturbation obtained from a classical dual-norm problem.For p-norm maximization, the dual exponent satisfies 1/p + 1/q = 1.
  • SAM computes the perturbed gradient using automatic differentiation and tractable Hessian-vector products, while the final effective approximation drops second-order terms.Including second-order terms degraded performance in an initial experiment, motivating further investigation.
  • The final algorithm samples a batch, computes its gradient and perturbation, evaluates the perturbed gradient, and updates weights with a standard optimizer such as SGD.The algorithm uses equation 3 for the SAM gradient approximation and repeats these steps until convergence.

3 EMPIRICAL EVALUATION

SAM improves generalization across image classification, large-scale ImageNet training, finetuning, and noisy-label settings, often achieving state-of-the-art results.

  • 3.1 IMAGE CLASSIFICATION FROM SCRATCH: SAM improves generalization across all evaluated CIFAR-10 and CIFAR-100 settings, including complex architectures with existing regularization.A WideResNet reaches 1.6% test error versus 2.2% without SAM, while PyramidNet with ShakeDrop reaches 10.3% error on CIFAR-100.
  • 3.1 IMAGE CLASSIFICATION FROM SCRATCH: SAM achieves 0.99% error on SVHN and 3.59% error on Fashion-MNIST with a simple WideResNet.
  • 3.1 IMAGE CLASSIFICATION FROM SCRATCH: 1.9 percentage points: SAM reduces ResNet-152 ImageNet top-1 error from 20.3% to 18.4%.SAM also permits longer training to improve accuracy, whereas standard training generally overfits from 200 to 400 epochs.
  • 3.2 FINETUNING: SAM uniformly improves finetuning performance and yields novel state-of-the-art errors of 0.30% on CIFAR-10, 3.92% on CIFAR-100, and 11.39% on ImageNet.
  • 3.3 ROBUSTNESS TO LABEL NOISE: SAM provides label-noise robustness on par with specialized state-of-the-art procedures, outperforming prior methods except MentorMix.Bootstrapping SAM achieves performance comparable to MentorMix, which is substantially more complex.

4 SHARPNESS AND GENERALIZATION THROUGH THE LENS OF SAM

The analysis examines SAM through batch-level perturbation structure, curvature, and sharpness measures. Smaller accelerator batch subsets improve generalization and make m-sharpness more predictive of generalization gaps.

  • Batch-level SAM: SAM updates are computed per batch or independently per accelerator, modifying the objective through separate perturbation maximizations.
  • Effect of m: Smaller values of m tend to produce better generalization, aligning with the need to parallelize training across accelerators.
  • Effect of m: As m decreases, m-sharpness exhibits better correlation with actual generalization gaps than the full-training-set measure.
  • Hessian spectrum: SAM converges to minima with lower curvature than standard SGD, as shown by the Hessian spectrum.
  • Hessian spectrum: At convergence, λmax is approximately 1.0 with SAM versus 24 without SAM, while λmax/λ5 reaches 2.6 versus 11.4.

5 RELATED WORK

Prior work links flat minima and loss-landscape geometry to generalization, motivating sharpness-based approaches such as SAM.

  • 5 RELATED WORK: Sharpness-based measures have been studied as indicators of generalization, with one large empirical study finding the highest correlation among 40 complexity measures.

6 DISCUSSION AND FUTURE WORK

The discussion highlights m-sharpness as a new lens for generalization and identifies methodological and theoretical directions for future work. It also presents a sharpness-based generalization bound under a stated perturbation assumption.

  • Future work: SAM motivates m-sharpness, which measures sharpness per data point rather than globally across the training set.
  • Future work: The authors suggest investigating whether SAM can replace Mixup in robust or semi-supervised methods, including MentorSAM.
  • Generalization bound: Theorem 2 states a generalization bound based on sharpness for any ρ > 0 and distribution D, with probability 1 − δ over the training set.
  • Generalization bound: The bound assumes Gaussian perturbations do not decrease test error, an expectation stated for final solutions but not necessarily for every parameter vector.
  • Generalization bound: The proof uses PAC-Bayesian bounds and selects among predefined prior standard deviations because the prior must be chosen before observing the training data.

B ADDITIONAL EXPERIMENTAL RESULTS

Additional experiments report that SAM reaches or improves state-of-the-art accuracy on SVHN and Fashion-MNIST using a simple WideResNet.

  • Results: 0.99% error rate on SVHN and 3.59% on Fashion-MNIST are reported for SAM with a simple WideResNet.
  • SVHN: The SVHN experiment uses all available training data, including 73,257 digits and 531,131 additional samples.
  • Results: Table 5 reports results for both SVHN and Fashion-MNIST.

C.1 HYPERPARAMETERS FOR EXPERIMENTS

The experiments use dataset- and model-specific training configurations, with hyperparameters selected by grid search for several datasets and fixed training schedules for ImageNet.

  • CIFAR, SVHN, and Fashion-MNIST: For CIFAR-10, CIFAR-100, SVHN, and Fashion-MNIST, learning rate and weight decay are selected by joint grid search before applying SAM.
  • ImageNet: ImageNet ResNet models are trained for 100, 200, or 400 epochs on 32 Google Cloud TPUv3 cores with batch size 4096.
  • Noisy labels: For noisy-label experiments, ρ is selected by grid search using accuracy on a validation set containing 10% of the usual CIFAR training samples.
  • Reported configurations: Tables 6 and 7 report the hyperparameters used for CIFAR-10, CIFAR-100, SVHN, and Fashion-MNIST results.
  • Finetuning: Finetuning models use public checkpoints, replace and randomly initialize the final dense layer, and train with weight decay 1e−5 and learning rate 0.016.

C.3 EXPERIMENTAL RESULTS WITH ρ = 0.05

The experiments evaluate SAM with ρ = 0.05 to assess whether performance remains strong without finely tuning this hyperparameter.

  • Sensitivity: ρ = 0.05 is used for the CIFAR and finetuning experiments reported in the sensitivity tables.
  • Sensitivity: The same ρ = 0.05 value is used for all ImageNet experiments.
  • Sensitivity: The experiment is motivated by testing whether SAM performs when ρ is not finely tuned.

C.4 ABLATION OF THE SECOND ORDER TERMS

The ablations examine whether SAM’s efficient approximations preserve useful update directions and which perturbation choices best support generalization. Discarding second-order terms remains effective, while Euclidean adversarial perturbations outperform random perturbations in the reported example.

  • During the first half of training, discarding second-order terms leaves update directions nearly unchanged, with cosine similarity very close to 1.The comparison uses WideResNet-40x2 on CIFAR-10 and evaluates first- versus second-order SAM updates.
  • The paper reports anecdotal evidence that SAM makes finetuning more robust to overtraining.
  • Near convergence, the similarity between first- and second-order updates weakens, but omitting second-order terms reaches lower test error in this example.The authors state that the reason for this outcome is unclear and requires follow-up analysis.
  • Adversarial perturbations outperform random perturbations, while p = 2 yields superior accuracy on the reported WideResNet CIFAR-10 comparison.The ablation compares p = ∞, p = 2, and random fixed-Euclidean-norm perturbations.

C.6 SEVERAL ITERATIONS IN THE INNER MAXIMIZATION

The experiments test whether SAM’s linearized inner maximization adequately estimates the loss-increasing perturbation. A single projected-gradient step is generally effective during training, but becomes less accurate near convergence.

  • Several iterations in the inner maximization: The study trains WideResNets on CIFAR datasets with multiple projected-gradient ascent steps to estimate maxϵ L(w + ϵ).The resulting loss increase is tracked during training, with test accuracy and estimated sharpness reported at the end across 20 runs.
  • Several iterations in the inner maximization: The reported diagnostics compare training and test error, first- and second-order variants, update similarity, perturbation norms, and inner-step counts.These comparisons are presented in Figures 4–7 and Table 11.
  • Several iterations in the inner maximization: For most training, one projected-gradient step provides a good approximation to the perturbation found with multiple inner-maximization steps.This is the approximation used in standard SAM.
  • Several iterations in the inner maximization: Near convergence, the one-step approximation weakens, and several projected-gradient-ascent iterations produce a better perturbation estimate.The difference is illustrated on CIFAR-10, where five inner steps find about 3% more maximum loss per batch than one step.
Loading 2010.01412v3…