SAM operates far from home: eigenvalue regularization as a dynamical phenomenon

TL;DR

This paper reveals that SAM regularizes eigenvalues dynamically throughout training, inducing a stabilization phenomenon related to the edge of stability, which impacts generalization and model behavior far from minima.

Contribution

It demonstrates that SAM's eigenvalue regularization is a dynamic process affecting the entire learning trajectory, not just near minima, and links this to the edge of stability phenomenon.

Findings

SAM stabilizes eigenvalues during training.

Theoretical prediction of eigenvalue behavior based on learning rate and SAM radius.

Practical models exhibit edge of stability dynamics influenced by SAM.

Abstract

The Sharpness Aware Minimization (SAM) optimization algorithm has been shown to control large eigenvalues of the loss Hessian and provide generalization benefits in a variety of settings. The original motivation for SAM was a modified loss function which penalized sharp minima; subsequent analyses have also focused on the behavior near minima. However, our work reveals that SAM provides a strong regularization of the eigenvalues throughout the learning trajectory. We show that in a simplified setting, SAM dynamically induces a stabilization related to the edge of stability (EOS) phenomenon observed in large learning rate gradient descent. Our theory predicts the largest eigenvalue as a function of the learning rate and SAM radius parameters. Finally, we show that practical models can also exhibit this EOS stabilization, and that understanding SAM must account for these dynamics far away from any minima.