Year: 2021

Abstract:
Despite the recent success in many applications, the high computational requirements of vision transformers limit their use in resource-constrained settings. While many existing methods improve the quadratic complexity of attention, in most vision transformers, self-attention is not the major computation bottleneck, e.g., more than 80% of the computation is spent on fully-connected layers. To improve the computational complexity of all layers, we propose a novel token downsampling method, called Token Pooling, efficiently exploiting redundancies in the images and intermediate token representations. We show that, under mild assumptions, softmax-attention acts as a high-dimensional low-pass (smoothing) filter. Thus, its output contains redundancy that can be pruned to achieve a better trade-off between the computational cost and accuracy. Our new technique accurately approximates a set of tokens by minimizing the reconstruction error caused by downsampling. We solve this optimization problem via cost-efficient clustering. We rigorously analyze and compare to prior downsampling methods. Our experiments show that Token Pooling significantly improves the cost-accuracy trade-off over the state-of-the-art downsampling. Token Pooling is simple and effective and can benefit many architectures with global attention. Applied to DeiT, NEW it achieves the same ImageNet top-1 accuracy using 42% fewer computations.

text: Despite the recent success in many applications, the high computational requirements of vision transformers limit their use in resource-constrained settings. While many existing methods improve the quadratic complexity of attention, in most vision transformers, self-attention is not the major computation bottleneck, e.g., more than 80% of the computation is spent on fully-connected layers. To improve the computational complexity of all layers, we propose a novel token downsampling method, called Token Pooling, efficiently exploiting redundancies in the images and intermediate token representations. We show that, under mild assumptions, softmax-attention acts as a high-dimensional low-pass (smoothing) filter. Thus, its output contains redundancy that can be pruned to achieve a better trade-off between the computational cost and accuracy. Our new technique accurately approximates a set of tokens by minimizing the reconstruction error caused by downsampling. We solve this optimization problem via cost-efficient clustering. We rigorously analyze and compare to prior downsampling methods. Our experiments show that Token Pooling significantly improves the cost-accuracy trade-off over the state-of-the-art downsampling. Token Pooling is simple and effective and can benefit many architectures with global attention. Applied to DeiT, NEW it achieves the same ImageNet top-1 accuracy using 42% fewer computations.

heading: 1 INTRODUCTION
text: Vision transformers (Dosovitskiy et al., 2020; Touvron et al., 2021; Liu et al., 2021; Heo et al., 2021; Zheng et al., 2021) have demonstrated state-of-the-art results in many vision applications, from image classification to segmentation. However, the high computational cost limits their use in resource-restricted, real-time, or low-powered applications. While most prior work in Natural Language Processing (NLP) improve the time-complexity of attention (Tay et al., 2020b; Ilharco et al., 2020), in vision transformers the main computation bottleneck is the fully-connected layers, as we show in §3.1. The computational complexity of these layers is determined by the number of tokens and their feature dimensionality. While reducing the dimensionality improves computational cost, it sacrifices model capacity and often significantly deteriorates the accuracy of the model. On the other hand, since images often contain mostly smooth surfaces with sparsely located edges and corners, they contain similar (and thus redundant) features. Moreover, we show that, under mild assumptions, softmax-attention is equivalent to low-pass filtering of tokens and thereby produces tokens with similar features, as empirically observed by Goyal et al. (2020) and Rao et al. (2021). This redundancy in representations suggests that we can reduce the number of tokens, i.e., downsampling, without a significant impact to the accuracy, achieving a better cost-accuracy trade-off than reducing feature dimensionality alone.
Downsampling is widely used in Convolutional Neural Network (CNN) architectures to improve computational efficiency, among other purposes. Given a grid of pixels or features, downsampling gradually reduces the grid dimensions via combining neighboring vertices on the grid. The prevailing max/average pooling and sub-sampling are examples of (spatially uniform) grid-downsampling that only uses locations on the grid to decide which vertices to combine. Such methods do not efficiently address non-uniformly distributed redundancy in images and features (Recasens et al., 2018; Marin et al., 2019). Unlike CNNs that require grid preservation, transformers allow a wider range of nonuniform data-aware downsampling layers, where a better operator can be designed.
We propose Token Pooling, a novel nonuniform data-aware downsampling operator for transformers efficiently exploiting redundancy in features. See the illustration and performance metric in Fig-
ures 1a & 1b. Motivated by nonuniform sampling and image compression (Marvasti, 2012; Unat et al., 2009; Belfor et al., 1994; Rabbani, 2002), we formulate token downsampling as an optimization problem that minimizes the reconstruction error caused by downsampling. We show that clustering algorithms, K-Means and K-Medoids, efficiently solve this problem, see illustration in Figure 1c. To the best of our knowledge, we are the first to use this formulation and simple clustering analysis for token donwsampling in transformers. We compare with various prior downsampling techniques, including grid-downsampling (Pan et al., 2021; Liu et al., 2021; Wang et al., 2021) and token pruning (Goyal et al., 2020; Rao et al., 2021). The proposed Token Pooling outperforms existing methods and provides the best trade-off between computational cost and classification accuracy.
Contributions. The paper makes the following contributions:
• We conduct an extensive study of prior downsampling techniques for visual transformers by comparing their computation-accuracy trade-offs.
• We analyze the computational cost of vision-transformer components and the limitations of the prior score-based downsampling methods. We also show that attention layers behave like low-pass filtering and thus produce redundant tokens.
• Motivated by the redundancy in images and features, we propose a novel token downsampling technique, Token Pooling, for transformers with global attention via error minimization and NEW achieve a significant improvement in the computation-accuracy trade-off.

heading: 2 RELATED WORK
text: In this section, we introduce vision transformers, and review existing methods that improve the efficiency of transformers including existing token-downsampling methods.

heading: 2.1 VISION TRANSFORMERS
text: Vision transformers (Dosovitskiy et al., 2020; Touvron et al., 2021; Heo et al., 2021; Liu et al., 2021; Pan et al., 2021) utilize the transformer architecture that is originally designed for NLP by Vaswani et al. (2017) and further popularized by Radford et al. (2018) and Devlin et al. (2019). In a high level, a vision transformer is a composition ofL transformer blocks that take a set of input tokens and return another set of output tokens. In vision, input tokens are features representing individual nonoverlapping image patches. To perform classification, a classification token is inserted to estimate the probabilities of individual classes. To achieve the state-of-the-art, ViT (Dosovitskiy et al., 2020) used pretraining on JFT-300M, a proprietary dataset much larger than standard ImageNet1k (Deng et al., 2009). Recently, DeiT (Touvron et al., 2021) achieved state-of-the-art results with advanced training on ImageNet1k only.
Let the set of tokens at depth l be F l = {f l0, . . . ,f lN} where f li ∈ RM is the feature values of the i-th token. A typical transformer block φ at depth l processes F l by a Multi-head Self-Attention (MSA) layer and a point-wise Multi-Layer Perceptron (MLP). Let matrix F ∈ RN×M be a row-wise concatenation of tokens F l. Then1,
φ(F ) = MLP(MSA(F )) such that (1)
MSA(F ) = [O1,O2, . . . ,OH ]W O (2)
where H is the number of heads, matrix WO ∈ RM×M is a learnable parameter of the block, [, ] is column-wise concatenation and Oh ∈ RN×d is the output of h-th attention head for d =M/H:
Oh = AhVh such that Ah = softmax(QhK>h /√ d) ∈ RN×N . (3)
Keys Kh, queries Qh and values Vh are linear projections of the input tokens (QKV projections):
Qh = FW Q h , Kh = FW K h , Vh = FW V h (4)
where WQh ∈ RM×d, WKh ∈ RM×d, W Vh ∈ RM×d are learnable linear transformations. Note, the number of tokens is not affected by the transformer blocks, i.e., |F l+1| = |F l|.

heading: 2.2 EFFICIENT TRANSFORMERS
text: Similar to many machine learning models, the efficiency of transformers can be improved via metaparameter search (Howard et al., 2017; Tan & Le, 2019), automated neural architecture search (Elsken et al., 2019; Tan et al., 2019; Wu et al., 2019), manipulating the input size and resolution of feature maps (Paszke et al., 2016; Howard et al., 2017; Zhao et al., 2018), pruning (LeCun et al., 1990), quantization (Jacob et al., 2018), and sparsification (Gale et al., 2019), etc. For example, Dosovitskiy et al. (2020) and Touvron et al. (2021) obtain a family of ViT and DeiT models, respectively, by varying the input resolution, the number of heads H , and the feature dimensionality M . Each of the models operates with a different computational requirement and accuracy. In the following, we review techniques developed for transformers.

heading: 2.2.1 EFFICIENT SELF-ATTENTION
text: The softmax-attention layer (3) has a quadratic time complexity w.r.t. the number of tokens, i.e., O(N2). In many NLP applications where every token represents a word or a character, N can be large, making attention a computation bottleneck (Dai et al., 2019; Rae et al., 2019). While many works improve the time complexity of attention layers, as we will see in §3.1, they are not the bottleneck in most current vision transformers.
The time complexity of an attention layer can be reduced by restricting the attention field-of-view and thus imposing sparsity on Ah in (3). This can be achieved using the spatial relationship between tokens in the image/text domain (Parmar et al., 2018; Ramachandran et al., 2019; Qiu et al., 2020; Beltagy et al., 2020; Child et al., 2019; Zaheer et al., 2020) or based on token values using localitysensitive hashing, sorting, compression, etc. (Kitaev et al., 2020; Vyas et al., 2020; Tay et al., 2020a; Liu et al., 2018; Wang et al., 2020; Tay et al., 2021). Prior works have also proposed attention mechanisms with lower time complexity, e.g., O(N) or O(N logN) (Katharopoulos et al., 2020;
Peng et al., 2021; Choromanski et al., 2021; Tay et al., 2021). Roy et al. (2021) cluster queries and NEW keys to sparsify attention matrices to speed-up the attention, but they do not downsample the tokens.
Note that the goal of these methods is to reduce the time complexity of the attention layer—the number of tokens remains the same across the transformer blocks. In contrast, our method reduces the number of tokens after attention has been computed. Thereby, we can utilize these methods to further improve the overall efficiency of transformers.
Recently, Wu et al. (2020) proposed a new attention-based layer that learns a small number of query vectors to extract information from the input feature map. Similarly, Wu et al. (2021) replace selfattention with a new recurrent layer that outputs a smaller number of tokens. In comparison, our method directly minimizes the token reconstruction error due to token downsampling. Also, our layer has no learnable parameters and can be easily incorporated into existing vision transformers.

heading: 2.2.2 DOWNSAMPLING METHODS FOR TRANSFORMERS
text: Grid-downsampling. The input tokens of the first vision transformer block are computed from image patches. Therefore, even though transformers treat tokens as a set, we can still associate a grid to the tokens using their initial locations on the image. The regular grid structure allows typical downsampling methods, such as max/mean pooling, uniform sub-sampling. For example, Liu et al. (2021), Heo et al. (2021), and Wang et al. (2021) use convolutions with stride to downsample the feature maps formed by the tokens.
Score-based token downsampling. In the area of NLP, Goyal et al. (2020) introduce the idea of dropping tokens. Their approach, called PoWER-BERT, is based on a measure of significance score, which is defined as the total attention given to a token from all other tokens. Specifically, the significance scores of all tokens in the l-th transformer block, sl ∈ RN , is computed by
sl = H∑ h=1 Alh > 1, (5)
where Alh is the attention weights of head h defined in (3). They only pass Kl tokens with the highest scores in sl to the next transformer block. The pruning is performed on all blocks.
PoWER-BERT is trained using a three-stage process. First, given a base architecture, they pretrain a model without pruning. In the second stage, a soft-selection layer is inserted after each transformer block, and the model is finetuned for a small number of epochs. Once learned, the number of tokens to keep, Kl, for each layer is computed from the soft-selection layers. Last, the model is finetuned again with the tokens pruned using the Kl from the second stage. See Goyal et al. (2020) for details.
Recently, Rao et al. (2021) proposed Dynamic-ViT that also uses scores to prune tokens. Unlike PoWER-BERT, which computes significance scores from attention weights, Rao et al. use a dedicated sub-network with learned parameters. The method requires knowledge distillation, GumbelSoftmax, and straight-through estimators on top of the DeiT training.
We will analyze the limitations of score-based methods in §3.2. Note that while downsampling NEW tokens can complicate some downstream tasks, e.g., semantic segmentation and pose estimation, many techniques Wu et al. (2020); Marin et al. (2019) have been proposed to deal with the problem.

heading: 3 ANALYSIS
text: This section addresses three questions. First, it identifies the computational bottleneck of vision transformers. Second, it discusses the limitations of score-based downsampling. Third, it analyzes how the softmax-attention affects the redundancy in tokens.

heading: 3.1 COMPUTATION ANALYSIS OF VISION TRANSFORMERS
text: We analyze the time complexity and computational costs (measured in flops) of commonly used vision transformers, namely ViT and DeiT. We breakdown the computation into four categories:
1For compactness, we omit layer norm and skip-connections, see Dosovitskiy et al. (2020) for details.
softmax-attention (3), QKV projections (4), O projection (2) and MLP (1). As shown in Table 1, in all these vision transformers, the main computational bottleneck is the fully-connected layers that spend over 80% of the total computation. In comparison, softmax-attention only takes less than 15%. Note that we explicitly break down the multi-head attention into the softmax-attention, QKV and O projections, as they have different time complexity, see Table 1. This decomposition reveals that the QKV and O projections spend most of the computations of the multi-head self-attention.

heading: 3.2 LIMITATIONS OF SCORE-BASED TOKEN DOWNSAMPLING
text: Existing score-based token downsampling methods like PoWER-BERT and Dynamic-ViT utilize scoring functions to determine the tokens to keep (or prune). They keep tokens with the top-K scores and discard the rest. Since these scoring functions are continuous with limited Lipschitz constants, tokens that are close in the feature space will be assigned similar scores. Therefore, these similar tokens will likely either be all kept or discarded, as illustrated in Figure 2a. As our experiments show, this redundancy (in the kept tokens) and severe information loss (in the pruned tokens) deteriorate the computation-accuracy trade-off of the score-based downsampling methods.

heading: 3.3 ATTENTION AS A LOW-PASS FILTER
text: Given a query vector q, a set of key vectors K = {k1, . . . ,kN}, the corresponding value vectors V = {v1, . . . ,vN} and a scalar α > 0, softmax-attention computes the output via
o(q) = 1
z(q) N∑ i=1 exp(α q · ki)vi where z(q) = N∑ i=1 exp(α q · ki). (6)
Note that we write o(q) to indicate that the output vector o is a function of the query q. If the query vector and all key vectors are normalized to have a fixed `2 norm, we can rewrite (6) as
o(q) = 1
z′(q) N∑ i=1 exp ( −α 2 ‖q − ki‖2 ) vi = 1 z′(q) ∫ exp ( −α 2 ‖q − k‖2 )( N∑ i=1 δ(k − ki)vi ) dk
= 1
z′(q) G
( q; 1
α
) ∗ S(q;K,V) (7)
where ∗ represents high-dimensional convolution, z′(q) = ∑ i exp ( −α2 ‖q − ki‖ 2 ) = G ( q; 1α ) ∗
S(q;K, 1) is the normalization scalar function,G ( q;σ2 ) = exp ( −‖q‖2 / 2σ2 )
is an isometric Gaussian kernel, and S(q;K,V) = ∑N i=1 δ(q − ki)vi is a high-dimensional sparse signal, which is composed of N delta functions located at ki with value vi. According to (7), given query vectors q1, . . . , qN , there are two conceptual steps to compute softmax-attention:
1. filter S(q;K,V) with a Gaussian kernel to get o(q), and 2. sample o(q) at coordinates q1, . . . , qN to get the output vectors o1, . . . ,oN .
Since Gaussian filtering is low-pass, o(q) is a smooth signal. Therefore, the output tokens of the attention layer, i.e., discrete samples of o(q), contain similar feature values. So, based on NyquistShannon sampling theorem (Shannon, 1949), there exists redundant information in the output, which is used by our methods to reduce computation without losing much of the important information.
Note that our analysis is based on the normalized query and key vectors, which can be achieved by inserting a normalizing layer before the softmax-attention layer without significantly affecting the performance of a transformer, as shown by Kitaev et al. (2020). It has also been empirically observed by Goyal et al. (2020) and Rao et al. (2021) that even without the normalization, transformers produce tokens with similar values. To demonstrate this, in all our experiments, we use standard multi-head attention without normalizing keys and queries. We conduct an ablation study with normalized keys and queries in Appendix F.

heading: 4 TOKEN POOLING
text: Pruning tokens inevitably loses information. In this section, we formulate a new token downsampling principle enabling strategical tokens selection that preserves the most information. Based on this principle, we formulate and discuss several Token Pooling algorithms.
Given a set of output tokens F = {f1, . . . ,fN} of a transformer block, our goal is to find a smaller set of tokens F̂ = {f̂1, . . . , f̂K} that after upsampling minimizes the reconstruction error of F . Specifically, the reconstruction of u(fi; F̂) of token fi is computed via interpolating the tokens in F̂ . The reconstruction error is then defined as
`(F , F̂) = ∑ fi∈F ‖fi − u(fi; F̂)‖2. (8)
To simplify our formulation and reduce computation, we use nearest-neighbor interpolation as u. As a consequence, the reconstruction error (8) becomes
`(F , F̂) = ∑ fi∈F min f̂j∈F̂ ‖fi − f̂j‖2, (9)
which is the asymmetric Chamfer divergence between F and F̂ (Barrow et al., 1977; Mechrez et al., 2018). The loss (9) can be minimized by the K-Means algorithm, i.e., clustering the tokens in F into K clusters.
Algorithm 1: Token Pooling input : tokens F = {f1, . . . ,fN}; target set size K; (optional) weightsW = {w1, . . . , wN} output: downsampled set F̂ = {f̂1, . . . , f̂K}
1 if k ≥ N then return F; 2 initialize cluster centers F̂ to be the K tokens from F with the highest weights ; 3 while not converged and max number of iterations is not reached do 4 for i ∈ {1, ..., N} do update cluster assignment z(i)← argminKj=1 ‖fi − f̂j‖ ; 5 for j ∈ {1, ...,K} do update cluster center f̂j according to the chosen clustering algorithm, that is either (11), (12), (15) or (18) ; 6 return weighted means of tokens in each cluster
The proposed Token Pooling layer is defined in Algorithm 1. It downsamples input tokens via clustering the tokens and returns the cluster centers (the average of the tokens in a cluster). As we have shown above, this operation directly minimizes the reconstruction error (9) caused by the downsampling. Intuitively, clustering the tokens provides a more accurate and diverse representation of the original set of tokens, compared to the top-K selection used by score-based downsampling methods, as shown in Figure 2b. Note that Token Pooling is robust to the initialization of cluster centers, as shown in Appendix B. Below, we provide details of the clustering algorithms. K-Means. We use the K-Means algorithm to minimize (9) via the following iterations:
a(i) ← argmin j ‖fi − f̂j‖ ∀i ∈ {1, 2, . . . , N} (10)
f̂j ← N∑ i=1 [a(i) = j]fi / N∑ i=1 [a(i) = j] ∀j ∈ {1, 2, . . . ,K} (11)
where [ ] is the Iverson bracket and a is the cluster assignment function. The overall algorithm complexity is O(TKNM) where T is the number of iterations. The vast majority of the computation is spent on the repetitive evaluation of the distances between tokens and centroids in step (10). K-Medoids. We can use the more efficient K-Medoids algorithm by replacing step (11) with:
n(j) ← arg min i:a(i)=j ∑ i′:a(i′)=j ‖fi − fi′‖2 and f̂j ← fn(j). (12)
These steps minimize objective (9) under the medoid constraint: F̂ ⊆ F . The advantage of the K-Medoids algorithm is its time complexity O(TKN + N2M), which is substantially lower in practice as we only compute the distances between tokens once. In our experience, it requires less than 5 iterations to converge. Apart from the distance matrix computation, the cost of the K-Medoids algorithm is negligible when compared with the cost of the other layers. Weighted clustering. Reconstruction error (9) treats every token equally; however, each token contributes differently to the final output of an attention layer. Thus, we also consider a weighted reconstruction error: `(F , F̂ ;w) = ∑ fi∈F minf̂j∈F̂ wi‖fi − f̂j‖
2 where w = [w1, . . . , wN ] are the positive weights corresponding to the individual tokens in F . Appendix A details the clustering algorithms for the weighted case. A good choice of the weights is the significance scores (5), i.e., wi = si. The significance score identifies the tokens that influence the current transformer block most and thus should be approximated more precisely.
Note, when applying Token Pooling, it is important to keep special tokens like the classification NEW token. If spatial locations on an image are need, e.g., in the case of semantic segmentation or pose estimation, techniques like (Wu et al., 2020) can be utilized to project tokens back to the image grid.

heading: 5 EXPERIMENTS
text: Our implementation is based on DeiT (Touvron et al., 2021) where we insert downsampling layers after each transformer block. To evaluate the pure effect of downsampling, we keep all metaparameters of DeiT, including the feature dimensionality, network depth, learning rate schedule, etc.
We do not use knowledge distillation. We use ImageNet1k (Russakovsky et al., 2015). Our cost and performance metrics are flops and top-1 accuracy. Appendix H reports throughput of our models. NEW Methods. We evaluate the following token downsampling methods: 1. Convolution downsampling. We implement uniform grid-downsampling via convolution with
stride 2, i.e., the tokens corresponding to the adjacent image patches are concatenated and mapped to RM . This design is used in Liu et al. (2021); Heo et al. (2021). See details in Appendix D.
2. PoWER-BERT. We implement PoWER-BERT on DeiT, following Goyal et al. (2020).
3. Random selection, a simple baseline randomly selecting K tokens without replacement.
4. Importance selection chooses K tokens by drawing from a categorical distribution without replacement with probabilities proportional to the significance score (5) of each token.
5. Token Pooling use K-Means or K-Medoids algorithms, or their weighted versions, WK-Means or WK-Medoids, respectively. The weights are the significance scores (5).
Selection of K = [K1, . . . ,KL]. To fairly compare our Token Pooling with PoWER-BERT and other baselines, all methods (except convolution downsampling) use the same number of retained tokens for downsampling layers. Appendix C details the selection of K.
Training protocol. (1) A base DeiT model (e.g., DeiT-S) is trained using the original training (Touvron et al., 2021). (2) We then finetune the base DeiT model using the second stage of PoWERBERT’s training and acquire K. (3) We further finetune the downsampling methods using K. (4) We also finetune the base DeiT model. Our protocol ensures a fair comparison such that all of the models are trained with the same number of iterations, the same learning rate schedule, etc.

heading: 5.1 MAIN RESULTS
text: First, we apply different downsampling methods on DeiT-S. As shown in Figure 3a, random selection achieves a similar trade-off as lowering feature dimensionality M . While convolution with stride is better than adjusting M at the low-compute regime, it fails in high-compute regimes. Importance selection improves upon random selection but is still outperformed by PoWER-BERT. Our Token Pooling (with weighted K-Medoids) achieves the best trade-off in all regimes.
Next, we apply Token Pooling to DeiT models with different M (DeiT-Ti, DeiT-e252, DeiT-e318, and DeiT-S). Figure 4 shows trade-off curves for each DeiT model. Token Pooling enables each of the models to achieve a better computation-accuracy trade-off than simply varying the feature dimensionality M . For each computational budget, we find the combination of M and K that gives the highest accuracy. The best balance is achieved by applying Token Pooling and selecting the best
1 2 3 4
Gflops
72
74
76
78
80
ac cu
ra cy
( % ) DeiT-e252
DeiT-e318 DeiT-S
Ours (on DeiT-S)
Ours (on DeiT-e318)
Ours (on DeiT-e252)
Figure 4: The figure shows the results when we apply Token Pooling to various DeiT architectures. Token Pooling consistently improves the computation-accuracy trade-off for all evaluated architectures. By utilizing both Token Pooling and architecture search, we can further improve the accuracy at a given flops budget. For example, at 1 Gflop, we should use Token Pooling on DeiT-e252 instead of DeiT-S.
1 1.5 2 2.5 3 3.5 4
Gflops
74
76
78
80
ac cu
ra cy
( %
)
Ours (WK-Medoids)
PoWER-BERT
Importance Selection
Random Selection
(a) Ours vs. methods using significance score
1 2 3 4
Gflops
74
76
78
80
ac cu
ra cy
( %
)
(Ours) WK-Medoids
(Ours) K-Medoids
(Ours) WK-Means
(Ours) K-Means
PoWER-BERT
(b) Variants of Token Pooling
Figure 5: Ablation studies of (a) downsampling methods using significance score, and (b) proposed Token Pooling using different clustering algorithms. The base model is DeiT-S for all methods.
M . We do the same for PoWER-BERT. Figure 3b shows the results of the proposed Token Pooling and state-of-the-art methods. We cite the results of Rao et al. (2021), d’Ascoli et al. (2021), and Pan et al. (2021). Our Token Pooling achieves the best accuracy across all evaluated compute regimes.
Finally, we compare the best computation-accuracy trade-off achieved by Token Pooling (with WKMedoids and varying M ) with standard DeiT models in Figure 1b. As can be seen, utilizing Token Pooling, we significantly improve the computation costs of DeiT-Ti by 42% and improve the top-1 accuracy by 3.3 points at the same flops. Similar benefits can be seen on DeiT-e252 and DeiT-e318.

heading: 5.2 ABLATION STUDIES
text: Figure 5a compares methods utilizing significance scores. As can be seen, using importance selection improves upon the simple random selection. By minimizing the reconstruction error (9), our method achieves better cost-accuracy trade-off. Figure 5b evaluates Token Pooling with different clustering algorithms. Weighted versions outperform regular versions of K-Means and K-Medoids. Due to the higher time complexity, K-Means is outperformed by K-Medoids (the curves are shifted toward the right). See the metrics and flops used by clustering in table format in Appendix G. All Token Pooling variants outperform the baseline, demonstrating the effectiveness of our method.
More ablation studies are in the appendices: clustering initialization (Appendix B), convolution downsampling (D), Token Pooling without finetuning (E), results with normalized keys and queries (F), and detailed information (K, flops, accuracy, clustering cost) of our models (G).

heading: 6 CONCLUSIONS
text: This paper provides two insights of vision transformers: first, their computational bottleneck is the fully-connected layers, and second, attention layers generate redundant representations due to the connection to Gaussian filtering. Token Pooling, our novel nonuniform data-aware downsampling operator, utilizes these insights and significantly improves the computation-accuracy balance of DeiT, compared to existing downsampling techniques. We believe that Token Pooling is a timely development and can be used with other techniques (e.g., architecture search, quantization) to spur the design of efficient vision transformers.

heading: A WEIGHTED CLUSTERING ALGORITHMS
text: Weighted K-Means minimizes the following objective w.r.t. F̂ = {f̂1, . . . , f̂K} ⊂ RM : `(F , F̂) = ∑ fi∈F min f̂j∈F̂ wi‖fi − f̂j‖2 (13)
The extension of the K-Means algorithm to the weighted case iterates the following steps:
a(i) ← argmin j ‖fi − f̂j‖ ∀i ∈ {1, 2, . . . , N}, (14)
f̂j ← ∑N i=1[a(i) = j]wifi∑N i=1[a(i) = j]wi
∀j ∈ {1, 2, . . . ,K}., (15)
Weighted K-Medoids optimizes objective (13) under the medoid constraint F̂ ⊂ F :
a(i) ← argmin j ‖fi − f̂j‖ ∀i ∈ {1, 2, . . . , N}, (16)
n(j) ← argmin i: z(i)=j ∑ i′: a(i′)=j ‖fi − fi′‖2 ∀j ∈ {1, 2, . . . ,K}, (17)
f̂j ← fn(j). (18)

heading: B CLUSTERING INITIALIZATION ABLATIONS
text: We examine the effect of the cluster center initialization. We compare our default initialization, which uses the tokens with top-K significance scores as initial cluster centers, with random initialization, which randomly selects tokens as initial cluster centers. As shown in Figure 6, Token Pooling is robust to the initialization methods.

heading: C TRAINING DETAILS & HYPERPARAMETERS
text: A Token Pooling layer has one parameter—the number of the output tokens K (or equivalently the NEW downsampling ratio)—which is the same as downsampling layers like max or average pooling. For max or average pooling layers, the sizes of the kernels and strides play the same role as K. We use the method proposed by Goyal et al. (2020) to automatically select Ks of all Token Pooling layers (see below for details). All other hyperparameters are inherited from DeiT (Touvron et al., 2021).
To fairly compare PoWER-BERT and the baseline methods with the proposed Token Pooling, all methods (except convolution downsampling) use the same target number of tokens for downsampling layers after each transformer block. Specifically, we run the second stage of the PoWER-BERT training (Goyal et al., 2020) for 30 epochs with various values of the token-selection weight parameter λ producing a family of models. The single parameter λ controls the overall efficiency of a model, and the numbers of retained tokens of all Token Pooling layers are selected automatically based on the choice of λ. Each of the resulted models has a different number of retained tokens at each of its L transformer blocks: K = (K1, . . . ,KL). Appendix G lists all combinations of automatically determined K. We then finetune these models using the third (last) stage of PoWER-BERT. Note that we apply the same process for PoWER-BERT, random selection, importance selection, and our Token Pooling, and we use the same K when comparing these methods.
We find that the DeiT models provided by Touvron et al. (2021) are under-fit, and their accuracy improves with additional training, see Figure 7. After the standard DeiT training, we restart the training. This ensures that downsampling models and DeiT with a similar number of training steps.

heading: D CONVOLUTION DOWNSAMPLING
text: As mentioned in the main paper, we enumerate combinations of layers to insert the convolution downsampling layer. We use 2x2 convolution with stride 2 like in Liu et al. (2021). To keep the feature dimensionality of DeiT the same (and evaluate the pure effect of the downsampling layers), the output feature dimensionality is the same as the input. With 196 tokens in DeiT-S model, we can include no more than 3 convolution downsampling layers as each layer reduces the number of token by a factor of 4. When using 3 layers at depths l1 < l2 < l3, we restrict l2− l1 = l3− l2. With these constraints, we enumerate all possible downsampling configurations. Each of the combinations produces a model with a different computation-accuracy trade-off, and we report the Pareto front, i.e., the best accuracy these models achieve at a given flop. See Figure 8.
Note that the convolution downsampling layer is equivalent to the patch-merging layer used by Liu NEW et al. (2021), Heo et al. (2021), and Wang et al. (2021) (except that we keep the output feature dimension the same to compare fairly with DeiT). It is also a generalized version of the commonly used average pooling layer. Average pooling can be implemented by convolving the feature map with a constant kernel of size n×n containing 1n2 and subsample (stride) every n pixels. By learning its kernel, the convolution downsampling layer is able to find a kernel suitable for the downsampling and the current task than average pooling.

heading: E RESULTS WITHOUT FINETUNING
text: Since Token Pooling minimizes the reconstruction error due to token downsampling, in this section, we evaluate the performance when we insert Token Pooling layers into a pre-trained model without finetuning. Figure 9 shows the results when we directly insert Token Pooling layers (using the same downsampling ratios in Figure 5b) into a pretrained DeiT-S. As can be seen, minimizing the reconstruction error, Token Pooling preserves information that enables the model to retain accuracy during token downsampling.
In Figure 9, we also show the results when we replace tokens to their cluster centers (WK-Means, carry, no finetuning and WK-Medoids, carry, no finetuning). Specifically, in addition to outputting K cluster centers, we count the number of tokens assigned to a cluster center and carry the count when we compute softmax-attention in the next transformer blocks. This operation preserves the attention weights, and since the models are not trained after inserting Token Pooling layers, it preserves the most accuracy. In practice, when the models are trained with Token Pooling layers, the carry operation is not important.

heading: F NORMALIZED KEY AND QUERY VECTORS
text: Our analysis of softmax-attention in §3.3 assumes the query and the key vectors are normalized to have a constant `2 norm. It has been observed by Kitaev et al. (2020) that normalizing key and query vectors does not change the performance of a transformer. Thus, for all experiments, we train models without the norm normalization. In this section, we verify this observation by training various DeiT, PoWER-BERT, and Token Pooling with normalized keys and queries. We found that the scalar α in the softmax attention layer (6) can affect the performance of a transformer with normalized keys and queries. As can be seen in Figure 10a, setting the scalar α = 1 slightly deteriorates the performance of the model. Instead of using a fixed α, we let the model learn the α for each layer. As can be seen in Figure 10b, learning α enables the resulting models to achieve similar cost-accuracy trade-off as the standard (unnormalized) models. With or without the normalization and the learnable α, the proposed Token Pooling significantly improves the cost-accuracy trade-off of DeiT and outperforms PoWER-BERT.

heading: G CLUSTERING ALGORITHM ABLATIONS
text: Tables 2–5 detail the results of PoWER-BERT and Token Pooling on the DeiT architectures that we tested (DeiT-S, DeiT-e318, and DeiT-e252). Figure 11 shows the ablation results with different backbone models. Table 2 details the results of the best cost-accuracy trade-off achieved by the proposed Token Pooling (using K-Medoids and WK-Medoids) and PoWER-BERT via varying token sparsity and feature dimensionality of DeiT.
Apart from the standard K-Means and K-Medoids, other clustering approaches could be used. Many methods are not suitable due to efficiency constraints. For example, normalized cut (Shi & Malik, 2000) uses expensive spectral methods. One prerequisite of using K-Means and K-Medoids is the number of clusters, K. While selecting K for each of the layers may be tedious and difficult, one can choose K via computational budgets and heuristics. In this work, we use the automatic search procedure proposed by Goyal et al. (2020) to determine K, see Appendix C.
Since both K-Means and K-Medoids require specifying the number of clusters K in advance, one may consider using methods automatically determining K. Nevertheless, such methods typically have other parameters, which are less interpretable than K. For example, mean-shift (Cheng, 1995) or quick-shift (Vedaldi & Soatto, 2008) require specifying the kernel size. From our experience, determining these parameters is challenging. Also, since the number of clusters (and hence the number of output tokens) is determined on the fly during inference, the computational requirement can fluctuate, making deployment of these models difficult.
H THROUGHPUT NEW
Throughput reflects the actual speed during inference time, and it is often measured by the number of processed images per second (fps). As a result, the value of throughput highly depends on the specific hardware, implementation quality, and the workload and the state (e.g., temperature) of the machines. Since it is usually difficult to control all these factors to have a fair comparison with other papers, in the main paper, we choose to report the theoretical computational cost (i.e., flops) that is known to be highly correlated with the energy consumption on device.
In Figure 12, we provide the throughput of our models, served as a reference on how our models actually perform given our implementation and the type of GPU used. Note that our PyTorch implementation of the clustering algorithms can be significantly improved, for example, via implementing as a CUDA kernel. To determine the throughput, we run the model inference several times with different batch sizes. We report the average throughput of 30 different runs using the best batch size. As can be seen, with the same accuracy, using Token Pooling increases the throughput of DeiT-e318 by 25% (1640 vs. 1310 fps), of DeiT-e252 by 22% (2190 vs. 1791 fps), and of DeiT-Ti by 14% (3220 vs. 2834 fps).

Review summaries:
In this paper, the authors propose a novel token-pooling method to reduce redundancies in tokens for recent vision transformers. They analyze the computation cost distribution of vision transformers and the limitations of grid-based & score-based token downsampling methods. They further formulate a reconstruction loss and optimize it with the token pooling layer. The experimental results based on DeiT show that token pooling improves accuracy while reducing computation cost by a large margin. This idea can be applied to other vision transformers as well.

This paper proposes a new token downsampling method for vision transformer, called Token Pooling, to prune redundant tokens efficiently, so as to achieve a better flop-accuracy trade-off. Specifically, token pooling is a nonuniform data-aware downsampling method, which uses cluster algorithms to aggregate information from tokens automatically. To keep important information, the authors also proposed minimizing the reconstruction loss caused by downsampling. The authors performed experiments on the ImageNet-1k dataset, showing that the proposed token pooling can significantly improve the flop-accuracy trade-off over the existing downsampling methods.