Source-linked AI summary
AdaViT: Adaptive Tokens for Efficient Vision Transformer
Hongxu Yin, Arash Vahdat, Jose Alvarez, Arun Mallya, Jan Kautz, Pavlo Molchanov
TL;DR
Vision transformers use fixed computation despite differing image complexity and patch informativeness. A-ViT reformulates adaptive computation to halt redundant spatial tokens at different depths without extra halting parameters or architecture changes. On ImageNet1K, it improves DeiT-Tiny throughput by 62% and DeiT-Small throughput by 38% with a 0.3% accuracy drop, while distributional priors stabilize training.
Problem
Vision transformers generally use fixed inference cost even though image complexity and patch informativeness vary.
Method
A-ViT adaptively halts spatial tokens at different depths, using existing model parameters and distributional prior regularization to guide halting.
Results
62% higher DeiT-Tiny throughput and 38% higher DeiT-Small throughput were achieved with only a 0.3% accuracy drop on ImageNet1K.
Takeaways & Limitations
A-ViT reduces vision-transformer inference cost while retaining nearly the reported baseline accuracy and requiring no extra halting sub-network.
Takeaways & Limitations
The work primarily evaluates image classification, leaving extensions such as video processing for future work.
Abstract
from arXiv · showhide
We introduce A-ViT, a method that adaptively adjusts the inference cost of vision transformer (ViT) for images of different complexity. A-ViT achieves this by automatically reducing the number of tokens in vision transformers that are processed in the network as inference proceeds. We reformulate Adaptive Computation Time (ACT) for this task, extending halting to discard redundant spatial tokens. The appealing architectural properties of vision transformers enables our adaptive token reduction mechanism to speed up inference without modifying the network architecture or inference hardware. We demonstrate that A-ViT requires no extra parameters or sub-network for halting, as we base the learning of adaptive halting on the original network parameters. We further introduce distributional prior regularization that stabilizes training compared to prior ACT approaches. On the image classification task (ImageNet1K), we show that our proposed A-ViT yields high efficacy in filtering informative spatial features and cutting down on the overall compute. The proposed method improves the throughput of DeiT-Tiny by 62% and DeiT-Small by 38% with only 0.3% accuracy drop, outperforming prior art by a large margin. Project page at https://a-vit.github.io/
1. Introduction
Vision transformers use fixed inference costs despite variation in image and patch complexity. A-ViT adapts computation by halting less informative spatial tokens at different depths, improving throughput with minimal accuracy loss.
- Motivation: Vision transformers typically incur fixed inference cost even when image complexity and patch informativeness vary.Simple images and background patches may require less computation than complex images and discriminative object regions.
- A-ViT: A-ViT halts different spatial tokens at different depths, reserving computation for discriminative tokens dynamically.This extends input-dependent inference from jointly halting all tokens to fine-grained spatial halting.
- A-ViT: A-ViT learns adaptive token halting from existing model parameters without adding parameters or compute for a separate halting module.The method uses the original architecture’s embedding dimensions for halting decisions.
- Results: 62% higher DeiT-Tiny throughput and 38% higher DeiT-Small throughput were achieved with only a 0.3% accuracy drop on ImageNet1K.The reported result demonstrates reduced inference cost for two vision transformer baselines.
- A-ViT: A-ViT introduces distributional prior regularization to guide halting toward a desired token-depth distribution and stabilize training.The regularizer targets computational budgets through the halting behavior of tokens.
2. Related Work
Prior transformer-efficiency methods modify sharing, attention span, exit depth, or token interactions. Related adaptive-inference approaches also use separate halting models, reinforcement learning, or conditional gating.
- Transformer efficiency: Transformer efficiency has been addressed through weight sharing, dynamic attention spans, early exits, pruning, and reduced token interactions.These approaches target different parts of transformer computation and architecture.
- Adaptive inference: Several adaptive-inference methods train separate networks with reinforcement learning to decide when computation should halt.The passage identifies dedicated halting models and high-variance reinforcement-learning signals as training challenges.
- Adaptive inference: Conv-AIG uses the Gumbel-softmax trick to learn conditional gating of residual blocks.This represents another route to input-dependent computation in residual architectures.
3. A-ViT
A-ViT adapts vision-transformer inference by assigning input-dependent halting scores to tokens and dropping each token at its own stopping depth. It uses existing token embeddings for halting, masks halted tokens, and regularizes the halting distribution to balance accuracy and compute.
- Parameter-free halting: Vision transformers support token dropping because their token embedding dimension remains consistent while the number of tokens can change across layers.This enables computational gains when halted tokens are removed.
- Adaptive token halting: A-ViT assigns each token an input-dependent halting score and stops tokens when cumulative halting exceeds 1 − ϵ.The small constant ϵ allows halting after one layer.
- Adaptive token halting: Halted tokens are masked from subsequent layers by zeroing their values and blocking their attention to other tokens.At inference, halted tokens are removed from computation to measure speedup.
- Parameter-free halting: A-ViT uses one existing MLP embedding dimension for halting, requiring no additional parameters or halting sub-network beyond shared scalars β and γ.The first embedding dimension performs well empirically, while changing its index does not affect original performance.
- Output and optimization: The class token uses a halting-probability-weighted mean-field output, while spatial tokens contribute through attention without aggregating image tokens.Each token maintains its own stopping accumulator and may halt at a different depth.
- Output and optimization: The training objective combines task loss with ponder loss to encourage an accuracy–efficiency trade-off, but its behavior is sensitive to the ponder-loss weight α_p.Larger α_p imposes a stronger penalty and encourages earlier token halting.
- Output and optimization: A distributional prior regularizes layerwise halting scores toward a target stopping-depth distribution while allowing per-image variation.The target is a Gaussian-shaped distribution centered at the expected stopping depth N_target, optimized with KL divergence.
4. Experiments
Experiments on ImageNet-1K show that A-ViT adaptively allocates token computation according to image and patch difficulty, improving efficiency while retaining accuracy. Analyses, comparisons, and ablations examine token behavior, class sensitivity, throughput, and the contributions of token-level halting and existing-parameter training.
- Qualitative results: A-ViT-T processes salient tokens more deeply while halting relatively irrelevant tokens, including background regions, during inference.Qualitative analyses report retention of discriminative features such as animal eyes, textures, and colors, while redundant computation is reduced.
- Token depth distribution: Halting scores increase during initial transformer layers, peak at intermediate depths, and decrease in deeper layers across 5K sampled validation images.Figure 4 analyzes average token halting scores over the 12 transformer layers.
- Sharp-halting baseline: A-ViT-T exceeds an 8-layer DeiT-T baseline by 1.4% top-1 accuracy at similar throughput.Informative tokens can still use deeper layers, preserving the expressivity of the original 12-layer network.
- Easy and hard samples: Images with homogeneous backgrounds require less computation, whereas images with informative features distributed across the image incur more computation despite correct classification.Hard samples contain visual information throughout the image, while easy samples are processed faster.
- Class-wise sensitivity: Adaptive inference leaves originally confident or uncertain samples largely unaffected and improves accuracy for visually dominant classes such as furniture and animals.Class-wise sensitivity is measured against the full fixed-computation model.
- Ablations: Token-level ACT reduces token depths by roughly 3 layers and yields 25% more FLOP reductions than conventional layer-wise ACT.The distributional prior guides convergence toward a target average depth, while using it alone causes an accuracy drop of more than 2%.
- Ablations: Using one existing embedding element for halting changes ImageNet1K top-1 accuracy by only 0.08% ± 0.04% for DeiT-T and 0.04% ± 0.03% for DeiT-S.An added two-layer halting network increases accuracy by 0.06% but adds 0.2M parameters and 12.6% inference-throughput overhead.
5. Limitations & Future Directions
The work primarily focuses on image classification and identifies video processing as an interesting extension because video inputs contain both spatial and temporal redundancy.
- The work primarily focuses on the classification task.
- Video processing is proposed as a potential extension because input tokens may contain both spatial and temporal redundancy.
6. Conclusions
A-ViT adjusts token computation according to input complexity, improving vision transformer throughput without extra parameters or transformer-block modifications. Its token-importance distributions vary by image and align with human perception.
- A-ViT adaptively adjusts token computation based on input complexity.
- A-ViT improves vision transformer throughput without imposing extra parameters or modifications of transformer blocks.
- Captured token-importance distributions vary across input images yet coincide with human perception.
- A-ViT outperforms prior dynamic approaches.
Appendix A - More Examples
Additional ImageNet-1K validation examples show A-ViT dynamically allocating token depth to informative regions while filtering complex backgrounds.
- A-ViT processes informative image regions while filtering out complex backgrounds.
- In examples with human faces and coats, A-ViT can ignore faces and focus computation on the coats.
- A-ViT can effectively process images even when the target object occupies a very small informative region.
Appendix B - Additional Details
The appendix details training recipes, latency measurement, and baseline implementation for evaluating A-ViT and competing dynamic inference methods.
- Training: Training follows the original DeiT repository hyperparameters apart from scaling constants and learning rate.
- Training: Adapting a static model to its adaptive counterpart requires minimal training-recipe changes, including dropout, momentum, and preprocessing settings.
- Latency: Latency is measured on an NVIDIA TITAN RTX 2080 GPU with PyTorch using batch size 64 and the median of 1,000 warmed-up measurements.
- SOTA baselines: Baseline experiments use DeiT recipes and checkpoints, DynamicViT's public implementation, and reimplemented adaptive methods from CNN and NLP literature.