11institutetext: Paper ID 167011institutetext: Alibaba Group
11email: {pichao.wang, xue.w, fan.w, ming.l, shuning.csn, lihao.lh,
<EMAIL_ADDRESS>

# KVT: $k$-NN Attention for Boosting Vision Transformers

Anonymous ECCV submission Paper ID 1670Alibaba Group
11email: {pichao.wang, xue.w, fan.w, ming.l, shuning.csn, lihao.lh,
<EMAIL_ADDRESS>
Pichao Wang The first two authors contribute equally.Paper ID 1670Alibaba
Group
11email: {pichao.wang, xue.w, fan.w, ming.l, shuning.csn, lihao.lh,
<EMAIL_ADDRESS>
Xue Wang††footnotemark:  Paper ID 1670Alibaba Group
11email: {pichao.wang, xue.w, fan.w, ming.l, shuning.csn, lihao.lh,
<EMAIL_ADDRESS>
Fan Wang Paper ID 1670Alibaba Group
11email: {pichao.wang, xue.w, fan.w, ming.l, shuning.csn, lihao.lh,
<EMAIL_ADDRESS>
Ming Lin Paper ID 1670Alibaba Group
11email: {pichao.wang, xue.w, fan.w, ming.l, shuning.csn, lihao.lh,
<EMAIL_ADDRESS>
Shuning Chang Paper ID 1670Alibaba Group
11email: {pichao.wang, xue.w, fan.w, ming.l, shuning.csn, lihao.lh,
<EMAIL_ADDRESS>
Hao Li Paper ID 1670Alibaba Group
11email: {pichao.wang, xue.w, fan.w, ming.l, shuning.csn, lihao.lh,
<EMAIL_ADDRESS>
Rong Jin Paper ID 1670Alibaba Group
11email: {pichao.wang, xue.w, fan.w, ming.l, shuning.csn, lihao.lh,
<EMAIL_ADDRESS>

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

##  1 Introduction

Traditional CNNs provide state of the art performance in vision tasks, due to
its ability in capturing locality and translation invariance, while
transformer [54] is the de-facto standard for natural language processing
(NLP) tasks thanks to its advantages in modelling long-range dependencies.
Recently, various vision transformers [16, 52, 71, 56, 25, 70, 42, 60, 53, 72]
have been proposed by building pure or hybrid transformer models for visual
tasks. Inspired by the transformer scaling success in NLP tasks, vision
transformer converts an image into a sequence of image patches (tokens), with
each patch encoded into a vector. Since self-attention in the transformer is
position agnostic, different positional encoding methods [16, 11, 14] have
been developed, and in [60, 8] their roles have been replaced by convolutions.
Afterwards, all tokens are fed into stacked transformer encoders for feature
learning, with an extra $CLS$ token [16, 52, 14] or global average pooling
(GAP) [42, 8] for final feature representation. Compared with CNNs,
transformer-based models explicitly exploit global dependencies and
demonstrate comparable, sometimes even better, results than highly optimised
CNNs [27, 48].

Albeit achieving its initial success, vision transformers suffer from slow
training. One of the key culprits is the fully-connected self-attention, which
takes all the tokens to calculate the attention map. The dense attention not
only neglects the locality of images patches, an important feature of CNNs,
but also involves noisy tokens into the computation of self-attention,
especially in the situations of cluttered background and occlusion. Both
issues can slow down the training significantly [12, 14]. Recent works [70,
60, 8] try to mitigate this problem by introducing convolutional operators
into vision transformers. Despite encouraging results, these studies fail to
resolve the problem fundamentally from the transformer structure itself,
limiting their success. In this study, we address the challenge by directly
attacking its root cause, i.e. the fully-connected self-attention.

To this end, we propose the $k$-NN attention to replace the fully-connected
attention. Specifically, we do not use all the tokens for attention matrix
calculation, but only select the top-$k$ similar tokens from the sequence for
each query token to compute the attention map. The proposed $k$-NN attention
not only naturally inherits the local bias of CNNs as the nearby tokens tend
to be more similar than others, but also builds the long range dependency by
choosing the most similar tokens from the entire image. Compared with
convolution operator which is an aggregation operation built on Ising model
[44] and the feature of each node is aggregated from nearby pixels, in the
$k$-NN attention, the aggregation graph is no longer limited by the spatial
location of nodes but is adaptively computed via attention maps, thus, the
$k$-NN attention can be regarded as a relieved version of local bias. The
similar idea is proposed in [76] where the $k$-NN attention is mostly
evaluated on NLP tasks. Despite the similarity in terms of the calculation of
top-$k$, our work focuses on the recent vision transformers, makes a deep
theoretical understanding and presents a thoroughly analysis by defining
several metrics. We verify, both theoretically and empirically, that $k$-NN
attention is effective in speeding up training and distilling noisy tokens of
vision transformers. Eleven different available vision transformer
architectures are adopted to verify the effectiveness of the proposed $k$-NN
attention.

##  2 Related Work

###  2.1 Self-attention

Self-attention [54] has demonstrated promising results on NLP related tasks,
and is making breakthroughs in speech and computer vision. For time series
modeling, self-attention operates over sequences in a step-wise manner.
Specifically, at every time-step, self-attention assigns an attention weight
to each previous input element and uses these weights to compute the
representation of the current time-step as a weighted sum of the past inputs.
Besides the vanilla self-attention, many efficient transformers [51] have been
proposed. Among these efficient transformers, sparse attention and local
attention are one of the main streams, which are highly related to our work.
Sparse attention can be further categorized into data independent (fixed)
sparse attention [9, 30, 1, 73] and content-based sparse attention [13, 46,
35, 49]. Local attention [43, 41, 42] mainly considers attending only to a
local window size. Our work is also content-based attention, but compared with
previous works [13, 46, 35, 49], our $k$-NN attention has its merits for
vision domain. For example, compared with routing transformer [46] that
clusters both queries and keys, our $k$-NN attention equals only clustering
keys by assigning each query as the cluster center, making the quantization
more continuous which is a better fitting of image domain; compared with
reformer [35] which adopts complex hashing attention that cannot guarantee
each bucket contain both queries and keys, our $k$-NN attention can guarantee
that each query has number $k$ keys for attention computing. In addition, our
$k$-NN attention is also a generalized local attention, but compared with
local attention, our $k$-NN attention not only enjoys the locality but also
empowers the ability of global relation mining.

###  2.2 Transformer for Vision

Transformer [54] is an effective sequence-to-sequence modeling network, and it
has achieved state-of-the-art results in NLP tasks with the success of BERT
[15]. Due to its great success, it has also be exploited in computer vision
community, and ‘Transformer in CNN’ becomes a popular paradigm [2, 59, 7, 81,
37, 38, 26, 4]. ViT [16] leads the other trend to use ‘CNN in Transformer’
paradigm for vision tasks [28, 36, 69, 63, 67]. Even though ViT has been
proved compelling in vision recognition, it has several drawbacks when
compared with CNNs: large training data, fixed position embedding, rigid patch
division, coarse modeling of inner patch feature, single scale, unstable
training process, slow speed training, easily fitting data and poor
generalization, shallow & narrow architecture, and quadratic complexity. To
deal with these problems, many variants have been proposed [75, 33, 58, 20,
31, 21, 45, 68, 79, 17, 57, 65, 80]. For example, DeiT [52] adopts several
training techniques and uses distillation to extend ViT to a data-efficient
version; CPVT [11] proposes a conditional positional encoding that is
adaptable to arbitrary input sizes; CvT [60], CoaT [64] and Visformer [8]
safely remove the position embedding by introducing convolution operations;
T2T ViT [71], CeiT [70], and CvT [60] try to deal with the rigid patch
division by introducing convolution operation for patch sequence generation;
Focal Transformer [66] makes each token attend its closest surrounding tokens
at fine granularity and the tokens far away at coarse granularity; TNT [25]
proposes the pixel embedding to model the inner patch feature; PVT [56], Swin
Transformer [42], MViT [18], ViL [74], CvT [60], PiT [29], LeViT [23], CoaT
[64], and Twins [10] adopt multi-scale technique for rich feature learning;
DeepViT [78], CaiT [53], and PatchViT [22] investigate the unstable training
problem, and propose the re-attention, re-scale and anti-over-smoothing
techniques respectively for stable training; to accelerate the convergence of
training, ConViT [14], PiT [29], CeiT [70], LocalViT [39] and Visformer [8]
introduce convolutional bias to speedup the training; conv-stem is adopted in
LeViT [23], EarlyConv [61], CMT [24], VOLO [72] and ScaledReLU [55] to improve
the robustness of training ViTs; LV-ViT [32] adopts several techniques
including MixToken and Token Labeling for better training and feature
generation; T2T ViT [71], DeepViT [78] and CaiT [53] try to train deeper
vision transformer models; T2T ViT [71], ViL [74] and CoaT [64] adopt
efficient transformers [51] to deal with the quadratic complexity; To further
exploit the capacities of vision transformer, OmniNet [50], CrossViT [6] and
So-ViT [62] propose the dense omnidirectional representations, coarse-fine-
grained patch fusion and cross co-variance pooling of visual tokens,
respectively. However, all of these works adopt the fully-connected self-
attention which will bring the noise or irrelevant tokens for computing and
slow down the training of networks. In this paper, we propose an efficient
sparse attention, called $k$-NN attention, for boosting vision transformers.
The proposed $k$-NN attention not only inherits the local bias of CNNs but
also achieves the ability of global feature exploitation. It can also speed up
the training and achieve better performance.

##  3 $k$-NN Attention

###  3.1 Vanilla Attention

For any sequence of length $n$, the vanilla attention in the transformer is
the dot product attention [54]. Following the standard notation, the attention
matrix $\bm{A}\in\mathbbm{R}^{n\times n}$ is defined as:

$\bm{A}=\textrm{softmax}\left(\frac{\bm{Q}\bm{K}^{\top}}{\sqrt{d}}\right),$

where $\bm{Q}\in\mathbbm{R}^{n\times d}$ denotes the queries while
$\bm{K}\in\mathbbm{R}^{n\times d}$ denotes the keys, and $d$ represents the
dimension. By multiplying the attention weights $\bm{A}$ with the values
$\bm{V}\in\mathbbm{R}^{n\times d}$, the new values $\hat{\bm{V}}$ are
calculated as:

$\hat{\bm{V}}=\bm{A}\bm{V}.$

The intuitive understanding of the attention is the weighted average over the
old ones, where the weights are defined by the attention matrix $\bm{A}$. In
this paper, we consider the $\bm{Q}$, $\bm{K}$ and $\bm{V}$ are generated via
the linear projection of the input token matrix $\bm{X}$:

$\bm{Q}=\bm{X}\bm{W}_{\bm{Q}},\ \bm{K}=\bm{X}\bm{W}_{\bm{K}},\
\bm{V}=\bm{X}\bm{W}_{V},$

where $\bm{X}\in\mathbbm{R}^{n\times d_{m}}$,
$\bm{W}_{\bm{Q}},\bm{W}_{\bm{K}},\bm{W}_{V}\in\mathbbm{R}^{d_{m}\times d}$ and
$d_{m}$ is the input token dimension.

One shortcoming with fully-connected self-attention is that irrelevant tokens,
even though assigned with smaller weights, are still taken into consideration
when updating the representation $\bm{V}$, making it less resilient to noises
in $\bm{V}$. This shortcoming motivates us to develop the $k$-NN attention.

###  3.2 $k$-NN Attention

Instead of computing the attention matrix for all the query-key pairs as in
vanilla attention, we select the top-$k$ most similar keys and values for each
query in the $k$-NN attention. There are two versions of $k$-NN attention, as
described below.

Slow Version: For the $i$-th query, we first compute the Euclidean distance
against all the keys, and then obtain its $k$-nearest neighbors
$\mathcal{N}_{i}^{k}$ and $\mathcal{N}_{i}^{v}$ from keys and values, and
lastly calculate the scaled dot product attention as:

$\bm{A}_{i}=\textrm{softmax}\left(\frac{\langle\bm{q}_{i},(\bm{k}_{j_{1}},...,\bm{k}_{j_{l}},...,\bm{k}_{j_{k}})\rangle}{\sqrt{d}}\right),\bm{k}_{j_{l}}\in\mathcal{N}_{i}^{k}.$

The shape of final attention matrix is $\bm{A}^{knn}\in{\mathbbm{R}^{n\times
k}}$, and the new values $\hat{\bm{V}}^{knn}$ is the same size of values
$\hat{\bm{V}}$. The slow version is the exact definition of $k$-NN attention,
but it is extremely slow because for each query it has to compute distances
for different $k$ keys.

Fast Version: As the computation of Euclidean distance against all the keys
for each query is slow, we propose a fast version of $k$-NN attention. The key
idea is to take advantage of matrix multiplication operations. Same as vanilla
attention, all the queries and keys are calculated by the dot product, and
then row-wise top-$k$ elements are selected for $\mathrm{softmax}$ computing.
The procedure can be formulated as:

$\hat{\bm{V}}^{knn}=\textrm{softmax}\left(\mathcal{T}_{k}\left(\frac{\bm{Q}\bm{K}^{\top}}{\sqrt{d}}\right)\right)\cdot\bm{V},$

where $\mathcal{T}_{k}\left(\cdot\right)$ denotes the row-wise top-k selection
operator:

$\left[\mathcal{T}_{k}(\bm{A})\right]_{ij}=\begin{cases}A_{ij}\quad&{A_{ij}\in\textrm{top-k}(\textrm{row
$j$})}\\\ -\infty\quad&\textrm{ otherwise.}\end{cases}$

###  3.3 Theoretical Analysis on $k$-NN Attention

In this section, we will show theoretically that despite its simplicity,
$k$-NN attention is powerful in speeding up network training and in distilling
noisy tokens. All the proof of the lemmas are provided in the supplementary.

Convergence Speed-up. Compared to CNNs, the fully-connected self-attention is
able to capture long range dependency. However, the price to pay is that the
dense self-attention model requires to mix each image patch with every other
patch in the image, which has potential to mix irrelevant information
together, e.g. the foreground patches may be mixed with background patches
through the self-attention. This defect could significantly slow down the
convergence as the goal of visual object recognition is to identify key visual
patches relevant to a given class.

To see this, we consider the model with only learnable parameters
$\bm{W}_{\bm{Q}}$, $\bm{W}_{\bm{K}}$ in attention layers and adopting Adam
optimizer [34]. According to Theorem 4.1 in [34], Adam’s convergence is
proportional to $\mathcal{O}\left(\alpha^{-1}(G_{\infty}+1)+\alpha
G_{\infty}\right)$, where $\alpha$ is the learning rate and $G_{\infty}$ is an
element-wise upper bound on the magnitude of the batch gradient111 Theorem 4.1
in [34] describes the upper bound for regrets (the gap on loss function value
between the current step parameters and optimal parameters). One can telescope
it to the average regrets to consider the Adam’s convergence.. Let $f_{i}$ be
the loss function corresponding to batch $i$. Via chain rule of derivative,
the gradient w.r.t the $\bm{W}_{Q}$ in a self-attention block can be
represented as
$\nabla_{\bm{W}_{\bm{Q}}}f_{i}=F_{i}(\hat{\bm{V}}^{knn})\cdot\frac{\partial\hat{\bm{V}}^{knn}}{\partial\bm{W}_{\bm{Q}}}$,
where $F_{i}(\hat{\bm{V}}^{knn})$ is a matrix output function. Since the
possible value of $\hat{\bm{V}}^{knn}$ is a subset of its fully-connected
counterpart, the upper bound of on the magnitude of
$F_{i}(\hat{\bm{V}}^{knn})$ is no larger than the full attention. We then
introduce the weighed covariate matrix of patches to characterize the scale of
$\frac{\partial\hat{\bm{V}}^{knn}}{\partial\bm{W}_{\bm{Q}}}$ in the following
lemma.

###### Lemma 1

(Informal) Let $\hat{\bm{V}}^{knn}_{l}$ be the $l$-th row of the
$\hat{\bm{V}}^{knn}$. We have

$\frac{\partial\hat{\bm{V}}_{l}^{knn}}{\partial\bm{W}_{Q}}\propto\textrm{Var}_{\bm{a}_{l}}(\bm{x})\textrm{
and
}\frac{\partial\hat{\bm{V}}_{l}^{knn}}{\partial\bm{W}_{K}}\propto\textrm{Var}_{\bm{a}_{l}}(\bm{x}),$

where $\textrm{Var}_{\bm{a}_{l}}(\bm{x})$ is the covariate matrix on patches
$\\{\bm{x}_{1},...,\bm{x}_{n}\\}$ with probability from $l$-th row of the
attention matrix.
The same is true for $\hat{\bm{V}}$ of the fully-connected self-attention.

Since $k$-NN attention only uses patches with large similarity, its
$\textrm{Var}_{\bm{a}_{l}}(\bm{x})$ will be smaller than that computed from
the fully-connected attention. As indicated in Lemma 1,
$\frac{\partial\hat{\bm{V}}^{knn}}{\partial\bm{W}_{\bm{Q}}}$ is proportional
to variance $\textrm{Var}_{\bm{a}_{l}}(\bm{x})$ and thus the scale of
$\nabla_{\bm{W}_{\bm{Q}}}f_{i}$ becomes smaller in k-NN attention. Similarly,
the scale of $\nabla_{\bm{W}_{\bm{K}}}f_{i}$ is also smaller in k-NN
attention. Therefore, the element-wise upper bound on batch gradient
$G_{\infty}$ in Adam analysis is also smaller for k-NN attention. For the same
learning rate, the k-NN attention yields faster convergence. It is
particularly significant at the beginning of training. This is because, due to
the random initialization, we expect a relatively small difference in
similarities between patches, which essentially makes self-attention behave
like “global average". It will take multiple iterations for Adam to turn the
"global average" into the real function of self-attention. In Table 2 and
Figure 2, we numerically verify the training efficiency of $k$-NN attention as
opposed to the fully-connected attention.

Noisy patch distillation. As already mentioned before, the fully-connected
self-attention model may mix irrelevant patches with relevant ones,
particularly at the beginning of training when similarities between relevant
patches are not significantly larger than those for irrelevant patches. $k$-NN
attention is more effective in identifying noisy patches by only considering
the top $k$ most similar patches. To formally justify this point, we consider
a simple scenario where all the patches are divided into two groups, the group
of relevant patches and the group of noisy patches. All the patches are
sampled independently from unknown distributions. We assume that all relevant
patches are sampled from distributions with the same shared mean, which is
different from the means of distributions for noisy patches. It is important
to know that although distributions for the relevant patches share the mean,
those relevant patches can look quite differently, due to the large variance
in stochastic sampling. In the following Lemma, we will show that the $k$-NN
attention is more effective in distilling noises for the relevant patches than
the fully-connected attention.

######  Lemma 2 (informal)

We consider the self-attention for query patch $l$. Let’s assume the patch
$\bm{x}_{i}$ are bounded with mean $\bm{\mu}_{i}$ for $i=1,2,...,n$, and
$\rho_{k}$ is the ratio of the noisy patches in all selected patches. Under
mild conditions, the follow inequality holds with high probability:

$\left\|\hat{\bm{V}}_{l}^{knn}-\bm{\mu}_{l}\bm{W}_{V}\right\|_{\infty}\leq\mathcal{O}(k^{-1/2}+c_{1}\rho_{k}),$

where $c_{1}$ is a positive number.

In the above lemma, the quantity
$\left\|\hat{\bm{V}}_{l}^{knn}-\bm{\mu}_{l}\bm{W}_{V}\right\|_{\infty}$
measures the distance between $\hat{\bm{V}}^{knn}_{l}$, represention vector
updated by the $k$-NN attention, and its mean $\mu_{l}\bm{W}_{\bm{V}}$. We now
consider two cases: the normal $k$-NN attention with appropriately chosen $k$,
and fully-connected attention with $k=n$. In the first case, with
appropriately chosen $k$, we should have most of the selected patches coming
from the relevant group, implying a small $\rho_{k}$. By combining with the
fact that $k$ is decently large, we expect a small upper bound for the
distance
$\left\|\hat{\bm{V}}_{l}^{knn}-\bm{\mu}_{l}\bm{W}_{V}\right\|_{\infty}$,
indicating that $k$-NN attention is powerful in distilling noise. For the case
of fully-connected attention model, i.e. $k=n$, it is clearly that
$\rho_{n}\approx 1$, leading to a large distance between transformed
representation $\hat{\bm{V}}_{l}$ and its mean, indicating that fully-
connected attention model is not very effective in distilling noisy patches,
particularly when noise is large.

Besides the instance with low signal-noise-ratio, the instance with a large
volume of backgrounds can also be hard. In the next lemma, we show that under
a proper choice of $k$, with a high probability the $k$-NN attention will be
able to select all meaningful patches.

######  Lemma 3 (informal)

Let $\mathcal{M}^{*}$ be the index set contains all patches relevant to query
$\bm{q}_{l}$. Under mild conditions, there exist $c_{2}\in(0,1)$ such that
with high probability, we have

$\sum_{i=1}^{n}\mathbbm{1}(\bm{q}_{l}\bm{k}_{i}^{\top}\geq\min_{j\in\mathcal{M}^{*}}\bm{q}_{l}\bm{k}_{j}^{\top})\leq\mathcal{O}(nd^{-c_{2}}).$

The above lemma shows that if we select the top $\mathcal{O}(nd^{-c_{2}})$
elements, with high probability, we will be able to eliminate almost all the
irrelevant noisy patches, without losing any relevant patches. Numerically, we
verify the proper $k$ gains better performance (e.g., Figure 1) and for the
hard instance $k$-NN gives more accurate attention regions. (e.g., Figure 4
and Figure 5).

##  4 Experiments for Vision Transformers

In this section, we replace the dense attention with $k$-NN attention on the
existing vision transformers for image classification to verify the
effectiveness of the proposed method. The recent DeiT [52] and its variants,
including T2T ViT [71], TNT [25], PiT [29], Swin [42], CvT [60], So-ViT [62],
Visformer [8], Twins [10], Dino [3] and VOLO [72], are adopted for evaluation.
These methods include both supervised methods [52, 71, 25, 29, 42, 60, 62, 8,
10, 72] and self-supervised method [3]. Ablation studies are provided to
further analyze the properties of $k$-NN attention.

###  4.1 Experimental Settings

We perform image classification on the standard ILSVRC-2012 ImageNet dataset
[47]. In our experiments, we follow the experimental setting of original
official released codes. For fair comparison, we only replace the vanilla
attention with proposed k-NN attention. Unless otherwise specified, the fast
version of $k$-NN attention is adopted for evaluation. To speed up the slow
version, we develop the CUDA version $k$-NN attention. As for the value $k$,
different architectures are assigned with different values. For DeiT [52], So-
ViT [62], Dino [3], CvT [60], TNT [25] PiT [29] and VOLO [72], as they
directly split an input image into rigid tokens and there is no information
exchange in the token generation stage, we suppose the irrelevant tokens are
easy to filter, and tend to assign a smaller $k$ compared with these
complicated token generation methods [71, 42, 8, 10]. Specifically, we assign
$k$ to approximate $\frac{n}{2}$ at each scale stage; for these complicated
token generation methods [71, 42, 8, 10], we assign a larger $k$ which is
approximately $\frac{2}{3}{n}$ or $\frac{4}{5}{n}$ at each scale stage.

###  4.2 Results on ImageNet

Table 1 shows top-$1$ accuracy results on the ImageNet-1K validation set by
replacing the dense attention with $k$-NN attention using eleven different
vision transformer architectures. From the Table we can see that the proposed
$k$-NN attention improves the performance from 0.2% to 0.8% for both global
and local vision transformers. It is worth noting that on ImageNet-1k dataset,
it is very hard to improve the accuracy after 85%, but our $k$-NN attention
can still consistently improve the performance even without model size
increase.

Table 1: The $k$-NN attention performance on ImageNet-1K validation set. "!" means we pretrain the model with 300 epochs and finetune the pretrained model for 100 epoch for linear eval, following the instructions of Dino training and evaluation; "$\rightarrow$ $k$-NN Attn" represents replacing the vanilla attention with proposed $k$-NN attention;$\rightarrow$ $k$-NN Attn-slow means adopting the slow version.  Arch. | Model | Input | Params | GFLOPs | Top-1
---|---|---|---|---|---
Transformers | DeiT-Tiny [52] | $224^{2}$ | 5.7M | 1.3 | 72.2%
(Supervised) | DeiT-Tiny [52] $\rightarrow$ k-NN Attn | $224^{2}$ | 5.7M | 1.3 | 73.0%
| DeiT-Tiny [52] $\rightarrow$ k-NN Attn-slow | $224^{2}$ | 5.7M | 1.3 | 73.0%
| So-ViT-7 [62] | $224^{2}$ | 5.5M | 1.3 | 76.2%
| So-ViT-7 [62] $\rightarrow$ k-NN Attn | $224^{2}$ | 5.5M | 1.3 | 77.0%
Transformers | Visformer-Tiny [8] | $224^{2}$ | 10M | 1.3 | 78.6%
(Supervised) | Visformer-Tiny [8] $\rightarrow$ k-NN Attn | $224^{2}$ | 10M | 1.3 | 79.0%
Transformers | CvT-13 [60] | $224^{2}$ | 20M | 4.6 | 81.6%
(Supervised) | CvT-13 [60] $\rightarrow$ k-NN Attn | $224^{2}$ | 20M | 4.6 | 81.9%
| DeiT-Small [52] | $224^{2}$ | 22M | 4.6 | 79.8%
| DeiT-Small [52] $\rightarrow$ k-NN Attn | $224^{2}$ | 22M | 4.6 | 80.1%
| TNT-Small [25] | $224^{2}$ | 24M | 5.2 | 81.5%
| TNT-Small [25] $\rightarrow$ k-NN Attn | $224^{2}$ | 24M | 5.2 | 81.9%
| VOLO-D1 [72] | $384^{2}$ | 27M | 22.8 | 85.2%
| VOLO-D1 [72] $\rightarrow$ k-NN Attn | $384^{2}$ | 27M | 22.8 | 85.4%
| Swin-Tiny [42] | $224^{2}$ | 28M | 4.5 | 81.2%
| Swin-Tiny [42] $\rightarrow$ k-NN Attn | $224^{2}$ | 28M | 4.5 | 81.3%
| T2T-ViT-t-19 [71] | $224^{2}$ | 39M | 9.8 | 82.2%
| T2T-ViT-t-19 [71] $\rightarrow$ k-NN Attn | $224^{2}$ | 39M | 9.8 | 82.7%
Transformer | Dino-Small [3]! | $224^{2}$ | 22M | 4.6 | 76.0%
(Self-supervised) | Dino-Small [3]! $\rightarrow$ k-NN Attn | $224^{2}$ | 22M | 4.6 | 76.2%
Transformers | Twins-SVT-Base [10] | $224^{2}$ | 56M | 8.3 | 83.2%
(Supervised) | Twins-SVT-Base [10] $\rightarrow$ k-NN Attn | $224^{2}$ | 56M | 8.3 | 83.4%
| PiT-Base[29] | $224^{2}$ | 74M | 12.5 | 82.0%
| PiT-Base[29] $\rightarrow$ k-NN Attn | $224^{2}$ | 74M | 12.5 | 82.6%
| VOLO-D3 [72] | $448^{2}$ | 86M | 67.9 | 86.3%
| VOLO-D3 [72] $\rightarrow$ k-NN Attn | $448^{2}$ | 86M | 67.9 | 86.5%

###  4.3 The Impact of Number $k$

The only parameter for $k$-NN attention is $k$, and its impact is analyzed in
Figure 1. As shown in the figure, for DeiT-Tiny, $k$ = 100 is the best, where
the total number of tokens $n$ = 196 (14 $\times$ 14), meaning that $k$
approximates half of $n$; for CvT-13, there are three scale stages with the
number of tokens $n_{1}$ = 3136, $n_{2}$ = 784 and $n_{3}$ = 196, and the best
results are achieved when the $k$ in each stage is assigned to 1600/400/100,
which also approximate half of $n$ in each stage; for Visformer-Tiny, there
are two scale stages with the number of tokens $n_{1}$ = 196 and $n_{2}$ = 49,
and the best results are achieved when $k$ in each stage is assigned to
150/45, as there are more than 21 conv layers for token generation and the
information in each token are already mixed, making it hard to distinguish the
irrelevant tokens, thus larger values of $k$ are desired; for PiT-Base, there
are three scale stages with the number of tokens $n_{1}$ = 961, $n_{2}$ = 256
and $n_{3}$ = 64, and the optimal values of $k$ also approximate the half of
$n$. Please note that, we do not perform exhaustive search for the optimal
choice of $k$, instead, a general rule as below is sufficient: $k\approx$
$\frac{n}{2}$ at each scale stage for simple token generation methods and
$k\approx\frac{2}{3}{n}$ or $\frac{4}{5}{n}$ for complicated token generation
methods at each scale stage.

(a) DeiT-Tiny

(c) Visformer-Tiny

(b) CvT-13

(d) PiT-Base

Figure 1: The impact of $k$ on DeiT-Tiny, Visformer-Tiny, CvT-13 and PiT-Base.

###  4.4 Convergence Speed of $k$-NN Attention

In Table 2, we investigate the convergence speed of $k$-NN attention. Three
methods are included for comparison, i.e. DeiT-Small [52], CvT-13 [60] and
T2T-ViT-t-19 [71]. From the Table we can see that the convergence speed of
$k$-NN attention is faster than full-connected attention, especially in the
early stage of training. These observations reflect that removing the
irrelevant tokens benefits the convergence of neural networks training.

Table 2: Ablation study on the convergence speed of $k$-NN attention. Epoch | Top-1 accuracy
---|---
DeiT-S | DeiT-S $\rightarrow$ k | CvT-13 | CvT-13 $\rightarrow$ k | T2T-ViT-t-19 | T2T-ViT-t-19 $\rightarrow$ k
10 | 29.1% | 31.3% | 51.4% | 54.2% | 0.52% | 0.68%
30 | 54.4% | 55.4% | 65.4% | 68.1% | 63.0% | 63.2%
50 | 60.9% | 62.0% | 68.1% | 70.5% | 73.8% | 74.4%
70 | 65.0% | 65.8% | 69.9% | 72.2% | 76.9% | 77.3%
90 | 67.7% | 68.2% | 71.0% | 73.0% | 78.4% | 78.6%
120 | 69.9% | 70.7% | 72.4% | 73.7% | 79.7% | 80.0%
150 | 72.4% | 72.4% | 74.4% | 74.9% | 80.7% | 80.9%
200 | 75.5% | 75.7% | 77.3% | 77.7% | 82.0% | 82.3%
300 | 79.8% | 80.0% | 81.6% | 81.9% | 81.3% | 81.7%

###  4.5 Other properties of $k$-NN attention

To analyze other properties of $k$-NN attention, four quantitative metrics are
defined as follows.
Layer-wise cosine similarity between tokens: following [22] this metric is
defined as:

$\textrm{CosSim({t})}=\frac{1}{n(n-1)}\sum_{i\neq
j}\frac{t^{T}_{i}t_{j}}{\lVert t_{i}\rVert\lVert t_{j}\rVert},$

where $t_{i}$ represents the $i$-th token in each layer and
$\lVert\cdot\rVert$ denotes the Euclidean norm. This metric implies the
convergence speed of the network.

Layer-wise standard deviation of attention weights: Given a token $t_{i}$ and
its softmax attention weight sfm($t_{i}$), the standard deviation of the
softmax attention weight std(sfm($t_{i}$)) is defined as the second metric.
For multi-head attention, the standard deviations over all heads are averaged.
This metric represents the degree of training stability.

Ratio between the norms of residual activations and main branch: The ratio
between the norm of the residual activations and the norm of the activations
of the main branch in each layer is defined as $\lVert f_{l}(t)\rVert/\lVert
t\rVert$, where $f_{l}(t)$ can be the attention layer or the FFN layer. This
metric denotes the information preservation ability of the network.

Nonlocality: following [14], the nonlocality is defined by summing, for each
query patch $i$, the distances $\left\|\delta_{ij}\right\|$ to all the key
patches $j$ weighted by their attention score $\bm{A}_{ij}$. The number
obtained over the query patch is averaged to obtain the nonlocality metric of
head $h$, which can the be averaged over the attention heads to obtain the
nonlocality of the whole layer $l$:

$D_{loc}^{l,h}:=\frac{1}{L}\sum_{ij}\bm{A}_{ij}^{h,l}\left\|\delta_{ij}\right\|,D_{loc}^{l}:=\frac{1}{N_{h}}\sum_{h}D_{loc}^{l,h},$

where $D_{loc}$ is the number of patches between the center of attention and
the query patch; the further the attention heads look from the query patch,
the higher the nonlocality.

(a) Layer-wise cosine similarity of tokens

(b) Layer-wise $s.t.d$ of attention weights

(c) Ratio of residual and main branch for attn

(d) Ratio of residual and main branch for ffn

Figure 2: The properties of $k$-NN attention. Blue and red dotted lines
represent the metrics for $k$-NN attetion and the original fully-connected
self-attention, respectively.

Comparisons of the four metrics on DeiT-tiny without distillation token are
shown in Figure 2 and Figure 3. From Figure 2 (a) we can see that by using
$k$-NN attention, the averaged cosine similarity is larger than that of using
dense self-attention, which reflects that the convergence speed is faster for
$k$-NN attention. Figure 2 (b) shows that the averaged standard deviation of
$k$-NN attention is smoother than that of fully-connected self-attention, and
the smoothness will help make the training more stable. Figure 2 (c) and (d)
show the ratio between the norms of residual activations and main branch are
consistent with each other for $k$-NN attention and dense attention, which
indicates that there is nearly no information lost in $k$-NN attention by
removing the irrelevant tokens. Figure 3 shows that, with k-NN attention,
lower layers tend to focus more on the local areas (with more lines being
pushed toward the bottom area in Figure 3), while the higher layers still
maintain their capability of extracting global information. Additionally, it
is also observed that the non-locality of different layers is spreading more
evenly, indicating that they can explore a larger variety of dependencies at
different ranges.

Figure 3: The nonlocality of DeiT-Tiny. It is plotted averaged over all the
images from training set of ImageNet-1k.

###  4.6 Comparisons with temperature in softmax

$k$-NN attention effectively zeros the bottom $N-k$ tokens out of the
attention calculation. How does this compare with introducing a temperature
parameter to softmax over the attention values? We compare our $k$-NN
attention with temperature $t$ in softmax as softmax(attn/$t$). The
performance over the $t$ is shown in Table 3. From the Table we can see that
small $t$ makes the training crash due to large value of attention values; the
performance increases a little bit to 72.5 (baseline 72.2) with $t$ assigned
to appropriate values. The $k$-NN attention is more robust compared with
temperature in softmax, and achieves much better performance, 73.0 ($k$-NN
attention) vs 72.5 (best performance for temperature in softmax).

Table 3: The Top-$1$ (%) over the temperature $t$ in softmax. $t$ | 0.05 | 0.1 | 0.25 | 0.75 | 2 | 4 | 8 | 16
---|---|---|---|---|---|---|---|---
Top-$1$ (%) | crash | crash | 72.0 | 72.5 | 72.5 | 72.5 | 72.5 | 72.1

dense

$k$-NN

input

head0

head1

head2

head3

head4

head5

Figure 4: Self-attention heads from the last layer.

input

dense attn

$k$-NN attn

(a)dog

(b)wheel

(c)ferret

(d)monitor

(e)hornbill

(f)chain

(g)crib

Figure 5: Visualization using Transformer Attribution [5].

###  4.7 Visualization

Figure 4 visualizes the self-attention heads from the last layer on Dino-Small
[3]. We can see that different heads attend to different semantic regions of
an image. Compared with dense attention, the $k$-NN attention filters out most
irrelevant information from background regions which are similar to the
foreground, and successfully concentrates on the most informative foreground
regions. Images from different classes are visualized in Figure 5 using
Transformer Attribution method [5] on DeiT-Tiny. It can be seen that the
$k$-NN attention is more concentrated and accurate, especially in the
situations of cluttered background and occlusion.

Table 4: Object detection and Segmentation results for Swin-Tiny and Twins-SVT-Base with/without $k$-NN attention on the COCO and ADE20K validation sets. All the models are pretrained on ImageNet-1k.  | COCO | ADE20K
---|---|---
Backbone | Method | mAP(box) | Method | mIoU
Swin-T | Mask R-CNN 3x | 46.0 | UPerNet | 44.5
Swin-T-$k$-NN | Mask R-CNN 3x | 46.2 | UPerNet | 44.7
Twins-SVT-Base | Mask R-CNN 1x | 45.2 | UPerNet | 47.4
Twins-SVT-Base-$k$-NN | Mask R-CNN 1x | 45.6 | UPerNet | 47.9

###  4.8 Object Detection and Semantic Segmentation

To verify the effects of $k$-NN attention on object detection and semantic
segmentation tasks, the widely-used COCO [40] and ADE20K [77] are adopted for
evaluation. We adopt Swin-Tiny [42] and Twins-SVT-Base [10] for comparisons
due to the well released codes, and the results are shown in Table 4. From the
Table we can see that by replacing the vanilla attention with our $k$-NN
attention, the performance increases with almost no overhead.

##  5 Conclusion

In this paper, we propose an effective $k$-NN attention for boosting vision
transformers. By selecting the most similar keys for each query to calculate
the attention, it screens out the most ineffective tokens. The removal of
irrelevant tokens speeds up the training. We theoretically prove its
properties in speeding up training, distilling noises without losing
information, and increasing the performance by choosing a proper $k$. Several
vision transformers are adopted to verify the effectiveness of the $k$-NN
attention.

## References

  * [1] Beltagy, I., Peters, M.E., Cohan, A.: Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150 (2020)
  * [2] Carion, N., Massa, F., Synnaeve, G., Usunier, N., Kirillov, A., Zagoruyko, S.: End-to-end object detection with transformers. In: ECCV (2020)
  * [3] Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., Joulin, A.: Emerging properties in self-supervised vision transformers. arXiv preprint arXiv:2104.14294 (2021)
  * [4] Chang, S., Wang, P., Wang, F., Li, H., Feng, J.: Augmented transformer with adaptive graph for temporal action proposal generation. arXiv preprint arXiv:2103.16024 (2021)
  * [5] Chefer, H., Gur, S., Wolf, L.: Transformer interpretability beyond attention visualization. In: CVPR (2021)
  * [6] Chen, C.F., Fan, Q., Panda, R.: Crossvit: Cross-attention multi-scale vision transformer for image classification. arXiv preprint arXiv:2103.14899 (2021)
  * [7] Chen, X., Yan, B., Zhu, J., Wang, D., Yang, X., Lu, H.: Transformer tracking. In: CVPR (2021)
  * [8] Chen, Z., Xie, L., Niu, J., Liu, X., Wei, L., Tian, Q.: Visformer: The vision-friendly transformer. arXiv preprint arXiv:2104.12533 (2021)
  * [9] Child, R., Gray, S., Radford, A., Sutskever, I.: Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 (2019)
  * [10] Chu, X., Tian, Z., Wang, Y., Zhang, B., Ren, H., Wei, X., Xia, H., Shen, C.: Twins: Revisiting spatial attention design in vision transformers. arXiv preprint arXiv:2104.13840 (2021)
  * [11] Chu, X., Zhang, B., Tian, Z., Wei, X., Xia, H.: Do we really need explicit position encodings for vision transformers? arXiv preprint arXiv:2102.10882 (2021)
  * [12] Cordonnier, J.B., Loukas, A., Jaggi, M.: On the relationship between self-attention and convolutional layers. In: ICLR (2019)
  * [13] Correia, G.M., Niculae, V., Martins, A.F.: Adaptively sparse transformers. In: EMNLP. pp. 2174–2184 (2019)
  * [14] d’Ascoli, S., Touvron, H., Leavitt, M., Morcos, A., Biroli, G., Sagun, L.: Convit: Improving vision transformers with soft convolutional inductive biases. arXiv preprint arXiv:2103.10697 (2021)
  * [15] Devlin, J., Chang, M.W., Lee, K., Toutanova, K.: Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 (2018)
  * [16] Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al.: An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 (2020)
  * [17] El-Nouby, A., Touvron, H., Caron, M., Bojanowski, P., Douze, M., Joulin, A., Laptev, I., Neverova, N., Synnaeve, G., Verbeek, J., et al.: Xcit: Cross-covariance image transformers. arXiv preprint arXiv:2106.09681 (2021)
  * [18] Fan, H., Xiong, B., Mangalam, K., Li, Y., Yan, Z., Malik, J., Feichtenhofer, C.: Multiscale vision transformers. arXiv preprint arXiv:2104.11227 (2021)
  * [19] Fan, J., Lv, J.: Sure independence screening for ultrahigh dimensional feature space. Journal of the Royal Statistical Society: Series B (Statistical Methodology) 70(5), 849–911 (2008)
  * [20] Fang, J., Xie, L., Wang, X., Zhang, X., Liu, W., Tian, Q.: Msg-transformer: Exchanging local spatial information by manipulating messenger tokens. arXiv preprint arXiv:2105.15168 (2021)
  * [21] Gao, P., Lu, J., Li, H., Mottaghi, R., Kembhavi, A.: Container: Context aggregation network. arXiv preprint arXiv:2106.01401 (2021)
  * [22] Gong, C., Wang, D., Li, M., Chandra, V., Liu, Q.: Improve vision transformers training by suppressing over-smoothing. arXiv preprint arXiv:2104.12753 (2021)
  * [23] Graham, B., El-Nouby, A., Touvron, H., Stock, P., Joulin, A., Jégou, H., Douze, M.: Levit: a vision transformer in convnet’s clothing for faster inference. arXiv preprint arXiv:2104.01136 (2021)
  * [24] Guo, J., Han, K., Wu, H., Xu, C., Tang, Y., Xu, C., Wang, Y.: Cmt: Convolutional neural networks meet vision transformers. arXiv preprint arXiv:2107.06263 (2021)
  * [25] Han, K., Xiao, A., Wu, E., Guo, J., Xu, C., Wang, Y.: Transformer in transformer. arXiv preprint arXiv:2103.00112 (2021)
  * [26] Han, L., Wang, P., Yin, Z., Wang, F., Li, H.: Exploiting better feature aggregation for video object detection. In: ACM MM (2020)
  * [27] He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: CVPR (2016)
  * [28] He, S., Luo, H., Wang, P., Wang, F., Li, H., Jiang, W.: Transreid: Transformer-based object re-identification. arXiv preprint arXiv:2102.04378 (2021)
  * [29] Heo, B., Yun, S., Han, D., Chun, S., Choe, J., Oh, S.J.: Rethinking spatial dimensions of vision transformers. arXiv preprint arXiv:2103.16302 (2021)
  * [30] Ho, J., Kalchbrenner, N., Weissenborn, D., Salimans, T.: Axial attention in multidimensional transformers. arXiv preprint arXiv:1912.12180 (2019)
  * [31] Huang, Z., Ben, Y., Luo, G., Cheng, P., Yu, G., Fu, B.: Shuffle transformer: Rethinking spatial shuffle for vision transformer. arXiv preprint arXiv:2106.03650 (2021)
  * [32] Jiang, Z., Hou, Q., Yuan, L., Zhou, D., Jin, X., Wang, A., Feng, J.: Token labeling: Training a 85.5% top-1 accuracy vision transformer with 56m parameters on imagenet. arXiv preprint arXiv:2104.10858 (2021)
  * [33] Jonnalagadda, A., Wang, W., Eckstein, M.P.: Foveater: Foveated transformer for image classification. arXiv preprint arXiv:2105.14173 (2021)
  * [34] Kingma, D.P., Ba, J.: Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 (2014)
  * [35] Kitaev, N., Kaiser, Ł., Levskaya, A.: Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451 (2020)
  * [36] Li, W., Liu, H., Ding, R., Liu, M., Wang, P.: Lifting transformer for 3d human pose estimation in video. arXiv preprint arXiv:2103.14304 (2021)
  * [37] Li, X., Hou, Y., Wang, P., Gao, Z., Xu, M., Li, W.: Transformer guided geometry model for flow-based unsupervised visual odometry. Neural Computing and Applications pp. 1–12 (2021)
  * [38] Li, X., Hou, Y., Wang, P., Gao, Z., Xu, M., Li, W.: Trear: Transformer-based rgb-d egocentric action recognition. arXiv preprint arXiv:2101.03904 (2021)
  * [39] Li, Y., Zhang, K., Cao, J., Timofte, R., Van Gool, L.: Localvit: Bringing locality to vision transformers. arXiv preprint arXiv:2104.05707 (2021)
  * [40] Lin, T.Y., Maire, M., Belongie, S., Hays, J., Perona, P., Ramanan, D., Dollár, P., Zitnick, C.L.: Microsoft coco: Common objects in context. In: European conference on computer vision. pp. 740–755. Springer (2014)
  * [41] Liu, P.J., Saleh, M., Pot, E., Goodrich, B., Sepassi, R., Kaiser, L., Shazeer, N.: Generating wikipedia by summarizing long sequences. In: ICLR (2018)
  * [42] Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., Guo, B.: Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint arXiv:2103.14030 (2021)
  * [43] Luong, M.T., Pham, H., Manning, C.D.: Effective approaches to attention-based neural machine translation. In: EMNLP (2015)
  * [44] Pai, S.: Convolutional neural networks arise from ising models and restricted boltzmann machines
  * [45] Rao, Y., Zhao, W., Liu, B., Lu, J., Zhou, J., Hsieh, C.J.: Dynamicvit: Efficient vision transformers with dynamic token sparsification. arXiv preprint arXiv:2106.02034 (2021)
  * [46] Roy, A., Saffar, M., Vaswani, A., Grangier, D.: Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics 9, 53–68 (2021)
  * [47] Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., et al.: Imagenet large scale visual recognition challenge. IJCV 115(3), 211–252 (2015)
  * [48] Tan, M., Le, Q.: Efficientnet: Rethinking model scaling for convolutional neural networks. In: ICML (2019)
  * [49] Tay, Y., Bahri, D., Yang, L., Metzler, D., Juan, D.C.: Sparse sinkhorn attention. In: ICML (2020)
  * [50] Tay, Y., Dehghani, M., Aribandi, V., Gupta, J., Pham, P., Qin, Z., Bahri, D., Juan, D.C., Metzler, D.: Omninet: Omnidirectional representations from transformers. arXiv preprint arXiv:2103.01075 (2021)
  * [51] Tay, Y., Dehghani, M., Bahri, D., Metzler, D.: Efficient transformers: A survey. arXiv preprint arXiv:2009.06732 (2020)
  * [52] Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., Jégou, H.: Training data-efficient image transformers & distillation through attention. arXiv preprint arXiv:2012.12877 (2020)
  * [53] Touvron, H., Cord, M., Sablayrolles, A., Synnaeve, G., Jégou, H.: Going deeper with image transformers. arXiv preprint arXiv:2103.17239 (2021)
  * [54] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, L., Polosukhin, I.: Attention is all you need. In: NIPS (2017)
  * [55] Wang, P., Wang, X., Luo, H., Zhou, J., Zhou, Z., Wang, F., Li, H., Jin, R.: Scaled relu matters for training vision transformers. arXiv preprint arXiv:2109.03810 (2021)
  * [56] Wang, W., Xie, E., Li, X., Fan, D.P., Song, K., Liang, D., Lu, T., Luo, P., Shao, L.: Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. arXiv preprint arXiv:2102.12122 (2021)
  * [57] Wang, W., Yao, L., Chen, L., Cai, D., He, X., Liu, W.: Crossformer: A versatile vision transformer based on cross-scale attention. arXiv preprint arXiv:2108.00154 (2021)
  * [58] Wang, Y., Huang, R., Song, S., Huang, Z., Huang, G.: Not all images are worth 16x16 words: Dynamic vision transformers with adaptive sequence length. arXiv preprint arXiv:2105.15075 (2021)
  * [59] Wang, Y., Xu, Z., Wang, X., Shen, C., Cheng, B., Shen, H., Xia, H.: End-to-end video instance segmentation with transformers. In: CVPR (2021)
  * [60] Wu, H., Xiao, B., Codella, N., Liu, M., Dai, X., Yuan, L., Zhang, L.: Cvt: Introducing convolutions to vision transformers. arXiv preprint arXiv:2103.15808 (2021)
  * [61] Xiao, T., Singh, M., Mintun, E., Darrell, T., Dollár, P., Girshick, R.: Early convolutions help transformers see better. arXiv preprint arXiv:2106.14881 (2021)
  * [62] Xie, J., Zeng, R., Wang, Q., Zhou, Z., Li, P.: So-vit: Mind visual tokens for vision transformer. arXiv preprint arXiv:2104.10935 (2021)
  * [63] Xu, T., Chen, W., Wang, P., Wang, F., Li, H., Jin, R.: Cdtrans: Cross-domain transformer for unsupervised domain adaptation. arXiv preprint arXiv:2109.06165 (2021)
  * [64] Xu, W., Xu, Y., Chang, T., Tu, Z.: Co-scale conv-attentional image transformers. arXiv preprint arXiv:2104.06399 (2021)
  * [65] Xu, Y., Zhang, Z., Zhang, M., Sheng, K., Li, K., Dong, W., Zhang, L., Xu, C., Sun, X.: Evo-vit: Slow-fast token evolution for dynamic vision transformer. arXiv preprint arXiv:2108.01390 (2021)
  * [66] Yang, J., Li, C., Zhang, P., Dai, X., Xiao, B., Yuan, L., Gao, J.: Focal self-attention for local-global interactions in vision transformers. arXiv preprint arXiv:2107.00641 (2021)
  * [67] Yin, Z., Wang, P., Wang, F., Xu, X., Zhang, H., Li, H., Jin, R.: Transfgu: A top-down approach to fine-grained unsupervised semantic segmentation. In: European conference on computer vision (2022)
  * [68] Yu, Q., Xia, Y., Bai, Y., Lu, Y., Yuille, A., Shen, W.: Glance-and-gaze vision transformer. arXiv preprint arXiv:2106.02277 (2021)
  * [69] Yu, Z., Li, X., Wang, P., Zhao, G.: Transrppg: Remote photoplethysmography transformer for 3d mask face presentation attack detection. arXiv preprint arXiv:2104.07419 (2021)
  * [70] Yuan, K., Guo, S., Liu, Z., Zhou, A., Yu, F., Wu, W.: Incorporating convolution designs into visual transformers. arXiv preprint arXiv:2103.11816 (2021)
  * [71] Yuan, L., Chen, Y., Wang, T., Yu, W., Shi, Y., Tay, F.E., Feng, J., Yan, S.: Tokens-to-token vit: Training vision transformers from scratch on imagenet. arXiv preprint arXiv:2101.11986 (2021)
  * [72] Yuan, L., Hou, Q., Jiang, Z., Feng, J., Yan, S.: Volo: Vision outlooker for visual recognition. arXiv preprint arXiv:2106.13112 (2021)
  * [73] Zaheer, M., Guruganesh, G., Dubey, A., Ainslie, J., Alberti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q., Yang, L., et al.: Big bird: Transformers for longer sequences. arXiv preprint arXiv:2007.14062 (2020)
  * [74] Zhang, P., Dai, X., Yang, J., Xiao, B., Yuan, L., Zhang, L., Gao, J.: Multi-scale vision longformer: A new vision transformer for high-resolution image encoding. arXiv preprint arXiv:2103.15358 (2021)
  * [75] Zhang, Z., Zhang, H., Zhao, L., Chen, T., Pfister, T.: Aggregating nested transformers. arXiv preprint arXiv:2105.12723 (2021)
  * [76] Zhao, G., Lin, J., Zhang, Z., Ren, X., Su, Q., Sun, X.: Explicit sparse transformer: Concentrated attention through explicit selection. arXiv preprint arXiv:1912.11637 (2019)
  * [77] Zhou, B., Zhao, H., Puig, X., Xiao, T., Fidler, S., Barriuso, A., Torralba, A.: Semantic understanding of scenes through the ade20k dataset. IJCV 127(3), 302–321 (2019)
  * [78] Zhou, D., Kang, B., Jin, X., Yang, L., Lian, X., Hou, Q., Feng, J.: Deepvit: Towards deeper vision transformer. arXiv preprint arXiv:2103.11886 (2021)
  * [79] Zhou, D., Shi, Y., Kang, B., Yu, W., Jiang, Z., Li, Y., Jin, X., Hou, Q., Feng, J.: Refiner: Refining self-attention for vision transformers. arXiv preprint arXiv:2106.03714 (2021)
  * [80] Zhou, J., Wang, P., Wang, F., Liu, Q., Li, H., Jin, R.: Elsa: Enhanced local self-attention for vision transformer. arXiv preprint arXiv:2112.12786 (2021)
  * [81] Zou, C., Wang, B., Hu, Y., Liu, J., Wu, Q., Zhao, Y., Li, B., Zhang, C., Zhang, C., Wei, Y., et al.: End-to-end human object interaction detection with hoi transformer. In: CVPR (2021)

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

Anonymous ECCV submission

Pichao WangThe first two authors contribute equally. Xue Wang††footnotemark:
Fan Wang Ming Lin Shuning Chang Hao Li Rong Jin

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

Similarities: Our method part is similar to EST in terms of the calculation of
top-$k$.
Differences:

  * •

Our paper is focused not only on the methodology part, but also the deep
understanding. There are many variants of Transformers in the NLP and vision
community now, but few of them provide a deep and thorough analysis of their
proposed methods. The proposed $k$-NN attention indeed happens to be similar
to EST, which was arxived 2 years ago and we were not aware of it when
conducting our research. In addition to applying the idea to transformers and
conducting extensive experiments as EST did, we provide theoretical
justifications about the idea, which we think is equally or more important
than the method itself, and helps with a more fundamental understanding.

  * •

The conclusion about how to select $k$ is different. In EST, it is found that
a small $k$ is better (8 or 16), but we find a larger $k$, namely,
$\geq\frac{1}{2}N$ is better ($N$ is the sequence length).

  * •

The motivations of these two papers are different: EST targets to get sparse
attention maps while ours aims to distill noisy patches.

  * •

Our paper focuses on vision transformers but EST focuses on NLP tasks, even
though EST applied it to the image captioning task. Since late 2020, vision
transformer backbones have become very popular, and $k$-NN attention deserves
a deeper analysis. Therefore, we apply the $k$-NN attention on 11 different
vision transformer backbones for empirical evaluations and find it simple and
effective for vision transformer backbones.

  * •

More analysis about the properties of $k$-NN attention in the context of
vision transformer backbones are provided in our paper. Besides the $k$
selection and convergence speed as EST presented, we also define several
metrics to facilitate the analysis, e.g. layer-wise cosine similarity between
tokens, layer-wise standard deviation of attention weights, ratio between the
norms of residual activation and main branch, and nonlocality. We also compare
it with temperature in $softmax$ and provide the visualizations.

In summary, our paper provides deeper understanding with comprehensive
analysis of the $k$-NN attention for vision transformers, which provides well-
grounded knowledge advancement.

##  7 Source codes of fast version $k$-NN attention in Pytorch

The source codes of fast version $k$-NN attention in Pytorch are shown in
Algorithm 1, and we can see that the core codes of fast version $k$-NN
attention is consisted of only four lines, and it can be easily imported to
any architecture using fully-connected attention.

Algorithm 1 Codes of fast version $k$-NN attention in Pytorch.

⬇

1class kNN-Attention(nn.Module):

2 def
__init__(self,dim,num_heads=8,qkv_bias=False,qk_scale=None,attn_drop=0.,proj_drop=0.,topk=100):

3 super().__init__()

4 self.num_heads=num_heads

5 head_dim=dim//num_heads

6 self.scale=qk_scale or head_dim**-0.5

7 self.topk=topk

8

9 self.qkv=nn.Linear(dim,dim*3,bias=qkv_bias)

10 self.attn_drop=nn.Dropout(attn_drop)

11 self.proj=nn.Linear(dim,dim)

12 self.proj_drop=nn.Dropout(proj_drop)

13

14 def forward(self,x):

15 B,N,C=x.shape

16
qkv=self.qkv(x).reshape(B,N,3,self.num_heads,C//self.num_heads).permute(2,0,3,1,4)

17 q,k,v=qkv[0],qkv[1],qkv[2] #B,H,N,C

18 attn=(q@k.transpose(-2,-1))*self.scale #B,H,N,N

19 # the core code block

20 mask=torch.zeros(B,self.num_heads,N,N,device=x.device,requires_grad=False)

21 index=torch.topk(attn,k=self.topk,dim=-1,largest=True)[1]

22 mask.scatter_(-1,index,1.)

23 attn=torch.where(mask>0,attn,torch.full_like(attn,float(’-inf’)))

24 # end of the core code block

25 attn=torch.softmax(attn,dim=-1)

26 attn=self.attn_drop(attn)

27 x=(attn@v).transpose(1,2).reshape(B,N,C)

28 x=self.proj(x)

29 x=self.proj_drop(x)

30

31 return x

##  8 Comparisons between slow version and fast version

We develop two versions of $k$-NN attention, one slow version and one fast
version. The $k$-NN attention is exactly defined by slow version, but its
speed is extremely slow, as for each query it needs to select different $k$
keys and values, and this procedure is very slow. To speedup, we developed the
CUDA version, but the speed is still slower than fast version. The fast
version takes advantages of matrix multiplication and greatly speedup the
computing. The speed comparisons on DeiT-Tiny using 8 V100 are illustrated in
Table 5.

method | time per iteration (second)
---|---
slow version (pytorch) | 8192
slow version (CUDA) | 1.55
fast version (pytorch) | 0.45
Table 5: The speed comparisons on DeiT-tiny for slow and fast version

##  9 Evaluations on CIFAR10 or CIFAR100.

As vision transformers are data-hungry, directly training vision transformer
backbones from scratch on small-size datasets such as CIFAR10 or CIFAR100
would yield much worse performances compared with ConvNets. Following the
paradigm and codes of the NIPS2021 paper “Efficient Training of Visual
Transformers with Small Datasets", we briefly conducted experiments on CIFAR10
and CIFAR100 using Swin-T and T2T-ViT-14 with $k$-NN attention as shown in
Table 6. Adding $k$-NN attention brings much larger performance gain in the
scratch training (ST) due to its faster convergence speed, while the gain in
the setting of ImageNet-1k pretraining and CIFAR finetuning (FT) is not as
large.

Model | C10 (ST) | C100 (ST) | C10 (FT) | C100 (FT)
---|---|---|---|---
Swin-T | 83.9 | 66.2 | 98.4 | 88.4
Swin-T$\rightarrow$ k-NN Attn | 84.5 | 67.1 | 98.6 | 88.7
T2T-VIT-14 | 87.6 | 68.0 | 98.5 | 87.7
T2T-VIT-14$\rightarrow$ k-NN Attn | 88.2 | 68.8 | 98.8 | 88.1
Table 6: Results on CIFAR10 and CIFAR100 (100 epochs).

##  10 Proof

Notations. Throughout this appendix, we denote $x_{i}$ as $i$-th element of
vector $\bm{x}$, $\bm{W}_{ij}$ as the element at $i$-th row and $j$-th column
of matrix $\bm{W}$, and $\bm{W}_{j}$ as the $j$-th row of matrix $\bm{W}$.
Moreover, we denote $\bm{x}_{i}$ as the $i$-th patch (token) of the inputs
with $\bm{x}_{i}=\bm{X}_{i}$.

Proof for Lemma 1 We first give the formal statement of Lemma 1.

######  Lemma 4 (Formal statement of Lemma 1)

Let $\hat{\bm{V}}^{knn}_{l}$ be the $l$-th row of the $\hat{\bm{V}}^{knn}$ and
$\textrm{Var}_{\bm{a}_{l}}(\bm{x})=\mathbbm{E}_{\bm{a}_{l}}[\bm{x}^{\top}\bm{x}]-\mathbbm{E}_{\bm{a}_{l}}[\bm{x}^{\top}]\mathbbm{E}_{\bm{a}_{l}}[\bm{x}]$
with $\mathbbm{E}_{\bm{a}_{l}}[\bm{x}]=\sum_{t=1}^{n}a_{lt}\bm{x}_{t}$. Then
for any $i,j=1,2,...,n$, we have

$\displaystyle\frac{\partial\hat{\bm{V}}_{l}}{\partial
W_{\bm{Q},ij}}=x_{li}\bm{W}_{\bm{K},j}^{\top}\textrm{Var}_{a_{l}}(\bm{x})\bm{W}_{\bm{V}}\propto\textrm{Var}_{a_{l}}(\bm{x})$

and

$\displaystyle\frac{\partial\hat{\bm{V}}_{l}}{\partial
W_{\bm{K},ij}}=x_{li}\bm{W}_{\bm{Q},j}^{\top}\textrm{Var}_{a_{l}}(\bm{x})\bm{W}_{\bm{V}}\propto\textrm{Var}_{a_{l}}(\bm{x}).$

The same is true for $\hat{\bm{V}}$ of the fully-connected self-attention.

###### Proof

Let’s first consider the derivative of $\hat{\bm{V}}_{l}$ over
$W_{\bm{Q},ij}$. Via some algebraic computation, we have

$\displaystyle\frac{\partial\hat{\bm{V}}_{l}}{\partial
W_{\bm{Q},ij}}=\frac{\partial(\bm{a}_{l}\bm{V})}{\partial
W_{\bm{Q},ij}}=\displaystyle\sum_{t=1}^{n}\displaystyle
a_{lt}\left(\frac{\partial\mathcal{T}_{l}^{knn}(t)}{\partial
W_{\bm{Q},ij}}-\displaystyle\sum_{k_{1}=1}^{n}a_{lk_{1}}\frac{\partial\mathcal{T}_{l}^{knn}(k_{1})}{\partial
W_{\bm{Q},ij}}\right)\bm{x}_{t}\bm{W}_{\bm{V}},$ (1)

where we denote $\mathcal{T}_{l}^{knn}(k)$ as follow for shorthand:

$\displaystyle\mathcal{T}_{l}^{knn}(k_{1})=\begin{cases}\bm{x}_{l}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{k_{1}}^{\top},\
&\textrm{if patch $k_{1}$ is selected in row $l$}\\\ -\infty,\
&\textrm{otherwise}\end{cases}$

Let denote set $\mathcal{S}\doteq\\{i:\textrm{patch $i$ is selected in row
$l$}\\}$ and then we consider the right-hand-side of (1).

$\displaystyle\eqref{eq:lem:proof:1}=$
$\displaystyle\sum_{t\in\mathcal{S}}^{n}a_{lt}\left(\frac{\partial\left(\bm{x}_{i}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{k}^{\top}\right)}{\partial
W_{\bm{Q},ij}}-\sum_{k_{1}\in\mathcal{S}}a_{lk_{1}}\frac{\partial\left(\bm{x}_{i}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{k_{1}}^{\top}\right)}{\partial
W_{\bm{Q},ij}}\right)\bm{x}_{t}\bm{W}_{\bm{V}}$ $\displaystyle=$
$\displaystyle\sum_{t\in\mathcal{S}}^{n}a_{lt}\left(x_{1i}\bm{x}_{t}\bm{W}_{\bm{K},j}-\sum_{k_{1}\in\mathcal{S}}a_{lk_{1}}x_{li}\bm{x}_{k_{1}}\bm{W}_{\bm{K},j}\right)\bm{x}_{t}\bm{W}_{\bm{V}}$
$\displaystyle=$
$\displaystyle\underbrace{\sum_{t\in\mathcal{S}}^{n}a_{lt}x_{li}\bm{x}_{t}\bm{W}_{\bm{K},j}\bm{x}_{t}\bm{W}_{\bm{V}}}_{(a)}-\underbrace{\sum_{t\in\mathcal{S}}^{n}a_{lt}\bm{x}_{t}\bm{W}_{\bm{V}}}_{(b)}\cdot\underbrace{\sum_{k_{1}\in\mathcal{S}}a_{lk_{1}}x_{li}\bm{x}_{k_{1}}\bm{W}_{\bm{K},j}}_{(c)}.$
(2)

Since $\bm{a}_{l}$ is the $l$-th row of the attention matrix, we have
$a_{lt}\geq 0$ and $\sum_{t}a_{lt}=1$. It is possible to treat terms $(a)$,
$(b)$ and $(c)$ as the expectation of some quantities over $t$ replicates with
probability $a_{lt}$. Then (2) can be further simplified as

$\displaystyle\eqref{eq:lem:proof:2}$
$\displaystyle=\mathbbm{E}_{\bm{a}_{l}}[x_{li}\bm{x}\bm{W}_{\bm{K},j}\cdot\bm{x}\bm{W}_{\bm{V}}]-\mathbbm{E}_{\bm{a}_{l}}[\bm{x}\bm{W}_{\bm{K},j}]\cdot\mathbbm{E}_{\bm{a}_{l}}[x_{li}\bm{x}\bm{W}_{\bm{V}}]$
$\displaystyle=x_{li}\left(\mathbbm{E}_{\bm{a}_{l}}[\bm{W}_{\bm{K},j}^{\top}\bm{x}^{\top}\cdot\bm{x}\bm{W}_{\bm{V}}]-\mathbbm{E}_{\bm{a}_{l}}[\bm{W}_{\bm{K},j}^{\top}\bm{x}^{\top}]\cdot\mathbbm{E}_{\bm{a}_{l}}[\bm{x}\bm{W}_{\bm{V}}]\right)$
$\displaystyle=x_{li}\bm{W}_{\bm{K},j}^{\top}\left(\mathbbm{E}_{\bm{a}_{l}}[\bm{x}^{\top}\bm{x}]-\mathbbm{E}_{\bm{a}_{l}}[\bm{x}^{\top}]\cdot\mathbbm{E}_{\bm{a}_{l}}[\bm{x}_{t}]\right)\bm{W}_{\bm{V}}$
$\displaystyle=x_{li}\bm{W}_{\bm{K},j}^{\top}\textrm{Var}_{\bm{a}_{l}}(\bm{x})\bm{W}_{\bm{V}},$
(3)

where the second equality uses the fact that $\bm{x}_{t}\bm{W}_{\bm{K},j}$ is
a scalar.

Combing (1)-(3), we have

$\displaystyle\frac{\partial\hat{\bm{V}}_{l}}{\partial
W_{\bm{Q},ij}}=x_{li}\bm{W}_{\bm{K},j}^{\top}\textrm{Var}_{\bm{a}_{l}}(\bm{x})\bm{W}_{\bm{V}}\propto\textrm{Var}_{\bm{a}_{l}}(\bm{x}).$
(4)

Due the symmetric on $\bm{Q}$ and $\bm{K}$, we can follow the similar
procedure to show

$\displaystyle\frac{\partial\hat{\bm{V}}_{l}}{\partial
W_{\bm{K},ij}}=x_{li}\bm{W}_{\bm{Q},j}^{\top}\textrm{Var}_{\bm{a}_{l}}(\bm{x})\bm{W}_{\bm{V}}\propto\textrm{Var}_{\bm{a}_{l}}(\bm{x}).$
(5)

Finally, by setting $k=n$, one may verify that equations (4) and (5) also hold
for fully-connected self-attention.

Proof for Lemma 2 Before given the formal statement of the Lemma 2, we first
show the assumptions.

Assumption 2

  1. 1.

The token $\bm{x}_{i}$ is the sub-gaussian random vector with mean
$\bm{\mu}_{i}$ and variance $(\sigma^{2}/d)I$ for $i=1,2,...,n$.

  2. 2.

$\bm{\mu}$ follows a discrete distribution with finite values
$\bm{\mu}\in\mathcal{V}$. Moreover, there exist $0<\nu_{1},0<\nu_{2}<\nu_{4}$
such that a) $\|\bm{\mu}_{i}\|=\nu_{1}$, and b)
$\bm{\mu}_{i}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{T}\bm{\mu}_{i}\in[\nu_{2},\nu_{4}]$
for all $i$ and
$|\bm{\mu}_{i}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{j}^{\top}|\leq\nu_{2}$
for all $\bm{\mu}_{i}\neq\bm{\mu}_{j}\in\mathcal{V}$.

  3. 3.

$\bm{W}_{V}$ and $\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}$ are element-wise
bounded with $\nu_{5}$ and $\nu_{6}$ respectively, that is,
$|\bm{W}_{V}^{(ij)}|\leq\nu_{5}$ and
$|(\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top})^{(ij)}|\leq\nu_{6}$, for all $i,j$
from 1 to $d$.

In Assumption 2 we ensure that for a given query patch, the difference between
the clustering center and noises are large enough to be distinguished.

######  Lemma 5 (formal statement of Lemma 2)

Let patch $\bm{x}_{i}$ be $\sigma^{2}$-subgaussian random variable with mean
$\bm{\mu}_{i}$ and there are $k_{1}$ patches out of all $k$ patches follows
the same clustering center of query $l$. Per Assumption 2, when $\sqrt{d}\geq
3(\psi(\delta,d)+\nu_{2}+\nu_{4})$, then with probability $1-5\delta$, we have

$\displaystyle\left\|\frac{\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{l}\bm{W}_{\bm{Q}}\bm{W}^{\top}_{k}\bm{x}_{i}\right)\bm{x}_{i}\bm{W}_{V}}{\sum_{j=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{l}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{j}\right)}-\bm{\mu}_{l}\bm{W}_{V}\right\|_{\infty}$
(6) $\displaystyle\leq
4\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)\sigma\nu_{5}\sqrt{\frac{2}{dk}\log\left(\frac{2d}{\delta}\right)}$
$\displaystyle+\left[8\exp\left(\frac{\nu_{2}-\nu_{4}+\psi(\delta,d)}{\sqrt{d}}\right)-\left(7+\exp\left(\frac{\nu_{2}-\nu_{4}+\psi(\delta,d)}{\sqrt{d}}\right)\right)\frac{k_{1}}{k}\right]\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty},$

where
$\psi(\delta,d)=2\sigma\nu_{1}\nu_{6}\sqrt{2\log\left(\frac{1}{\delta}\right)}+2\sigma^{2}\nu_{6}\log\left(\frac{d}{\delta}\right)$.

###### Proof

Without loss of generality, we assume the first $k$ patch are the top-$k$
selected patches. From Assumption 2.1, we can decompose
$\bm{x}_{i}=\bm{\mu}_{i}+\bm{h}_{i}$, $i=1,2,...,k$, where $\bm{h}_{i}$ is the
sub-gaussian random vector with zero mean. We then analyze the numerator part.

$\displaystyle\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{l}\bm{W}_{\bm{Q}}\bm{W}^{\top}_{k}\bm{x}_{i}\right)\bm{x}_{i}\bm{W}_{V}$
$\displaystyle=$
$\displaystyle\overbrace{\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{\mu}_{l}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}\right)\bm{\mu}_{i}\bm{W}_{v}}^{(a)}+\overbrace{\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{l}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)\bm{h}_{i}\bm{W}_{v}}^{(b)}$
$\displaystyle+$
$\displaystyle\overbrace{\sum_{i=1}^{k}\left[\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{l}\bm{W}_{\bm{Q}}\bm{W}^{\top}_{k}\bm{x}_{i}\right)-\exp\left(\frac{1}{\sqrt{d}}\bm{\mu}_{l}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}\right)\right]\bm{\mu}_{i}\bm{W}_{v}}^{(c)}.$
(7)

Below we will bound $(a)$, $(b)$ and $(c)$ separately.

Upper bound for $(a)$. Let denote index set
$\mathcal{S}_{1}=\\{i:\bm{\mu}_{1}=\bm{\mu}_{i},\ i=1,2,...,k\\}$. We then
have

$\displaystyle\left\|(a)-\sum_{i\in\mathcal{S}_{1}}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)\bm{\mu}_{1}\bm{W}_{V}\right\|_{\infty}$
$\displaystyle\leq$
$\displaystyle(k-|\mathcal{S}_{1}|)\max_{i}\left\\{\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)\right\\}\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty}$
$\displaystyle\leq$
$\displaystyle(k-k_{1})\exp\left(\frac{\nu_{2}}{\sqrt{d}}\right)\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty},$
(8)

where last inequality is from the Assumption 2.2.

Upper bound for $(b)$. Since each dimension in $\bm{h}_{l}$ is the i.i.d
random vector with zero mean variance $\sigma^{2}/d$ based on Assumption 2.1,
we can use Hoeffding Inequality to derive the following result holds with
probability $1-\delta$.

$\displaystyle\left\|\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)\bm{h}_{i}\bm{W}_{V}\right\|_{\infty}\leq\sigma\nu_{5}\sqrt{\frac{2(k_{1}U_{1}^{2}+(k-k_{1})U_{2}^{2})}{d}\log\left(\frac{2d}{\delta}\right)},$
(9)

where

$U_{1}=\max_{i\in\mathcal{S}_{1}}\left\\{\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)\right\\}$

and

$U_{2}=\max_{i\notin\mathcal{S}_{1}}\left\\{\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)\right\\}.$

We then build the upper bound for $U_{1}$ and $U_{2}$. Since
$\bm{x}_{i}=\bm{\mu}_{i}+\bm{h}_{i}$ for $i=1,2,...,k$, we have

$\displaystyle|\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}|\leq|\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}|+|\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{h}_{i}^{\top}|$
(10)
$\displaystyle+|\bm{h}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}|+|\bm{h}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{h}_{i}^{\top}|$

Via Assumption 2.3 and Hoeffding Inequality, with probability $1-4\delta$, the
follow results hold.

$\displaystyle|\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{h}_{i}^{\top}|$
$\displaystyle\leq\sigma\nu_{1}\nu_{6}\sqrt{2\log\left(\frac{1}{\delta}\right)}$
(11)
$\displaystyle|\bm{h}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}|$
$\displaystyle\leq\sigma\nu_{1}\nu_{6}\sqrt{2\log\left(\frac{1}{\delta}\right)}$
(12) $\displaystyle\
|\bm{h}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{h}_{i}^{\top}|$
$\displaystyle\leq 2\sigma^{2}\nu_{6}\log\left(\frac{d}{\delta}\right)$ (13)

and we denote
$\psi(\delta,d)=2\sigma\nu_{1}\nu_{6}\sqrt{2\log\left(\frac{1}{\delta}\right)}+2\sigma^{2}\nu_{6}\log\left(\frac{d}{\delta}\right)$
for shorthand and then we have

$\displaystyle U_{1}$
$\displaystyle\leq\exp\left[\frac{1}{\sqrt{d}}\left(\nu_{2}+\psi(\delta,d)\right)\right]$
$\displaystyle U_{2}$
$\displaystyle\leq\exp\left[\frac{1}{\sqrt{d}}\left(\nu_{4}+\psi(\delta,d)\right)\right]$

As a result, with a probability $1-5\delta$, we have:

$\displaystyle\|(b)\|_{\infty}\leq\sigma\nu_{5}\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)\sqrt{\frac{2(k_{1}\exp\left(\frac{2\nu_{2}}{\sqrt{d}}\right)+(k-k_{1})\exp\left(\frac{2\nu_{4}}{\sqrt{d}}\right))}{d}\log\left(\frac{2d}{\delta}\right)}.$
(14)

Upper bound for $(c)$.

$\displaystyle\|(c)\|_{\infty}$
$\displaystyle\leq\left|\sum_{i=1}^{k}\left[\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)-\exp\left(\frac{1}{\sqrt{d}}\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}\right)\right]\right|\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty}$

and it implies

$\displaystyle\frac{\|(c)\|_{\infty}}{\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty}}$
$\displaystyle\leq\left|\sum_{i\notin\mathcal{S}_{1}}\left[\left(\exp\left(\frac{1}{\sqrt{d}}\left(\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}-\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}\right)\right)-1\right)\right]\right|\exp\left(\frac{\nu_{2}}{\lambda}\right)$
$\displaystyle+\left|\sum_{i\in\mathcal{S}_{1}}\left[\left(\exp\left(\frac{1}{\sqrt{d}}\left(\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}-\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}\right)\right)-1\right)\right]\right|\exp\left(\frac{\nu_{5}}{\lambda}\right).$
(15)

Combine (15) with (11)-(13) and we have with probability $1-4\delta$:

$\displaystyle\|(c)\|_{\infty}$
$\displaystyle\leq\left|\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)-1\right|\left[(k-k_{1})\exp\left(\frac{\nu_{2}}{\sqrt{d}}\right)+k_{1}\exp\left(\frac{\nu_{5}}{\sqrt{d}}\right)\right]\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty}.$
(16)

From (7), (8), (14) and (16), with probability $1-5\delta$, we have

$\displaystyle\left\|\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)\bm{x}_{i}\bm{W}_{V}-\sum_{i\in\mathcal{S}_{1}}\exp\left(\frac{1}{\sqrt{d}}\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}\right)\bm{\mu}_{1}\bm{W}_{V}\right\|_{\infty}$
$\displaystyle\leq\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)\left[\tau_{1}(k,k_{1})\|W_{v}\bm{\mu}_{1}\|_{\infty}+\sigma\nu_{5}\sqrt{\frac{2\tau_{2}(k,k_{1})}{d}\log\left(\frac{2d}{\delta}\right)}\right],$
(17)

where
$\tau_{1}(k,k_{1})=(k-k_{1})\exp\left(\frac{\nu_{2}}{\sqrt{d}}\right)+k_{1}\exp\left(\frac{\nu_{5}}{\sqrt{d}}\right)$
and
$\tau_{2}(k,k_{1})=k_{1}\exp\left(\frac{2\nu_{2}}{\sqrt{d}}\right)+(k-k_{1})\exp\left(\frac{2\nu_{4}}{\sqrt{d}}\right)$.

Now we consider the upper bound the denominator part.

$\displaystyle\left|\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)-\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}\right)\right|$
$\displaystyle\leq$
$\displaystyle\overbrace{\left|\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}\right)\left[\exp\left(\frac{1}{\sqrt{d}}\left(\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}-\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}\right)\right)-1\right]\right|}^{(g_{1})}.$
(18)

Via Assumption 2.2, (11)-(13) and the definition of $\psi(\delta,d)$, with
probability $1-5\delta$ the follow results hold.

$\displaystyle(g_{1})$
$\displaystyle\leq\exp\left(\frac{\nu_{4}}{\sqrt{d}}\right)\left|\sum_{i=1}^{k}\left[\exp\left(\frac{1}{\sqrt{d}}\left[\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}-\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}\right]\right)-1\right]\right|$
$\displaystyle\leq\exp\left(\frac{\nu_{4}}{\sqrt{d}}\right)\left|\sum_{i=1}^{k}\left[\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)-1\right]\right|$
$\displaystyle\leq
k\exp\left(\frac{\nu_{4}}{\sqrt{d}}\right)\left[\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)-1\right].$
(19)

Combining (18), (19) and Assumption 2.2, we have

$\displaystyle\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)$
$\displaystyle\geq$
$\displaystyle\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}\right)-k\exp\left(\frac{\nu_{4}}{\sqrt{d}}\right)\left(\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)-1\right)$
$\displaystyle\geq$ $\displaystyle
k\exp\left(-\frac{\nu_{2}}{\sqrt{d}}\right)-k\exp\left(\frac{\nu_{4}}{\sqrt{d}}\right)\left(\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)-1\right).$
(20)

When $\sqrt{d}\geq 3(\psi(\delta,d)+\nu_{2}+\nu_{4})$, one may verify

$\displaystyle\exp\left(\frac{\psi(\delta,d)+\nu_{2}+\nu_{4}}{\sqrt{d}}\right)-\exp\left(\frac{\nu_{2}+\nu_{4}}{\sqrt{d}}\right)\leq\exp\left(\frac{1}{3}\right)-1\approx
0.39<1$ $\displaystyle\Rightarrow$
$\displaystyle\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)-1-\exp\left(-\frac{\nu_{2}+\nu_{4}}{\sqrt{d}}\right)<0$
$\displaystyle\Rightarrow$
$\displaystyle\exp\left(-\frac{\nu_{2}}{\sqrt{d}}\right)-\exp\left(\frac{\nu_{4}}{\sqrt{d}}\right)\left(\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)-1\right)>0$
$\displaystyle\Rightarrow$ $\displaystyle
k\exp\left(-\frac{\nu_{2}}{\sqrt{d}}\right)-k\exp\left(\frac{\nu_{4}}{\sqrt{d}}\right)\left(\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)-1\right)>0.$
(21)

Finally, via (17), we have

$\displaystyle\left\|\frac{\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)}{\sum_{j=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{j}^{\top}\right)}\bm{x}_{i}\bm{W}_{V}-\bm{\mu}_{1}\bm{W}_{V}\right\|_{\infty}\cdot\sum_{j=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{j}^{\top}\right)$
$\displaystyle=$
$\displaystyle\left\|\sum_{j=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)\bm{x}_{i}\bm{W}_{V}-\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{j}^{\top}\right)\bm{\mu}_{1}\bm{W}_{V}\right\|_{\infty}$
$\displaystyle\leq$
$\displaystyle\left\|\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)\bm{x}_{i}\bm{W}_{V}-\sum_{i\in\mathcal{S}_{1}}\exp\left(\frac{1}{\sqrt{d}}\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}\right)\bm{\mu}_{1}\bm{W}_{V}\right\|_{\infty}$
$\displaystyle+$
$\displaystyle\left\|\sum_{i\notin\mathcal{S}_{1}}\exp\left(\frac{1}{\sqrt{d}}\bm{\mu}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{\mu}_{i}^{\top}\right)\bm{\mu}_{1}\bm{W}_{V}\right\|_{\infty}$
$\displaystyle\leq$
$\displaystyle\exp\left(\frac{\psi(\delta,d)}{\lambda}\right)\left[2\tau_{1}(k,k_{1})\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty}+\sigma\nu_{5}\sqrt{\frac{2\tau_{2}(k,k_{1})}{d}\log\left(\frac{2d}{\delta}\right)}\right]$
$\displaystyle-$ $\displaystyle
k_{1}\exp\left(\frac{\nu_{4}}{\sqrt{d}}\right)\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty}+(k-k_{1})\exp\left(\frac{\nu_{2}}{\sqrt{d}}\right)\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty}.$
(22)

Per (22) and (21), we have

$\displaystyle\left\|\frac{\sum_{i=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{i}^{\top}\right)}{\sum_{j=1}^{k}\exp\left(\frac{1}{\sqrt{d}}\bm{x}_{1}\bm{W}_{\bm{Q}}\bm{W}_{\bm{K}}^{\top}\bm{x}_{j}^{\top}\right)}\bm{x}_{i}\bm{W}_{V}-\bm{\mu}_{1}\bm{W}_{V}\right\|_{\infty}$
$\displaystyle\leq\left(k\exp\left(-\frac{\nu_{2}}{\sqrt{d}}\right)-k\exp\left(\frac{\nu_{4}}{\sqrt{d}}\right)\left(\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)-1\right)\right)^{-1}$
$\displaystyle\cdot\left(\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)\left[2\tau_{1}(k,k_{1})\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty}+\sigma\nu_{5}\sqrt{\frac{2\tau_{2}(k,k_{1})}{d}\log\left(\frac{2d}{\delta}\right)}\right]\right.$
$\displaystyle\left.\quad\quad+(k-k_{1})\exp\left(\frac{\nu_{2}}{\sqrt{d}}\right)\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty}-k_{1}\exp\left(\frac{\nu_{4}}{\sqrt{d}}\right)\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty}\right)$
$\displaystyle\leq\frac{\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)\sigma\nu_{5}\sqrt{\exp\left(\frac{-2\nu_{4}}{\sqrt{d}}\right)\frac{2\tau_{2}(k,k_{1})}{dk^{2}}\log\left(\frac{2d}{\delta}\right)}}{1+\exp\left(-\frac{\nu_{4}+\nu_{2}}{\sqrt{d}}\right)-\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)}$
$\displaystyle+\frac{\exp\left(\frac{\psi(\delta,d)-\nu_{4}}{\sqrt{d}}\right)\frac{2\tau_{1}(k,k_{1})}{k}-\frac{k_{1}}{k}\exp\left(\frac{\nu_{4}}{\sqrt{d}}\right)+(1-\frac{k_{1}}{k})\exp\left(\frac{\nu_{4}}{\sqrt{d}}\right)}{1+\exp\left(-\frac{\nu_{4}+\nu_{2}}{\sqrt{d}}\right)-\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)}\cdot\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty}$
$\displaystyle\leq\left[8\exp\left(\frac{\nu_{2}-\nu_{4}+\psi(\delta,d)}{\sqrt{d}}\right)-\left(7+\exp\left(\frac{\nu_{2}-\nu_{4}+\psi(\delta,d)}{\sqrt{d}}\right)\right)\frac{k_{1}}{k}\right]\|\bm{\mu}_{1}\bm{W}_{V}\|_{\infty}$
$\displaystyle+4\exp\left(\frac{\psi(\delta,d)}{\sqrt{d}}\right)\sigma\nu_{5}\sqrt{\frac{2}{dk}\log\left(\frac{2d}{\delta}\right)},$
(23)

where second inequality uses $\nu_{4}\geq\mu_{2}$ and last inequality uses the
definition of $\tau_{1}(k,k_{1})$, $\tau_{2}(k,k_{1})$ and $\sqrt{d}\geq
3(\psi(\delta,d)+\nu_{2}+\nu_{4})$.

Proof of Lemma 3

Without loss of generality, we only consider the case with first query
patches. In the $k$-NN attention scheme, we first use the dot-product product
to compute the similarity between query and each key-patches and then use the
$\mathrm{softmax}$ to normalize the similarities. We make the following
assumption to facility our analysis.

Assumption 3.1 There exists $\bm{\beta}^{*}\in\mathbbm{R}^{1\times n}$ and
$\bm{\beta}^{*}\in\Delta$ such that
$\bm{q}_{1}=\bm{\beta}\bm{K}+\bm{\epsilon}$, where $\bm{\epsilon}$ is filled
with random variable follows $\mathcal{N}(0,\sigma^{2})$ for some $\sigma>0$.

To see the connection between the Assumption 3.1 with attention scheme, we
consider the follow problem.

$\displaystyle\min_{\bm{\beta}\in\Delta}\|\bm{K}^{\top}\bm{\beta}^{\top}-\bm{q}_{1}^{\top}\|_{2}^{2},$
(24)

If $\bm{K}$ is normalized with zero columns mean and we apply the exponential
gradient method on the initial solution $\bm{\beta}_{0}=\frac{1}{n}\bm{e}$
with step length $1/\sqrt{d}$, the one step updated solution $\bm{\beta}_{1}$
is

$\displaystyle\bm{\beta}_{1}=\frac{\exp(\frac{1}{\sqrt{d}}\bm{K}\bm{q}_{1}^{\top})}{\sum_{i=1}^{k}\exp(\frac{1}{\sqrt{d}}\bm{k}_{i}\bm{q}_{1}^{\top})}.$
(25)

The above equation (25) is just the attention scheme used standard transformer
type model. Based on Assumption 3.1, we can treat $\bm{\beta}_{1}$ as an
approximation of underlying true parameters $\bm{\beta}^{*}$.

On the other hand, it is commonly believed that only part of patches are
correlated with the query patch (i.e., with non-zero similarity weights.) and
it would be ideal if we could use a computational cheap method to eliminate
the irrelevant patches. In this paper, we consider the top-k selection scheme.
To see the rationality of the top-$k$ selection, we consider augmenting (24)
with $L_{2}$ regularization on $\bm{\beta}$.

$\displaystyle\min_{\bm{\beta}\in\Delta}\|\bm{K}^{\top}\bm{\beta}^{\top}-\bm{q}_{1}^{\top}\|_{2}^{2}+\frac{\lambda}{2}\|\bm{\beta}\|_{2}^{2},$
$\displaystyle\Rightarrow$
$\displaystyle\bm{K}(\bm{K}^{\top}\bm{\beta}^{\top}+\bm{q}_{1}^{\top})+\lambda\bm{\beta}^{\top}+\bm{\lambda_{1}}+\bm{e}\lambda_{2}=0,$
$\displaystyle\Rightarrow$
$\displaystyle\bm{\beta}^{\top}=(\bm{K}\bm{K}^{\top}+\lambda\bm{I})^{-1}\left(\bm{K}\bm{q}_{1}^{\top}+\bm{\lambda_{1}}+\bm{e}\lambda_{2}\right),$

where we use the KKT optimal condition, and $\bm{\lambda}_{1}$, $\lambda_{2}$
are Lagrange multipliers to make sure $\bm{\beta}\in\Delta$. If $\lambda$ is
large enough and $\bm{\beta}>0$, we will have

$\displaystyle\bm{\beta}^{\top}\approx\frac{1}{\lambda}\left(\bm{K}\bm{q}_{1}^{\top}+\bm{e}\lambda_{2}\right)\propto\bm{K}\bm{q}_{1}^{\top}$
(26)

The above result indicates that we may selection the important elements in
$\bm{\beta}$ (e.t., with large magnitude) by considering rankness in vector
$\bm{K}\bm{q}_{1}^{\top}$.

We then discussion the correctness of the top-k selection with the following
regularity assumptions on $\bm{K}$ and $\bm{q}_{1}$.

Assumption 3.2

  1. 1.

$\bm{K}$ is normalized with row zero mean. Let
$\bm{\Sigma}=\bm{K}\bm{K}^{\top}$ and
$\bm{Z}=\bm{\Sigma^{-1/2}}\bm{K}^{\top}$. We assume there exist some
$c,c_{4}>1$ and $C_{1}>0$ such that the following inequality

$\mathbbm{P}\left(\lambda_{\max}(\tilde{p}^{-1}\tilde{\bm{Z}}\tilde{\bm{Z}}^{\top})>c_{4}\textrm{
and
}\lambda_{\max}(\tilde{p}^{-1}\tilde{\bm{Z}}\tilde{\bm{Z}}^{\top})<\frac{1}{c_{4}}\right)<\exp(-C_{1}d)$

holds for any $\tilde{p}\times d$ submatrix $\tilde{\bm{Z}}$ of $\bm{Z}$ with
$cn<\tilde{p}\leq n$.

  2. 2.

$\textrm{var}(\bm{q}_{1})=\mathcal{O}(1)$ and for some $\kappa\geq 0$ and
$c_{5},c_{6}>0$,

$\min_{i:\beta_{i}^{*}>0}\beta_{i}\geq\frac{c_{5}}{d^{\kappa}}\textrm{ and
}\min_{i:\beta_{i}^{*}>0}\textrm{cov}(\beta_{i}^{-1}\bm{q}_{1}^{\top},\bm{k}_{i}^{\top})\geq
c_{6}$

  3. 3.

There exist some $\tau\in[0,1)$ and $c_{7}>0$ such that

$\lambda_{\max}(\bm{\Sigma})\leq c_{7}d^{\tau}.$

######  Lemma 6 (formal statement of Lemma )

Let’s assume only be $s$ keywords are relevant to the query $l$. Under
Assumption 3.1 and 3.2, when $2\kappa+\tau<1$, with probability
$1-\mathcal{O}(s\exp(-Cd^{1-2\kappa}/\log d))$, we have

$\displaystyle\sum_{i=1}^{n}\mathbbm{1}(\bm{q}_{l}\bm{k}_{i}^{\top}\geq\min_{i\in\mathcal{M}^{*}}\bm{q}_{l}\bm{k}_{i}^{\top})\leq
cnd^{2\kappa+\tau-1},$ (27)

where $\mathcal{M}^{*}=\\{i:\textrm{keyword $i$ is relevant to the query
$l$.}\\}$ , and $\tau$, $\kappa$, $c$ and $C$ are positive constants.

###### Proof

Our strategy is the similar to the proof of Theorem 1 in [19].

Based on equation (44) in [19], we have

$\displaystyle\bm{K}\bm{K}^{\top}=n\bm{\Sigma}^{1/2}\tilde{\bm{U}}^{\top}\textrm{diag}(\mu_{1},...,\mu_{d})\tilde{\bm{U}}\bm{\Sigma}^{1/2},$
(28)

where $\mu_{1}$,…, $\mu_{d}$ are $d$ eigenvalues of
$p^{-1}\bm{Z}\bm{Z}^{\top}$, $\tilde{\bm{U}}=(I_{d},\bm{0})_{d\times
n}\bm{U}$, and $\bm{U}$ is uniformally distributed on the orthogonal group
$O(n)$. To facility our further analysis we denote
$\bm{\omega}=\bm{q}_{l}\bm{K}^{\top}$. By definition of $\bm{\omega}$ and per
Assumption 3.1, we have

$\displaystyle\bm{\omega}^{\top}=\bm{K}\bm{q}_{l}^{\top}=\bm{K}\bm{K}^{\top}\bm{\beta}^{\top}+\bm{K}\bm{\epsilon}^{\top}\doteq\bm{\xi}+\bm{\eta}.$
(29)

We then separately study $\bm{\xi}$ and $\bm{\eta}$.

Analysis on $\bm{\xi}$. We first bound $\bm{\xi}$ from above. Since
$\\{\mu_{i}\\}$ is the eigenvalues of $n^{-1}\bm{Z}\bm{Z}^{\top}$, we have

$\displaystyle\textrm{diag}(\mu_{1}^{2},...,\mu_{d}^{2})\leq\left[\lambda_{\max}(n^{-1}\bm{Z}\bm{Z}^{\top})\right]^{2}I_{d}$
(30)

and

$\tilde{\bm{U}}\bm{\Sigma}\tilde{\bm{U}}^{\top}\leq\lambda_{\max}(\bm{\Sigma})I_{d}$

There lead to

$\displaystyle\|\bm{\xi}\|^{2}\leq
n^{2}\lambda_{\max}(\bm{\Sigma})\left[\lambda_{\max}(n^{-1}\bm{Z}\bm{Z}^{\top})\right]^{2}\cdot\bm{\beta}^{\top}\bm{\Sigma}^{1/2}\tilde{\bm{U}}^{\top}\tilde{\bm{U}}\bm{\Sigma}^{1/2}\bm{\beta}.$
(31)

Let $Q$ belongs to the orthogonal group $O(n)$ such that
$\bm{\Sigma}^{1/2}\bm{\beta}=\|\bm{\Sigma}^{1/2}\bm{\beta}\|Q\bm{e}_{1}$.
Then, it follows from Lemma 1 in [19] that

$\displaystyle\bm{\beta}^{\top}\bm{\Sigma}^{1/2}\tilde{\bm{U}}^{\top}\tilde{\bm{U}}\bm{\Sigma}^{1/2}\bm{\beta}=\|\bm{\Sigma}^{1/2}\bm{\beta}\|\langle
Q^{\top}\bm{S}Q\bm{e}_{1},\bm{e}_{1}\rangle\overset{(d)}{=}\|\bm{\Sigma}^{1/2}\bm{\beta}\|\langle\bm{S}\bm{e}_{1},\bm{e}_{1}\rangle,$
(32)

where we use the symbol $\overset{(d)}{=}$ to denote being identical in
distribution. By part 3 in Assumption 3.2,
$\|\bm{\Sigma}^{1/2}\bm{\beta}\|^{2}=\bm{\beta}^{\top}\bm{\Sigma}\bm{\beta}\leq\textrm{var}(\bm{y})=\mathcal{O}(1)$,
and thus via Lemma 4 in [19], we have for some $C>0$,

$\displaystyle\mathbbm{P}\left(\bm{\beta}^{\top}\bm{\Sigma}^{1/2}\tilde{\bm{U}}^{\top}\tilde{\bm{U}}\bm{\Sigma}^{1/2}\bm{\beta}>(d/n)\right)\leq\mathcal{O}\left(\exp(-Cd)\right).$
(33)

Combining with $\lambda_{\max}(\bm{\Sigma})=O(d^{\tau})$ and
$\mathbbm{P}(\lambda_{\max}(n^{-1}\bm{Z}\bm{Z}^{\top})>c_{1})\leq\exp\left(-C_{1}d\right)$
by parts 1 and 3 in Assumption 3.2 along with union bound, we have

$\displaystyle\mathbbm{P}\left(\|\bm{\xi}\|^{2}\geq\mathcal{O}(d^{1+\tau}n)\right)\leq\mathcal{O}(\exp(-Cd)).$
(34)

We then consider the lower bound on $\xi_{i}$ for $i\in\mathcal{M}_{*}$. By
(28), we have

$\displaystyle\xi_{i}=n\bm{e}_{i}^{\top}\bm{\Sigma}^{1/2}\tilde{\bm{U}}^{\top}\textrm{diag}(\mu_{1},...,\mu_{d})\tilde{\bm{U}}\bm{\Sigma}^{1/2}\bm{\beta}.$
(35)

Note that $\|\bm{\Sigma}^{1/2}\bm{e}_{i}\|=\sqrt{\textrm{var}(\bm{X}_{i})}=1$,
$\|\bm{\Sigma}^{1/2}\bm{\beta}\|=\mathcal{O}(1)$. By part 2 of Assumption 3.2,
there exists some $c>0$ such that

$\displaystyle\langle\bm{\Sigma}^{1/2}\bm{\beta},\bm{\Sigma}^{1/2}\bm{e}_{i}\rangle=\beta_{i}\textrm{cov}(\beta_{i}^{-1}\bm{q}_{1}^{\top},\bm{k}_{i})\geq\frac{c}{d^{\kappa}}.$
(36)

Thus, there exists $Q$ in orthogonal group $O(n)$ such that
$\bm{\Sigma}^{1/2}\bm{e}_{i}=Q\bm{e}_{1}$ and

$\displaystyle\bm{\Sigma}^{1/2}\bm{\beta}=\langle\bm{\Sigma}^{1/2}\bm{\beta},\bm{\Sigma}^{1/2}\bm{e}_{i}\rangle
Qe_{1}+\mathcal{O}(1)Q\bm{e}_{2}.$ (37)

Since $(\mu_{1},....,\mu_{d})^{\top}$ is independent of $\tilde{\bm{U}}$ by
Lemma 1 in [19] and the uniform distribution on the orthogonal group $O(n)$ is
invariant under itself, it follows that

$\displaystyle\xi_{i}\overset{(d)}{=}\langle\bm{\Sigma}^{1/2}\bm{\beta},\bm{\Sigma}^{1/2}\bm{e}_{i}\rangle
R_{1}+\mathcal{O}(n)R_{2}\doteq\xi_{i,1}+\xi_{i,2},$ (38)

where
$\bm{R}=(R_{1},R_{2},...,R_{n})^{\top}=\tilde{\bm{U}}^{\top}\textrm{diag}(\mu_{1},...,\mu_{d})\tilde{\bm{U}}^{\top}\bm{e}_{1}$.
We will bound the above two terms $\xi_{i,1}$ and $\xi_{i,2}$ separately. One
can verify

$\displaystyle
R_{1}\geq\bm{e}_{1}^{\top}\tilde{\bm{U}}^{\top}\lambda_{\min}(n^{-1}\bm{Z}\bm{Z}^{\top})I_{d}\tilde{\bm{U}}\bm{e}=\lambda_{\min}(n^{-1}\bm{Z}\bm{Z}^{\top})\langle\bm{S}\bm{e}_{1},\bm{e}_{1}\rangle,$
(39)

and thus by part 1 of assumption 3.2, Lemma 4 in [19], and union bound, we
have for some $c,C>0$,

$\displaystyle\mathbbm{P}(R_{1}<cd/n)\leq\mathcal{O}(\exp(-Cd)).$ (40)

Therefore we have for some $c>0$,

$\displaystyle\mathbbm{P}(\xi_{i,1}\leq
cd^{1-\kappa})\leq\mathcal{O}(\exp(-Cd))$ (41)

Similarly, we can also show that

$\displaystyle\mathbbm{P}(\|\bm{R}\|^{2}\geq
O(d/n))\leq\mathcal{O}(\exp(-Cd)).$ (42)

Via some analysis, we can show that the distribution of
$\tilde{\bm{R}}=(R_{2},...,R_{n})^{\top}$ is invariant under the orthogonal
group $O(n-1)$. Then, it follows that
$\tilde{\bm{R}}\overset{(d)}{=}\|\tilde{\bm{R}}\|\bm{W}/\|\bm{W}\|$, where
$\bm{W}=(W_{1},...,W_{n-1})^{\top}\sim\mathcal{N}(0,I_{n-1})$, independent of
$\|\tilde{\bm{R}}\|$. Thus, we have

$\displaystyle
R_{2}\overset{(d)}{=}\|\tilde{\bm{R}}\|\frac{W_{1}}{\|\bm{W}\|}.$ (43)

And via the Lemma 5 in [19], we can show

$\displaystyle\mathbbm{P}(|\xi_{i,2}|\geq
c\sqrt{d}|W|)\leq\mathcal{O}(\exp(-Cd)),$ (44)

where $W$ is $\mathcal{N}(0,1)$-distributed random variable. We then pick
$x_{d}=c\sqrt{2C}d^{1-\kappa}/\sqrt{\log d}$. Then, by standard tail bound, we
have

$\displaystyle\mathbbm{P}(c\sqrt{d}|W|\geq
x_{d})\leq\mathcal{O}(\exp(-Cd^{1-2\kappa}/\log d)),$ (45)

Then

$\displaystyle\mathbbm{P}(|\xi_{i,2}|\geq
x_{d})\leq\mathcal{O}(\exp(-Cd^{1-2\kappa}/\log d)).$ (46)

It implies that for $i$ with $\beta_{i}^{*}>0$, we have

$\displaystyle\mathbbm{P}\left(\xi_{i,1}-|\xi_{i,2}|\leq
cd^{1-\kappa}\right)\leq\mathcal{O}(\exp(-Cd^{1-2\kappa}/\log d))$
$\displaystyle\Rightarrow$ $\displaystyle\mathbbm{P}\left(\xi_{i}\leq
cd^{1-\kappa}\right)\leq\mathcal{O}(\exp(-Cd^{1-2\kappa}/\log d)).$

Next, we examine term
$\bm{\eta}=(\eta_{1},...,\eta_{n})^{\top}=\bm{K}\bm{\epsilon}^{\top}$.
Clearly, we have

$\displaystyle\bm{K}^{\top}\bm{K}=\bm{Z}\bm{\Sigma}\bm{Z}^{\top}\leq\bm{Z}\lambda_{\max}(\bm{\Sigma})I_{n}\bm{Z}^{\top}=n\lambda_{\max}(\bm{\Sigma})\lambda_{\max}(n^{-1}\bm{Z}\bm{Z}^{\top})I_{d}.$
(47)

Then, it follows that

$\displaystyle\|\bm{\eta}\|^{2}=\bm{\epsilon}\bm{K}^{\top}\bm{K}\bm{\epsilon}^{\top}\leq
n\lambda_{\max}(\bm{\Sigma})\lambda_{\max}(n^{-1}\bm{Z}\bm{Z}^{\top})\|\bm{\epsilon}\|^{2}.$
(48)

From Assumption 3.1, we know that $\\{\epsilon_{i}^{2}/\sigma^{2}\\}$ are
i.i.d. $\chi_{1}^{2}$-distributed random variables. Thus there exist $c,C>0$
such that

$\displaystyle\mathbbm{P}(\|\bm{\epsilon}\|^{2}>cd\sigma^{2})\leq\exp(-Cd)$
(49)

Along with parts 1 and 3 of Assumption 3.2, we have

$\displaystyle\mathbbm{P}(\|\bm{\eta}\|^{2}>O(d^{1+\tau}n))\leq\mathcal{O}(\exp(-Cd)).$
(50)

We then bound $|\eta_{i}|$ from above. Given that
$\bm{\eta}=K\bm{\epsilon}^{\top}\sim\mathcal{N}(0,\sigma^{2}KK^{\top})$. Hence
$\eta_{i}|_{K=\bm{K}}\sim\mathcal{N}(0,\textrm{var}(\eta_{i}|_{K=\bm{K}}))$
with
$\textrm{var}(\eta_{i}|K=\bm{K})=\sigma^{2}\bm{e}_{i}^{\top}\bm{K}\bm{e}_{i}$.

Let $\mathcal{E}$ be the event $\\{\textrm{var}(\eta_{i}|\bm{K})\leq cd\\}$
for some c>0\. Then, using the same argument in the previous proof. we can
show that

$\displaystyle\mathbbm{P}(\mathcal{E}^{c})\leq\mathcal{O}(\exp(-Cd)).$ (51)

Condition on the event $\mathcal{E}$, for all $x>0$, we have

$\displaystyle\mathbbm{P}(|\eta_{i}|>x|\bm{K})\leq\mathbbm{P}(\sqrt{cd}|W|>x)$
(52)

, where $W$ is $\mathcal{N}(0,1)$ random variable. Via union bound, we have

$\displaystyle\mathbbm{P}(|\eta_{i}|>x)\leq\mathcal{O}(\exp(-Cd))+\mathbbm{P}(\sqrt{cd}|W|>x)$
(53)

By setting $x=\sqrt{2cC}d^{1-\kappa}/\sqrt{\log d}$, we have

$\displaystyle\mathbbm{P}(\sqrt{cn}|W|>x)\leq\mathcal{O}(\exp(-Cd^{1-2\kappa}/\log
d))$ (54)

Then

$\displaystyle\mathbbm{P}(|\eta_{i}|>o(d^{1-\kappa}))\leq\mathcal{O}(\exp(-Cd^{1-2\kappa}/\log
d))$ (55)

Finally, we could reach

$\displaystyle\mathbbm{P}(\min_{i:\beta^{*}_{i}>0}\xi_{i}-|\eta_{i}|\leq
c_{1}d^{\kappa})\leq\mathcal{O}(s\exp(-Cd^{1-2\kappa}/\log d))$
$\displaystyle\Rightarrow\mathbbm{P}(\min_{i:\beta_{i}^{*}>0}\omega_{i}\leq
c_{1}d^{\kappa})\leq\mathcal{O}(s\exp(-Cd^{1-2\kappa}/\log d))$

Therefore, with probability $1-\mathcal{O}(s\exp(-Cd^{1-2\kappa}/\log d))$,
the magnitudes of $\omega_{i}$ with $\beta_{i}^{*}>0$ are uniformly at least
of order $d^{1-\kappa}$ and for some $c>0$, we have

$\displaystyle\sum_{i=1}^{n}\mathbbm{1}(\omega_{k}\geq\min_{i:\beta^{*}_{i}>0}\omega_{i})\leq\frac{cp}{d^{1-2\kappa-\tau}}\leq
cnd^{2\kappa+\tau-1}$ $\displaystyle\Rightarrow$
$\displaystyle\sum_{i=1}^{n}\mathbbm{1}(\bm{q}_{l}\bm{k}^{\top}_{i}\geq\min_{i\in\mathcal{M}_{*}}\bm{q}_{l}\bm{k}^{\top}_{i})\leq
cnd^{2\kappa+\tau-1}.$