Stabilizing Sharpness-aware Minimization Through A Simple Renormalization Strategy
This addresses a stability problem in optimization for machine learning practitioners, offering an incremental improvement to SAM with minimal computational cost.
The paper tackles the issue of sharpness-aware minimization (SAM) getting stuck at saddle points, which degrades performance, by proposing Stable SAM (SSAM), a simple renormalization strategy that stabilizes training and extends the effective learning rate regime, demonstrating improved performance on several datasets.
Recently, sharpness-aware minimization (SAM) has attracted much attention because of its surprising effectiveness in improving generalization performance. However, compared to stochastic gradient descent (SGD), it is more prone to getting stuck at the saddle points, which as a result may lead to performance degradation. To address this issue, we propose a simple renormalization strategy, dubbed Stable SAM (SSAM), so that the gradient norm of the descent step maintains the same as that of the ascent step. Our strategy is easy to implement and flexible enough to integrate with SAM and its variants, almost at no computational cost. With elementary tools from convex optimization and learning theory, we also conduct a theoretical analysis of sharpness-aware training, revealing that compared to SGD, the effectiveness of SAM is only assured in a limited regime of learning rate. In contrast, we show how SSAM extends this regime of learning rate and then it can consistently perform better than SAM with the minor modification. Finally, we demonstrate the improved performance of SSAM on several representative data sets and tasks.