Revisiting Sharpness-Aware Minimization: A More Faithful and Effective Implementation

TL;DR

This paper revisits Sharpness-Aware Minimization (SAM), providing a new interpretation of its effectiveness, identifying limitations, and proposing an improved method called XSAM that explicitly estimates the maximum sharpness, leading to better generalization.

Contribution

The paper introduces XSAM, a novel implementation of SAM that explicitly estimates sharpness, improves approximation accuracy, and enhances training effectiveness with minimal computational cost.

Findings

XSAM outperforms existing SAM variants in experiments.

Explicit sharpness estimation improves approximation accuracy.

XSAM maintains low computational overhead.

Abstract

Sharpness-Aware Minimization (SAM) enhances generalization by minimizing the maximum training loss within a predefined neighborhood around the parameters. However, its practical implementation approximates this as gradient ascent(s) followed by applying the gradient at the ascent point to update the current parameters. This practice can be justified as approximately optimizing the objective by neglecting the (full) derivative of the ascent point with respect to the current parameters. Nevertheless, a direct and intuitive understanding of why using the gradient at the ascent point to update the current parameters works superiorly is still lacking. Our work bridges this gap by proposing a novel and intuitive interpretation. We show that the gradient at the single-step ascent point, \uline{when applied to the current parameters}, provides a better approximation of the direction from the current parameters toward the maximum within the local neighborhood than the local gradient. This improved approximation thereby enables a more direct escape from the maximum within the local neighborhood. Nevertheless, our analysis further reveals two issues. First, the approximation by the gradient at the single-step ascent point is often inaccurate. Second, the approximation quality may degrade as the number of ascent steps increases. To address these limitations, we propose in this paper eXplicit Sharpness-Aware Minimization (XSAM). It tackles the first by explicitly estimating the direction of the maximum during training, while addressing the second by crafting a search space that effectively leverages the gradient information at the multi-step ascent point. XSAM features a unified formulation that applies to both single-step and multi-step settings and only incurs negligible computational overhead. Extensive experiments demonstrate the consistent superiority of XSAM against existing counterparts.