[Paper Review] Adaptive Token Sampling For Efficient Vision Transformers

Mohsen Fayyaz, Soroush Abbasi Koohpayegani+6 more|arXiv (Cornell University)|Nov 30, 2021

Cell Image Analysis Techniques4 citations

TL;DR

This paper proposes Adaptive Token Sampling (ATS), a differentiable, parameter-free module that dynamically reduces the number of tokens in Vision Transformers based on input image content. By using attention weights from the classification token to score and sample significant patches via inverse transform sampling, ATS reduces GFLOPs by up to 2× without accuracy drop on ImageNet, Kinetics-400, and Kinetics-600, enabling plug-and-play efficiency gains on pre-trained models.

ABSTRACT

While state-of-the-art vision transformer models achieve promising results in image classification, they are computationally expensive and require many GFLOPs. Although the GFLOPs of a vision transformer can be decreased by reducing the number of tokens in the network, there is no setting that is optimal for all input images. In this work, we therefore introduce a differentiable parameter-free Adaptive Token Sampler (ATS) module, which can be plugged into any existing vision transformer architecture. ATS empowers vision transformers by scoring and adaptively sampling significant tokens. As a result, the number of tokens is not constant anymore and varies for each input image. By integrating ATS as an additional layer within the current transformer blocks, we can convert them into much more efficient vision transformers with an adaptive number of tokens. Since ATS is a parameter-free module, it can be added to the off-the-shelf pre-trained vision transformers as a plug and play module, thus reducing their GFLOPs without any additional training. Moreover, due to its differentiable design, one can also train a vision transformer equipped with ATS. We evaluate the efficiency of our module in both image and video classification tasks by adding it to multiple SOTA vision transformers. Our proposed module improves the SOTA by reducing their computational costs (GFLOPs) by 2X, while preserving their accuracy on the ImageNet, Kinetics-400, and Kinetics-600 datasets.

Motivation & Objective

• To address the high computational cost of Vision Transformers, which scale quadratically with sequence length.
• To enable dynamic, content-aware token reduction that varies per input image, avoiding fixed-ratio pruning.
• To design a parameter-free, differentiable module that can be added to pre-trained models without retraining.
• To reduce GFLOPs while preserving or improving accuracy across image and video classification benchmarks.
• To provide a plug-and-play solution for efficient inference on edge devices without architectural retraining.

Proposed method

• ATS uses attention weights from the classification token to compute significance scores for all input tokens.
• It applies inverse transform sampling over the cumulative distribution of these scores to select a subset of significant tokens.
• The selected tokens are softly downsampled using a weighted combination of attention weights and values to preserve feature representation.
• The module is differentiable and parameter-free, allowing integration into any Vision Transformer block without additional training.
• It can be inserted into existing pre-trained models as a plug-and-play layer, or jointly trained with the backbone.
• The method enables progressive token pruning across multiple stages, with the number of tokens varying per image based on complexity.

Experimental results

Research questions

• RQ1 Can a parameter-free, differentiable module dynamically reduce the number of tokens in Vision Transformers based on input content?
• RQ2 Does adaptive token sampling preserve accuracy while reducing computational cost compared to fixed-ratio pruning?
• RQ3 Can ATS be applied as a plug-and-play module to pre-trained Vision Transformers without fine-tuning?
• RQ4 How does progressive token sampling across multiple stages affect efficiency and accuracy?
• RQ5 Does the use of attention-based significance scoring outperform learnable scoring networks in terms of efficiency and generalization?

Key findings

• ATS reduces GFLOPs by up to 2× in DeiT-S and other SOTA Vision Transformers while maintaining or improving top-1 accuracy on ImageNet.
• The DeiT-S+ATS model achieves 79.7% top-1 accuracy with only 2.9 GFLOPs, outperforming the baseline DeiT-S (78.9%) at the same FLOP level.
• Integrating ATS into later stages (e.g., stage 3) yields higher accuracy (78.5%) than earlier stages (73.1%), due to more reliable attention scores.
• Multi-stage integration of ATS (stages 3–11) achieves 79.2% top-1 accuracy with 2.6 GFLOPs, demonstrating progressive efficiency gains.
• When added to EViT-DeiT-S, ATS further reduces GFLOPs from 3.0 to 2.5 without accuracy drop, proving compatibility and superiority over static pruning.
• Visualizations confirm that the model samples more tokens for complex, cluttered images and fewer for simple ones, with more focused attention on salient regions.