Bandit-based Attention Mechanism in Vision Transformers

Vision Transformers (ViT) have demonstrated remarkable performance on many computer vision tasks. However, their high computational cost and quadratic complexity pose chal-lenges for deployment in resource-constrained environments. The core of Vision Transformers is the self-attention mecha-nism which aggregates information from different image re-gions, or patches. In a conventional ViT, processing involves attention to all patches, creating a substantial computational bottleneck and extended training times. We hypothesize that applying soft attention to all patches may be unnecessary and instead focusing on relevant and significant patches (hard attention) would be sufficient. To address this, we introduce a module within the Vision Transformer that allows the at-tention mechanism to selectively process only the essential patches. We propose a novel bandit-based attention mecha-nism that leverages the idea of exploration and exploitation. The extensive experimentation across various datasets il-lustrates that the proposed bandit attention-based ViT not only achieves superior performance compared to the existing state-of-the-art vision transformer models but also results in greater throughput and lower computational time in the training as well as the inference. The code is publicly avail-able at https://github.com/aquorio15/bandit_wacv