References & Citations
Computer Science > Computer Vision and Pattern Recognition
Title: KVT: k-NN Attention for Boosting Vision Transformers
(Submitted on 28 May 2021 (v1), last revised 22 Jul 2022 (this version, v3))
Abstract: Convolutional Neural Networks (CNNs) have dominated computer vision for years, due to its ability in capturing locality and translation invariance. Recently, many vision transformer architectures have been proposed and they show promising performance. A key component in vision transformers is the fully-connected self-attention which is more powerful than CNNs in modelling long range dependencies. However, since the current dense self-attention uses all image patches (tokens) to compute attention matrix, it may neglect locality of images patches and involve noisy tokens (e.g., clutter background and occlusion), leading to a slow training process and potential degradation of performance. To address these problems, we propose the $k$-NN attention for boosting vision transformers. Specifically, instead of involving all the tokens for attention matrix calculation, we only select the top-$k$ similar tokens from the keys for each query to compute the attention map. The proposed $k$-NN attention naturally inherits the local bias of CNNs without introducing convolutional operations, as nearby tokens tend to be more similar than others. In addition, the $k$-NN attention allows for the exploration of long range correlation and at the same time filters out irrelevant tokens by choosing the most similar tokens from the entire image. Despite its simplicity, we verify, both theoretically and empirically, that $k$-NN attention is powerful in speeding up training and distilling noise from input tokens. Extensive experiments are conducted by using 11 different vision transformer architectures to verify that the proposed $k$-NN attention can work with any existing transformer architectures to improve its prediction performance. The codes are available at \url{this https URL}.
Submission history
From: Pichao Wang [view email][v1] Fri, 28 May 2021 06:49:10 GMT (6266kb,D)
[v2] Wed, 12 Jan 2022 00:53:35 GMT (6786kb,D)
[v3] Fri, 22 Jul 2022 23:18:16 GMT (14790kb,D)
Link back to: arXiv, form interface, contact.