Explicit Eigenvalue Regularization Improves Sharpness-Aware Minimization

TL;DR

This work analyzes Sharpness-Aware Minimization (SAM) through the lens of the top Hessian eigenvalue $\lambda_1(\nabla^2 f(x))$ and derives a third-order stochastic differential equation (SDE) that reveals a complex mix of second- and third-order terms governing SAM dynamics. It shows that, while perturbation–eigenvector alignment can implicitly reduce sharpness, practical alignment is often weak, motivating Eigen-SAM, which explicitly aligns the perturbation with the leading Hessian eigenvector to directly regularize the top eigenvalue. The authors prove a PAC-Bayes-based generalization connection to the top eigenvalue, formalize the third-order SDE and its implications, and demonstrate through simulations and large-scale experiments that Eigen-SAM improves generalization and reshapes the Hessian spectrum across datasets and architectures. The proposed approach provides both a deeper theoretical understanding of SAM and a practical method to enhance sharpness control with manageable computational overhead.

Abstract

Sharpness-Aware Minimization (SAM) has attracted significant attention for its effectiveness in improving generalization across various tasks. However, its underlying principles remain poorly understood. In this work, we analyze SAM's training dynamics using the maximum eigenvalue of the Hessian as a measure of sharpness, and propose a third-order stochastic differential equation (SDE), which reveals that the dynamics are driven by a complex mixture of second- and third-order terms. We show that alignment between the perturbation vector and the top eigenvector is crucial for SAM's effectiveness in regularizing sharpness, but find that this alignment is often inadequate in practice, limiting SAM's efficiency. Building on these insights, we introduce Eigen-SAM, an algorithm that explicitly aims to regularize the top Hessian eigenvalue by aligning the perturbation vector with the leading eigenvector. We validate the effectiveness of our theory and the practical advantages of our proposed approach through comprehensive experiments. Code is available at https://github.com/RitianLuo/EigenSAM.

Figures (5)

• Figure 1: Alignment and top eigenvalue for a 6-layer CNN model trained on CIFAR-10. The left panel shows the trend of alignment during SAM training; the shaded area represents the 95% confidence interval. The right panel displays the trend of the top eigenvalue over the course of training.
• Figure 2: Training dynamics of discrete SAM, second-order SDE, and third-order SDE during training. Metrics include training loss, test loss, test accuracy, parameter norm, gradient norm, and the top Hessian eigenvalue. Each plot illustrates how each approach affects loss dynamics and key stability metrics.
• Figure 3: Left: Sensitivity analysis of $\alpha$; the blue lines indicate the confidence interval. Right: Spectrum of the Hessian at the end of training.
• Figure 4: Comparison of training loss and test loss metrics across algorithms.
• Figure 5: The effect of the number of Hessian-vector product steps in Algorithm \ref{['algorithm:1']} (power iteration) on the alignment of the estimated vector with the top eigenvalue. The dataset is CIFAR-100, and the models are ResNet18 and WideResNet-28-10 at mid-training stage (100th epoch).

Theorems & Definitions (15)

• Theorem 3.1
• Theorem 4.1
• Corollary 4.1.1
• Definition B.1
• Lemma B.1
• Lemma B.2
• Lemma B.3
• proof
• Theorem B.4
• proof