C-SWAP: Explainability-Aware Structured Pruning for Efficient Neural Networks Compression

TL;DR

This paper tackles the challenge of efficient neural network compression by enabling high-sparsity structured pruning without retraining. It proposes C-SWAP, a causal-aware, one-shot structured pruning algorithm that uses a global causal effect $\xi_n$ and class-wise predicates $\pi_n^{(k)}$ to categorize neurons as Neutral, Critical, or Detrimental and prune non-critical units progressively. Across CNNs, vision transformers, and a semantic segmentation task, C-SWAP demonstrates superior pruning curves and SAUCE performance, achieving substantial reductions in parameters and FLOPs with minimal accuracy loss and without fine-tuning. The work highlights the role of explainable AI in guiding compression and shows scalability to complex architectures, while noting computational costs for per-unit causality and the potential for group-wise extensions in future work.

Abstract

Neural network compression has gained increasing attention in recent years, particularly in computer vision applications, where the need for model reduction is crucial for overcoming deployment constraints. Pruning is a widely used technique that prompts sparsity in model structures, e.g. weights, neurons, and layers, reducing size and inference costs. Structured pruning is especially important as it allows for the removal of entire structures, which further accelerates inference time and reduces memory overhead. However, it can be computationally expensive, requiring iterative retraining and optimization. To overcome this problem, recent methods considered one-shot setting, which applies pruning directly at post-training. Unfortunately, they often lead to a considerable drop in performance. In this paper, we focus on this issue by proposing a novel one-shot pruning framework that relies on explainable deep learning. First, we introduce a causal-aware pruning approach that leverages cause-effect relations between model predictions and structures in a progressive pruning process. It allows us to efficiently reduce the size of the network, ensuring that the removed structures do not deter the performance of the model. Then, through experiments conducted on convolution neural network and vision transformer baselines, pre-trained on classification tasks, we demonstrate that our method consistently achieves substantial reductions in model size, with minimal impact on performance, and without the need for fine-tuning. Overall, our approach outperforms its counterparts, offering the best trade-off. Our code is available on GitHub.

Figures (12)

• Figure 1: C-SWAP, an explainability-aware structured pruning framework for DNNs. Unlike conventional pipelines, C-SWAP uses a causal-based, one-shot progressive pruning strategy that removes neurons while preserving model performance without costly re-training. Its progressive approach scales effectively to complex architectures such as transformers.
• Figure 2: Neurons and paths for a toy MLP. Critical path conveys the relevant information. Detrimental and Neutral neurons are not essential.
• Figure 3: Simple MLP with a residual connection (dotted). Structured pruning forces to remove all the weights in red to remove the crossed neuron.
• Figure 4: Summary of models and datasets used in our experiments.
• Figure 5: Validation accuracy as a function of percentage of parameters removed for four architectures. Dashed line represents original model performance.

Theorems & Definitions (2)

• Definition 1: Global Causal Effect
• Definition 2: Neutral, Critical and Detrimental neurons