Token Selection is a Simple Booster for Vision Transformers.

Area of Science:

• Computer Vision
• Machine Learning
• Artificial Intelligence

Background:

• Vision Transformers (ViTs) achieve state-of-the-art results in visual recognition, largely due to their self-attention mechanism.
• Self-attention models global dependencies among image patches (tokens) but faces training difficulties and attention over-smoothing, limiting feature learning.
• Existing research often focuses on complex self-attention modifications to address these challenges.

Purpose of the Study:

• To explore simple approaches to enhance the vanilla self-attention mechanism in Vision Transformers.
• To address the limitations of low diversity in token selection caused by attention over-smoothing.
• To develop a drop-in module that boosts the performance of various ViT backbones.

Main Methods:

• Investigated the token selection behavior of self-attention in Vision Transformers.
• Developed simple methods to improve the selectivity and diversity of self-attention in token selection.
• Introduced a novel token selector module designed as a drop-in component for existing ViT architectures.

Main Results:

• The proposed token selector module enhances selectivity and diversity in self-attention.
• ViTs equipped with the token selector achieved 84.6% top-1 accuracy on ImageNet with 25M parameters.
• Scaling up to 81M parameters further improved performance to 86.1% top-1 accuracy.
• The token selector consistently boosted performance across various transformer-based models for image classification, semantic segmentation, and NLP tasks.

Conclusions:

• Simple approaches can effectively unlock the potential of vanilla self-attention in Vision Transformers.
• The developed token selector module is a versatile and effective enhancement for transformer-based models.
• This work demonstrates significant performance gains in computer vision and natural language processing tasks through improved token selection.