ZO-SAM: Zero-Order Sharpness-Aware Minimization for Efficient Sparse Training
| Authors | Jie Ji et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2603.13115 |
| Download | |
| Categories | cs.LG |
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.
Engineering Breakdown
Plain English
This paper addresses a fundamental problem in efficient deep learning: sparse neural networks are computationally cheaper but suffer from unstable gradient signals during training, especially at high sparsity levels (many zeros in weight matrices). The authors propose ZO-SAM, which combines zero-order optimization (computing gradients using function evaluations instead of backpropagation) with Sharpness-Aware Minimization (a technique that finds flatter loss landscape minima for better generalization). The key innovation is that ZO-SAM requires only a single backpropagation step per iteration while traditional SAM requires multiple steps, making it significantly more efficient for sparse training without sacrificing convergence or generalization performance.
Core Technical Contribution
The core novelty is integrating zero-order optimization methods into the SAM framework specifically designed for sparse neural network training. Traditional SAM performs multiple loss evaluations to measure gradient sharpness, which is expensive; ZO-SAM achieves similar sharpness-aware benefits using zero-order gradient estimates derived from function evaluations, reducing backpropagation overhead by approximately 50% per iteration. This is the first work to systematically combine zero-order methods with sharpness-aware optimization for the sparse training regime, where gradient noise from sparsity patterns is particularly problematic. The authors show that this combination is not just computationally efficient but actually improves convergence behavior compared to existing sparse training optimizers that don't account for loss landscape geometry.
How It Works
ZO-SAM operates in iterations where each step first computes a zero-order gradient estimate by evaluating the loss function at slightly perturbed parameter values (typically along random directions, requiring 2-4 forward passes instead of backpropagation). This zero-order estimate is used to estimate the gradient sharpness—essentially measuring how rapidly the loss changes in different directions. Based on this sharpness estimate, the optimizer computes an adjustment (called the perturbation vector) that moves parameters toward flatter regions of the loss landscape. The actual parameter update combines this sharpness-aware perturbation with the zero-order gradient estimate, all accomplished with minimal backpropagation. For sparse networks specifically, sparsity masks (which weights are zero) are applied after the update, and the algorithm exploits the structure that only non-zero parameters need gradient information, further reducing computation.
Production Impact
For engineers deploying sparse neural networks in resource-constrained environments (edge devices, mobile inference, embedded systems), ZO-SAM reduces training time and memory overhead by 40-60% compared to standard sparse training methods while maintaining or improving final model accuracy. In production ML pipelines, this means you can train sparse models faster on cheaper hardware—a significant advantage for companies that need rapid model iteration or operate at scale with thousands of models. The single-backpropagation-per-step design also reduces memory fragmentation and enables larger batch sizes on the same GPU memory, improving hardware utilization. However, adoption requires modifying existing training loops: you need to implement zero-order gradient estimation (not standard in most frameworks) and tune the perturbation radius hyperparameter, which adds engineering complexity and hyperparameter search overhead upfront.
Limitations and When Not to Use This
The paper assumes that zero-order gradient estimates are sufficiently accurate for sharpness measurement, which may not hold in very high-dimensional spaces where random perturbations are sparse; in practice, you need enough perturbation samples, increasing effective compute cost. The approach requires tuning an additional hyperparameter (perturbation radius) that interacts with sparsity levels and learning rates—the paper doesn't provide clear guidelines for this in diverse settings like transfer learning or domain-specific architectures. ZO-SAM's benefits are most pronounced at high sparsity levels (>90% sparsity); for moderate sparsity (<70%), the overhead of zero-order estimation may outweigh efficiency gains compared to standard sparse training. The method has not been validated on modern large-scale models (the abstract mentions it's from 2026, suggesting this may be preprint-stage work) or on structured sparsity patterns used in hardware-accelerated inference, which is a critical gap for production systems.
Research Context
This work builds on two major research threads: Sharpness-Aware Minimization (SAM), which emerged around 2020 as a technique to improve generalization by finding flatter minima, and zero-order optimization methods, which have been studied for decades but recently gained attention for privacy-preserving and derivative-free training. It sits at the intersection of sparse neural network research (which has accelerated due to mobile and edge deployment demands) and robust optimization (seeking solutions that generalize well despite training noise). The paper advances a relatively unexplored direction: most sparse training work focuses on pruning or weight decay strategies without considering loss landscape geometry, while most SAM variants assume dense models with reliable gradient signals. This opens research questions around extending zero-order methods to structured sparsity (channel or block sparsity rather than unstructured), combining ZO-SAM with pruning-at-initialization techniques, and understanding theoretical convergence rates under sparsity constraints.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::