[Paper Review] DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification

Yongming Rao, Wenliang Zhao+4 more|arXiv (Cornell University)|Jun 3, 2021

Advanced Neural Network ApplicationsComputer Science31 references310 citations

TL;DR

DynamicViT introduces hierarchical, input-dependent token pruning in vision transformers using a lightweight prediction module and attention masking, achieving substantial FLOPs reduction with minimal accuracy loss.

ABSTRACT

Attention is sparse in vision transformers. We observe the final prediction in vision transformers is only based on a subset of most informative tokens, which is sufficient for accurate image recognition. Based on this observation, we propose a dynamic token sparsification framework to prune redundant tokens progressively and dynamically based on the input. Specifically, we devise a lightweight prediction module to estimate the importance score of each token given the current features. The module is added to different layers to prune redundant tokens hierarchically. To optimize the prediction module in an end-to-end manner, we propose an attention masking strategy to differentiably prune a token by blocking its interactions with other tokens. Benefiting from the nature of self-attention, the unstructured sparse tokens are still hardware friendly, which makes our framework easy to achieve actual speed-up. By hierarchically pruning 66% of the input tokens, our method greatly reduces 31%~37% FLOPs and improves the throughput by over 40% while the drop of accuracy is within 0.5% for various vision transformers. Equipped with the dynamic token sparsification framework, DynamicViT models can achieve very competitive complexity/accuracy trade-offs compared to state-of-the-art CNNs and vision transformers on ImageNet. Code is available at https://github.com/raoyongming/DynamicViT

Motivation & Objective

• Motivate accelerating vision transformers by exploiting sparsity in informative image patches.
• Propose a dynamic token sparsification framework that prunes tokens progressively and input-dependently.
• Develop end-to-end trainable prediction modules with differentiable pruning via Gumbel-Softmax and attention masking.
• Demonstrate substantial FLOPs reduction and throughput gains across multiple backbone transformers on ImageNet.

Proposed method

• Insert lightweight prediction modules at multiple transformer blocks to estimate per-token importance.
• Compute local-global embeddings from token features to predict drop/keep probabilities for each token.
• Use Gumbel-Softmax to sample a binary keep/drop mask while preserving differentiability for training.
• Apply attention masking to remove interactions involving pruned tokens during self-attention, keeping computation uniform during training.
• Train with a combination of cross-entropy loss, distillation losses against a teacher backbone, KL divergence, and a ratio-constrained pruning loss.
• During inference, prune a fixed number of tokens per stage according to learned scores to achieve target keeping ratios.

Experimental results

Research questions

• RQ1 Can vision transformers be accelerated by pruning uninformative tokens without substantial accuracy loss?
• RQ2 How can we train a dynamic token pruning mechanism end-to-end within a transformer framework?
• RQ3 What are the effects of hierarchical, input-dependent token pruning on model efficiency and accuracy across different backbones?

Key findings

• Hierarchical token sparsification can prune up to 66% of input tokens, reducing 31%–37% FLOPs while increasing throughput by over 40% with only ~0.5% accuracy drop across backbones.
• DynamicViT achieves competitive complexity/accuracy trade-offs on ImageNet compared to state-of-the-art CNNs and vision transformers.
• Dynamic pruning behaves sensibly, preserving tokens around the image center and objects while pruning peripheral areas, and shows interpretability as pruning progresses.
• The approach provides a viable alternative to width scaling for model scaling, achieving similar or better efficiency with dynamic token sparsification.
• Larger models (DeiT-B and 384x384 inputs) benefit from DynamicViT, with substantial FLOPs reductions and modest accuracy drops.
• Ablations show the effectiveness of dynamic, learned pruning over static or random token removal strategies.