MPruner: Optimizing Neural Network Size with CKA-Based Mutual Information Pruning

TL;DR

A new pruning algorithm is developed that leverages mutual information through vector similarity, allowing it to incorporate global information from the neural network for more precise and efficient layer-wise pruning.

Abstract

Determining the optimal size of a neural network is critical, as it directly impacts runtime performance and memory usage. Pruning is a well-established model compression technique that reduces the size of neural networks while mathematically guaranteeing accuracy preservation. However, many recent pruning methods overlook the global contributions of individual model components, making it difficult to ensure that a pruned model meets the desired dataset and performance requirements. To address these challenges, we developed a new pruning algorithm, MPruner, that leverages mutual information through vector similarity. MPruner utilizes layer clustering with the Centered Kernel Alignment (CKA) similarity metric, allowing us to incorporate global information from the neural network for more precise and efficient layer-wise pruning. We evaluated MPruner across various architectures and configurations, demonstrating its versatility and providing practical guidelines. MPruner achieved up to a 50% reduction in parameters and memory usage for CNN and transformer-based models, with minimal to no loss in accuracy.

Figures (5)

• Figure 1: Overview of MPruner
• Figure 2: Three Phases of MPruner: (1) Analysis Stage with CKA Similarity Score, (2) Optimization Stage with Multi-Layer Cluster Pruning, and (3) Recovery Stage with Retraining.
• Figure 3: Clustering Heatmaps for Bert.
• Figure 4: Result of Wanda Sun24, and Integration of Wanda and MPruner.
• Figure 5: Clustering Heatmaps for CodeT5.