Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Improving SAM Requires Rethinking its Optimization Formulation

Wanyun Xie, Fabian Latorre, Kimon Antonakopoulos, Thomas Pethick, Volkan Cevher

TL;DR

BiSAM replaces SAM's inner max over a surrogate loss with a bilevel formulation that directly targets the 0-1 misclassification loss $L^{01}$. The outer objective minimizes a differentiable upper bound (e.g., cross-entropy) while the inner maximization uses a differentiable lower-bound surrogate $Q_{\phi,\mu}$, enabling efficient, scalable optimization with the same computational footprint as SAM. Empirically, BiSAM yields consistent improvements over SAM on CIFAR-10/100 and ImageNet, and extends effectively to NLP and fine-tuning tasks, including integration with ASAM and ESAM variants. This principled reframing offers sharper perturbations and broader applicability, with code available at the cited GitHub repository.

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.

Improving SAM Requires Rethinking its Optimization Formulation

TL;DR

BiSAM replaces SAM's inner max over a surrogate loss with a bilevel formulation that directly targets the 0-1 misclassification loss . The outer objective minimizes a differentiable upper bound (e.g., cross-entropy) while the inner maximization uses a differentiable lower-bound surrogate , enabling efficient, scalable optimization with the same computational footprint as SAM. Empirically, BiSAM yields consistent improvements over SAM on CIFAR-10/100 and ImageNet, and extends effectively to NLP and fine-tuning tasks, including integration with ASAM and ESAM variants. This principled reframing offers sharper perturbations and broader applicability, with code available at the cited GitHub repository.

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.
Paper Structure (22 sections, 2 theorems, 19 equations, 4 figures, 12 tables, 3 algorithms)

This paper contains 22 sections, 2 theorems, 19 equations, 4 figures, 12 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries and Problem Setup
  3. Preliminaries: Surrogate loss
  4. SAM and its limitation
  5. Example.
  6. Bilevel Sharpness-aware Minimization (BiSAM)
  7. Recall the example.
  8. Experiments
  9. Image classification
  10. Incorporation with variants of SAM
  11. Adaptive BiSAM.
  12. Efficient BiSAM.
  13. Fine-tuning on image classification
  14. NLP Fine-tuning
  15. Noisy labels task
  16. ...and 7 more sections

Key Result

Lemma 3.1

Let $\phi(x)$ be a lower bound of the 0-1 step function $\mathbb{I}\{x > 0\}$. For each $j \in [K]$, let $F_{w+\epsilon}(x_i, y_i)_j = f_{w+\epsilon}(x_i)_j - f_{w+\epsilon}(x_i)_{y_i}$ and let $\mu > 0$. It holds that

Figures (4)

  • Figure 1: Plot of suggested lower bounds.
  • Figure 2: Number of misclassified samples under perturbation of Resnet56 on CIFAR-10.
  • Figure : Adaptive BiSAM (A-BiSAM)
  • Figure : Efficient BiSAM (E-BiSAM)

Theorems & Definitions (5)

  • Lemma 3.1
  • Remark 3.2
  • proof
  • Theorem 3.1
  • Remark 3.2