Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Jie Ji, Gen Li, Kaiyuan Deng, Fatemeh Afghah, Xiaolong Ma

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.

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

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.
Paper Structure (20 sections, 16 equations, 7 figures, 6 tables, 1 algorithm)

This paper contains 20 sections, 16 equations, 7 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Sparse Training Methods
  4. Zero-Order Optimization
  5. Sharpness-Aware Minimization
  6. Method
  7. Sparse Training Formulation
  8. Sharpness-Aware Minimization (SAM)
  9. Zero-Order Optimization (ZO)
  10. Algorithm
  11. Experiments
  12. Accuracy Improvement for State-of-the-art Sparse Training Methods
  13. Loss Surface Visualization
  14. Convergence Rate Comparison
  15. Feature Map Comparison
  16. ...and 5 more sections

Figures (7)

  • 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.
  • Figure 3: Overview of perturbation step in ZO-SAM.
  • Figure 4: Loss surface improvement at various sparsity levels.
  • Figure 5: Convergence Speed Comparison at sparsity level 90%.
  • ...and 2 more figures