Leveraging Extragradient for Effective Sharpness-Aware Minimization in Deep Learning

Abstract: Generalization remains a pivotal challenge in deep learning, where traditional optimizers like Stochastic Gradient Descent (SGD) often converge to sharp minima, leading to overfitting and reduced performance on unseen data. Building on Sharpness-Aware Minimization (SAM), for seeking flat minima associated with improved generalization, we propose the Extragradient-Inspired Sharpness-Aware Minimization (EISAM), a novel optimizer that enhances generalization via the extragradient technique. EISAM uses a two-step update process: a prediction step investigating the geometry of the loss landscape and a perturbation step that refines updates with a base optimizer. This approach achieves better generalization performance than SAM. Crucially, EISAM reduces sensitivity to the perturbation radius, enhancing robustness, and simplifying the tuning across diverse settings. Extensive experiments on benchmark datasets demonstrate that EISAM consistently outperforms SGD, Adaptive Moment Estimation (Adam), and SAM in test accuracy and training efficiency across various architectures. Theoretical analysis further confirms that EISAM tightens the generalization bound by steering parameters toward flatter minima with reduced curvature. Accompanied by a thorough hyperparameter analysis, EISAM offers practical tuning guidance, establishing it as a robust, scalable, and broadly applicable optimization solution that advances both the theory and practice in deep learning.

• The paper introduces EISAM, a two-step extragradient update that drives convergence to flatter minima while improving generalization and reducing hyperparameter sensitivity.
• It employs a prediction step followed by sharpness-aware correction, effectively exploring the loss landscape as demonstrated by Hessian eigenvalue analysis.
• Empirical results on datasets like CIFAR, ImageNet, and COCO show that EISAM offers stability, reduced computational overhead, and superior performance in multiple deep learning tasks.

Leveraging Extragradient for Effective Sharpness-Aware Minimization in Deep Learning

Introduction

The prevailing optimization paradigm for deep neural networks, such as SGD and Adam, tends toward convergence at sharp minima—regions of high loss curvature—known to be correlated with poor generalization to unseen data. Sharpness-Aware Minimization (SAM) provides a generalizable improvement by minimizing loss sharpness within a prescribed neighborhood, but it incurs computational overhead and exposes heightened hyperparameter sensitivity, particularly to the perturbation radius $\rho$. This work introduces Extragradient-Inspired Sharpness-Aware Minimization (EISAM), which integrates the extragradient method's two-step prediction-correction update into sharpness-aware optimization, thereby improving generalization, stability, and hyperparameter robustness across a spectrum of deep learning tasks (2607.06151).

Methodological Contributions

The EISAM Algorithm

EISAM introduces a two-step update process:

1. Prediction Step: Compute an intermediate point $$\mathbf{y}_t = \mathbf{w}_t - s \nabla F_S(\mathbf{w}_t)$$ using a prediction step size $s$, emulating an extragradient lookahead.
2. Sharpness-Aware Perturbation: At $\mathbf{y}_t$, compute the adversarial perturbation $$\boldsymbol{\epsilon}_t = \rho {\nabla F_S(\mathbf{y}_t)}/{\|\nabla F_S(\mathbf{y}_t)\|}$$.
3. Correction Step: Perform the update $$\mathbf{w}_{t+1} = \mathbf{y}_t - \eta \nabla F_S(\mathbf{y}_t + \boldsymbol{\epsilon}_t)$$ with learning rate $\eta$.

This procedure differs from SAM in its explicit geometric exploration before sharpness-aware correction, guiding optimization toward flatter minima, improving robustness to $\rho$, and reducing overall hyperparameter tuning burden.

Figure 2: Schematic of EISAM parameter update, showing the prediction step and subsequent sharpness-aware correction step.

Motivation and Loss Geometry

The extragradient mechanism's primary theoretical advantage is the explicit consideration of local curvature and landscape geometry. Visualization of optimizer trajectories on the Rosenbrock function demonstrates that EISAM follows more direct and stable paths to minima across a range of $\rho$ values, in contrast to SGD and SAM, which either meander or exhibit sensitivity to hyperparameters.

Figure 3: 3D contour visualization of optimization trajectories on the Rosenbrock function for SGD, Adam, EISAM, and SAM at various $\rho$.

Theoretical Analysis

Generalization and Stability Bounds

Rigorous excess risk and generalization bounds are provided for EISAM under strong convexity, smoothness, and Hessian Lipschitz assumptions. The stability analysis confirms that EISAM's prediction step tightens the generalization error bound, introducing dependence on $$\mathbf{y}_t = \mathbf{w}_t - s \nabla F_S(\mathbf{w}_t)$$0 (where $$\mathbf{y}_t = \mathbf{w}_t - s \nabla F_S(\mathbf{w}_t)$$1 is the prediction step and $$\mathbf{y}_t = \mathbf{w}_t - s \nabla F_S(\mathbf{w}_t)$$2 the strong convexity constant), an improvement over SAM's corresponding factor. This demonstrates EISAM's heightened resistance to hyperparameter-induced variance and its inherent tendency to favor flatter minima.

Curvature Analysis

Empirical eigenvalue analysis of the Hessian during training on CIFAR-100 reveals that EISAM persistently converges to regions of significantly lower curvature ($$\mathbf{y}_t = \mathbf{w}_t - s \nabla F_S(\mathbf{w}_t)$$3) compared to vanilla SGD and even SAM, consistent across training and testing regimes.

Figure 1: Trends of the maximum Hessian eigenvalue ($$\mathbf{y}_t = \mathbf{w}_t - s \nabla F_S(\mathbf{w}_t)$$4) during training, demonstrating improved curvature properties of EISAM minima.

Figure 4: Visualization of training trajectories in the loss landscape; EISAM finds flatter minima with lower average curvature.

Empirical Evaluation

Classification, Detection, and Segmentation

EISAM is evaluated on CIFAR-10, CIFAR-100, ImageNet-1K, COCO2017, LVIS, and ISIC2018, encompassing classification, object detection, and segmentation tasks with varied architectures (ResNet, ViT, U-Net).

• Classification: EISAM consistently yields higher test accuracy and lower loss relative to SGD, Adam, SAM, and their variants across all architectures and datasets. The improvement is robust even under strong regularization (e.g., CutMix).
• Object Detection: On COCO and LVIS, EISAM achieves higher AP scores with randomly initialized detectors and in the long-tailed regime, reflecting strong generalizability and optimization stability.
• Segmentation: On ISIC2018, EISAM achieves the highest validation accuracy, with competitive composite segmentation scores.

Figure 6: Training loss and test accuracy curves for ResNet18/ResNet101 on CIFAR-100 with CutMix, highlighting rapid and stable convergence of EISAM.

Hyperparameter Sensitivity and Robustness

Comprehensive sweeps over $$\mathbf{y}_t = \mathbf{w}_t - s \nabla F_S(\mathbf{w}_t)$$5 and the prediction step $$\mathbf{y}_t = \mathbf{w}_t - s \nabla F_S(\mathbf{w}_t)$$6 reveal that EISAM displays markedly reduced sensitivity in both training and test metrics, as compared to SAM. Notably, when $$\mathbf{y}_t = \mathbf{w}_t - s \nabla F_S(\mathbf{w}_t)$$7, EISAM degenerates to SAM, recovering its sensitivity profile.

Figure 5: EISAM hyperparameter sensitivity to $$\mathbf{y}_t = \mathbf{w}_t - s \nabla F_S(\mathbf{w}_t)$$8; stability is observed across the full parameter range.

Figure 7: EISAM hyperparameter sensitivity to prediction step $$\mathbf{y}_t = \mathbf{w}_t - s \nabla F_S(\mathbf{w}_t)$$9; generalization remains robust for a broad spectrum of $s$0 values.

Figure 8: Joint sensitivity analyses for $s$1 and $s$2; EISAM exhibits stable optima with lower test loss variance compared to SAM.

Sharpness and Loss Landscape Visualizations

EISAM minima are characterized by broader, lower-curvature basins in the loss landscape, as visualized by 2D contour projections. These findings, aligned with excess risk theory and empirical Hessian spectra, indicate that EISAM's optimization trajectory strongly biases toward solutions with improved expected generalization.

Figure 9: Loss landscape and trajectory visualization show EISAM achieves lower curvature minima than SAM and SGD.

Computational Efficiency

Although EISAM, like all SAM variants, requires two forward-backward passes per iteration, comparative analysis shows its computational overhead (wall-clock time, memory, FLOPs) is negligible relative to SAM, ASAM, GSAM, and FSAM. The practical impact is further reduced owing to accelerated convergence and reduced need for parameter search.

Implications and Future Directions

EISAM's results substantiate the efficacy of extragradient-inspired updates for flat minima discovery in deep networks, with both tighter theoretical bounds and superior empirical outcomes. The lowered sensitivity to perturbation radius represents a significant advance in optimizer usability and deployability, especially in regimes characterized by highly nonconvex or ill-conditioned landscapes (e.g., large-scale transformers, long-tailed detection, and biomedical segmentation).

Potential future directions motivated by this work include:

• Extending extragradient-inspired sharpness minimization to large-scale pretraining and self-supervised scenarios.
• Integrating EISAM with advanced scheduler, label smoothing, and adversarial data regimes.
• Investigating EISAM in LLM and RL training settings, where sharpness-aware dynamics may interact with reward structure.

Conclusion

EISAM advances the field of sharpness-aware deep learning optimization by integrating an extragradient-inspired two-step mechanism, yielding strong theoretical guarantees, superior empirical generalization, and robustness to critical hyperparameters. Given its scalability, stability, and compatibility with a diversity of architectures and tasks, EISAM represents a substantive improvement over SAM and its variants, establishing itself as a highly effective practitioner- and researcher-oriented optimizer for deep learning applications.

(2607.06151)