Adaptive MLP Pruning for Large Vision Transformers

TL;DR

This paper introduces an adaptive pruning method for large vision transformers that significantly reduces parameters and computational costs while maintaining performance, by evaluating neuron importance with a novel label-free criterion.

Contribution

The proposed AMP method adaptively prunes MLP modules in vision transformers using a new importance evaluation, outperforming existing pruning techniques without needing predefined compression ratios.

Findings

Achieves roughly 40% reduction in parameters and FLOPs.

Maintains near-original performance after pruning.

Outperforms other pruning methods without fine-tuning.

Abstract

Large vision transformers present impressive scalability, as their performance can be well improved with increased model capacity. Nevertheless, their cumbersome parameters results in exorbitant computational and memory demands. By analyzing prevalent transformer structures, we find that multilayer perceptron (MLP) modules constitute the largest share of the model's parameters. In this paper, we propose an Adaptive MLP Pruning (AMP) method to substantially reduce the parameters of large vision transformers without obvious performance degradation. First, we adopt Taylor based method to evaluate neuron importance of MLP. However, the importance computation using one-hot cross entropy loss ignores the potential predictions on other categories, thus degrading the quality of the evaluated importance scores. To address this issue, we introduce label-free information entropy criterion to fully model the predictions of the original model for more accurate importance evaluation. Second, we rank the hidden neurons of MLP by the above importance scores and apply binary search algorithm to adaptively prune the ranked neurons according to the redundancy of different MLP modules, thereby avoiding the predefined compression ratio. Experimental results on several state-of-the-art large vision transformers, including CLIP and DINOv2, demonstrate that our method achieves roughly 40\% parameter and FLOPs reduction in a near lossless manner. Moreover, when the models are not finetuned after pruning, our method outperforms other pruning methods by significantly large margin. The source code and trained weights are available at https://github.com/visresearch/AMP.