Source-linked AI summary
AdaViT: Adaptive Vision Transformers for Efficient Image Recognition
Lingchen Meng, Hengduo Li, Bor-Chun Chen, Shiyi Lan, Zuxuan Wu, Yu-Gang Jiang, Ser-Nam Lim
TL;DR
Vision transformers are computationally expensive because their cost grows with patches, attention heads, and blocks, while image complexity varies across inputs. AdaViT adaptively selects patches, heads, and blocks per image using a lightweight decision network jointly optimized with the backbone. On ImageNet, it reports more than 2× efficiency improvement with only a 0.8% classification-accuracy drop and supports multiple efficiency/accuracy trade-offs.
Problem
Vision transformers have high computational costs that grow with patches, attention heads, and transformer blocks, although easy images may not require full computation.
Method
AdaViT learns per-input policies for patch selection, head activation, and block skipping using a lightweight decision network optimized jointly with the transformer backbone.
Results
More than 2× efficiency improvement with only a 0.8% drop of classification accuracy is reported on ImageNet.
Takeaways & Limitations
AdaViT provides efficiency/accuracy trade-offs across computational budgets and offers quantitative and qualitative insights into redundancy in vision transformers.
Takeaways & Limitations
The experiments consider the ViT architecture, while extension to other vision-transformer variants is described as straightforward.
Abstract
from arXiv · showhide
Built on top of self-attention mechanisms, vision transformers have demonstrated remarkable performance on a variety of vision tasks recently. While achieving excellent performance, they still require relatively intensive computational cost that scales up drastically as the numbers of patches, self-attention heads and transformer blocks increase. In this paper, we argue that due to the large variations among images, their need for modeling long-range dependencies between patches differ. To this end, we introduce AdaViT, an adaptive computation framework that learns to derive usage policies on which patches, self-attention heads and transformer blocks to use throughout the backbone on a per-input basis, aiming to improve inference efficiency of vision transformers with a minimal drop of accuracy for image recognition. Optimized jointly with a transformer backbone in an end-to-end manner, a light-weight decision network is attached to the backbone to produce decisions on-the-fly. Extensive experiments on ImageNet demonstrate that our method obtains more than 2x improvement on efficiency compared to state-of-the-art vision transformers with only 0.8% drop of accuracy, achieving good efficiency/accuracy trade-offs conditioned on different computational budgets. We further conduct quantitative and qualitative analysis on learned usage polices and provide more insights on the redundancy in vision transformers.
1. Introduction
Vision transformers offer strong image recognition through self-attention, but their computational cost grows with patches, heads, and blocks. AdaViT addresses this by allocating computation adaptively per image, reporting over 2× efficiency improvement with a 0.8% accuracy drop on ImageNet.
- Self-attention enables vision transformers to model long-range dependencies, supporting strong performance across vision tasks.
- Computational cost is high and grows quadratically with the number of patches as attention heads and transformer blocks increase.
- Image complexity varies, so cluttered or occluded images may require more patches and attention components than easy object-centric images.
- AdaViT learns instance-specific policies for retaining patches, activating self-attention heads, and skipping transformer blocks.
- The framework controls target computation through γ ∈(0, 1], supporting efficiency/accuracy trade-offs under different computational budgets.
- More than 2× inference-efficiency improvement is achieved with only a 0.8% drop of classification accuracy on ImageNet.
2. Related Work
Prior vision-transformer work improves recognition and efficiency mainly through fixed architectural designs. AdaViT instead learns input-specific computation policies, complementing static architecture approaches.
- Vision-transformer variants have explored recognition performance, training efficiency, and inference efficiency across computer-vision tasks.
- LeViT and Swin Transformer use multi-stage architectures with down-sampling to improve inference efficiency.
- Existing efficient variants use fixed architectures for all samples, whereas AdaViT adaptively allocates computation according to input-specific redundancy.
3. Approach
AdaViT adaptively selects patches, attention heads, and transformer blocks for each image, using decision networks to reduce computation while maintaining classification accuracy. Its policies operate within a vision-transformer backbone and can be trained for target computation budgets.
- Adaptive Vision Transformer: AdaViT inserts a lightweight decision network before each transformer block to produce input-specific usage policies.The policies determine computation throughout the backbone rather than using one fixed strategy for every image.
- Vision Transformer Preliminaries: Vision transformers model patch dependencies with stacked multi-head self-attention and feed-forward networks, but their cost grows with patches, heads, and blocks.Self-attention cost also grows quadratically with the number of patches.
- Adaptive Vision Transformer: The framework selects which patch embeddings to keep, which attention heads to activate, and which transformer blocks or sublayers to skip.Patch, head, and block decisions are produced for each input image.
- Decision Network and Optimization: The decision networks produce sigmoid probabilities for patch, head, and block usage, and target budgets γp, γh, γb specify desired percentages to keep.Binary keep/discard actions are sampled with a differentiable Gumbel-Softmax relaxation during training.
- Patch Selection: Patch selection removes embeddings whose decisions deactivate them, while the class token is always retained as the whole-image representation.This reduces the sequence processed by subsequent transformer blocks.
- Head Selection: Head selection supports partial deactivation through an identity attention map or full deactivation by removing the head and reducing the MSA output embedding size.Full deactivation saves more computation at the same deactivation percentage but may cause more classification errors.
4. Experiment
Experiments on ImageNet evaluate AdaViT's efficiency, accuracy, learned usage policies, component choices, and computational allocation. AdaViT achieves strong efficiency/accuracy trade-offs while adaptively varying patches, heads, and blocks across inputs and network stages.
- Main Results: 81.1% Top-1 accuracy at 3.9 GFLOPs per image gives AdaViT more than 2× efficiency than T2T-ViT with approximately 0.8% accuracy drop.AdaViT also reports better classification performance with less computational cost than comparable ResNets and vision transformers.
- Main Results: AdaViT obtains 48.1% and 9.6% higher accuracy than random and random+ baselines, respectively, at a similar cost of 3.9 GFLOPs per image.Varying γp, γh, and γb allows AdaViT to cover a wide range of efficiency/accuracy trade-offs.
- Ablation Study: Replacing learned patch, head, or block policies with random policies reduces accuracy by clear margins under similar computational budgets.AdaViT achieves 31.9%, 23.7%, and 16.4% higher accuracy than random patch, head, and block selection, respectively.
- Ablation Study: Using only patch, head, or block selection still covers a wide range of accuracy/efficiency trade-offs and outperforms the corresponding Random+ baselines.This supports the effectiveness of each selection component independently.
- Ablation Study: At 50% head deactivation, partial deactivation reaches 81.7% accuracy and 6.9 GFLOPs, while full deactivation reaches 80.3% and 5.1 GFLOPs.Partial deactivation skips attention-map computation, whereas full deactivation removes the entire head and its output to the FFN; full deactivation becomes more accurate as head count increases.
- Computational Allocation: AdaViT allocates more computation in earlier network stages for patches, while head and block selection retain relatively more computation in the last few blocks.The authors hypothesize that later layers are more responsible for final prediction.
- Computational Allocation: AdaViT allocates more computation to difficult classes with complex scenes and less computation to classes requiring fewer resources.Figures report computational-resource distributions across classes and show selected patches at different blocks.
5. Conclusion
AdaViT adaptively selects patches, self-attention heads, and transformer blocks per input using lightweight decision networks jointly optimized with the backbone. It improves efficiency while retaining accuracy and reveals redundancy in vision transformers.
- A lightweight decision network is attached to each transformer block and optimized jointly with the backbone end-to-end.
- More than 2× efficiency improvement is achieved with only a small accuracy drop compared with state-of-the-art vision transformers.
- The method covers a wide range of efficiency/accuracy trade-offs and provides quantitative and qualitative insight into redundancy in vision transformers.
A. Qualitative Results
AdaViT allocates computation adaptively across images, using less for easy object-centric images and more for hard images with cluttered backgrounds. Its patch-selection policies progressively reduce selected patches through the network.
- AdaViT uses less computation on easy object-centric images and more computation on hard, complex images with cluttered backgrounds.
- Patch-selection visualizations show that the number of selected patches gradually decreases through the network.
B. Compatibility to Other Backbones
AdaViT is model-agnostic and is evaluated with DeiT-small as its backbone. It achieves better efficiency/accuracy trade-offs than standard DeiT variants and consistently outperforms Random+, while retaining the same layerwise allocation trend.
- AdaViT is model-agnostic and can be applied to different vision transformer backbones.
- With DeiT-small as backbone, AdaViT achieves better efficiency/accuracy trade-offs than standard DeiT variants.
- AdaViT consistently outperforms its Random+ baseline by large margins in the DeiT-small comparison.
- With DeiT-small, patch-selection policies also keep more computation in earlier layers and allocate progressively less throughout the network.