[Paper Review] Vision Transformer Pruning

Mingjian Zhu, Yehui Tang+1 more|arXiv (Cornell University)|Apr 17, 2021

Advanced Neural Network ApplicationsComputer Science44 references53 citations

TL;DR

The paper introduces Vision Transformer Pruning (VTP), a method that learns dimension-wise importance scores with L1 sparsity to prune MHSA and MLP projections in ViT models, achieving substantial parameter and FLOPs reductions with minimal accuracy loss.

ABSTRACT

Vision 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 to mobile devices. Here we present a vision transformer pruning approach, which identifies the impacts of dimensions in each layer of transformer and then executes pruning accordingly. By encouraging dimension-wise sparsity in the transformer, important dimensions automatically emerge. A great number of dimensions with small importance scores can be discarded to achieve a high pruning ratio without significantly compromising accuracy. The pipeline for vision transformer pruning is as follows: 1) training with sparsity regularization; 2) pruning dimensions of linear projections; 3) fine-tuning. The reduced parameters and FLOPs ratios of the proposed algorithm are well evaluated and analyzed on ImageNet dataset to demonstrate the effectiveness of our proposed method.

Motivation & Objective

• Motivate and enable practical deployment of vision transformers on edge devices by reducing storage, memory, and computation.
• Propose a principled pruning framework that identifies and removes unimportant feature dimensions in transformer projections.
• Show that sparsity-inducing training yields salient dimensions automatically and enables substantial compression with limited accuracy loss.
• Provide empirical validation on ImageNet-1K and ImageNet-100 demonstrating effective pruning and speedups.

Proposed method

• Introduce learnable importance scores for the dimensions of linear projections in MHSA and MLP blocks.
• Relax discrete pruning decisions with real-valued importance scores and enforce sparsity using an L1 penalty.
• Train with sparsity regularization to obtain near-zero importance scores, then threshold to obtain a binary prune mask.
• Apply pruning across all MHSA and MLP components, rewire pruned projections, and fine-tune the pruned model.
• Evaluate compression in terms of parameter count, FLOPs, and accuracy on ImageNet-1K and ImageNet-100.

Experimental results

Research questions

• RQ1 Can dimension-wise pruning via learnable importance scores meaningfully reduce parameters and FLOPs in vision transformers without large accuracy losses?
• RQ2 How does the sparsity-regularized training regime influence the emergence of important versus prunable dimensions in ViT?
• RQ3 What are the trade-offs between pruning rate, model size, computational cost, and accuracy on standard vision benchmarks?
• RQ4 Is the proposed VTP approach effective on large-scale datasets like ImageNet-1K as well as smaller subsets like ImageNet-100?

Key findings

• The method yields substantial reductions in parameters and FLOPs with only modest accuracy loss on ImageNet-based benchmarks.
• Pruning up to 40% of dimensions can preserve most of the baseline accuracy while achieving notable FLOPs reduction.
• Pruning performance scales with the sparsity level and is consistent across Imagenet-100 and Imagenet-1K.
• The approach provides a simple baseline for pruning vision transformers and demonstrates practical deployment potential.