Stability Analysis of Sharpness-Aware Minimization

Sharpness-aware minimization (SAM) is a training method that seeks to find flat minima in deep learning, resulting in state-of-the-art performance across various domains. Instead of minimizing the loss of the current weights, SAM minimizes the worst-case loss in its neighborhood in the parameter space. In this paper, we investigate the convergence instability of SAM near a saddle point. Using the qualitative theory of dynamical systems, we explain how SAM becomes stuck in the saddle point and theoretically prove that the saddle point can become an attractor under SAM dynamics. Additionally, we show that this convergence instability can also occur in stochastic dynamical systems by establishing the diffusion of SAM. We prove that SAM diffusion is worse than that of vanilla gradient descent in terms of saddle point escape. Finally, we demonstrate that often overlooked training tricks, momentum and batch-size, might be important to mitigate the convergence instability and achieve high generalization performance. Our theoretical and empirical results are thoroughly verified through experiments on several well-known optimization problems and benchmark tasks.

Sharpness-aware minimization (SAM) is a powerful method used to train deep learning models. However, SAM sometimes suffers from instability because it gets trapped in saddle points. We found that SAM is surprisingly worse at escaping these traps than standard training methods. To understand why this happens, we analyzed SAM's behavior using the mathematics of moving systems. We mapped out exactly how SAM gets pulled into these saddle points and gets stuck. Fortunately, we discovered that simple, often-overlooked training adjustments can fix this issue. By tweaking the momentum of the learning process and changing how much data the model processes at once (the batch size), we can effectively push SAM out of these traps. Our findings explain why this leading training method sometimes fails and provide practical fixes. Applying these simple adjustments allows researchers to prevent their models from getting stuck, unlocking more reliable and high-performing artificial intelligence models.