ZO-SAM: Zero-Order Sharpness-Aware Minimization for Efficient Sparse Training

TL;DR

ZO-SAM introduces a zero-order optimization method integrated with sharpness-aware minimization to enhance sparse neural network training efficiency, stability, and robustness while reducing computational costs.

Contribution

It presents a novel zero-order SAM framework that halves backpropagation costs and improves convergence and robustness in sparse training scenarios.

Findings

Reduces backpropagation cost by 50% compared to traditional SAM.

Enhances training stability and convergence speed in sparse models.

Improves robustness under distribution shifts.

Abstract

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.