Preserving Deep Representations in One-Shot Pruning: A Hessian-Free Second-Order Optimization Framework

Authors: Ryan Lucas, Rahul Mazumder
Primary area: optimization
Keywords: Neural Network Pruning, Structured Pruning, Optimization, Hessian-free Optimization

Abstract:
We present SNOWS, a one-shot post-training pruning framework aimed at reducing the cost of vision network inference without retraining. Current leading one-shot pruning methods minimize layer-wise least squares reconstruction error which does not take into account deeper network representations. We propose to optimize a more global reconstruction objective. This objective accounts for nonlinear activations deep in the network to obtain a better proxy for the network loss. This nonlinear objective leads to a more challenging optimization problem---we demonstrate it can be solved efficiently using a specialized second-order optimization framework. A key innovation of our framework is the use of Hessian-free optimization to compute exact Newton descent steps without needing to compute or store the full Hessian matrix. A distinct advantage of SNOWS is that it can be readily applied on top of any sparse mask derived from prior methods, readjusting their weights to preserve deep feature representations. SNOWS obtains state-of-the-art results on various one-shot pruning benchmarks including residual networks and Vision Transformers (ViT/B-16 and ViT/L-16, 86m and 304m parameters respectively). Our open-source implementation is available at https://github.com/mazumder-lab/SNOWS.

Decision: accept (poster)
Reviewer scores: 8, 5, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1730043475890

thread_type:
official_review

<thread_object>
creator: reviewer_yaxu

<content>
summary:
This paper proposes SNOWS, a one-shot pruning framework aimed at reducing vision network inference costs without retraining. Unlike traditional approaches that prioritize layer-wise error minimization, SNOWS optimizes a global objective that captures deeper, nonlinear network representations. To efficiently solve the resulting complex optimization problem, SNOWS leverages a Hessian-free second-order method that computes Newton descent steps without needing the full Hessian matrix. SNOWS also enhances any pre-existing sparse mask by refining weights to better leverage nonlinearities in deep features. The framework demonstrates state-of-the-art results across multiple one-shot pruning benchmarks, including residual networks and Vision Transformers. Results are shown for CIFAR10/100 and Mini-Imagenet.

strengths:
**===After rebuttal, I raised score from 5 to 6=====**
+ A method with a new pruning strategy at the level of global reconstruction shows promising results on common datasets CIFAR and Mini-Imagenet.
+ The proposed framework is reportedly efficient where it does not require storing the full Hessian matrix, allowing it to perform running with large neural networks.
+ Providing more details of pruning with ViT models gives some new insights.

weaknesses:
**Concerns**
+ The baselines seem to be out-of-date, it would be helpful to specify additional recent baselines from 2024. More baselines should be considered for example [1,2].

+ Considering specific aspects of the method that could be condensed or moved to an appendix, and particular experimental analyses or comparisons should be expanded:  Accuracy and latency speed-up with different network architectures: Res50, MobileNet, ViT.

+ Section 5, "Experimental setup" is written too long, it contains also the results and a description of results that is not "setup". It should be separated into a new paragraph or subsection. More breaks between paragraphs should be added for better readability.

+ Considering particular efficiency metrics the authors could report (e.g., runtime, memory usage) and specific methods (in Table 3) they should compare against to demonstrate their efficiency claims.

**Other suggestions**
+ Mentioning that experiments on ImageNet-1k (lines 102, 444)  might be misleading since the fact that the authors only use the Mini-ImageNet with 10,000 samples compared to ImageNet-1k (1.2M samples). Replacing ImageNet-1k with MiniImagenet-1k is more appropriate.

+ There are two "Experimental setup" parts which are the same level as the network part (CNN running, Vision transformers, ...), which might cause confusion, suggesting making it a subsection for each type of network.

+ The so-called "Figure 3" (Line 472-484) should be a "Table 3". The text referred to that table should be revised accordingly.

+ For better readability, I suggest rather "minimize (7)" --> "minimize Eq. 7", Taylor expansion in (7) --> "Taylor expansion in Eq. 7", each call to (5) --> "each call to Eq. 5", etc.

+ Equations are written in a separate line with a center and should be enumerated, for example for line 214, line 382, and 384.

+ Although it is lookable, Figure 4 should make the text inside the figure larger for better readability.

[1] FALCON: FLOP-Aware Combinatorial Optimization for Neural Network Pruning, AISTATS 2024

[2] LayerMerge: Neural Network Depth Compression through Layer Pruning and Merging, ICML 2024

questions:
Suggest that the authors define the MP metric/method when it is first used and provide a citation if it refers to a specific published method.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730533615708

thread_type:
official_review

<thread_object>
creator: reviewer_etde

<content>
summary:
(+) The authors propose optimizing a more global reconstruction objective referred to as SNOW, which accounts for non-linear activations deep in the network to obtain a better proxy for the network loss.

(+) This nonlinear objective leads to a more challenging optimization problem - the authors demonstrate that it can be solved efficiently using Hessian-free optimization.

(+) SNOWS demonstrates its effectiveness with state-of-the-art results on various one-shot pruning benchmarks, including  CNNs and Vision Transformers (ViTs).

strengths:
(+) SNOW can be seamlessly applied to any sparse mask from prior methods, refining weights to leverage the nonlinearities in the deep feature representations of CNNs and ViTs.

(+) The updated rule (Newton) and the illustrative examples in the Supplementary section are well articulated.

weaknesses:
(-) The motivations for the improved layer-wise model (Eq. 3) need clarification. Why are the subsequent operations $f^{l:l+K}$ necessary-are they intended to minimize residuals? Please clarify the motivation behind SNOW by explaining the concept of Fig. 2 and including additional captions.

(-) While SNOW’s performance with increasing K iterations is well documented, computational comparisons with baselines are lacking. For consistent comparisons, all experimental results should explicitly include the number of K-iterations in captions of tables and figures. Moreover, I recommend the authors add an ablation study on the time complexity of K-iterations to the experimental result tables.

questions:
Please see the above weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730676857046

thread_type:
official_review

<thread_object>
creator: reviewer_96mw

<content>
summary:
This paper proposes a novel pruning method called SNOWS (Stochastic Newton Optimal Weight Surgeon), which is an adaptation of pruning methods like the Optimal Brain Surgeon method, that uses the second-order information of the weights to do the pruning.
However, SNOWS does not calculate the full hessian but merely the hessian products, and those too are not all calculated (approximated) at once and the optimization is performed using stochastic newton descent.

strengths:
The following are the strengths of the paper:
1. The idea contributed by this paper seems to be novel and is very interesting. 

2. The theoretical aspect of the paper seems sound.

3. The proposed method seems universally applicable to convolution and ViT-based models. 

4. The writing of the paper is quite clear.

weaknesses:
The paper's introduction makes some glaring errors in defining concepts. For example, it describes one-shot pruning methods as pruning methods that do not require retraining after pruning. However, this is not accurate. Pruning can be done iteratively or one-shot; however, some finetuning might be required after one-shot pruning, as shown by [1], [2], and many other methods covered by [3]. 

With the above definition for one-shot in mind, [4] achieves unstructured pruning with more than 90% sparsity while not dropping any clean performance for multiple architectures. The proposed SNOWS method should also be evaluated against [4] for unstructured pruning performance. 


Lastly, the major motivation of the proposed work is gain in speed when performing pruning (as the entire hessian need not be calculated), however these no exists latency comparison to other pruning methods. 

Improving on these weaknesses might require a major rewrite, and thus the recommendation.

**References**

[1] Sanh, Victor, Thomas Wolf, and Alexander Rush. "Movement pruning: Adaptive sparsity by fine-tuning." Advances in neural information processing systems 33 (2020): 20378-20389.

[2] Sehwag, Vikash, et al. "Hydra: Pruning adversarially robust neural networks." Advances in Neural Information Processing Systems 33 (2020): 19655-19666.

[3] Hoefler, Torsten, et al. "Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks." Journal of Machine Learning Research 22.241 (2021): 1-124.

[4] Hoffmann, Jasper, et al. "Towards improving robustness of compressed cnns." ICML Workshop on Uncertainty and Robustness in Deep Learning (UDL). Vol. 2. 2021.

questions:
Q1- Since this method is able to achieve semi-structured sparsity, how does its latency compare to unstructured pruning methods? A comprehensive latency study would be very helpful.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730722701447

thread_type:
official_review

<thread_object>
creator: reviewer_xymy

<content>
summary:
The following paper presents SNOWS, a one-shot post-training pruning technique that performs pruning at a layer-wise level while globally optimizing joint layers within a deep neural network to include nonlinear activation functions within. While it resulted in a more challenging optimization problem due to nonlinearities, it is possible to use the Hessian-free approach integrated with a customized Conjugate Gradient (CG) method that efficiently solves second-order approximation during gradient computation while simultaneously allowing reduced memory requirements. Experiments on various Computer Vision tasks under both CNN-based architecture and Vision transformer highlighted the effectiveness of SNOWS.

strengths:
1. Besides achieving performance improvement compared to other one-shot pruning methods, the proposed method also improves the classification performance in image classification tasks when built on top of existing sparse mask selection algorithms for N:M pruning, denoting its practicability for accelerated training.
2. Highly appreciate the fact that the source code is available for inspection.

weaknesses:
1. There are no conclusions and limitations section in the paper, and the writing in the methods section can be more concise in Section 4, which mainly focuses on explaining the SNOW method. This could leave more space to explain the experiment section in the main paper properly along with the results. Overall, the organization of this paper needs a lot of improvement and it feels like the paper is not ready for publication in the current form. 
2. Typos:
- "Like retraining approaches" -> "Like retraining, approaches" (Section 1)
- "network parammeters $\mathbf{W} = (\mathbf{W}^1, ..., \mathbf{W}^1)$ should have been $\mathbf{W} = (\mathbf{W}^1, ..., \mathbf{W}^L)$ (Section 3)

questions:
See the weakness section.

===EDIT: Rating is raised from 3 to 5=====
</content>
</thread_object>

</thread>
<thread>
creation date:
1730722756610

thread_type:
official_review

<thread_object>
creator: reviewer_tgxm

<content>
summary:
The paper argues it is important to consider the effect of pruning of the whole network instead of just layerwise when performing one-shot post-training pruning. They use a Hessian-free second-order method to optimize their formulation. The results show significant improvement in image-classification tasks on resnets and vits.

strengths:
1. The formulation is neat, and the use of tools from optimization literature (Hessian-free conjugate gradients) is elegant.
2. The results show significant improvements over OBC and pruning using the layerwise objective.
3. Writing is clear and the main claims are supported by the experiments.

weaknesses:
1. The emphasis on using true Hessian information is not supported in the experiments. This is not a crucial point but in most scenarios in deep learning, Fisher matrix-based approximation of Hessian [a] (or even Ada-grad [b] or Adam) works well. It would be worth providing that baseline to better support the proposed method.

- [a] Martens, James. "New insights and perspectives on the natural gradient method." Journal of Machine Learning Research 21.146 (2020): 1-76. 
- [b] Duchi, John, Elad Hazan, and Yoram Singer. "Adaptive subgradient methods for online learning and stochastic optimization." Journal of machine learning research 12.7 (2011).

questions:
1. There are a few cases where other pruning methods fail (very low accuracy, Fig. 3, Tab. 1), could you comment on the potential reasons?
2. Does retraining improves the performance further after pruning using the proposed method?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732310906770

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Revision Changes

comment:
We would like to thank the reviewers for their detailed feedback and comments. We have provided a revised manuscript incorporating your feedback. We detail the changes below.

**Formatting and organization**

We thank the reviewers for their helpful formatting and structuring suggestions. We have made your formatting changes such as equation references and figure legends, as well as fixing typos. Moreover, we have significantly shortened Section 4 on the SNOWS method, and in its place we have added references to additional experiments (such as run-time and retraining experiments) in Section 5, and added a new Section 6 in the main text on the limitations and conclusions of our work. Moreover, due to the additional content, we have improved the structure of the supplementary section, dividing it into "Primary Methods and Analysis", "Ablation Studies", and "Implementation and Experimental Details".

**Comparison to Fisher approximation**

We have provided new experiments demonstrating the utility of incorporating the true Hessian information versus the popular Fisher approximation. We observe that the second-order descent with the Fisher approximation can take a long time to converge. For large networks, this becomes prohibitive, whereas our Hessian-free method typically can be safely terminated after a few seconds per layer. This is discussed in Section A.2.2.

**Analysis on the SNOWS loss function**

We have provided some additional analysis on the merit of the SNOWS loss function in Section A.1.2, comparing to the regular layer-wise approach, at the request of Reviewer EtDe. To summarize, SNOWS better preserves the input to future layers, during the pruning of a given layer. We thank the reviewer for prompting this further analysis with their question.

**Effect of the parameter $K$ on run-time** 

Several reviewers asked for ablations showing the effect of the parameter $K$ on the run-time and memory consumption of our algorithm. This is discussed in the new Section A.2.4. As we mentioned in the main text, the parameter $K$ induces a trade-off between run-time and performance, with larger values of $K$ performing better but being slower to run. In Section A.3.7, we provide some additional run-time experiments comparing to CHITA [1], showing that for moderate-sized $K$, SNOWS performs better and scales much better in terms of architecture size. We reiterate that our method avoids computing, approximating or storing the full network Hessian, and even avoids storing layer-wise Hessians, which makes its complexity much more manageable with respect to architecture size. 

**Updating benchmarking**

We have added FALCON, a new benchmark from 2024 AISTATS, to our unstructured sparsity benchmarks (Figure 3). We decided not to include the experiments from [3] in our final manuscript, though we addressed this concern with reviewer 96mw, and we hope our response to the reviewer clarifies this issue. Similarly, to Reviewer YAxu, while we see LayerMerge [4] as a good structured pruning technique, it should not be considered a one-shot pruning method, since the method has close to 10\% accuracy following pruning (cf. Figure 3 of their paper), and requires extensive retraining to recover performance (they use 180 epochs of retraining on ImageNet-1k).

**Retraining experiments**

At the request of several reviewers, we provide experiments in Section A.3.4 showing SNOWS can benefit from even limited retraining (e.g. 5 epochs on ImageNet-1k). We compare to retraining performance of FALCON [2].

**Latency analysis**

A challenge in our evaluation has been the measurement of inference time speed-ups for N:M sparsity patterns. Unfortunately, there is limited native support in PyTorch for such measurements. Current hardware acceleration is particularly constrained, with only the 2:4 sparsity pattern having support through NVIDIA's sparse tensor cores. To address these limitations, we have report comprehensive end-to-end FLOPs reduction across several of our tested architectures and sparsity patterns in Section A.3.6. We tested this end-to-end to ensure our measurements account for all network operations and layers. 

[1] Benbaki, R., Chen, W., Meng, X., Hazimeh, H., Ponomareva, N., Zhao, Z., & Mazumder, R. (2023). Fast as CHITA: Neural Network Pruning with Combinatorial Optimization. In Proceedings of the 40th International Conference on Machine Learning.

[2] Meng, X., Chen, W., Benbaki, R., & Mazumder, R. (2024). FALCON: FLOP-Aware Combinatorial Optimization for Neural Network Pruning. In Proceedings of The 27th International Conference on Artificial Intelligence and Statistics (AISTATS).

[3] Hoffmann, J., Agnihotri, S., Saikia, T., & Brox, T. (2021). Towards improving robustness of compressed CNNs. In Proceedings of the ICML Workshop on Uncertainty and Robustness in Deep Learning.

[4] LayerMerge: Neural Network Depth Compression through Layer Pruning and Merging, International Conference on Machine Learning, 2024.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a one-shot post-training pruning method, that involves layer-wise pruning within a global  optimization. It is an adaptation of previous second order methods which, however, avoids the calculation of the full Hessian by using Hessian-free conjugate gradients.
The proposed method, SNOWS, is interesting and presented within a theoretically sound framework and provides good results that support the theoretical claims. While there are slight remaining concerns on the organization of the paper, the reviewers are overall positive, stressing the theoretic contribution of the work.