On Statistical Properties of Sharpness-Aware Minimization: Provable Guarantees

TL;DR

This paper provides a theoretical analysis of Sharpness-Aware Minimization (SAM), demonstrating its statistical advantages and flatter solutions in both convex and non-convex settings, supported by numerical experiments.

Contribution

It offers the first theoretical explanation of SAM's statistical properties and generalization benefits, especially in non-convex problems, complementing prior empirical findings.

Findings

SAM achieves smaller prediction error than Gradient Descent under certain conditions.

SAM solutions are less sharp, indicating flatter minima.

Numerical experiments validate theoretical predictions across various scenarios.

Abstract

Sharpness-Aware Minimization (SAM) is a recent optimization framework aiming to improve the deep neural network generalization, through obtaining flatter (i.e. less sharp) solutions. As SAM has been numerically successful, recent papers have studied the theoretical aspects of the framework and have shown SAM solutions are indeed flat. However, there has been limited theoretical exploration regarding statistical properties of SAM. In this work, we directly study the statistical performance of SAM, and present a new theoretical explanation of why SAM generalizes well. To this end, we study two statistical problems, neural networks with a hidden layer and kernel regression, and prove under certain conditions, SAM has smaller prediction error over Gradient Descent (GD). Our results concern both convex and non-convex settings, and show that SAM is particularly well-suited for non-convex problems. Additionally, we prove that in our setup, SAM solutions are less sharp as well, showing our results are in agreement with the previous work. Our theoretical findings are validated using numerical experiments on numerous scenarios, including deep neural networks.