A Single-Step, Sharpness-Aware Minimization is All You Need to Achieve Efficient and Accurate Sparse Training

Authors: jie  ji
gen  li
jingjing  fu
fatemeh  afghah
linke  guo
xiaoyong  yuan
xiaolong  ma
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 03:35:20

Abstract: Sparse training stands as a landmark approach in addressing the considerable training resource demands imposed by the continuously expanding size of Deep Neural Networks (DNNs). However, the training of a sparse DNN encounters great challenges in achieving optimal generalization ability despite the efforts from the state-of-the-art sparse training methodologies. To unravel the mysterious reason behind the difficulty of sparse training, we connect the network sparsity with neural loss functions structure, and identify the cause of such difficulty lies in chaotic loss surface. In light of such revelation, we propose $S^{2} - SAM$, characterized by a **S**ingle-step **S**harpness_**A**ware **M**inimization that is tailored for **S**parse training. For the first time, $S^{2} - SAM$ innovates the traditional SAM-style optimization by approximating sharpness perturbation through prior gradient information, incurring *zero extra cost*. Therefore,  $S^{2} - SAM$ not only exhibits the capacity to improve generalization but also aligns with the efficiency goal of sparse training. Additionally, we study the generalization result of $S^{2} - SAM$ and provide theoretical proof for convergence. Through extensive experiments, $S^{2} - SAM$ demonstrates its universally applicable plug-and-play functionality, enhancing accuracy across various sparse training methods. Code available at https://github.com/jjsrf/SSAM-NEURIPS2024.

Main paper: 
 introduction: 
 The arrival of the Artificial General Intelligence (AGI) [1] era has urged an ever-expanding realm of artificial intelligence, bringing significant growth in deep neural networks (DNNs) depth and intricacy. The efficient training of DNN has thus emerged as an uttermost imperative, demanding immediate and concerted efforts.To train an overparameterized, large-scale DNN efficiently while preserving its high accuracy, stateof-the-art literature has introduced sparse training [2] [3] [4] [5] [6] as a straightforward solution that reduces both parameter footprint and computation cost. With the presence of a large portion of zeros in the network, both forward and backward computation for training can be saved by skipping those zeros, as well as reducing the memory consumption since the zeros are not necessary to be stored. However, training a sparse neural network is difficult since the optimization under sparse regime readily converges to stationary points with a sub-optimal generalization accuracy (i.e., saddle points) [7] . Due to the difficulty of directly (i.e., without the time-consuming linear interpolation) assessing if the sparse solution is in the saddle point, finding workaround approaches for training a better sparse neural network is critical, especially for the practical usage in efficient learning [8] [9] [10] [11] [12] [13] [14] [15] [16] [17] .To address the above difficulty, many efforts have been made. For instance, the Static Sparse Training (SST) such as the Lottery Ticket Hypothesis (LTH) [20] , SNIP [21] , GraSP [18] and SynFlow [22] The loss surface visualization for training a sparse neural network using ResNet-32 on CIFAR-10. We select two representative sparse training methods [18, 3] and incorporate different levels of sparsity. We also quantify the loss surface behavior using coefficient Ra [19] to evaluate sharpness. With increased sparsity, Ra becomes larger, indicating sharper and steeper surface. determine a static sparse pattern at training initialization. On the other hand, Dynamic Sparse Training (DST) such as SET [2] , RigL [23] , and MEST [3] iteratively updates the sparse topology during training to find a better sparse model. However, those methods either suffer from suboptimal generalization ability due to heuristic nature of their methodology or metric settings, or experience difficulties in parameter setting and dynamic sparsity scheduling. With the increase of the sparsity, these phenomenon become more severe.We identify that the cause of the sparse network learning difficulties lies in the variable learning dynamics, which is closely related to the network topology. When incorporating sparsity into a neural network, the effective structure of the original network becomes narrower. According to the study of neural loss function structure [24] , a wide network resulted in flat minima and wide regions of apparent convexity, which helps prevent the chaotic behavior that occurs during training. Evidently, higher sparsity indicates a narrower structure, which suggests more chaotic behavior is expected during training, thus degrading the accuracy. We perform different types of sparse training with different levels of sparsity and plot the loss surface in Figure 1 to demonstrate the convexity of the training behavior. According to the sharpness of the loss basin, we observe a transition from a smooth and close-to-convex surface to a steep one as sparsity is introduced and increased. Since higher sparsity levels lead to more chaotic behavior, in landscapes that are narrow and sharp with a large Ra value, we might expect to encounter more chaotic training behavior that degrades generalization ability.It seems that achieving coexistence of sparsity and good generalization ability during training is challenging. Therefore, we raise the following question: is there a simple, effective method that can improve generalization ability of training a sparse neural network, without sacrificing efficiency (sparsity) and incurring zero extra cost? We believe the answer to the above question lies in the sharpness of the loss surface. Inspired by the Sharpness-Aware Minimization (SAM) [25] technique that finds flatter minima that have uniformly low loss in nearby regions, we argue that such technique is especially suitable for sparse training since the steep loss surface induced by sparsity can be directly mitigated. However, to leverage sharpness for better generalization, SAM uses an additional full training step (i.e., forward and backward propagation) to quantify and evaluate loss on the summation of current weights and a constrained perturbation, which roughly doubles the computation of training. Such extra costs contradict the goal of sparse training. Although some literature [26, 27] have been proposed to reduce the computation cost for SAM, the computation is still significant, posing a roadblock to the wide application of SAM to the realm of sparse training.In this paper, we propose a novel approach to achieve sharpness-aware minimization with zero extra cost, tailored for the sparse training regime to maintain efficiency while improving generalization ability. It is the first time that a Single-step Sharpness-Aware Minimization is proposed for Sparse training (S 2 -SAM). Different from the traditional two-step computation regime of SAM, S 2 -SAM uses a unique single-step approach that leverages sharpness and trains weights with one training step. Specifically, S 2 -SAM uses the weight gradients from the prior step to approximate the perturbation to the weights, thus solving the sharpness evaluation without performing an extra full training step. Therefore, S 2 -SAM incurs zero extra cost to achieve sharpness-aware training, which aligns with the efficiency goal of sparse training. We also study the generalization result of S 2 -SAM and provide theoretical proof for convergence. We demonstrate that S 2 -SAM provides a straightforward plug-andplay functionality on variety of sparse training methods, significantly boosting their accuracy. Our contributions are summarized as follows:• We identify that the difficulty of training a sparse neural network lies in the increasingly chaotic and steep loss surface when sparsity is introduced and increased. SAM is an optimization technique designed to enhance neural network generalization and mitigate overfitting. It minimizes the maximum loss in a neighborhood around the current parameters, as opposed to solely focusing on the loss at the current point. This approach identifies flatter minima that have uniformly low loss in nearby regions, ultimately contributing to improved generalization performance.Specifically, consider a family of models parameterized by w ∈ W ⊆ R d ; L is the loss function, and S denotes the training dataset. SAM aims to minimize the following upper bound of the PAC-Bayesian generalization error: for any ρ > 0,EQUATIONTo solve the above minimax problem, at each iteration t, SAM updates the following steps:EQUATIONwhere 1/p + 1/q = 1, ρ > 0 is a hyperparameter, λ > 0 is the parameter for weight decay, and η t > 0 is the learning rate. By setting p = q = 2 and introducing an intermediate variable u t , we have:EQUATIONEQUATION 
the proposed s 2 -sam method: 
 As shown in Equation (2) , SAM needs to compute the gradient twice at each iteration, involving additional computation costs. Further, the two-step gradient computation is not parallelizable, which presents challenges for deployment in large-scale training scenarios. Thus, we propose a new algorithm that only needs to compute the gradient once in each iteration. Different from prior efficient SAM works [26] [27] [28] aiming to reduce the computation by introducing periodically SAM steps or data selection, our proposed framework S 2 -SAM achieves zero extra computation cost while maintaining the improved generalization ability of sparse network training.In Figure 2 , we demonstrate that the perturbation on the current weights is approximated by the weight gradients g t-1 from the prior step. The rational behind such design is that the gradient direction (i.e., w t-2 to w t-1 ) from prior step is optimizing the prior loss to a considerable extent, thus it can be used to represent a sharp direction among all perturbation directions at t -1. Since the loss surface of a sparse network is too chaotic, a perturbation with relatively high degree of sharpness (i.e., the prior gradient) can still find a neighborhood with near-maximum loss. The following equations specifying g t-1 information are used to substitute the ∇L S (w t-1 ) step in Equation (2) . Thus, we only need to compute the gradient once in each iteration:EQUATIONEQUATIONEQUATION)Remark 1. The parameter ρ can vary in terms of iteration t : ρ t = c/t, where c > 0 is a constant.In this section, we study the generalization result of S 2 -SAM. First, we give some notations, where most of them are followed by [29, 30] . Let w S = A(S) be a solution that generated by a random algorithm A based on dataset S. Recall that problem (8)EQUATIONis called empirical risk minimization in the literature, and the true risk minimization is given by min w∈WEQUATIONWe define its optimal solution: w * ∈ arg min w∈W F (w). Then the excess risk bound (ERB) is defined as E A,S [F (w S )] -F (w * ) . (10) It has been shown that the ERB can be upper bounded by optimization error and generalization error [29, 30] . We notice that there are several works [29, 30] studying the generalization result of SGD for non-convex setting under different conditions, such as bounded stochastic gradient ∥∇ w ℓ(y, f (w; x))∥ ≤ G and decaying learning rate η t ≤ c t with a constant c > 0, where t is the optimization iteration. In this paper, we are not interested in establishing a fast rate in ERB under different conditions, but we want to explore the generalization ability of S 2 -SAM with the fewest possible modifications when building a bridge between theory and practice. For example, weight decay is a widely used trick when training deep neural networks. With the use of weight decay, the empirical risk minimization in practice becomes min w∈W F S (w) := F S (w) + λ 2 ∥w∥ 2 . Then, we define some notations as follows. Specifically, letEQUATIONFollowed by [30] , we use the following decomposition of testing error:EQUATION) where the upper bound is the optimization error plus the generalization error.Next, we present some notations and assumptions that will be used in the convergence analysis. Throughout this paper, we also make the following assumptions for solving the problem (8) .Assumption 1. Assume the following conditions hold: (i) The stochastic gradient of F S (w) is unbiased, i.e., E (x,y) [∇ℓ(y, f (w; x))] = ∇F S (w), and the variance of stochastic gradient is bounded, i.e., there exists a constant σ 2 > 0, such thatE (x,y) ∥∇ℓ(y, f (w; x)) -∇F S (w)∥ 2 = σ 2 .(ii) F S (w) is smooth with an L-Lipchitz continuous gradient, i.e., it is differentiable and there exists a constant L > 0 such that ∥∇F S (w) -∇F S (u)∥ ≤ L∥w -u∥, ∀w, u ∈ W.Assumption 2. There exists a constant µ > 0 such that 2µ (F S (w) -F S (w * S )) ≤ ∥∇F S (w)∥ 2 , ∀w ∈ W, where w * S ∈ min w∈W F S (w) is a optimal solution (PL condition [31] ). Now consider that A = S 2 -SAM. We define the gradient update rule G ℓ,η as followsEQUATIONEQUATIONThen we have the following lemma, which is similar to Lemma 2.5 and Lemma 4.2 in [29] that use recursive definition through variable G ′ and w′ t . Lemma 1. Assume that ℓ(y, f (w, x)) is L-smooth and B-Lipschitz. Let w t+1 = G (w t ) and another sequence w ′ t+1 = G ′ (w ′ t ), then w t+1 -w ′ t+1 = (1 + ηL -ηλ) ∥w t -w ′ t ∥ + 2ηρ, G = G ′ (1 -ηλ) ∥w t -w ′ t ∥ + 2ηB, G ̸ = G ′Please see the proof of Lemma 1 in Appendix B.Theorem 1. Under Assumption 1, assume that ℓ(y, f (w, x)) is L-smooth and B-Lipschitz, suppose F S (w) satisfies Assumption 2 and min w∈W F S (w) ≤ F S (w * S ) + λ 2 ∥w t ∥ 2 with λ = 2L, where w t is the intermediate solution of A, thenE R,A,S [F (w R )] -E S [F (w * )] ≤ F S (w 0 ) µηT + η(L + λ)σ 2 2µ + 6B + 1 nwhere A is SGD.The w * S ∈ min w∈W F S (w) is an optimal solution, and we show that the generalization error is bounded. Based on Lemma 1, the proof of Theorem 1 is derived in Appendix C. 
experimental results: 
 In this section, we carry out comprehensive experiments to demonstrate how S 2 -SAM improves sparse training performance. We test S 2 -SAM on CIFAR-10/100 [34] with ResNet-32 [32] and VGG-19 [35] , and we also perform experiments on ImageNet-1K [36] and ImageNet-C [37] based on ResNet-50 [38] .Following the recent developments in sparse training techniques, we apply S 2 -SAM on static sparse training such as LTH [20] , SNIP [21] , GraSP [18] , as well as dynamic sparse training methods such as SET [2] , DSR [33] , RigL [23] , MEST [3] , CHEX [10] and Chase [39] . We apply S 2 -SAM to the official codes or published implementations to show performance gains. All of our experiments are performed on NVIDIA 4× A6000 GPUs. We repeat training experiments for 3 times and report the mean and standard deviation of the accuracy. For training throughput evaluation, we adopt the original settings of each baseline method and record the throughput on 4× A6000 GPUs. We stress that for all experiments, we only use one setting for S 2 -SAM, which indicates an easy implementation that has the potential for wide applicability. We also find that S 2 -SAM works better on high sparsity. For all sparse training with a 98% sparsity ratio, S 2 -SAM improves the original baseline accuracy by an average of 0.77% on CIFAR-10 and 1.12% on CIFAR-100, which shows S 2 -SAM is especially suitable for challenging tasks.ImageNet-1K. We evaluate S 2 -SAM on ImageNet-1K dataset. The results are shown in Table 2 . We report the top-1 accuracy of the sparse training results, as well the overall training cost with FLOPs. We set hyper-parameter ρ = 1.0 for all experiments. Compared to the baseline methods, our result shows higher test accuracy at the same computational cost. Similarly, S 2 -SAM also shows better gain on a higher sparsity ratio for most of the sparse training baselines, which indicates the consistency of using our methods on different scales of datasets. To make a fair comparison with different training recipes, we slightly scale up our training epochs to have the same or less overall training FLOPs to compare with longer training baselines, and we find out that S 2 -SAM shows stable improvement on those methods.In this section, we evaluate how S 2 -SAM performs on structured sparse training methods. Note that all previous baselines use unstructured sparsity where the zeros are scattered in the network. Such unstructured sparse training shows the potential of training a sparse model, but cannot fully represent training acceleration ability. Therefore, we apply S 2 -SAM to two representative structured sparse trainings, CHEX [10] and Chase [39] which dynamically train sparse networks with channel-level sparsity. The results are shown in Table 3 . From the results, we can see that the average accuracy improvement is around 0.5%. Considering structured (channel-wise) sparsity is one challenging sparsity type to achieve good accuracy, these results indicate S 2 -SAM works very well on structured sparse training methods.In this part, we plot the loss surface of the sparse training with and without S 2 -SAM. The comparisons are demonstrated in Figure 3 , where we perform our experiments with three representative sparse training methods SNIP [21] , GraSP [18] and MEST [3] at 90%, 95% and 98% sparsity ratios on S 2 -SAM is a highly efficient approach to optimizing the sharpness of the loss surface, which incurs zero extra cost for achieving its functionality. Compared to the traditional SAM, S 2 -SAM uses fewer computations and has better training performance on GPUs. In Since sparse training approaches have potentially high usage in practical scenarios, the system's robustness against perturbation (e.g., inclement weather conditions for image/video tasks) is usually critical. We perform experiments to evaluate the sparse model (80% sparsity) robustness against perturbations in Table 5 . Our intuition is that the model trained with its sharpness optimized has a wider loss basin, which indicates higher endurance on perturbations since the loss won't change much when it is located in a wide and flat region. We adopt ImageNet-C [37] which contains a test set with the same images as ImageNet-1K but with nineteen types of corruptions applied with five different levels of severity. We train the sparse networks using different methods on ImageNet-1K, and then test their accuracy on ImageNet-C. We report test accuracy for both datasets. We can see that when the sparse model is trained using the original method, the accuracy of ImageNet-C is significantly lower than the accuracy of the original ImageNet-1K test set (around 35% lower). With S 2 -SAM, the test accuracy on ImageNet-C shows promising improvement. The robust accuracy improves by an average of 3.23%. We also apply S 2 -SAM to the dense model training. The results are shown in Table 6 . We test on two datasets CIFAR-10 and ImageNet-1K with two additional networks MobileNet-V2 [41] and EfficientNet-B0 [42] . For ResNet-32 and VGG-19 on CIFAR-10, we follow the settings in [3] and train for 160 epochs, and for MobileNet-V2, we use the 1.0 width version and train for 350 epochs for better convergence [43] . For ResNet family on ImageNet-1K, we follow the setting in [3] and train for 150 epochs, and we train EfficientNet-B0 for 350 epochs for better convergence [44] . From the results, we can see that S 2 -SAM still improves the accuracy of the dense networks. We find out that the less the parameter count of the network, the better effectiveness of S 2 -SAM it achieves. This phenomenon proves our finding that a narrower network is harder to train, and S 2 -SAM is an effective solution for better generalization. We must also stress that S 2 -SAM is focusing on sparse neural network training, and we will leave the study of S [20, 45, 46] uses iterative magnitude-based pruning (IMP) to find a subnetwork that can be trained from scratch without losing accuracy. SNIP [21] , GraSP [18] , SynFlow [22] determine a static sparse pattern at training initialization by obtaining gradient information with a few iterations of dense training. FISH [47] acquires fixed subnetwork by pre-computing a sparse mask using Fisher information.Dynamic Sparse Training Dynamic sparse training starts with a randomly selected sparse network structure and adapts it throughout the training process in an effort to find a better sparse structure. Sparse Evolutionary Training (SET) [2] prunes small magnitude weights and grows back randomly at the end of each training epoch. Deep R [48] uses a combination of stochastic parameter updates and dynamic sparse parameterization for training. Dynamic Sparse Reparameterization (DSR) [33] proposes to redistribute parameters between layers during training. Sparse Networks from Scratch (SNFS) [49] creates the sparse momentum algorithm, which finds the layers and weights that effectively reduce the error by using exponentially smoothed gradients (momentum). In RigL [23] , the gradients of all the weights are computed when the model needs to be updated to grow new connections. Top-KAST [40] proposes a scalable and performant sparse-to-spars DST framework for maximum efficacy. Powerpropagation [50] suggests a novel neural network weight parameterization that largely preserves low-magnitude parameters from learning. ITOP [51] investigates the underlying DST mechanism and finds that the advantages of DST result from a time-based search for all potential factors. MEST [3] designs a memory-economic sparse training framework targeting accurate and fast execution on edge devices. By co-training dense and sparse models, AD/AC [52] suggests a technique that, at the conclusion of training, produces precise sparse-dense model pairings. Chase [39] dynamically translates the unstructured sparsity into channel-level sparsity to achieve direct speedup on GPU.Sharpness-Aware Minimization Sharpness-Aware Minimization was first introduced in [25] . This optimization technique is designed to identify flatter minima characterized by consistently low loss in neighboring regions, aiming to enhance the model generalization capability during training. Nevertheless, the computational cost of Sharpness-Aware Minimization is doubled due to its two-step gradient computation regime, presenting challenges for deployment in large-scale training scenarios.To reconcile such, ESAM [26] introduces two novel and efficient training strategies: stochastic weight perturbation and sharpness-sensitive data selection, enhancing the efficiency of the SAM process without compromising its generalization performance. LookSAM [27] only periodically calculates the inner gradient ascent to significantly reduce the additional training cost of SAM. SAF [28] introduces a novel trajectory loss based on KL-divergence to measure the rate of change in training loss along the model update trajectory and replace the SAM sharpness measure. CrAM [53] optimizes over the compression projection applied to the intermediate model at every training step, which results in a compressible models that can be pruned in an on-shot manner after training. However, all those methods need extra computation cost, and not target on sparse training method. 
conclusion: 
 In this paper, we propose a novel Single-step Sharpness-Aware Minimization that is tailored for Sparse training (S 2 -SAM), which revolutionizes the originally computation-intensive sharpness-aware optimization into a highly efficient tool with zero extra cost. In light of the improved optimization trajectory in the loss surface, S 2 -SAM successfully enhances the accuracy of the sparse network training, as well as the robustness of the sparse model in practical scenarios. S 2 -SAM offers seamless plug-and-play functionality, showcasing its potential for widespread applicability in the evolving landscape of efficient training. The research is inherently scientific, and we anticipate no adverse societal impact stemming from its findings. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We report the training resources we use in the experimental settings.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: We discuss the broader impacts in the conclusion. The research is inherently scientific, and we anticipate no adverse societal impact stemming from its findings.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our method doesn't have high risk for misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 
licenses for existing assets: 
 Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We credit every referenced assets with citations.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [Yes] Justification: This paper introduces new code for sharpness-aware minimization, and we will release the code after the paper is accepted. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Proof. We consider two cases.(EQUATION(EQUATIONProof. By Theorem 3.2 of [29] , we haveEQUATIONwhere (a) uses Lemma 1; (b) uses λ = 2L; (c) uses ∆ t0 = 0; (d) uses ηL < 1. Then, by Lemma 3.11 of [29] and ℓ is (L + λ)-smooth, we haveEQUATIONwhere the last inequality holds by selecting t 0 = 1 and λ = 2L. Then, by Theorem 2.2 of [29]EQUATIONwhich is the generalization error. Next, we will bound the optimization error. Since F S is (L + λ)smooth and satisfies µ-PL condition, then follow the standard analysis, we have Since F S is (L + λ)-smooth, we haveEQUATIONEQUATIONwhere the last inequality is due to ℓ(y, f (w, x)) is B-Lipschitz. Therefore, we getEQUATION)Since F S satisfies µ-PL condition, thenEQUATIONThen by ( 12), ( 17) and ( 21), we getEQUATIONBy using the conditions that min w∈W F S (w) ≤ F S (w * S ) + λ 2 ∥w t ∥ 2 , definitions (11) , we getEQUATION) Therefore, we have the following inequality by ( 22) and ( 23):EQUATION) By Lemma 5.1 of [29] , (24) implies the ERB isEQUATIONBy settingη = O(1/ √ n) and T = O(n) then E R,A,S [F (w R )] -F (w * ) ≤ O(1/ √ n), where A is S 2 -SAM.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.Justification: Our research conforms with the NeurIPS Code of Ethics in every respect.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 


Decision: This paper proposes the application of Sharpness-Aware Minimization (SAM) to sparse training to enhance its generalization performance. To maintain the low cost of sparse training, the paper introduces a variant of SAM that requires only a single gradient computation by approximating sharpness perturbation using prior gradient information. Extensive empirical results are presented across various sparse training approaches and evaluation metrics. All reviewers agree that this paper makes a significant contribution to the field of sparse training.

Reviews: review: paper_summary: paper_summary: This paper studies the challenges of training sparse neural networks directly and identifies one of the contributing factors, i.e., the chaotic loss surface. Consequently, it proposes a new method, i.e., Single-step Sharpness-Aware Minimization (S2-SAM), tailored specially to train sparse networks. S2-SAM is based on SAM for dense neural networks training with the main difference that it uses just one gradient computation (thus, more efficient), while SAM uses two gradient computations. Experimental results show unanimously that S2-SAM improves the performance of all sparse training methods studied.
strength_weakness: strength_weakness: * I don’t see major weak points, except that it is not clear when the source code will be made available for easy reproducibility
questions: questions: Q1) It seems that S2-SAM works with sparse networks, but not too well with dense networks. As sparser the network is as more impactful S2-SAM is on the overall performance (lines 252-253 and the majority of the results). Can you prepare a systematic study (experiment) to quantify and illustrate better this behavior (e.g., by varying the sparsity level from 0 to 100% in small steps, or alternatives…)?

Q2) Could you present in an Appendix how the loss surface visualisations from Figures 1 and 3 have been computed? I see that you cite [6], but overall this seems to be an approximation method to visualize the loss surface of a very high dimensional space. If so, this shall be properly acknowledged to avoid bringing in readers’ mind inaccurate ideas. 

Q3) It seems that on ImageNet, which is a much more challenging dataset, your proposed method together with MEST or RigL outperforms the dense baseline (without S2-SAM). Can you present the results on this latter case also? Do you have any idea why this behaviour is different on CIFAR 10/100? Probably, this deserves a longer qualitative discussion.

Q4) The whole empirical validation has been performed on convolutional neural networks (not on all neural network architectures), and the “chaotic” loss surface is probably just one (not the only one) of the reasons which can help understand the training behaviors of sparse neural networks. Can you please identify better in the paper the boundaries of your work to avoid inaccurate conclusions in the readers’ mind?

Q5) (minor) To expand the broadness of the experimental results. In Table 1, would it be possible to add also the 80% sparsity level?In Table 2, would it be possible to add also a SET 5x?
limitations: limitations: n/a
paper_summary: paper_summary: This article introduces S2-SAM (Single-step Sharpness-Aware Minimization), an innovative sharpness-aware optimization method tailored specifically for sparse training with zero extra computational cost.
strength_weakness: strength_weakness: 1. S2-SAM seems to be a zero-extra-cost variety of SAM that is implemented in sparse cases. However, when the sparsity is high, the extra cost of SAM can be ignored.
questions: questions: 1. What is the topological initialization of Table 2? Why some of them are non-uniform and some of them are uniform?

2. The datasets in the article are all computer vision domains. Have you done some experiments on NLP domain?
limitations: limitations: This paper doesn't include a limitation section
paper_summary: paper_summary: The authors of this paper posit that sparse training is difficult due to a chaotic loss landscape as opposed to standard training of a dense network. In order to address this problem, they propose to perform sparse training with a Sharpness Aware Minimization approach. In order to do so efficiently, they leverage the gradient of the previous iteration to identify the SAM perturbation (S^2-SAM) instead of performing a second forward-backward step like SAM. They show that this approximation of SAM also converges to an optimal solution theoretically. 
Extensive empirical evaluations show that the proposed method is effective for sparse training and can also improve robustness of the sparse networks.
strength_weakness: strength_weakness: 1. The paper starts from the premise that sparse training is difficult in comparison to standard dense training. However, sparse training does not necessarily generalize poorly, in fact some methods are able to train sparse networks that outperform their dense counterparts as shown by Jin et. al. [1] In fact, Jin et. al. claim that pruning can behave as a regularizer, enabling better generalization. 
Similarly, Renda et. al. [2] have shown that Learning Rate Rewinding (LRR), a sparsifying algorithm similar to LTs, can find sparse networks that outperform their dense counterparts.   
Do these findings suggest that sparse networks can also be found without performing S^2-SAM. Does LRR also find better loss landscapes for sparse networks and if so, then is S^2-SAM necessary? It would also be beneficial to compare with LRR.
2. Does performing SAM instead of the proposed S^2-SAM have a stronger effect on sparse networks i.e. do sparse networks generalize better when trained with SAM instead of S^2-SAM, no matter the training cost. It would be nice to have this comparison to shed light on the robustness of the proposed training method.
questions: questions: Does the proposed method largely behaves as a regularizer and hence improves generalization of sparse networks? And if so, are there other such regularization methods that could potentially offer similar benefits other than the proposed S^2-SAM. For example, it was suggested by Paul et. al. [3] that as long as the loss landscape is linearly connected for LTs, sparse networks generalize well. Hence, would S2-SAM allow for faster training in this case or potentially pruning a larger fraction of parameters in one iteration in comparison to Iterative Magnitude Pruning?

[1] Jin, Tian, et al. "Pruning’s effect on generalization through the lens of training and regularization." Advances in Neural Information Processing Systems 35 (2022).

[2] Renda, Alex, Jonathan Frankle, and Michael Carbin. "Comparing Rewinding and Fine-tuning in Neural Network Pruning." International Conference on Learning Representations.

[3] Paul, Mansheej, et al. "Unmasking the Lottery Ticket Hypothesis: What's Encoded in a Winning Ticket's Mask?." The Eleventh International Conference on Learning Representations.
limitations: limitations: See above
paper_summary: paper_summary: This paper presents a sparse training method, $S^2$-SAM, which applies sharpness-aware minimization to sparse training. The authors demonstrate that sparsity during training leads to a sharper (to use the authors' words, "more chaotic") loss surface, something that can be mitigated by a variant of sharpness-aware training. Additionally, the authors avoid the extra gradient computation of SAM by reusing the gradient from the previous step as a proxy. The authors demonstrate theoretically that the error from $S^2$-SAM is bounded, and empirically by testing their algorithm with a number of sparse training methods, as well as with dense training.
strength_weakness: strength_weakness: The main experimental weakness of the paper is restricting the method to sparsities no higher than 90 on ImageNet and not doing a full ablation of the components of  $S^2$-SAM. It would have been interesting to see the high-sparsity results, as it seems that $S^2$-SAM would have been quite effective there (and if not, why not)? Likewise, an ablation would have been helpful to understand how much we lose by using the previous-gradient version of SAM (the authors do try $S^2$-SAM on dense models, but compare to regular SGD, not regular SAM). Comparing with the original SAM paper, the numbers actually seem pretty promising on dense models, but this is not included in this paper.

I found the sharpness figures confusing. The overall intuition is clear, but the actual process by which the figures were obtained is not.
questions: questions: How much test accuracy is lost by using the previous-gradient estimate of SAM as compared to the normal SAM?
limitations: limitations: The limitations were adequately addressed.

Comments: title: Rebuttal acknowledgement
Thank you for your support
Thank you for the rebuttal
Thank you for your support
Response to rebuttal
comment: Dear authors,

Thank you for considering my comments and for the well prepared rebuttal. I will keep my original score (accept).

Best wishes,
Dear Reviewer UgUE,

We sincerely appreciate the time and effort you’ve invested in providing thoughtful and constructive feedback on our submission. We're delighted that you view our work as a valuable contribution to SAM and sparse training. We hope our responses have fully addressed your concerns.

If you have any further questions, we would be glad to follow up.

Best regards,

The Authors
Thank you to the authors for the rebuttal and additional experiments. I remain convinced that this paper will be a useful addition to our understanding of SAM and sparse training, and I keep my (positive) score.
Dear Reviewer Z2wG,

We wanted to express our sincere gratitude for raising our score. Thank you for your support and constructive comments! We will include all the updates in our revision.

Best regards,

The Authors
I thank the authors for a detailed explanation and providing additional experiments to highlight the effectiveness of S^2-SAM. I do believe that their proposed method will be useful for sparse training. I have increased my score.