Sharpness-Aware Minimization: General Analysis and Improved Rates
This work addresses theoretical gaps in a popular optimization method for improving generalization in machine learning, but it is incremental as it builds on existing SAM frameworks.
The paper tackled open questions about the convergence properties of Sharpness-Aware Minimization (SAM) in non-convex settings by providing a unified analysis under a flexible update rule and relaxed noise assumptions, achieving convergence guarantees for non-convex and Polyak-Lojasiewicz functions with experiments showing practical effectiveness in deep neural network training.
Sharpness-Aware Minimization (SAM) has emerged as a powerful method for improving generalization in machine learning models by minimizing the sharpness of the loss landscape. However, despite its success, several important questions regarding the convergence properties of SAM in non-convex settings are still open, including the benefits of using normalization in the update rule, the dependence of the analysis on the restrictive bounded variance assumption, and the convergence guarantees under different sampling strategies. To address these questions, in this paper, we provide a unified analysis of SAM and its unnormalized variant (USAM) under one single flexible update rule (Unified SAM), and we present convergence results of the new algorithm under a relaxed and more natural assumption on the stochastic noise. Our analysis provides convergence guarantees for SAM under different step size selections for non-convex problems and functions that satisfy the Polyak-Lojasiewicz (PL) condition (a non-convex generalization of strongly convex functions). The proposed theory holds under the arbitrary sampling paradigm, which includes importance sampling as special case, allowing us to analyze variants of SAM that were never explicitly considered in the literature. Experiments validate the theoretical findings and further demonstrate the practical effectiveness of Unified SAM in training deep neural networks for image classification tasks.