PaperlayerView exact revision on arXiv

Source-linked AI summary

KVT: k-NN Attention for Boosting Vision Transformers

Pichao Wang, Xue Wang, Fan Wang, Ming Lin, Shuning Chang, Hao Li, Rong Jin

arXiv:2106.00515v3cs.CV

TL;DR

Vision transformers’ dense self-attention can overlook locality, process noisy tokens, and slow training. This paper introduces k-NN attention, which selects the top-k keys for each query to combine local bias with long-range correlation. Across eleven architectures, it improves ImageNet-1K top-1 accuracy by 0.2% to 0.8% without increasing model size.

  • Problem

    Fully connected self-attention in vision transformers can neglect image locality and include noisy tokens, contributing to slow training and potential performance degradation.

  • Method

    k-NN attention replaces dense attention by selecting the top-k similar keys for each query to compute the attention map.

  • Results

    Across eleven vision transformer architectures, k-NN attention improves ImageNet-1K top-1 accuracy by 0.2% to 0.8% without increasing model size.

  • Takeaways & Limitations

    The method retains local bias, explores long-range correlations, filters irrelevant tokens, and speeds up training in vision transformers.

  • Takeaways & Limitations

    The exact slow implementation of k-NN attention is extremely slow because it computes distances for different k keys separately for each query.

Abstract

from arXiv · show

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}.

1 Introduction

Vision transformers model long-range dependencies but fully connected self-attention can neglect image locality, include noisy tokens, and slow training. The paper proposes k-NN attention, selecting similar tokens to retain local and global information while filtering irrelevant tokens.

  • 1 Introduction: Vision transformers’ fully connected attention can neglect locality and include noisy tokens, slowing training especially with cluttered backgrounds or occlusion.The paper identifies fully connected self-attention as the structural source of these problems.
  • 1 Introduction: k-NN attention replaces dense attention by selecting the top-k similar tokens for each query when computing the attention map.This directly targets the use of all tokens in dense attention.
  • 1 Introduction: Selecting similar tokens preserves CNN-like local bias while also capturing long-range correlations across the image.The selected tokens are drawn from the entire image rather than restricted to spatial neighbors.
  • 1 Introduction: The method is evaluated with eleven vision transformer architectures to assess its effectiveness across existing designs.The paper presents k-NN attention as applicable to multiple transformer architectures.

2 Related Work

Prior work applies transformers and efficient-attention ideas to vision, addressing issues such as training stability, position modeling, scale, and quadratic complexity. However, existing vision transformer variants still use fully connected self-attention, motivating k-NN attention as a content-based sparse alternative with local bias and global feature exploitation.

  • 2 Related Work: Efficient attention includes fixed sparse, content-based sparse, and local attention, while k-NN attention is content-based and tailored to the vision domain.Unlike routing transformer, it clusters keys rather than both queries and keys.
  • 2 Related Work: Vision transformer research includes hybrid and pure transformer models, with ViT representing the CNN-in-Transformer direction for visual tasks.These models convert images into patch-token sequences and process them with transformer encoders.
  • 2 Related Work: Existing variants address data efficiency, positional encoding, patch modeling, multiscale learning, training stability, and quadratic complexity through diverse architectural techniques.Examples include distillation, conditional position encoding, pixel embeddings, multiscale designs, and efficient transformers.
  • 2 Related Work: Despite these advances, the reviewed vision transformer methods retain fully connected self-attention, which introduces irrelevant tokens and slows network training.The paper positions k-NN attention as an efficient sparse-attention response to this remaining issue.

3 k-NN Attention

k-NN attention replaces dense self-attention with top-k similar keys and values per query, preserving long-range matching while reducing irrelevant-token involvement. The paper presents fast and slow variants and theoretically links selective attention to faster convergence and noise distillation.

  • 3.2 k-NN Attention: Top-k selection replaces all query-key comparisons, producing an attention matrix of size n×k instead of n×n.The slow version exactly computes nearest neighbors but is extremely slow; the fast version uses matrix multiplication followed by row-wise top-k selection.
  • 3.2 k-NN Attention: The fast version computes query-key dot products and masks all but the row-wise top-k entries before softmax.This retains the matrix-multiplication structure of vanilla attention while restricting aggregation to selected tokens.
  • 3.3 Theoretical Analysis on k-NN Attention: Selecting similar tokens reduces attention variance and gradient scales, which the Adam analysis links to faster convergence at the same learning rate.The paper expects the benefit to be especially significant early in training, when random initialization makes attention resemble global averaging.
  • 3.3 Theoretical Analysis on k-NN Attention: With an appropriate k, k-NN attention preferentially selects relevant patches, lowering the noise ratio and the distance between updated representations and their means.For fully connected attention, k=n makes the noisy-patch ratio approximately one, weakening noise distillation when noise is large.
  • 3.3 Theoretical Analysis on k-NN Attention: Selecting the top O(nd−c2) similarities can eliminate almost all irrelevant patches without losing relevant ones with high probability.The paper reports numerical verification through better performance and more accurate attention regions for hard instances.

4 Experiments for Vision Transformers

The experiments replace vanilla attention with k-NN attention across eleven vision transformer architectures and evaluate ImageNet performance and the effect of k.

  • 4.2 Results on ImageNet: Replacing dense attention with k-NN attention improves ImageNet top-1 accuracy by 0.2%–0.8% across eleven vision transformers without increasing model size.The evaluation covers supervised and self-supervised architectures, including both global and local vision transformers.
  • 4.3 The Impact of Number k: For simple token-generation methods, the best k is generally approximately half the token count at each scale stage.Reported optima include k = 100 for DeiT-Tiny and 1600/400/100 for CvT-13.
  • 4.3 The Impact of Number k: For Visformer-Tiny, larger k values are preferred because extensive convolutional token mixing makes irrelevant tokens harder to distinguish.The reported stage-wise choices are k = 150/45 for token counts n1 = 196 and n2 = 49.

4.4 Convergence Speed of k-NN Attention

The section evaluates k-NN attention through convergence speed and four quantitative metrics, finding faster early training and changes in token similarity, attention-weight stability, information preservation, and nonlocality.

  • 4.4 Convergence Speed of k-NN Attention: k-NN attention converges faster than fully connected attention, especially during early training, in comparisons with DeiT-Small, CvT-13, and T2T-ViT-t-19.The authors attribute this observation to removing irrelevant tokens, which benefits neural-network training convergence.
  • 4.4 Convergence Speed of k-NN Attention: The analysis also measures residual-to-main-branch activation norms for information preservation and attention nonlocality across layers.Nonlocality weights distances from query patches by attention scores, with larger values indicating attention farther from the query.
  • 4.4 Convergence Speed of k-NN Attention: k-NN attention produces higher averaged layer-wise token cosine similarity than dense self-attention, reflecting faster convergence.Layer-wise cosine similarity is defined from token representations and is described as implying network convergence speed.
  • 4.4 Convergence Speed of k-NN Attention: Its attention-weight standard deviation is smoother than fully connected self-attention, a property associated with more stable training.The metric averages standard deviations across attention heads and represents training stability.
  • 4.4 Convergence Speed of k-NN Attention: More evenly distributed nonlocality across layers indicates exploration of dependencies at a larger variety of ranges.The comparisons are reported for DeiT-Tiny without a distillation token using Figures 2 and 3.

4.6 Comparisons with temperature in softmax

The comparison finds k-NN attention more robust than softmax temperature scaling, achieving higher Top-1 performance, while broader evaluations report gains with almost no overhead.

  • 4.6 Comparisons with temperature in softmax: 73.0 Top-1 (%) for k-NN attention exceeds 72.5 for the best softmax temperature setting.The baseline reaches 72.2, while small temperatures can cause training to crash.
  • 4.6 Comparisons with temperature in softmax: k-NN attention is more robust than temperature scaling in softmax.
  • 4.7 Visualization: Replacing vanilla attention with k-NN attention improves object detection and segmentation performance with almost no overhead.The evaluation uses Swin-Tiny and Twins-SVT-Base on COCO and ADE20K validation sets.

5 Conclusion

The paper concludes that k-NN attention improves vision transformers by selecting similar keys, removing ineffective tokens, and supporting faster training and noise distillation.

  • 5 Conclusion: k-NN attention selects the most similar keys for each query to calculate attention and screen out ineffective tokens.
  • 5 Conclusion: The paper theoretically analyzes faster training, noise distillation without information loss, and performance gains from choosing a proper k.
  • 5 Conclusion: Several vision transformers are used to verify the effectiveness of k-NN attention.

Supplemental Material: KVT: k-NN Attention for Boosting Vision

The supplemental material identifies the paper’s authors and their Alibaba Group affiliation.

  • Supplemental Material: KVT: k-NN Attention for Boosting Vision: The paper lists Pichao Wang, Xue Wang, Fan Wang, Ming Lin, Shuning Chang, and Hao Li among its authors.
  • Supplemental Material: KVT: k-NN Attention for Boosting Vision: The listed author affiliation is Alibaba Group.
  • Supplemental Material: KVT: k-NN Attention for Boosting Vision: The supplemental material provides Alibaba Group contact email addresses for the authors.

1 Differences with the arXiv paper: Explicit Sparse Transformer: Concentrated Attention Through Explicit Selection (EST)

The paper distinguishes its vision-transformer focus and noise-distillation motivation from EST, while emphasizing theoretical analysis and broader empirical evaluation.

  • Differences with the arXiv paper: The paper states that larger k values, namely ≥1⁄2N, performed better than the small k values of 8 or 16 reported for EST.
  • Differences with the arXiv paper: The paper aims to distill noisy patches, whereas EST targets sparse attention maps.
  • Differences with the arXiv paper: Unlike EST’s NLP focus, this paper applies k-NN attention to vision transformers and evaluates 11 transformer backbones.
  • Differences with the arXiv paper: The authors report theoretical justifications, deeper analysis, metric-based comparisons, and visualizations beyond applying top-k selection.

2 Source codes of fast version k-NN attention in Pytorch

Algorithm 1 implements fast k-NN attention in a compact PyTorch module that can be imported into architectures using fully connected attention. The implementation projects inputs to Q, K, and V, constructs a top-k mask, suppresses unselected attention entries, and computes the output with matrix multiplication.

  • Fast k-NN attention is implemented in a compact PyTorch module that can be imported into architectures using fully connected attention.The core implementation is described as consisting of only four lines.
  • The forward pass reshapes the input projections into query, key, and value tensors across attention heads.
  • A non-gradient mask marks the selected top-k keys for each query and assigns negative infinity to unselected attention entries.
  • The masked attention weights are multiplied by the values and reshaped to produce the output tensor.

3 Comparisons between slow version and fast version

The paper develops slow, CUDA, and fast versions of k-NN attention, with the fast version designed to accelerate computation through matrix multiplication. The supplied passages identify a speed-comparison table but do not report its numerical values.

  • The fast version accelerates k-NN attention by exploiting matrix multiplication, whereas the slow version selects different keys and values for each query.The CUDA version is described as still slower than the fast version.
  • The comparison is conducted on DeiT-Tiny, as indicated by the supplied speed-comparison table caption.

4 Evaluations on CIFAR10 or CIFAR100.

The paper evaluates k-NN attention on CIFAR10 and CIFAR100 using Swin-T and T2T-ViT-14. The supplied passages state that adding k-NN attention brings larger improvements but do not provide the reported values.

  • Experiments evaluate k-NN attention on CIFAR10 and CIFAR100 with Swin-T and T2T-ViT-14.The experiments use 100 training epochs.
  • The authors state that adding k-NN attention brings much larger performance improvements on these small datasets.

5 Proof

The proofs analyze how top-k attention affects gradients and token selection under stated assumptions. They relate gradients to attention-weighted token variance and establish probabilistic bounds for selecting relevant keys while limiting irrelevant ones.

  • Gradients with respect to query and key projections are proportional to the attention-weighted variance of input tokens.The result is stated for k-NN attention and also holds for fully connected self-attention when k=n.
  • The formal bounds rely on assumptions about sub-Gaussian token noise, discrete clustering means, and bounded projection matrices.
  • The analysis decomposes selected patches into shared clustering centers and sub-Gaussian noise, then bounds signal, noise, and approximation terms separately.
  • The top-k approximation can identify important elements by ranking the entries of Kq^T when their magnitudes are large.
  • Under the stated assumptions, top-k selection retains relevant keys with high probability while bounding the number of selected keys meeting the query-similarity threshold.The bound is given with probability 1 − O(s exp(−Cd^(1−2κ)/log d)).
Loading 2106.00515v3…