Post-Training Neural Network Pruning using Graph Curvature

Shuhang Tan, Jayson Sia, Paul Bogdan, Radoslav Ivanov

Abstract:

This paper provides a fresh view of the neural network (NN) pruning problem, through the lens of graph theory. To achieve effective pruning, we aim to identify the main NN data flows and the corresponding NN connections that are most (and least) important for the performance of the full model. Unlike the standard approach to NN data flow analysis, which is based on information theory, we employ the notion of graph curvature, specifically Ollivier-Ricci curvature (ORC). The ORC has been successfully used to identify important graph edges in various domains such as road traffic analysis, biological and social networks. In particular, edges with negative ORC are considered bottlenecks and as such are critical to the graph’s overall connectivity, whereas positive-ORC edges are not essential. We use this intuition for the case of NNs to: 1) construct a graph induced by the NN structure and introduce the notion of neural curvature (NC) based on the ORC; 2) calculate curvatures based on activation patterns for a set of input examples; 3) demonstrate that NC can
indeed be used to rank edges according to their importance for the overall NN functionality. We evaluate our method through pruning experiments on a variety of small-to-medium-size models trained on three image datasets, namely MNIST, CIFAR-10 and CIFAR-100. The results indicate that our method can identify a larger number of unimportant edges as compared to existing pruning methods. (Code: https://github.com/SH-Tan/Post-Training-NN-Pruning-using-Graph-Curvature)

Submission Type: Long submission (more than 12 pages of main content)

Changes Since Last Submission:

1. Adding the code link at the end of the abstract.
2. Claim that this work focuses on small- to medium-sized models in the abstract.
3. Adding discussion about the expensive computation and potential solutions in the second-to-last paragraph in the introduction.
4. Adding a new bullet point in the experiment summary in Section 6.1 about exploring how to mitigate the expensive computation.
5. Adding Acknowledgement section.

Code: https://github.com/SH-Tan/Post-Training-NN-Pruning-using-Graph-Curvature

Assigned Action Editor: ~Nadav_Cohen1

Submission Number: 6972