A Token Selection Method for Effective Token Pruning in Vision Transformers

Jaeyeon Lee, Dong-Wan Choi

http://doi.org/10.5626/JOK.2024.51.6.567

The self-attention-based models, vision transformers, have recently been employed in the field of computer vision. While achieving excellent performance in a variety of tasks, the computation costs increase in proportion to the number of tokens during inference, which causes a degradation in inference speed. Especially when deploying the model in real-world scenarios, many limitations could be encountered. To address this issue, we propose a new token importance measurement, which can be obtained by modifying the structure of multi-head self-attention in vision transformers. By pruning less important tokens through our method during inference, we can improve inference speed while preserving performance. Furthermore, our proposed method, which requires no additional parameters, exhibits better robustness without fine-tuning and demonstrates that it can maximize performance when integrated with existing token pruning methods.