TLDR: C-SWAP is a novel one-shot structured pruning framework that uses explainable deep learning to compress neural networks. It identifies and removes non-critical neurons based on their causal effect on model predictions, enabling significant model size reduction with minimal performance loss and no need for fine-tuning. The method outperforms existing baselines on various CNN and Vision Transformer architectures for classification and extends to semantic segmentation.
Deep neural networks, while incredibly powerful in computer vision, often come with a significant drawback: their large size and computational demands. This makes them challenging to deploy on devices with limited resources, such as those found at the ‘edge’ of a network. To address this, researchers have developed various compression techniques, with pruning being a prominent one. Pruning aims to reduce the size and complexity of a neural network by removing redundant parts, like weights, neurons, or even entire layers.
Structured pruning is particularly effective because it removes whole computational units, leading to faster inference times and reduced memory overhead. However, traditional structured pruning methods can be computationally expensive, often requiring iterative retraining and optimization. While one-shot pruning methods, which apply pruning directly after initial training, offer a more efficient alternative, they frequently suffer from a considerable drop in model performance.
A new framework called C-SWAP, which stands for Explainability-Aware Structured Pruning, has been introduced to tackle this performance trade-off. Developed by Baptiste Bauvin, Loïc Baret, and Ola Ahmad from Thales’ cortAIx Lab, C-SWAP offers a novel one-shot pruning approach that leverages explainable deep learning to efficiently compress neural networks without sacrificing performance.
The core of C-SWAP lies in its causal-aware pruning strategy. It identifies cause-effect relationships between a model’s predictions and its internal structures. This allows the framework to progressively prune the network, ensuring that only non-essential structures are removed. C-SWAP categorizes each neuron (or channel in convolutional neural networks) into one of three types: critical, neutral, or detrimental. Critical neurons are those that significantly enhance predictions, detrimental ones degrade performance, and neutral ones have an insignificant impact. By preserving critical neurons and removing neutral or detrimental ones, C-SWAP maintains model performance.
Unlike conventional methods that might rank neurons globally before pruning, C-SWAP integrates its causal analysis directly into a progressive pruning process. This means it can efficiently iterate through neurons, systematically removing irrelevant ones as it goes, without the need for computationally intensive re-ranking steps. This approach makes C-SWAP as computationally efficient as traditional ranking-based explainable AI pruning methods, but with superior results.
Experiments conducted on various neural network architectures, including Convolutional Neural Networks (ResNet-18, ResNet-50, MobileNetV2) and Vision Transformers (ViT), pre-trained on classification tasks (CIFAR10, ImageNet), demonstrated C-SWAP’s effectiveness. The method consistently achieved substantial reductions in model size with minimal impact on performance, all without requiring fine-tuning after pruning. It outperformed other baseline pruning techniques, offering a better balance between compression and accuracy.
To quantify the effectiveness of different pruning criteria, the researchers also introduced a new metric called Sparsity AUC Estimator (SAUCE), which measures the area under the accuracy curve as pruning ratios increase. C-SWAP consistently showed superior SAUCE values across all tested models and datasets.
Beyond classification, C-SWAP also proved its versatility by successfully adapting to semantic segmentation tasks, using mean Intersection-over-Union (IoU) as its scoring function. This highlights its potential for broader application in complex computer vision problems.
Also Read:
- Optimizing Vision-Language Models: A New Approach to Visual Token Pruning
- Neuronal Group Communication: A New Framework for Efficient and Interpretable AI Models
While C-SWAP marks a significant advancement in efficient structured pruning, the authors acknowledge certain limitations. Computing per-unit causal relevance can be computationally demanding for very wide layers. Future research directions include exploring group-wise relevance scores and extending the method to other complex tasks like object detection. For more details, you can read the full research paper here.
