Deep learning models, despite their impressive achievements, suffer from high computational costs and memory requirements, limiting their usability in resource-constrained environments. Sparse neural networks significantly alleviate these constraints by dramatically reducing parameter count and computational overhead. However, existing sparse training methods often experience chaotic and noisy gradient signals, severely hindering convergence and generalization performance, particularly at high sparsity levels. To tackle this critical challenge, we propose Zero-Order Sharpness-Aware Minimization (ZO-SAM), a novel optimization framework that strategically integrates zero-order optimization within the SAM approach. Unlike traditional SAM, ZO-SAM requires only a single backpropagation step during perturbation, selectively utilizing zero-order gradient estimations. This innovative approach reduces the backpropagation computational cost by half compared to conventional SAM, significantly lowering gradient variance and effectively eliminating associated computational overhead. By harnessing SAM's capacity for identifying flat minima, ZO-SAM stabilizes the training process and accelerates convergence. These efficiency gains are particularly important in sparse training scenarios, where computational cost is the primary bottleneck that limits the practicality of SAM. Moreover, models trained with ZO-SAM exhibit improved robustness under distribution shift, further broadening its practicality in real-world deployments.
Figures & tables
Figure 1: (a) Comparison of Gradient Variance during sparse training at sparsity 90% and 95% between SGD and ZO-SAM; (b) Loss landscape of dense training, sparse training at sparsity 90% and 95%.
Figure 2: Illustration of the optimization mechanism of ZO-SAM.
Datasets
CIFAR-10
CIFAR-100
Pruning ratio
90%
95%
98%
90%
95%
98%
ResNet-32
94.58 (Dense)
74.89 (Dense)
LT [ 14 ]
92.31
91.06
88.78
68.99
65.02
57.37
LT+ ZO-SAM (ours)
92.69 ± 0.05 (0.38 ↑ )
91.67 ± 0.10 (0.61 ↑ )
89.46 ± 0.15 (0.68 ↑ )
69.44 ± 0.06 (0.45 ↑ )
65.59 ± 0.13 (0.57 ↑ )
57.87 ± 0.14 (0.50 ↑ )
SNIP [ 22 ]
92.59 ± 0.10
91.01 ± 0.21
87.51 ± 0.31
68.89 ± 0.45
65.02 ± 0.69
57.37 ± 1.43
SNIP+ ZO-SAM (ours)
93.38 ± 0.12 (0.79 ↑ )
91.69 ± 0.17 (0.68 ↑ )
88.22 ± 0.21 (0.71 ↑ )
69.54 ± 0.23 (0.65 ↑ )
65.90 ± 0.49 (0.88 ↑ )
58.29 ± 0.47 (0.92 ↑ )
Table 1: Test accuracy (%) of pruned ResNet-32 on CIFAR-10/100.
Figure 3: Overview of perturbation step in ZO-SAM.
Figure 4: Loss surface improvement at various sparsity levels.
Model
DeiT-Small
DeiT-Tiny
Dense
79.90
72.20
Pruning ratio
70%
50%
RigL [ 11 ]
77.99
70.02
RigL+ZO-SAM
79.16 (1.17 ↑ )
70.79 (0.77 ↑ )
MEST [ 42 ]
77.15
69.69
MEST+ZO-SAM
78.04 (0.89 ↑ )
70.41 (0.72 ↑ )
Table 2: Test accuracy (%) of pruned DeiT on ImageNet-1K
Figure 5: Convergence Speed Comparison at sparsity level 90%.
Figure 6: Convergence Speed Comparison at sparsity level 95%.
CIFAR-10, MEST, Acc = 90%
SGD
ESAM
LookSAM(k=5)
GSAM
ZO-SAM
sp=0.9, Epoch#
104
75
79
84
70
sp=0.95, Epoch#
131
92
94
113
88
Table 3: Epochs to reach 90% accuracy on CIFAR-10 (lower is better).
Method
Sparsity
ResNet-32
WRN-28-10
CIFAR-10
CIFAR-100
Image/sec
CIFAR-10
CIFAR-100
Image/sec
SGD
-
94.58
74.89
5673.95
96.41
81.63
752.30
SAM
90%
93.77
72.64
2704.84 (47.67%)
95.37
80.14
354.94 (47.18%)
ESAM [ 9 ]
90%
93.67
72.49
2297.23 (40.48%)
95.31
80.23
305.76 (40.64%)
LS(k=5) [ 25 ]
90%
93.42
72.01
3980.62 (70.14%)
94.89
79.30
542.63 (72.13%)
LS(k=10) [ 25 ]
90%
93.03
71.62
4272.22 (75.29%)
94.17
78.84
593.02 (78.83%)
Table 4: Test accuracy (%) and computation efficiency (image/sec) of ResNet-32 and WideResNet-28-10 under 90% sparsity levels on CIFAR-10 and CIFAR-100. Comparison includes SGD, original SAM, efficient SAM variants, and our proposed methods.
Figure 7: Comparison of feature maps for different layers between Original DST and ZO-SAM under sparsity 90%.
Methods
CIFAR-10 (%)
CIFAR-10-C (%)
Δ↓
SNIP [ 22 ]
92.59
59.60
32.99
SNIP + ESAM [ 9 ]
93.40 (0.81 ↑ )
61.19 (1.59 ↑ )
32.21
SNIP + LS(k=5) [ 25 ]
93.31 (0.72 ↑ )
60.41 (0.81 ↑ )
32.90
SNIP + LS(k=10) [ 25 ]
92.97 (0.38 ↑ )
60.08 (0.48 ↑ )
32.89
SNIP + GSAM [ 48 ]
93.71 (1.12 ↑ )
61.65 (2.05 ↑ )
32.1
SNIP + ZO-SAM
93.38 (0.79 ↑ )
62.70 (3.10 ↑ )
30.68
Table 5: CIFAR-10 vs. CIFAR-10-C under 90% sparsity with SNIP.
Datasets
CIFAR-10
CIFAR-100
Pruning ratio
90%
95%
98%
90%
95%
98%
ResNet-50
95.23 (Dense)
78.42 (Dense)
SNIP [ 22 ]
92.65 ± 0.12
90.86 ± 0.17
87.02 ± 0.35
73.14 ± 0.32
69.25 ± 0.37
62.49 ± 0.32
SNIP+ ZO-SAM (ours)
92.98 ± 0.09 (0.33 ↑ )
91.47 ± 0.12 (0.61 ↑ )
87.92 ± 0.14 (0.90 ↑ )
73.64 ± 0.23 (0.50 ↑ )
69.97 ± 0.49 (0.72 ↑ )
63.79 ± 0.47 (1.30 ↑ )
GraSP [ 40 ]
92.47 ± 0.17
91.32 ± 0.15
88.77 ± 0.24
73.28 ± 0.21
70.29 ± 0.13
62.12 ± 0.33
GraSP+ ZO-SAM (ours)
92.89 ± 0.17 (0.42 ↑ )
91.85 ± 0.12 (0.53 ↑ )
89.47 ± 0.11 (0.70 ↑ )
73.79 ± 0.16 (0.51 ↑ )
70.98 ± 0.24 (0.69 ↑ )
63.27 ± 0.20 (1.15 ↑ )
6.Table: Test accuracy (%) of pruned ResNet-50 on CIFAR-10/100.
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) is an optimization method that improves the generalization performance of machine learning models. Despite its effectiveness, it has not been widely adopted in real-world applications due to its high computational cost. Existing approaches reduce this overhead by approximating or reusing gradients in the perturbation step, but still assume sequential perturbation and update steps. In this study, we introduce AsyncSAM, a variant of SAM that enables concurrent model perturbation and parameter updates across multiple compute resources. By adjusting the perturbation batch size in a system-aware manner, our asynchronous scheme fully overlaps perturbation with updates, eliminating the overhead and allowing even CPUs to perform perturbations without extra time. We also provide theoretical analysis showing that AsyncSAM achieves the same order of convergence as conventional SAM. Extensive experiments demonstrate that AsyncSAM delivers superior generalization performance and the highest training throughput among state-of-the-art SAM variants. For example, in ViT fine-tuning on CIFAR-100, it achieves accuracy comparable to SAM while maintaining training time nearly identical to SGD.
Zeroth-order (ZO) optimization offers a memory-efficient alternative to backpropagation for fine-tuning large language models. However, existing ZO methods typically sample training examples uniformly, overlooking their evolving importance and leaving avoidable estimation variance. To address this limitation, we propose Sample-Adaptive Zeroth-Order Optimization (SAZO), which learns sample importance online and adaptively allocates queries to accelerate convergence. SAZO extracts gradient-norm information from existing directional differences to estimate sample importance without access to true gradients. Furthermore, to track evolving importance from sparse minibatch feedback, SAZO maintains per-example exponential moving averages that aggregate observations across updates. The method requires no additional model evaluations and only O(N) additional scalar state compared with MeZO. We prove that SAZO asymptotically matches the leading convergence term of the minimum-variance sampling oracle, and achieves a strictly faster convergence rate than standard methods. Extensive experiments across four model families demonstrate improved average downstream performance and up to 6× faster convergence to the same training loss in terms of query count.
Haiyuan Wan, Yanyi Li, Weicheng Lin +1
School of Physics, Peking University · School of Intelligence Science and Technology, Peking University