Source-linked AI summary
DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification
Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, Cho-Jui Hsieh
TL;DR
Vision transformers require efficient acceleration, and their predictions depend on only a subset of informative tokens. DynamicViT progressively prunes tokens with lightweight, input-dependent prediction modules and end-to-end attention masking, reducing computation while retaining accuracy. On ImageNet, pruning 66% of tokens reduces GFLOPs by 31% ∼37% and improves throughput by over 40%, with accuracy drops within 0.5%.
Problem
Vision transformers are competitive across vision tasks but need acceleration for more suitable real-time deployment.
Method
DynamicViT uses lightweight prediction modules to generate input-specific binary token masks, prunes tokens hierarchically, and trains them end-to-end with Gumbel-Softmax and attention masking.
Results
66% token pruning reduces GFLOPs by 31% ∼37% and improves throughput by over 40%, while accuracy drops within 0.5% across vision transformers.
Takeaways & Limitations
DynamicViT achieves competitive complexity/accuracy trade-offs against state-of-the-art CNNs and vision transformers on ImageNet.
Takeaways & Limitations
The method always keeps the class token instead of applying the decision mask to it.
Abstract
from arXiv · showhide
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
1 Introduction
Vision transformers are desirable for real-time applications, and their patch-token representation enables dynamic sparsification because many tokens contribute little to prediction. DynamicViT progressively prunes uninformative tokens with lightweight prediction modules while preserving accuracy and improving efficiency.
- Vision transformers need acceleration to become more suitable for real-time applications.
- Patch-based inputs allow vision transformers to prune less important tokens, unlike convolutional networks whose pixel removal does not directly accelerate convolution.
- DynamicViT uses lightweight prediction modules to produce input-specific binary masks and hierarchically remove uninformative tokens across multiple layers.
- Gumbel-Softmax and attention masking support end-to-end optimization of the prediction module with the transformer backbone.
- 66% token pruning reduces GFLOPs by 31% ∼37% and raises throughput by over 40%, while accuracy drops within 0.5% across vision transformers.
2 Related Work
Related work establishes vision transformers as competitive alternatives to CNNs across multiple vision tasks and surveys established approaches for accelerating deep and transformer models.
- Vision transformers: Vision transformers extend self-attention-based architectures to image classification, object detection, semantic segmentation, and 3D analysis.
- Vision transformers: ViT applies transformers directly to non-overlapping image patches, while DeiT develops training techniques for vision transformers without requiring large-scale pre-training.
- Model acceleration: Deep-model acceleration methods include quantization, pruning, low-rank factorization, knowledge distillation, and transformer-specific approaches.
3 Dynamic Vision Transformers
DynamicViT combines hierarchical token prediction with attention masking to prune input-dependent tokens while preserving batch-friendly computation. Its modules predict which tokens to keep, progressively update the mask, and support end-to-end training and inference-time token reduction.
- Framework: DynamicViT inserts prediction modules between transformer blocks to generate token keep/drop probabilities and sparsify tokens hierarchically.The backbone can use vision transformers such as ViT, DeiT, or LV-ViT, with sparsification performed at selected locations.
- Hierarchical sparsification: Tokens are progressively removed through a binary decision mask, and a dropped token remains excluded from later computation.The mask is initialized with ones and updated stage by stage; the class token is always kept in practice.
- Prediction modules: The prediction module combines local token features with a global feature aggregated from existing tokens before predicting drop/keep probabilities.The local feature represents an individual token, while the global feature represents whole-image context.
- End-to-end optimization: Gumbel-Softmax makes sampling binary keep decisions differentiable, enabling the prediction modules and transformer backbone to be optimized end to end.The sampled mask represents kept tokens, and its expectation equals the predicted probabilities.
- Attention masking: Attention masking preserves a uniform token count during training while eliminating interactions from pruned tokens with other tokens.Simply zeroing pruned tokens is insufficient because they can still influence other tokens through Softmax attention.
- Inference: During inference, DynamicViT keeps exactly floor(ρ^sN) tokens at stage s by ranking tokens with predicted keeping probabilities.This enables dynamic pruning of less informative tokens and reduces inference computation.
4 Experimental Results
Experiments show that DynamicViT reduces computation and inference cost while preserving accuracy across vision-transformer backbones, and achieves favorable ImageNet complexity/accuracy trade-offs against competing models. Analyses further examine scaling, token selection, losses, sparsification strategies, and larger inputs or models.
- Main results: 31% ∼37% lower computational cost and 43% ∼54% higher runtime throughput accompany only −0.2% ∼−0.5% performance influence across three vision-transformer base models.The experiments use three hierarchical sparsification stages and measure throughput on an NVIDIA RTX 3090 GPU.
- Comparisons with state-of-the-art models: DynamicViT exhibits favorable complexity/accuracy trade-offs at three ImageNet complexity levels and outperforms the compared LV-ViT series.DynamicViT-LV-M/0.7 also beats EfficientNet-B5 and NFNet-F0 in the reported comparison.
- Model scaling: DynamicViT can match DeiT-Ti computational complexity while achieving around 4.3% higher ImageNet top-1 accuracy.
- Visualizations: Hierarchical sparsification progressively drops uninformative tokens and focuses on image objects, suggesting improved interpretability.The visualization shows the original image and sparsification results after three stages.
- Effects of different losses: Distillation and KL losses consistently further boost performance, although their improvements are not very significant.
- Sparsification strategies: At comparable complexity, dynamic token sparsification achieves the best accuracy among the evaluated strategies, while progressive sparsification outperforms one-stage sparsification.The comparison fixes accelerated models at 2.9G FLOPs.
- Accelerating larger models: For DeiT-S with 384 × 384 input, complexity decreases by over 50% with only 1.3% accuracy drop.The method also works on DeiT-B and larger-input models; the accuracy drop becomes less significant with larger feature maps.
5 Conclusion
DynamicViT accelerates vision transformers by dynamically pruning less important tokens with hierarchical prediction modules and end-to-end training strategies. It prunes 66% of input tokens with less than 0.5% accuracy drop, while extension beyond image classification remains an open direction.
- DynamicViT uses lightweight prediction modules and customized binary masks to prune less important tokens hierarchically for each input.Gumbel-Softmax and attention masking support end-to-end training with the transformer backbone.
- 66% token pruning improves efficiency while keeping the accuracy drop below 0.5% across different transformer backbones.
- The paper focuses on image classification, with video classification and dense prediction identified as possible extension directions.
A Implementation Details
Experiments use ImageNet with standard single-crop top-1 evaluation, while models are initialized from pretrained backbones and jointly trained with prediction modules.
- ImageNet training uses 1.28M images, with top-1 accuracy measured on 50k validation images using single-crop evaluation.
- The prediction modules use three sparsification stages with geometric keeping ratios [ρ, ρ^2, ρ^3].
- Each prediction stage uses identical LayerNorm, linear, GELU, and Softmax blocks to predict token-keeping probabilities.
- Backbones are initialized from pretrained vision transformers and jointly trained with prediction modules for 30 epochs.
- Training uses separate learning-rate settings for prediction modules and backbone models, with backbone weights fixed during the first 5 epochs.
B More Analysis
The analysis finds that progressive sparsification and auxiliary losses support the complexity–performance trade-off, while overly aggressive token removal causes substantial accuracy loss.
- Progressive sparsification: Progressive sparsification outperforms single-shot sparsification, with additional stages improving performance until gains beyond three stages become insignificant.The main experiments therefore use three progressive sparsification stages.
- Ablation on the distillation loss and ratio loss: The ratio loss encourages the desired acceleration rate, while distillation losses improve performance after sparsification.
- Smaller keeping ratio: ρ < 0.7 causes significant accuracy drops while reducing fewer FLOPs, because only 22% or 13% of tokens remain at the final stage for ρ = 0.6 or 0.5.
- More visual results: The method is reported to work well on randomly sampled ImageNet validation images from different categories.