Efficient Sharpness-aware Minimization for Improved Training of Neural Networks

TL;DR

This paper introduces ESAM, an efficient variant of SAM, which maintains generalization benefits while reducing computational costs by approximately 60% through novel training strategies.

Contribution

ESAM proposes two efficient training strategies, StochasticWeight Perturbation and Sharpness-Sensitive Data Selection, to significantly reduce SAM's computational cost without sacrificing performance.

Findings

ESAM reduces computational cost from 100% to 40% compared to SAM.

ESAM maintains or improves test accuracy on CIFAR and ImageNet.

Theoretical analysis supports the effectiveness of the proposed strategies.

Abstract

Overparametrized Deep Neural Networks (DNNs) often achieve astounding performances, but may potentially result in severe generalization error. Recently, the relation between the sharpness of the loss landscape and the generalization error has been established by Foret et al. (2020), in which the Sharpness Aware Minimizer (SAM) was proposed to mitigate the degradation of the generalization. Unfortunately, SAM s computational cost is roughly double that of base optimizers, such as Stochastic Gradient Descent (SGD). This paper thus proposes Efficient Sharpness Aware Minimizer (ESAM), which boosts SAM s efficiency at no cost to its generalization performance. ESAM includes two novel and efficient training strategies-StochasticWeight Perturbation and Sharpness-Sensitive Data Selection. In the former, the sharpness measure is approximated by perturbing a stochastically chosen set of weights in each iteration; in the latter, the SAM loss is optimized using only a judiciously selected subset of data that is sensitive to the sharpness. We provide theoretical explanations as to why these strategies perform well. We also show, via extensive experiments on the CIFAR and ImageNet datasets, that ESAM enhances the efficiency over SAM from requiring 100% extra computations to 40% vis-a-vis base optimizers, while test accuracies are preserved or even improved.