[Paper Review] KVT: k-NN Attention for Boosting Vision Transformers

Pichao Wang, Xue Wang+5 more|arXiv (Cornell University)|May 28, 2021

Advanced Neural Network Applications62 references4 citations

TL;DR

This paper proposes k-NN attention, a lightweight, differentiable mechanism that replaces the dense self-attention in Vision Transformers by selecting only the top-k most similar keys for each query, thereby enhancing local inductive bias and filtering noisy tokens. The method accelerates training and improves performance across 11 Vision Transformer architectures without architectural modifications.

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{https://github.com/damo-cv/KVT}.

Motivation & Objective

• Address the slow training and noise sensitivity of Vision Transformers caused by dense self-attention that aggregates all tokens.
• Overcome the limitation of fully-connected attention in capturing local image structure and handling cluttered backgrounds or occlusions.
• Propose a simple yet effective alternative to convolutional inductive bias that preserves global long-range modeling while improving inductive bias.
• Demonstrate that k-NN attention can be universally applied to any Vision Transformer architecture to improve accuracy and training speed.
• Theoretically and empirically validate that k-NN attention reduces noise and speeds up convergence without architectural changes.

Proposed method

• Replace standard dense self-attention with k-NN attention, where for each query token, only the top-k most similar key tokens are used to compute attention scores.
• Compute similarity between query and key vectors using dot-product attention, then select the k nearest keys based on cosine or dot-product similarity.
• The attention map is computed only over the k selected keys, reducing computational cost and filtering out irrelevant or noisy tokens.
• The k-NN mechanism is differentiable and end-to-end trainable, enabling integration into any Vision Transformer without architectural changes.
• The method naturally inherits local inductive bias since nearby tokens are more likely to be selected, even without explicit convolutional operations.
• The k-NN mechanism adapts the aggregation graph dynamically based on feature similarity, allowing long-range dependencies while suppressing distant, dissimilar tokens.

Experimental results

Research questions

• RQ1 Can replacing dense self-attention with k-NN attention improve training speed and model robustness in Vision Transformers?
• RQ2 Does k-NN attention effectively filter out noisy tokens such as background clutter and occlusions during training?
• RQ3 To what extent does k-NN attention preserve or enhance local inductive bias compared to standard self-attention?
• RQ4 Can k-NN attention be universally applied across diverse Vision Transformer architectures without architectural modifications?
• RQ5 How does k-NN attention compare to convolutional inductive bias in terms of performance and training dynamics?

Key findings

• k-NN attention significantly accelerates training across 11 Vision Transformer architectures, reducing convergence time without compromising final accuracy.
• The method improves generalization by filtering out noisy tokens such as background clutter and occluded regions, leading to more robust feature learning.
• k-NN attention achieves state-of-the-art or competitive performance on ImageNet-1K, with consistent improvements across multiple backbone architectures.
• Theoretical analysis shows that k-NN attention maintains strong generalization guarantees by controlling the influence of outliers and noisy features.
• Empirical results confirm that k-NN attention reduces the variance of attention maps and improves attention map sparsity, leading to more focused and meaningful attention heads.
• The method is universally applicable: it can be plugged into any Vision Transformer with no architectural changes, and it consistently improves performance across different models and datasets.