[Paper Review] Sharpness-Aware Minimization for Efficiently Improving Generalization

Pierre Foret, Ariel Kleiner+2 more|arXiv (Cornell University)|Oct 3, 2020

Advanced Neural Network ApplicationsComputer Science60 references104 citations

TL;DR

SAM improves generalization by optimizing not only the training loss but also the sharpness of the loss landscape, yielding better test performance across CIFAR, ImageNet, and transfer tasks.

ABSTRACT

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

Motivation & Objective

• Motivate the link between loss landscape geometry and generalization in overparameterized models.
• Propose a practical optimization objective that minimizes both loss value and loss sharpness.
• Demonstrate that optimizing for sharpness yields improved generalization across diverse datasets and architectures.

Proposed method

• Introduce SAM: min_w max_{||epsilon||_p <= rho} L_S(w+epsilon) + lambda||w||^2.
• Derive an efficient gradient approximation by differentiating through the inner maximization and using a first-order Taylor expansion.
• Compute the perturbation hat{epsilon}(w) as a rho-scaled normalization of the training loss gradient (p=2 case).
• Update w using SGD on the SAM objective with gradients evaluated at w+hat{epsilon}(w).
• Provide pseudo-code (Algorithm 1) and discuss practical implementation details, including parallelization and second-order term ablation.
• Analyze m-sharpness (sub-batch perturbations) and Hessian spectra to connect sharpness with generalization.

Experimental results

Research questions

• RQ1 Does incorporating a sharpness term into the training objective improve generalization across standard vision benchmarks?
• RQ2 How does SAM perform relative to SGD on CIFAR-10/100, ImageNet, and finetuning tasks?
• RQ3 What is the impact of the neighborhood size rho and the m-sharpness variant on performance and generalization?
• RQ4 Does SAM confer robustness to label noise comparable to state-of-the-art noisy-label methods?

Key findings

• SAM consistently improves generalization across CIFAR-10/100, ImageNet, and finetuning tasks compared to standard SGD.
• On CIFAR-10/100, SAM achieves state-of-the-art-like results for several models and augmentations (e.g., WRN, Shake-Shake, PyramidNet, and combinations), e.g., 1.6% error on CIFAR-10 with certain settings and 11.3% on CIFAR-100 with AA on some configurations.
• SAM provides robustness to label noise on par with specialized noisy-label methods, often outperforming baselines without explicit noisy-label strategies.
• The m-sharpness variant shows that smaller m (per-accelerator sub-batches) yields better generalization and stronger correlation with actual generalization gaps than full-batch sharpness.
• SAM trained ResNet variants on ImageNet show improved top-1 and top-5 accuracy; e.g., ResNet-152 top-1 error reduced from 20.3% (non-SAM) to 18.4% (with SAM) at 400 epochs.
• Hessian analysis confirms SAM converges to minima with substantially lower curvature (e.g., lambda_max ~1.0 with SAM vs ~24 without) and a flatter spectrum.