Sharpness-Aware Minimization (SAM) enhances generalization by minimizing the maximum training loss within a predefined neighborhood around the parameters. However, its practical implementation approximates this as gradient ascent(s) followed by applying the gradient at the ascent point to update the current parameters. This practice can be justified as approximately optimizing the objective by neglecting the (full) derivative of the ascent point with respect to the current parameters. Nevertheless, a direct and intuitive understanding of why using the gradient at the ascent point to update the current parameters works superiorly is still lacking. Our work bridges this gap by proposing a novel and intuitive interpretation. We show that the gradient at the single-step ascent point, \uline{when applied to the current parameters}, provides a better approximation of the direction from the current parameters toward the maximum within the local neighborhood than the local gradient. This improved approximation thereby enables a more direct escape from the maximum within the local neighborhood. Nevertheless, our analysis further reveals two issues. First, the approximation by the gradient at the single-step ascent point is often inaccurate. Second, the approximation quality may degrade as the number of ascent steps increases. To address these limitations, we propose in this paper eXplicit Sharpness-Aware Minimization (XSAM). It tackles the first by explicitly estimating the direction of the maximum during training, while addressing the second by crafting a search space that effectively leverages the gradient information at the multi-step ascent point. XSAM features a unified formulation that applies to both single-step and multi-step settings and only incurs negligible computational overhead. Extensive experiments demonstrate the consistent superiority of XSAM against existing counterparts.
Figures & tables
Figure 1: (a) Visualization of the local loss surface of single-step SAM 2 2 2 Data is collected at the first iteration of the 150th epoch in training ResNet-18 on CIFAR-100. on the hyperplane spanned by the gradient g0 at the current parameter ϑ0 and the gradient g1 at the single-step ascent point ϑ1 . ϑ0 is set as the origin, the Y -axis is defined along the direction of g0 , and the X -axis is aligned with the component of g1 perpendicular to g0 . The visualized arrows of gradients are set to have length ρ . We see that g1@ϑ0 (i.e., g1 applied to ϑ0 ) points clearly closer to the direction from ϑ0 toward the maximum within the local neighborhood than g0 . The targeted direction is roughly from the origin to the upper-right corner in the figure. The loss along g1@ϑ0 (i.e., L(ϑ0+ρm⋅g1/∥g1∥) ) is higher than that along g0 (i.e., L(ϑ0+ρm⋅g0/∥g0∥) ), for sufficiently large ρm . (b) A simulation of multi-step SAM on a 2D test function. The approximation quality by the SAM gradient may get worse as the number of ascent steps increases . g2@ϑ0 inferiorly identifies the direction from ϑ0 toward the maximum within the neighborhood (the upper-left high-loss region in yellow) than g1@ϑ0 .
Algorithm 1 XSAM
Figure 3: (a) Training trajectory comparisons on 2D test function. (b)-(c) Test accuracy comparisons of ResNet-18 trained on CIFAR-100 in single-step and multi-step ( k=3 ) settings with varying ρ .
Figure 4: XSAM robustness to the α∗ update frequency.
Appendix figures & tables23 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Visualization of loss surface during training: VGG-11 trained on CIFAR-100.
Figure 6: Visualization of loss surface during training: ResNet-18 trained on CIFAR-100.
Figure 7: Visualization of loss surface during training: ViT-Ti trained on CIFAR-100.
Figure 8: Visualization of loss surface during training: ResNet-18 trained on CIFAR-10.
Figure 9: Visualization of loss surface during training: DenseNet-121 trained on CIFAR-10.
Figure 10: (a) Visualization of the local loss surface of single-step SAM. The visualization procedure follows the same steps as in Figure 1(a) . Data is collected at the first iteration of the 100th epoch in training ResNet-18 on CIFAR-100. We see that g1@ϑ0 (i.e., g1 applied to ϑ0 ) points clearly closer to the direction from ϑ0 toward the maximum within the local neighborhood than g0 . he targeted direction is roughly from the origin to the upper-right corner in the figure. The loss along g1@ϑ0 (i.e., L(ϑ0+ρm⋅g1/∥g1∥) ) is higher than that along g0 (i.e., L(ϑ0+ρm⋅g0/∥g0∥) ), for sufficiently large ρm . (b) A simulation of multi-step SAM on a test function. The gradient at the multi-step ascent point, when applied to the current parameters, may be an inferior approximation of the direction toward the maximum .
Figure 11: Visualization of L(ϑ0+ρmg1/∣∣g1∣∣)−L(ϑ0+ρmg0/∣∣g0∣∣) during training.
Figure 12: Visualization of L(ϑ0+ρmg1/∣∣g1∣∣)−L(ϑ0+ρmg0/∣∣g0∣∣) during training.
CIFAR-10
CIFAR-100
Tiny-ImageNet
VGG-11
SGD
SAM
WSAM
XSAM
SGD
SAM
WSAM
XSAM
SGD
SAM
WSAM
XSAM
Epoch
200
200
200
Batch size
125
125
125
Initial learning rate
0.05
0.05
0.05
Momentum
0.9
0.9
0.9
Weight decay
1×10−3
1×10−3
1×10−3
Appendix
Table 5: Hyperparameter details for Results in Table 2 .
ρ =0.10
ρ =0.20
ρ =0.30
SAM
WSAM
XSAM
SAM
WSAM
XSAM
SAM
WSAM
XSAM
Epoch
200
200
200
Batch size
125
125
125
Initial learning rate
0.05
0.05
0.05
Momentum
0.9
0.9
0.9
Weight decay
1×10−3
1×10−3
1×10−3
Appendix
Table 6: Hyperparameter details for Results in Figure 3(b) . Note that, in this experiment, α for WSAM adopts the average value of the dynamic α∗ in the corresponding XSAM. We see from the results that such WSAM already clearly outperforms SAM.
ρ =0.04
ρ =0.08
ρ =0.12
SAM
WSAM
XSAM
SAM
WSAM
XSAM
SAM
WSAM
XSAM
Epoch
200
200
200
Batch size
125
125
125
Initial learning rate
0.05
0.05
0.05
Momentum
0.9
0.9
0.9
Weight decay
1×10−3
1×10−3
1×10−3
Appendix
Table 8: Hyperparameter details for Results in Figure 3(c) . Note that, in this experiment, α for WSAM adopts the average value of the dynamic α∗ in the corresponding XSAM. We see from the results that such WSAM already clearly outperforms SAM.
SAM
ASAM
XSAM
XSAM+ASAM
Test Accuracy
80.93±0.11
81.11±0.06
81.24±0.07
81.68 ± 0.11
Appendix
Table 10: Test accuracy of SAM variants and their combinations with XSAM.
Method
ResNet-18
DenseNet-121
SAM
80.93 ± 0.11
83.81 ± 0.02
ASAM
81.11 ± 0.06
83.99 ± 0.25
VaSSO
80.84 ± 0.15
83.78 ± 0.25
WSAM
80.95 ± 0.19
83.91 ± 0.15
XSAM
81.24 ± 0.07
83.96 ± 0.10
XSAM + ASAM
81.68 ± 0.11
84.06 ± 0.21
Appendix
Table 11: Comparison on CIFAR-100 with ResNet-18 and DenseNet-121. All baseline methods are carefully tuned for optimal performance. XSAM uses the same ρ as SAM, as in the paper.
Method
ρ=0.04
ρ=0.08
ρ=0.12
SAM
80.79 ±0.41
80.75 ±0.27
79.72 ±0.33
LSAM
81.00 ±0.21
81.20 ±0.24
81.16 ±0.04
LSAM+
80.56 ±0.20
80.77 ±0.04
80.21 ±0.27
MSAM
81.04 ±0.06
81.12 ±0.17
80.93 ±0.11
MSAM+
80.72 ±0.16
81.16 ±0.05
81.16 ±0.05
XSAM
81.23 ±0.06
81.36 ±0.08
81.29 ±0.06
Appendix
Table 12: Results on CIFAR-100 using ResNet-18 in multi-step ( k=3 ) setting.
Method
k=1
k=2
k=4
LSAM
83.81 ± 0.02
83.82 ± 0.28
83.40 ± 0.17
MSAM
83.81 ± 0.02
83.67 ± 0.23
83.74 ± 0.18
XSAM
83.96 ± 0.10
84.05 ± 0.04
84.02 ± 0.31
Appendix
Table 13: Results on DenseNet-121 with CIFAR-100 with different k . ρ=ρ∗/k with ρ∗ for single-step.
Method
k=1
k=2
k=4
LSAM
96.59 ± 0.06
96.66 ± 0.03
96.72 ± 0.07
MSAM
96.59 ± 0.06
96.78 ± 0.05
96.80 ± 0.07
XSAM
96.74 ± 0.04
96.81 ± 0.06
96.81 ± 0.11
Appendix
Table 14: Results on ResNet-18 with CIFAR-10 with different k . ρ=ρ∗/k with ρ∗ for single-step.
Table 15: Performance on CIFAR-10-C/CIFAR-100-C with ResNet-18 and DenseNet-121.
Figure 13: More visualizations of the dynamic estimations of α .
Figure 14: (a) Robustness analysis of XSAM with respect to ρm . (b)Training statistics of XSAM. (c) The norms of gi during training.
SGD
SAM
XSAM
λ1
78.79
36.15
33.92
λ1/λ5
2.26
1.89
1.59
Appendix
Table 16: Hessian spectrum of ResNet-18 using SGD, SAM, and XSAM on CIFAR-10.
Figure 15: Loss landscape visualizations of ResNet-18 on CIFAR-10 with SGD, SAM, and XSAM.
Figure 16: Visualization of the average sharpness of the loss landscape at the convergence point.
Figure 17: Comparison of various gradient scale strategies.
Recent progress in optimization research has highlighted the sharpness of the loss landscape as a key factor in narrowing the generalization gap. Motivated by this insight, Sharpness-Aware Minimization (SAM) was proposed as a training strategy that enhances generalization. Despite the promising performance, SAM suffers from its twice computational cost due to its core algorithm requiring an extra gradient computation during the perturbation step. To overcome this limitation, we introduce Exponential Moving Average Sharpness-Aware Minimization (EMASAM), a computationally efficient variant of SAM. EMASAM does not require the loss gradient in the perturbation step. Instead, EMASAM defines the perturbation direction based on the discrepancy between the main model and the EMA shadow model. This perturbation travels away from the stable average position toward the less stable area, acting as a softer yet cheaper alternative to SAM's worst-case scenario perturbation. Moreover, since EMASAM's perturbation does not rely on noisy mini-batch gradients, it mitigates the gradient-induced instability inherent in SAM. Hence, EMASAM eliminates the need for an extra backpropagation while also preserving the generalization ability of the SAM-style training. Several experiments have been performed and confirm the efficiency and robustness of our method.
Sharpness-Aware Minimization (SAM) improves generalization by minimizing the worst-case loss within a fixed parameter-space radius neighborhood. SAM and its variants mainly rely on a first-order linearized surrogate, while flat minima are inherently a second-order (curvature) notion.We revisit this mismatch and propose Loss-Equated SAM (LE-SAM), which inverts the traditional SAM mechanism that fixed perturbation radius with a fixed loss-space budget,effectively removing gradient-norm-dominated learning signals and shifting optimization toward curvature-dominated terms. Extensive experiments across diverse benchmarks and tasks demonstrate the strong generalization ability of LESAM that consistently outperforms SAM and even its variants, achieving the state-of-the-art performance.
Jinping Wang, Qinhan Liu, Zhiwu Xie +1
Wenzhou-Kean University · International Frontier Interdisciplinary Research Institute, Wenzhou-Kean University
The sharpness-aware minimization (SAM) algorithm and its variants, including gap guided SAM (GSAM), have been successful at improving the generalization capability of deep neural network models by finding flat local minima of the empirical loss in training. Meanwhile, it has been shown theoretically and practically that increasing the batch size or decaying the learning rate avoids sharp local minima of the empirical loss. In this paper, we consider the GSAM algorithm with increasing batch sizes or decaying learning rates, such as cosine annealing or linear learning rate, and theoretically show its convergence. Moreover, we numerically compare SAM (GSAM) with and without an increasing batch size and conclude that using an increasing batch size { achieves a lower worst-case ℓ∞ adaptive sharpness} than compared with using a constant batch size and learning rate.