Bandit-based Attention Mechanism in Vision Transformers

The following repo contains the code for the paper (Bandit-based Attention Mechanism in Vision Transformers) which was accepted at WACV-2025. Our approach aim to use the idea behind Multi-arm bandit problem especially the concept of exploration and exploitation used in UCB to intelligently keep tokens in ViT. Extensive experiments show that with our approach we are able to get better accuaracy and inference speed(throughout) with minimal drop in performance.