Adaptive Sharpness-Aware Pruning for Robust Sparse Networks

Thank you for your thoughtful comments. We hope that we have addressed your points satisfactorily below.

Weaknesses

1. In response to the first weakness, we agree that there are similarities between our method and SAM, and that there has been much work studying SAM in various settings. SAM was introduced as a method to improve (in-distribution) validation performance for (dense) models, while our goal is to improve OOD robustness of pruned models. We both study SAM in this novel context and introduce a new method that outperforms vanilla SAM. While we found that SAM can improve the robustness of pruned models (see Table 6 and Appendix Table 17), it was not intended to be used for pruning so it should be unsurprising that we can build on top of it to produce a better sharpness-aware pruning paradigm. The key difference in our method is that we introduce a novel adaptive weight perturbation, designed to be compatible with network pruning, which helps not only with standard validation performance, but also improves robust performance. Our adaptive perturbations penalize neurons that will be pruned differently than neurons which will be preserved post pruning, thus allowing the network to be optimally primed for pruning.

2. With regards to the second point, as described in appendix A.5, we choose our maximum $\rho$ based on the recommended values for $\rho$ in line with SAM and ASAM (Foret et al. 2020; Kwon et al., 2021). In our earliest experiments, we ran simple experiments where we evaluated the validation performance after one or two epochs (of training; no pruning involved) for varying $\rho_{min}$ and $\rho_{max}$. We tried minimum values of 0.01 and 0.05 and maximum values of 0.05, 0.1 and 0.5 (this formulation only considered perturbation balls, not other transformed shapes, and so we considered values close to the recommended $\rho=0.05$ in the SAM paper). We found here that the best $\rho_{min}$ and $\rho_{max}$ values were 0.01 and 0.05, respectively. When we added in the perturbation region transformations from ASAM, we only updated our maximum $\rho$ value to 2.0, in line with the recommended $\rho$ hyperparameter in the ASAM paper. We have added the following table to section A.5 of our paper (Table 7 in the appendix).

Remaining hyperparameters (i.e. learning rate, learning rate schedule, weight decay, warmup period, etc.) were fixed based on a previous SOTA paradigm with released code (HALP).

3. Finally, thank you for pointing out the two related works. We have cited them in our revised paper. Zhou et al. similarly analyzes SAM’s effectiveness as a pre-pruning optimizer, allowing for improved compression in certain regimes. Our work could perhaps be further refined by using the methods in this paper to vary the range of ($\rho_{min}$, $\rho_{max}$) based on the density of the model. As well, Peste et al. employs a similar approach to our paper - they modify the SAM objective to be more amenable to pruning and identify which weights would be preserved or pruned and optimize accordingly. CrAM primarily focuses on unstructured sparsity and does not evaluate OOD robustness, so the results are not immediately comparable to ours. However, it would be interesting to investigate under which conditions AdaSAP and CrAM perform well: one-shot vs gradual pruning, structured vs. unstructured sparsity patterns, standard validation vs OOD robustness performance.

We have cited these two works near the end of page 3:

Some work builds upon SAM, including improving the efficiency of the method (Liu et al., 2022; Du et al., 2022), modifying the optimization for improved performance (Kwon et al., 2021), applying it for improved model compression in NLP applications (Na et al., 2022) and as a pre-pruning optimizer (Zhou et al., 2023), and modifying the objective to help with unstructured pruning (Peste et al., 2022).

Questions

1. Yes, we would generally expect these results to be invariant to sharpness measure. We have some limited results measuring the sharpness via the top eigenvalue of the Hessian. We compare AdaSAP$_L$ to SGD with HALP pruning in the following table and observe that AdaSAP results in lower sharpness directly before pruning, directly after pruning, and after finetuning. We have added these results to Section B.5 and Table 11 in the Appendix.

Thank you for your thoughtful comments.

Questions

Q1: We think your first question may be answered by Table 14 in the updated Appendix. This table shows the loss and accuracy right before (i.e. post Adaptive Weight Perturbation Step / warmup) and right after pruning. Our AdaSAP method generally reaches the lowest loss and highest accuracy both right before and right after pruning. This indicates that the model is primed well for pruning, and that the pruning procedure minimally impacts the performance. Please let us know if we can clarify these results further.

(We also note that there was an omission in the initial version of this table which has since been updated - each grouped set of rows corresponds to a different pruning ratio. The first set of rows correspond to models with 20% of parameters preserved, or 80% pruned away; the second set is models with 43% of parameters preserved; and the third set is models with 76% of parameters preserved)

Q2: We thank you for pointing out the typo here. You are correct in your analysis of Equation 2 - terminology was accidentally left from an alternative characterization. Equation 2 should be $\rho_i = \rho_{max} - \frac{s_i - s_{min}}{s_{max} - s_{min}} (\rho_{max} - \rho_{min})$. We appreciate this feedback and have updated the equation (as well as the reference in Algorithm 1).

Weaknesses

W1: We would love to improve the presentation of our paper based on your recommendations. We read carefully through the paper and found one instance in Equation 1 in which epsilon was only implicitly defined; we will change the feasible set in this equation to explicitly define epsilon: the maximization will be over $\epsilon:(\|T_{\mathbf{w}}^{-1}\epsilon_i\| \leq \mathbf{\rho_i})_{i=1}^I$. However, we could not find any more examples of variables being used much before their description. We would appreciate if you could provide any examples where this occurs, or any other specific writing that lacks clarity. If you’re referring to the places where we define terms in Equations 1 and 2 after we state the equations, please let us know if defining those terms before the equations would improve readability.

W2: On $\rho_{min}$ and $\rho_{max}$, as described in Appendix A.5, we choose our maximum $\rho$ based on the recommended values for $\rho$ in line with SAM and ASAM (Foret et al. 2020; Kwon et al., 2021). In our earliest experiments, we ran simple experiments where we evaluated the validation performance after one or two epochs (of training; no pruning involved) for varying $\rho_{min}$ and $\rho_{max}$. We tried minimum values of 0.01 and 0.05 and maximum values of 0.05, 0.1, and 0.5 (this formulation only considered perturbation balls, not other transformed shapes, and so we considered values close to the recommended $\rho=0.05$ in the SAM paper). We found here that the best $\rho_{min}$ and $\rho_{max}$ values were 0.01 and 0.05, respectively. When we added in the perturbation region transformations from ASAM, we only updated our maximum $\rho$ value to 2.0, in line with the recommended $\rho$ hyperparameter in the ASAM paper. While we did not conduct an in-depth ablation of these hyperparameters, we found that extensive hyperparameter tuning was not needed to obtain our results and we simply chose consistent values that worked well from the beginning. We hope this can help clarify your question and have added this information to Appendix A.5 and Table 7 of the paper.

W3: With regards to your last point, we have updated the final paper to avoid text wrapping around the main tables. To the best of our knowledge, ICLR guidelines say nothing about text wrapping for figures and algorithms, but can further update the formatting as needed or to aid in clarity.

We hope our responses above can help clarify the points you made and we would appreciate if there is any additional information we can add for you to consider revising your rating.

Thank you for your thoughtful comments.

1. We would like to highlight that the adaptive nature of our approach is what makes it particularly suited to pruning. With the knowledge that some of the neurons will eventually be pruned, we regularize the neurons with the lowest importance (i.e. penalize sharpness) so that they have a smaller impact on the loss when they are pruned. Without this adaptivity, our method reduces to Sharpness-Aware Minimization (Foret et al. 2020) and has no particular compatibility with pruning algorithms.

2. We thank you for pointing out those works. Neither of these papers focus on pruned models, and while they do improve the robustness of dense models, the objective of our work is to specifically improve the robustness of pruned models. Our aim is to address the observed (see Liebenwein et al., 2021; Hooker et al., 2019) degradation of robustness during model pruning and reduce this degradation. While we already had cited “Understanding the robustness in vision transformers" in our paper, we have now also added "Robustifying token attention for vision transformers" and clarified that these works pursue robustness in dense models, unlike our work which focuses on pruned models. We have highlighted the difference between these two papers in the last sentence on page 1:

Despite significant progress in this area, most existing robustness work focuses on dense networks (Guo et al., 2023; Zhou et al., 2022).

3. We believe that this work is the first to introduce a method which has the primary aim of improving the OOD robustness of pruned models. Robust pruning methods (discussed in Section 3.4) are the closest to our approach, but differ since these methods requires train-time access to the types of data corruptions observed at test time (which are often adversarial attacks in this setting), which we do not want to assume (Sehwag et al., 2020; Vemparala et al., 2021; Zhao & Wressnegger, 2022). Some additional works (Section 2, paragraph “Robustness”) have observed, but not attempted to address, reduced OOD robustness in pruned models (Liebenwein et al., 2021; Hooker et al., 2019) or analyzed adversarial robustness (Sehwag et al., 2020; Fu et al., 2021). This gap in the literature is what motivated the formulation of our method. If you have other relevant comparison points that we have missed, we would be happy to provide a more detailed analysis and include the references in our paper.

4. We have some additional results on applying AdaSAP to other pruning criteria. In Table 5 in the paper (and an expanded set of results in Table 12 in the appendix) we use AdaSAP with the Taylor pruning criteria and show that it improves performance.

We agree that the object detection results are limited. We will attempt to run additional baselines before the end of the rebuttal period, but we may be unable to obtain results in time.

Thank you for your thoughtful comments. Please let us know if any of the following points are still unclear.

Questions

Addressing your question about our comparison points: we ran magnitude, Taylor pruning, and HALP following hyperparameter specifications from the corresponding papers in the latter two. For all other comparison methods, we took numbers from their papers. We have clarified this in the paper (see the quote from Section 4.1 "Baselines" below).

Weaknesses

1. We understand that the magnitude pruning results are unintuitive, however we believe that our experimental results are correct and consistent. The main point we want to emphasize is that AdaSAP, in combination with a particular pruning metric, should primarily be compared against a model trained with SGD (or another basic optimizer) and the same pruning metric. Therefore, magnitude is primarily intended as a comparison point with AdaSAP$_P$. We have attempted to clarify our results in the final paper in line with the following points. First, some of the discrepancy may come from variation in training procedures. This is unavoidable due to different setups across previous works. We used comparable training procedures when we ran both AdaSAP and magnitude pruning for our experiments, since AdaSAP utilizes L2-norm based magnitude pruning and we wanted to provide a direct comparison of our method to its closest baseline. Our goal in comparison to all other methods is primarily to demonstrate their lack of relative robustness out of the box. We would like to re-emphasize that many of these methods may be compatible with AdaSAP, which in turn could boost their performance. We demonstrate that AdaSAP + Taylor pruning outperforms Taylor pruning alone, and that AdaSAP + HALP outperforms HALP alone. Second, magnitude pruning is a strong baseline and results are often relatively close to SOTA methods. Therefore it should be unsurprising that small variations in training setups may lead to the type of results we observe. Finally, as a minor point, the comparisons to magnitude pruning in the original papers were slightly different: Greg compares to a L1-based magnitude method, whereas we use L2 and ABCPruner does not directly compare to magnitude pruning.
   We thank you for pointing out this discrepancy and we have clarified in the paper that magnitude pruning was intended primarily as a direct comparison to AdaSAP, since we did not rerun many other comparison methods from scratch.

The updated text in the paper in Section 4.1 under "Baselines" now reads as follows:

We primarily focus on comparing to two SOTA baselines which we re-run in our experiments, Taylor pruning for parameter-based pruning (Molchanov et al., 2019), and HALP for latency-based pruning (Shen et al., 2022a). Taylor pruning assigns a score to each neuron based on a Taylor series approximation of how much the loss would change if the neuron were removed. HALP is a pruning method that optimizes for improving latency. We additionally run results on magnitude pruning (Han et al., 2015b), to demonstrate the improvements of AdaSAP$_P$ over its closest baseline. We also compare to the results cited in several other pruning methods such as GR-eg (Wang et al., 2020b), EagleEye (Li et al., 2020), ABCPruner (Lin et al., 2020), and SMCP (Humble et al., 2022).

2. With regards to the misspellings, typos, and formatting, we appreciate your careful read and have fixed them in the final paper.
3. We agree that GPU hour statistics are beneficial to add. We were able to obtain some of these results and have augmented Section 4.6 with the following information:

One limitation is that AdaSAP requires twice the training time due to the two backward passes necessary for optimization. A 90 epoch pruning and finetuning run on 8 V100 GPUs to reduce a ResNet50 network to 0.2× its size takes takes 32 hours and 12 minutes for AdaSAP$_P$ while SGD with magnitude pruning takes 16 hours and 53 minutes. Recent strategies (Du et al., 2022; Liu et al., 2022) could be used to reduce this additional overhead.