[Paper Review] Visual Transformer Pruning.

Mingjian Zhu, Kai Han+2 more|arXiv (Cornell University)|Apr 17, 2021

Image Enhancement TechniquesComputer Science27 references23 citations

TL;DR

This paper proposes a channel pruning method for Vision Transformers that uses sparsity regularization to identify and remove unimportant channels, enabling high compression with minimal accuracy drop. The approach—training with regularization, pruning low-impact channels, and fine-tuning—achieves significant reductions in parameters and FLOPs on ImageNet with strong accuracy retention.

ABSTRACT

Visual transformer has achieved competitive performance on a variety of computer vision applications. However, their storage, run-time memory, and computational demands are hindering the deployment on mobile devices. Here we present an visual transformer pruning approach, which identifies the impacts of channels in each layer and then executes pruning accordingly. By encouraging channel-wise sparsity in the Transformer, important channels automatically emerge. A great number of channels with small coefficients can be discarded to achieve a high pruning ratio without significantly compromising accuracy. The pipeline for visual transformer pruning is as follows: 1) training with sparsity regularization; 2) pruning channels; 3) finetuning. The reduced parameters and FLOPs ratios of the proposed algorithm are well evaluated and analyzed on ImageNet dataset to demonstrate its effectiveness.

Motivation & Objective

• To address the high computational and memory demands of Vision Transformers that hinder mobile deployment.
• To identify and remove redundant channels in Vision Transformers without compromising performance.
• To develop a pruning pipeline that maintains high accuracy while achieving high compression ratios.
• To evaluate the effectiveness of the method on the ImageNet dataset.

Proposed method

• Apply sparsity regularization during training to encourage channel-wise sparsity in the Vision Transformer.
• Identify and remove channels with small coefficients after training, based on their importance scores.
• Perform structured pruning by eliminating entire channels rather than individual weights.
• Fine-tune the pruned model to recover accuracy lost during pruning.
• Use a three-stage pipeline: (1) sparse training, (2) channel pruning, (3) fine-tuning.
• Evaluate the final model's performance in terms of parameter reduction and FLOPs savings on ImageNet.

Experimental results

Research questions

• RQ1 Can sparsity regularization effectively identify unimportant channels in Vision Transformers for pruning?
• RQ2 What is the maximum pruning ratio achievable while maintaining competitive accuracy on ImageNet?
• RQ3 How does the three-stage pruning pipeline (train-prune-fine-tune) compare to end-to-end training in terms of efficiency and accuracy?
• RQ4 To what extent can parameter and FLOPs counts be reduced without significant performance degradation?

Key findings

• The proposed pruning method achieves significant reductions in model parameters and FLOPs on the ImageNet dataset.
• High pruning ratios are attainable with minimal drop in accuracy due to the identification of important channels via sparsity regularization.
• The three-stage pipeline—sparse training, pruning, and fine-tuning—effectively preserves model performance after compression.
• The method enables efficient deployment of Vision Transformers on resource-constrained mobile devices by reducing computational and memory demands.