Improving SAM Requires Rethinking its Optimization Formulation

TL;DR

This paper proposes a new formulation of Sharpness-Aware Minimization (SAM) as a bilevel optimization problem called BiSAM, using a 0-1 loss relaxation to improve perturbation strength and enhance model performance.

Contribution

It introduces BiSAM, a novel SAM variant reformulated as a bilevel optimization problem with a 0-1 loss surrogate, leading to stronger perturbations and better results.

Findings

BiSAM outperforms original SAM and variants in experiments.

BiSAM maintains similar computational complexity to SAM.

The code for BiSAM is publicly available.

Abstract

This paper rethinks Sharpness-Aware Minimization (SAM), which is originally formulated as a zero-sum game where the weights of a network and a bounded perturbation try to minimize/maximize, respectively, the same differentiable loss. To fundamentally improve this design, we argue that SAM should instead be reformulated using the 0-1 loss. As a continuous relaxation, we follow the simple conventional approach where the minimizing (maximizing) player uses an upper bound (lower bound) surrogate to the 0-1 loss. This leads to a novel formulation of SAM as a bilevel optimization problem, dubbed as BiSAM. BiSAM with newly designed lower-bound surrogate loss indeed constructs stronger perturbation. Through numerical evidence, we show that BiSAM consistently results in improved performance when compared to the original SAM and variants, while enjoying similar computational complexity. Our code is available at https://github.com/LIONS-EPFL/BiSAM.