Sharpness-Aware Minimization: General Analysis and Improved Rates

2025-03-04

AI Summary

Sharpness-Aware Minimization (SAM) has long suffered from restrictive convergence guarantees in non-convex optimization—relying on strong bounded-variance assumptions, fixed sampling schemes, and opaque normalization mechanisms. This work proposes Unified SAM, a general theoretical framework that for the first time establishes convergence for SAM under relaxed noise assumptions, eliminating dependence on bounded gradient variance. It further introduces the arbitrary sampling paradigm into SAM analysis—enabling flexible strategies such as importance sampling—and systematically characterizes the distinct convergence behaviors of normalized versus unnormalized updates. Theoretically, Unified SAM is proven to converge under both generic non-convex and Polyak–Łojasiewicz (PL) conditions. Empirically, it consistently outperforms standard SAM in both generalization performance and convergence speed on image classification benchmarks.

Technology Category

• Search and Optimization: Non-convex Optimization
• Computer Vision: Learning & Optimization for CV
• Machine Learning: Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Responsible Web: Machine-in-the-loop, human agency and autonomy

Abstract

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.

Problem

• Analyzes convergence properties of Sharpness-Aware Minimization (SAM) in non-convex settings.
• Explores benefits of normalization and relaxed noise assumptions in SAM updates.
• Provides convergence guarantees for SAM under various step sizes and sampling strategies.

Innovation

• Unified SAM with flexible update rule
• Convergence under relaxed noise assumptions
• Arbitrary sampling paradigm analysis