Adaptive Token Sampling For Efficient Vision Transformers (2111.15667v3)

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.

• The paper introduces a parameter-free, differentiable Adaptive Token Sampler (ATS) module that adaptively selects the most informative tokens based on attention weights, optimizing computational efficiency without requiring retraining.
• Incorporating ATS into vision transformer models like DeiT-S reduces GFLOPs by approximately 37% on ImageNet with negligible accuracy loss, demonstrating superior efficiency trade-offs compared to existing methods.
• The modularity and non-dependency on retraining make ATS a practical solution for deploying efficient vision transformers on resource-constrained devices and inspire future dynamic computational models.

Adaptive Token Sampling For Efficient Vision Transformers

Vision transformers have emerged as a promising advancement in the field of computer vision, often outperforming traditional convolutional neural networks (CNNs) in tasks like image classification. However, their deployment is often hindered by significant computational demands, particularly as their computational cost grows quadratically with the number of tokens. The paper "Adaptive Token Sampling For Efficient Vision Transformers" addresses this bottleneck by introducing a novel Adaptive Token Sampler (ATS) module designed to optimize the efficiency of vision transformers.

Overview of the Adaptive Token Sampler (ATS)

The core proposal of this paper is the integration of the ATS module into existing vision transformers. This module is parameter-free and differentiable, facilitating seamless implementation into off-the-shelf models without necessitating additional parameters. The primary function of ATS is to adaptively sample the most informative tokens based on the significance assigned through the attention weights of the classification token. The significance scores are calculated using the classification token's attention matrix entries, adjusted by the magnitude of the associated values in the self-attention mechanism.

ATS uses inverse transform sampling over a cumulative distribution function derived from these scores to select tokens, rather than employing traditional fixed sampling strategies. This adaptive approach ensures that the number of tokens processed varies in accordance with the specific input, thereby efficient managing of computational resources. For complex inputs, more tokens are selected, while simpler inputs necessitate fewer tokens, thereby reducing redundant computation without compromising accuracy.

Key Results and Evaluation

The paper thoroughly evaluates the efficacy of the ATS module across multiple state-of-the-art vision transformer architectures, including DeiT, CvT, and PS-ViT on image classification tasks using the ImageNet dataset. The results demonstrate that incorporating ATS leads to a reduction in GFLOPs by approximately 37% in the case of DeiT-S, with negligible impact on classification accuracy. Additionally, the module achieves similar savings on video classification tasks, with tests conducted on datasets like Kinetics-400 and Kinetics-600, thereby validating its broader application.

The adaptive sampling shows significantly better accuracy-GFLOPs trade-offs compared to other token reduction strategies, such as DynamicViT and EViT, which require additional learnable parameters. The non-dependency on re-training when adapting to different deployment scenarios further underscores the utility of the ATS module, making it a practical solution for reducing computational expenditure on edge devices.

Implications and Future Directions

The Adaptive Token Sampler presents substantial implications for the deployment of vision transformers in real-world applications, where computational efficiency is crucial. The modularity of ATS ensures that it can be integrated into pre-trained models, enabling quick adaptation and deployment across different computational environments without extensive retraining.

From a theoretical perspective, the approach opens avenues for future research into dynamic computational models that better utilize resources by considering the complexity of input data. Potential future developments may explore extending this adaptive methodology to other domains, such as natural language processing or audio processing applications, where transformer architectures are also prevalent.

In conclusion, the paper presents a meaningful advancement in the optimization of vision transformers, addressing the critical challenge of computational inefficiency without compromising performance. As researchers continue to seek methods to balance performance with resource utilization, the ideas put forth in this paper are likely to inspire further innovations and adaptations in the design and deployment of AI models.