X-SAM: Boosting Sharpness-Aware Minimization with Dominant-Eigenvector Gradient Correction

TL;DR

X-SAM enhances Sharpness-Aware Minimization by correcting gradients along the dominant eigenvector of the Hessian, leading to better regularization, convergence, and generalization in training deep models.

Contribution

We introduce X-SAM, a novel method that explicitly aligns gradient updates with the Hessian's dominant eigenvector to improve SAM's effectiveness.

Findings

X-SAM achieves faster convergence in training.

X-SAM improves model generalization performance.

Theoretical analysis confirms X-SAM's convergence properties.

Abstract

Sharpness-Aware Minimization (SAM) aims to improve generalization by minimizing a worst-case perturbed loss over a small neighborhood of model parameters. However, during training, its optimization behavior does not always align with theoretical expectations, since both sharp and flat regions may yield a small perturbed loss. In such cases, the gradient may still point toward sharp regions, failing to achieve the intended effect of SAM. To address this issue, we investigate SAM from a spectral and geometric perspective: specifically, we utilize the angle between the gradient and the leading eigenvector of the Hessian as a measure of sharpness. Our analysis illustrates that when this angle is less than or equal to ninety degrees, the effect of SAM's sharpness regularization can be weakened. Furthermore, we propose an explicit eigenvector-aligned SAM (X-SAM), which corrects the gradient via orthogonal decomposition along the top eigenvector, enabling more direct and efficient regularization of the Hessian's maximum eigenvalue. We prove X-SAM's convergence and superior generalization, with extensive experimental evaluations confirming both theoretical and practical advantages.