Paper: Adaptive Sharpness-Aware Pruning for Robust Sparse Networks (ICLR-2024)

Reviewer comment (Novelty):
As mentioned by the author, SAM as a post-pruning training can help pruning has already been studied by several works, this work seems to be a modification or extension of using sharpness-aware neuron-level perturbation to enhance the pruning. So there exists a lack of novelty and clarification on the difference between this method and the vanilla SAM method, as well as the motivation for proposing AdaSAP optimization to surpass common SAM.

Author response:
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.

Reviewer comment (Reproducibility):
There is no discussion of the hyperparameter tuning used in the method since it is especially critical for SAM-type methods.

Author response:
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).

| $\rho_{min}$ | $\rho_{max}$ | Epoch 2 Val Loss | Epoch 2 Val Acc |
| ---|---|---|---|
| $0.01$ | $0.05$ | $5.54$ | $8.63$ |
| $0.01$ | $0.1$ | $5.59$ | $7.11$ |
| $0.05$ | $0.1$ | $5.73$ | $5.69$ |
| $0.05$ | $0.5$ | $6.92$ | $0.1$ |

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).

Reviewer comment (Evaluation):
In Section 4.5 and Table 4, will the results be consistent if we use some other way to measure the sharpness like Hessian Top-eigenvalue or trace?

Author response:
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.

| Optimizer | Pre-pruning | Post-pruning | Post-finetuning |
|---|---|---|---|
| SGD | $35.91$ | $93.82$ | $8.17$ |
| AdaSAP$_L$ | $25.23$ | $16.79$ | $7.66$ |

Reviewer comment (Writing):
Writing needs to be improved especially within section 3.1 where mathematical preliminaries need to be set well. Also, variables are being used much before their description/introduction.

Author response:
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.

Reviewer comment (Experiments):
Lack of ablation with respect to $\rho_{min}, \rho_{max}$.

Author response:
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.

| $\rho_{min}$ | $\rho_{max}$ | Epoch 2 Val Loss | Epoch 2 Val Acc |
| ---|---|---|---|
| $0.01$ | $0.05$ | $5.54$ | $8.63$ |
| $0.01$ | $0.1$ | $5.59$ | $7.11$ |
| $0.05$ | $0.1$ | $5.73$ | $5.69$ |
| $0.05$ | $0.5$ | $6.92$ | $0.1$ |

Reviewer comment (Presentation):
Lack of adherence to ICLR template. Text is wrapped around images, tables, and algorithms.

Author response:
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.

Reviewer comment (Evaluation):
In each of the cases of AdaSAP pruning procedure, can you please report the accuracy numbers for the models post Adaptive Weight Perturbation Step (warmup)? It would be interesting to note the difference in numbers pre and post pruning.

Author response:
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)

Reviewer comment (Theory):
When computing $\alpha_i$ (equation 2), did you intend to write $\rho_{max}$ instead on $s_{max}$ in the first term? The expression looks odd otherwise. If yes, can you explain the choice of using $\rho_i = - \alpha_i s_i$ ? It is unclear why would you require a negative sign. If no, can you please explain the choice behind defining $\alpha_i$ in such a way?

Author response:
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).

Reviewer comment (Novelty):
Although the idea to improve both sparsity and robustness sounds good, the method to improve robustness via weight perturbation is indeed not novel. Is there any special designs closely related to pruning?

Author response:
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.

Reviewer comment (Novelty):
While the proposed AdaSAP method demonstrates notable effectiveness in concurrently optimizing sparsity and robustness, it is better for the authors to include a discussion of other relevant works that aim to achieve such dual objectives. This would provide readers with a better understanding of where AdaSAP stands in this field and allow for comparisons with existing literature.

Author response:
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.

Reviewer comment (Experiments):
In the experiments on image classification, evaluating AdaSAP on additional pruning criterias would provide a more comprehensive assessment the effectiveness and generality of AdaSAP.

Author response:
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.

Reviewer comment (Experiments):
The experiments and analysis on object detection are inadequate with only one model. More detection models should be evaluated. Furthermore, comparisons with other pruning methods are needed to further verify efficacy.

Author response:
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.

Reviewer comment (Writing):
Typos and grammar mistakes.  
- Table 1 caption: various -> various  
- HALP (NeurIPS’23 (Shen et al., 2022b)) -> NeurIPS'22  
- Section 4.5, first sentence misses a period.  
These small glitches, not very serious though, make me feel this paper is sort of rushed out.

Author response:
2. With regards to the misspellings, typos, and formatting, we appreciate your careful read and have fixed them in the final paper.

Reviewer comment (Experiments):
In Table 1, for the results of comparison methods, are they from the authors' rerun with the same base unpruned model, or cited from the original papers?

Author response:
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).