Token Pooling in Vision Transformers
Authors: Dmitrii Marin
Jen-Hao Rick Chang
Anurag Ranjan
Anish Prabhu
Mohammad Rastegari
Oncel Tuzel
Venue: ICLR 2022 Submitted (2022)
Keywords: Transformer
Pooling
Downsampling
Efficiency

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 a simple and effective operator that can benefit many architectures. Applied to DeiT, it achieves the same ImageNet top-1 accuracy using 42% fewer computations.

Under review as a conference paper at ICLR 2022
TOKEN POOLING IN VISION TRANSFORMERS
Anonymous authors
Paper under double-blind review
ABSTRACT
Despite the recent success in many applications, the high computational require- ments of vision transformers limit their use in resource-constrainee settings. While many existing methods improve the quadratic complexity of attention, 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 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 high-dimensional low-pass (smoothing) filter. Thus, its output contains redundancy that can be pruned to achieve better trade-off between the computational cost and accuracy. Our new technique accu- rately approximates a set of tokens by minimizing the reconstruction enror caused by downsampling. We solve this optimization problem via cost-efficient cluster- ing. We rigorously analyze and compare prior downsampling methods. Our experiments show that Token Pooling significantly improves the cost-accuracy trade-off over the state-of-the-ar downsampling. Token Pooling simple and ef- fective and can benefit many architectures with global attention. Applied to DeiT, it achieves the same ImageNet top-1 accuracy using 42% fewer computations.
NEW
INTRODUCTION
Vision transformers (Dosovitskiy et al., 2020; Touvron et al., 2021; Liu et al., 2021; Heo et al., 2021: Zheng 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 resturee-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, assumptions, softmax-attention is equivalent to low-pass filtering of tokens and thereby produces tokens with similar features, 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., down- sampling. without a significant impact to the accuracy, achieving a better cost-accuracy trade-off than reducing feature dimensionality alone.
show that, under mild
Downsampling is widely used in Convolutional Neural Network (CNN) architectures 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 prevail- ing max/average pooling and sub-sampling that only uses locations on the grid to decide which vertices to combine. Such methods do not effi- ciently address con-uniformly distributed redundancy in images and features (Recasens al., 2018; Marin et al., 2019). Unlike CNNs that require grid preservation, transformers allow a wider range of nonuniform data-aware downsampling layers, where better operator be designed.
examples of (spatially uniform) fridddowwssmppling
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-

---

Under review as a conference paper at ICLR 2022
typical transformer
with proposed Token Pooling
+0.7
80
-34%
DeiT-e318
79
Token Poaling
+1.6
tramsformer block
transformer block
-27%
DeiT-e252
77
76
+3.3
transfomerr block
Token Pooling
75
DefT (various dimensions)
tramsformer block
74
-42%
DoiT-Ti
Proposed Token Pooling
0.7
1.5
2.5
Gflops
(a) Proposed Token Pooling
(b) Accuracy computation
(c) Token Pooling cluster analysis token representations.
Figure 1: (a) We propose Token Pooling, novel token downsampling method, for visual trans- formers. (b) The proposed method achieves a state-of-the-arr trade-off between accuracy and com-
putation. (c) Token Pooling uses cluster analysis to aggregate information from individual tokens automatically. We show the input images and the token clusters at the 6-th layer of DeiT-S.
ures & 1b. Motivated by nonuniform sampling and image compression (Marvasti, 2012; Unat et al., 2009; Belfor et al., 1994; Rabbani, 2002), we formulate token downsampling mization 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 Ic. To the best of our knowledge, we are the first to use this formulation and simple cluster- ing analysis for token donwsampling 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 al., 2021). The proposed Token Pooling outperforms exist- ing methods and provides the best trade-off between computational cost and classification accuracy.
opti-
Contributions. The paper makes the following contributions:
We conduct an extensive study of prior downsampling techniques for visual transformers by com- paring their computation-accuracy trade-offs.
We analyze the computational cost of tision-tranffrmerr components and the limitations of the prior boore 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 achieve significant improvement in the computation-accural trade-off.
NEW
2 RELATED WORK
In this section, we introduce vision transformers, and review existing methods that improve the efficiency of transformers including existing token-downsampling methods.
2

---

Under review as a conference paper at ICLR 2022
2.1 VISION TRANSFORMERS
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 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 composition of I. transformer blocks that take set of input tokens and
level, a vision transformer return another set of output tokens. In vision, imput tokens are features representing individual non- overlapping image patches. To perform classification, a classification token is inserted to estimate the probabilities of individual classes. To achieve the state-of-the-an ViT (Dosovitskiy et al., 2020) used pretraining on JFT-300M, a proprietary dataset much larger than standard ImageNetIk (Deng et al., 2009). Recently, DeiT (Touvron et al., 2021) achieved state-of fiffthe-an art results with advanced training on ImageNetlk only.
J
Let the set of tokens at depth 1 be F = (f f'} where f R.M. is the fea- ture values of the i-th token. A typical transformer block 6 at depth I processes Fi by a Multi-head Self-Attention (MSA) layer and a point-wise Multi-Layer Perceptron (MLP). Let matrix F E IRNXM be row-wise concatenation of tokens Fr, Then
MLP
O-projection
softmax- attention
QKV-
Φ(F) (LP(MSA(F)) such that MSA(F) [0,, O2, ....., On]
projection
(2)
where H is the number of heads, matrix W E MRM is learnable parameter of the block, [.] is column-wise concatenation and Oj E IR Nxd is the output of /-th attention head for d = M/H:
O₂ = A such that A = softmax Q K /Vd) E NNN.
s
(3)
Keys KR, queries Qn and values Vn are linear projections of the input tokens (QKV projections):
Qx FW
K₁ FWK Vn = FWW
where 8º Mxd, w E Mxd W E Mxd are learnable linear transformations. Note, the number tokens is not affected by the transformer blocks, ie., | FT|.
2.2 EFFICIENT TRANSFORMERS
Similar to many machine learning models, the efficiency of transformers can be improved via meta- parameter 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 resolu- tion of feature maps (Paszke al., 2016; Howard al., 2017; Zhao al., 2018), pruning (LeCun et al., 1990), quantization (Jacob et al., 2018), and sparsification (Gale et al., 2019), etc. For exam- ple, Dosovitskiy et al. (2020) and Touvron 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.
2.2.1 EFFICIENT SELF-ATTENTION
The softmax-altention layer (3) has a quadratic time complexity w.r.t. the number of tokens, i.e., O(N3). In many NLP applications where every token represents word or a character, N can be large, making attention a computation bottleneck (Dai et al., 2019; Rae 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 A₁ 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 locality- sensitive 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., ONN) O(N (NNoo N) (Katharopoulos et al., 2020;
3

---

utilize these methods to Under review as a conference paper at ICLR 2022
Peng et al., 2021: Choromanski al., 2021; Tay et al., 2021). Roy et al. (2021) cluster queries and keys to sparsify attention matrices to speed-up the attention, but they do not downsample the tokens.
NEW
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, further improve the overall efficiency of transformers.
Recently, Wu et al. (2020) proposed new attention- based layer that learns small number of query vectors to extract information from the input feature map. Similarly, Wu et al. (2021) replace self- attention with new recurrent layer that outputs smaller number of tokens. In comparison, our method directly minimizes the token reconstruction eлrоr due to token downsampling. Also, our layer has no learnable parameters and can be easily incorporated into existing vision transformers.
2.2.2 DOWNSAMPLING METHODS FOR TRANSFORMERS
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 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 al. (2021), and Wang al. (2021) convolutions with stride to downsample the feature maps formed by the tokens.
Score-based token downsampling. In the area of NLP, Goyal al. (2020) introduce the idea of dropping tokens. Their approach, called POWER-BERT, is based on a measure of significance score, which defined as the total attention given to token from all other tokens. Specifically, the significance scores of all tokens in the I-th transformer block, s' E I is computed by
H
1,
(5)
where A is the attention weights of head h defined in (3). They only pass K₁ tokens with the highest scores in s to the next transformer block. The pruning is performed all blocks.
PoWER-BERT is trained using a three-stage process. First, given a base architecture, they pretrain a model without pruming. 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. Ki, for each layer is computed from the soft-selection layers. Last, the model is finetuned again with the tokens pruned using the K2 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 dedi- cated sub-networt with learned parameters. The method requires knowledge distillation, Gumbel- Softmax, and trraigh-tthoouhh estimators on top of the DeiT training.
We will analyze the limitations of score-based methods in $3.2. Note that while downsampling tokens can complicate some downstream tasks, e.g., semantic segmentation and pose estimation, many techniques Wu al. (2020); Marin et (2019) have been proposed deal with the problem.
NEW
ANALYSIS
This section addresses three questions. First, it identifies the computational bottleneck of vision transformers. Second, it discusses the limitations score-based downsampling. Third, analyzes how the softmax-attention affects the redundancy in tokens.
3.1 COMPUTATION ANALYSIS OF VISION TRANSFORMERS
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:
For compactness, omit layer norm and skip-connsttions, see Dosovitskiy al. (2020) for details.
4

---

Under review as a conference paper at ICLR 2022
Layer
Complexity
Computation (10" Flops)
VIT-B/384
ViT-B
DeiT-S
DeiT-Ti
577)
197)
197)
(N 197)
M)
6.18
0.72
0.36
0.18
QKV projections
OLL
12.25
4.18
1.05
0.26
projection
OLN
4.08
1.39
0.35
0.09
Multi-linear Perceptron
(
32.67
11.15
2.79
0.70
Total
MLLM NJ
55.5
17.6
4.6
1.3
8%
7%
22%
24%
21%
VIT-B/384
VIT-B
DeiT-S Flops
DefI-Ti
Flops
Flops
Flops
56%
11%
4%
63%%
8%
61%
14%
MLP
Attention
QKV projections
projection
other
Table 1: Time complexity and computation breakdown of ViT (Dosovitskiy et al., 2020) and DeiT (Touvron al., 2021). T, is the number of transformer blocks, N is the number of input tokens (patches), and M is the feature dimensionality. All models take input images of size 224 x 224 except ViT-B/384, which uses 384 x 384. The softmax-attention layers constitute fraction (15% or less) of the total compute, whereas fully-comnected layers (MLP and projections) spend over 80%.
selected tokens
selected cluster centers
scare
score
token f
token
(a) Score-based downsampling
(b) Proposed downsampling
Figure 2: Score-based downsampling methods (Goyal et al., 2020; Rao et al., 2021) vs. our method. In the figure, the x-axis represents the token values (in one dimension), and the y-axis represents their scores. Suppose four tokens are to be selected. higher scores. Since the scoring function is continuous, all tokens the left lobe will be selected, resulting in redundancy and information loss in the right lobe. (b) The proposed method first forms four clusters to approximate the set of tokens, then selects the cluster centers. Thus, the output tokens are more accurate representation of the original token set than the score-based methods.
(a) Score-based methods select tokens with
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 the fully-connected layers that spend over 80% of the total computation. In comparison, euttmax-attention only takes less than 15%. Note that we explicitly break down the multi-head attention into the softmax-attention QKV and O projections, they have different time complexity, see Table 1. This decomposition reveals that the QKV and O projections spend most of the computations the multi-head self-attention.
3.2 LIMITATIONS OF SCORE-BASED TOKEN DOWNSAMPLING
Existing score-based token downsampling methods like POWER-BERT and Dynamic-Vis 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 the feature space will be assigned similar scores. Therefore, these similar tokens will likely either be all kept discarded, experiments show, this redundancy (in the kept tokens) and severe information loss (in the pruned tokens) deteriorate the computation-aacuracy trade-off of the score-hased downsampling methods.
illustrated in Figure 2a. As
5

---

Under review as a conference paper at ICLR 2022
3.3 ATTENTION AS LOW-PASS FILTER
Given a query vector q. set of key vectors K = [ki,..., kN the corresponding value vectors {Vi,... UN and scalar a 0, softmax-attention computes the output via
N
N
( = z(g) Σ
exp(a q: k,)
where
2 - Σ expla q kt).
(6)
Note that we write o(q) to indicate that the output vector o is function of the query q. If the query vector and all key vectors are normalized have a fixed g norm, we can rewrite (6) as
((q) z'
Σ
N
exp
a -211
kill²)
vs
1 2'(4)
/
1
f exp
2 q -k1²)
Σ 8(k ks)b,
dk
q: = S(q; K,
(7)
(q)
where represents dimbnsiinn convolution, z'(g) Σ; exp(-2 10--k,, k;113) G( :)) S(q: K,, 1) the normalization scalar function, G(q;
exp(-1lqll/200) an isometric Gaus- sian kernel, and S(q: IC, V) Σ 8(q kj) 0, is high-dimensional sparse signal, which is composed of delta functions located at k; with value v; According to (7), given query vectors Q)...... QN. there are two conceptual steps to compute softmax-attention:
1. filter S(q: Kc, V) with a Gaussian kernel to get (Q)). and
2. sample o(4) at coordinates q1....: QN to get the output vectors 01,..., ON-
Since Gaussian filtering is low-pass, o()) is smooth signal. Therefore, the output tokens of the attention layer, i.e., discrete samples o(q). contain similar feature values. So, based on Nyquist- Shannon sampling theorem (Shannon, 1949), there exists redundant information in the output, which 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 be achieved by inserting a normalizing layer before the softmax-attention layer without significantly affecting the performance of transformer, observed by Goyal et al. (2020) and Rao et al. (2021) that even without the normalization, trans- formers produce tokens with similar values. To demonstrate this, in all experiments, we standard multi-head attention without normalizing keys and queries. We conduct an ablation study with normalized keys and queries in Appendix F.
shown by Kitaev et al. (2020). It has also been empirically
4
TOKEN POOLING
Pruning tokens inevitably loses information. In this section, we formulate a new token downsam- pling 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 = (f1...,fN' } of a transformer block, our goal to find smaller } that after upsampling minimizes the reconstruction error of F.
set of tokens f = (fi...., fxx Specifically, the reconstruction of F) of token fi is computed via interpolating the tokens in F. The reconstruction error is then defined as
((F
Σ - f.EF
(8)
To simplify our formulation and reduce computation, we use nearest-neighhor interpolation as u. As consequence, the reconstruction error (8) becomes
((F))) = Σ min lff. f2
F.EF feF
which is the asymmetric Chamfer divergence between F Fnd f (Barrow et al., 1977; Mechrez al., 2018). The loss (9) can be minimized by the K-Means algorithm, i.e.. clustering the tokens in J into K clusters.

---

Under review as a conference paper at ICLR 2022
Algorithm 1: Token Pooling
input : tokens F {fi.... fN }; target set size K; (optional) weights W , WN output: downsampled set F {fi...., if 2 N then return F:
initialize cluster centers F to be the K tokens from F with the highest weights while not converged and max number of iterations is not reached do
for E .I, N do update cluster assignment z(i) <- arg min Il fill: for j (1,..,.. ..., K do update cluster center f, according to the chosen clustering algorithm, that is either (11), (12), (15) or (18);
return weighted means of tokens in each cluster
The proposed Token Pooling layer is defined in Algorithm It downsamples input tokens via clustering the tokens and returns the cluster centers (the average of the tokens in cluster). As we have shown above, this operation directly minimizes the reconstruction error (9) caused by the downsampling. Intuitively, clustering the tokens provides more accurate and diverse representation of the original set of tokens, compared to the top-K selection used by score- score-based downsampling methods, shown in Figure 2b. Note that Token Pooling is robust 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) arg min Ilf: fyll
Vi E (1.2..... .)
(10)
N f Σ (a(1)) = i)f;
I(a(1)) =1
Aj (1,2,... K)
(11)
where |] is the Iverson bracket and is the cluster assignment function. The overall algorithm com- plexity is O(TKNM) where the number of iterations. The vast majority the computation is spent on the repetitive evaluation of the distances between tokens and centroids step (10).
K-Medoids. We can use the more efficient K-Medoids algorithm by replacing step (11) with: and
n(j) f arg min :=(()==
Σ lIf. fvll
f, + fnG)-
(12)
P:a(F'=}
These steps minimize objective (9) under the medoid constraint: f CCF.
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 ex- perience, it requires less than 5 iterations converge. Apart from the distance matrix computation, the cost 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 eиror: ((F,;; w) Et freF min e W; fl f2 where the positive weights corresponding 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., wy st. The significance score identifies the tokens that influence the current transformer block most and thus should be approximated more precisely.
[W],.... WN] are
Note, when applying Token Pooling, it is important keep special tokens like the classification 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 utilized to project tokens back to the image grid.
NEW
EXPERIMENTS
Our implementation is based on DeiT (Touvron al., 2021) where insert downsampling lay- ers after each transformer block. To evaluate the pure effect of downsampling, we keep all meta- parameters of DeiT, including the feature dimensionality, network depth, learning rate schedule, etc.

---

Under review as a conference paper at ICLR 2022
80
80
78
accuracy
76
accuracy
Ours (WK Medoid)
Ours (WK-M) (K-Medaids)
Ours (K-Medoid)
POWER-BERT
POWER- BERT
Importance Selection
Dynamic- ViT (Rao et al. 2021)
Random Selection
IIVT (Pan et al. 2021)
Convolution
68
ConViT (d Ascoli et al. 2021)
DeiT (Ti S)
DefT (Ti - S)
66
3
4
0.5
1
1.5
2.5
3.5 4 4.5
5.5
Gflops
Gflops
(a) Models based on DeiT-S
(b) Comparison with current SOTA
Figure 3: Main results. (a) shows the accuracy when we apply different downsampling methods to DeïT-S. More is in Appendix G. (b) shows comparison between the proposed method with the state-of-the-ae thr downsampling methods. The results of our method and PoWER-BERT BERT are acquired by varying K and the base architecture among DeiT-Ti, DeiT-e252, DeiT-e318, and DeiT-S.
We do not use knowledge distillation. We use ImageNetlk (Russakovsky et al., 2015). Our cost and performance metries are flops and top-1 accuracy. Appendix H reports throughput of our models.
NEW
Methods. We evaluate the following token downsampling methods:
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 IR M. This design used Liu et al. (2021); Heo al. (2021). See details Appendix D.
POWER BERT. We implement PoWER BERT on DeiT, following Goyal al. (2020).
Random selection, simple baseline randomly selecting K tokens without replacement.
4. Importance selection chooses K tokens by drawing from categorical distribution without placement with probabilities proportional to the significance score (5) of each token.
Token Pooling use K-Means or K-Medoids algorithms, or their weighted versions, WK-Means or WK-Medoids, respectively. The weights the significance scores (5).
Selection of K [K....,, Kz). To fairly compare our Token Pooling with POWER-BER' and other baselines, all methods (except convolution downsampling) use the same number of retained tokens for downsampling layers. Appendix C details the selection K.
Training protocol. (1) A base DeiT model (e.g., DeiT-S) trained using the original training (Touvron et al., 2021). (2) We then finetune the base DeiT model using the second stage of PoWER- BERT's training and acquire K. (3) We further finetune 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 iterations, the same learning rate schedule, etc.
5.1
MAIN RESULTS
First, we apply different downsampling methods on DeiT-S. As shown in Figure 3a, random se- lection achieves similar trade-off as lowering feature dimensionality M. While convolution with stride better than adjusting M the low-compute regime, it fails high-compute regimes. Im- portance selection improves upon random selection but is still outperformed by PoWER-BEET Our Token Pooling (with weighted K-Medoids) achieves the best trade-off all regimes.
Next, apply Token Pooling to DeÎT models with different M (DeiT-Ti, DeiT-e252, DeiT-e318. and DeiT-S). Figure shows trade-off curves for each DeiT model. Token Pooling enables each of the models to achieve better computation-accunt accuuaac 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 achieved by applying Token Pooling and selecting the best
8

---

Under review as a conference paper at ICLR 2022
Figure 4: The figure shows the results when we apply Token Pooling to various DeiTT architec- tures. Token Pooling consistently improves the computation-accuracyy 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 exam- ple, Gflop, should use Token Pooling DeiT-e252 instead of DeiT-S.
DeTT-S
80
DeTT-e318
accuracy (%)
78
Delf-ch c252
76
DeiT S)
Ours (on DeiT-S)
74
Ours ton DeiT-o318)
72
Ours (on DeiT +252)
2
4
Gflops
80
80
78
78
accuracy
(Ours) WK-Medoids
Recuiia
76
Ours (WMModdodds)
76
(Ours) A-Meedids
POWER-BERT
(Ours) WK Means
74
Importance Selection
74
(Qurs) K- K-Means
Random Selection
POWER-BERT
1.5
2
2.5
3
3.5
2
Gflloss
Gflops
(a) Ours vs. methods using significance score
(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-BERY 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-accume Medoids and varying M) with standard DeiT models in Figure 1b. As 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 DeiT-e252 and DeiT-e318.
trade-off achieved by Token Pooling (with WK-
5.2 ABLATION STUDIES
Figure 5a compares methods utilizing significance scores. As can be seen, using importance selec- tion 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).
6 CONCLUSIONS
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 downsam- pling operator, utilizes these insights and significantly improves the computation-accural balance of DeiT, compared to existing downsampling techniques. We believe that Token Pooling is timely development and can be used with other techniques (e.g., architecture search, quantization) to spur the design of efficient vision transformers.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Harry G Barrow, Jay M Tenenbaum, Robert C Bolles, and Helen C Wolf. Parametric correspon- dence and chamfer matching: Two new techniques for image matching. Technical report, SRI International Menlo Park CA Artificial Intelligence Center, 1977.
Ricardo AF Belfor, Mare PA Hesp, Reginald L Lagendijk, and Jan Biemond. Spatially adaptive subsampling of image sequences. Transactions on Image Processing, 3:492-500, 1994.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv:2004.05110, 2020.
Yizong Cheng. Mean shift, mode seeking, and clustering. IEEE Transactions on Pattern Analysis and Machine Intelligence (PAMI), 17(8):790-799, 1995.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint
2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, al. Rethinking attention with performers. International Conference on Learning Representations (ICLR), 2021.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-X XL: Attentive language models beyond fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 2978-298888 2988-298,,77 2019.
Stéphane d' Ascoli, Hugo Touvron, Matthew Leavitt, Ari Morcos, Giulio Biroli, and Levent Sagun.
Con ViT: Improving vision transformers with soft convolutional inductive biases. International Conference on Machine Learning (ICML), 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In IEEE Conference on Computer Vision and Patlern Recognition (CVPR), pp. 248-255, 2009.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 4171-4186, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghni,, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An im- age is worth 16x16 words: Transformers for image recognition scale. In International Confer- ence on Learning Representations (ICLR), 2020.
Thomas Elsken, Jan Hendrik Metzen, and Frank Hutter. Neural architecture search: A survey, The Journal of Machine Learning Research, 1997) 2019.
Trevor Gale, Erich Elsen, and Sara Hooker. The state of sparsity in deep neural networks. arXiv preprint wwXXiv:902.009575 2019.
Saurabh Goyal, Anamitra Roy Choudhury, Saurabh Raje, Venkatesan Chakaravarthy, Yogish Sab- harwal, and Ashish Verma. POWER-BERT Accelerating BERT inference via progressive word- vector elimination. In International Conference on Machine Learning (ICML), pp. 3690-3699, 2020.
Byeongho Heo, Sangdoo Yun, Dongyoon Han, Sanghyuk Chun, Junsuk Choe, and Seong Joon Oh. Rethinking spatial dimensions of vision transformers. In IEEE International Conference Computer Vision (ICCV), 2021.
Andrew G Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. MobileNets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint arXiv: 1704.04861, 2017.
10

---

Under review as a conference paper at ICLR 2022
Gabriel Ilharco, Cesar Ilharco, Iulia Ture, Tim Dettmers, Felipe Ferreira, and Kenton Lee. High performance natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: Tutorial Abstracts, pp- 24-27, 27, 2020.
Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient
inference. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2704-2713, 2018.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers RNNs: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning (ICML), pp. 5156-5165, 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference Learning Representations (ICLR), 2020.
Yann LeCun, John S Denker, and Sara Solla. Optimal brain damage. In Advances in Neural Information Processing Systems (NeurIPS), pp. 998-605, 1990.
Peter J. Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. International Conference on Learning Representations (ICLR), 2018.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows, October 2021.
Dmitrii Marin, Zijian He, Peter Vajda, Priyam Chatterjee, Sam Tsai, Fei Yang, and Yuri Boykov. Efficient segmentation: Learning downsampling near semantic boundaries. In IEEE International Conference on Computer Vision (ICCV), pp. 2131-2141, 2019.
Farokh Marvasti. Nonumiform Sampling: Theory and Practice Springer Science Business Media, 2012.
Roey Mechrez, Itamar Talmi, Firas Shama, and Lihi Zelnik-Manor. Maintaining natural image statistics with the contextual loss. In Asian Conference on Computer Vision (ACCV), pp. 427- 443, 2018.
Zizheng Pan, Bohan Zhuang, Jing Liu, Haoyu He, and Jianfei Cai. Scalable visual transformers with hierarchical pooling. IEEE International Conference Computer Vision (ICCV), 2021.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference Machine Learning (ICML), pp. PMLR, 2018.
Adam Paszke, Abhishek Chaurasia, Sangpil Kim. and Eugenio Culurciello. ENet: A deep neu- ral network architecture for real-time semantic segmentation, arXiv preprint arXiv:1606.00 02147, 2016.
Hao Peng. Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. In International Conference on Learning Representations (ICLR), 2021.
Jiezhong Qiu, Hao Ma, Omer Levy, Wen-tau Yib, Sinong Wang, and Jie Tang. Blockwise self- attention for long document understanding. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: Findings, pp. 2555-2565, 2020.
Majid Rabbani. JPEG2000: Image compression fundamentals, standards and practice. Journal of
Electronic Imaging, 11(2):286, 2002.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improv- ing language understanding by generative pre-training. 2018. URL https:
- amazonaws. http://ww...c.ooaaaaaaaaaaaaaaee Baanguage-unsueee brterrggggggggggrrrrmmmmammmaarr
/

---

Under review as a conference paper at ICLR 2022
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap- Compressive transformers for long-range sequence modelling. In International Conference on Learning Representations (ICLR), 2019.
Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jon Shlens.
Stand-alone self-attention in vision models. 32, 2019.
Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsieh. Dynamic Efficient vision transformers with dynamic token sparsification. ArXiv, abs/2106.02034, 2021.
Adria Recasens, Petr Kellnhofer, Simon Stent, Wojeiech Matusik, and Antonio Torralba. Learn- ing to zoom: a saliency- based sampling layer for neural networks. In European Conference on Computer Vision (ECCV), pp. 51-66, 2018.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguis- tics, 9:53-68, 2021.
Olga Russakovsky, Jia Deng. Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. ImageNet large scale visual recognition challenge. International Journal of Computer Vision (IJV)), (15())3111-255) 2015.
Claude E Shannon. Communication in the presence of noise. Proceedings of the IRE, 37())10021, jan 1949. doi: 109//rooo. 1110///pro..99999222992900 1949.232969, URL https: http:://ccooogoo.. .1.00.1199991 11t.0.111199///1////.// 1949. 232969.
Jianbo Shi and Jitendra Malik. Normalized cuts and image segmentation. IEEE Transactions on Pattern Analysis and Machine Intelligence (PAMI), 22(8):888 90588990551110 2000.
Mingxing Tan and Quoc Le. EfficientNet: Rethinking model scaling for convolutional neural net- works. In International Conference on Machine Learning (ICML), pp. 6105-6114, 14, 2019.
Mingxing Tan, Bo Chen, Ruoming Pang, Vijay Vasudevan, Mark Sandler, Andrew Howard, and Quoc V Le. MnasNet: Platform-aware neural architecture search for mobile. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2820-2828220022 2828. 2019.
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. In International Conference Machine Learning (ICML), pp. 9438 9447, 2020a.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: survey. arXiv preprint ar arXiv:2009.067:5 2020b.
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention in transformer models. International Conference on Machine Learning (ICML), 2021.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. pp- 10347-10357, 2021.
Didem Unat, Theodore Hromadka, and Scott B Baden. An adaptive sub-sampling method for in- memory compression of scientific data. In Data Compression Conference, pp. 262-271. IEEE, 2009.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz. Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Informa- tion Processing Systems (NeurIPS), 30:5998-6008 2017.
Andrea Vedaldi and Stefano Soatto. Quick shift and kernel methods for mode seeking. In European Conference on Computer Vision (ECCV), pp. 705-718, 2008.
Apoorv Vyas, Angelos Katharopoulos, and François Fleuret. Fast transformers with clustered atten- tion. Advances in Neural Information Processing Systems (NeurIPS), 33, 2020.
12

---

Under review as a conference paper at ICLR 2022
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04700 2020.
Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pyramid vision transformer: A versatile backbone for dense prediction with- out convolutions. In IEEE International Conference on Computer Vision (ICCV), pp. 568-578, October 2021.
Bichen Wu, Xiaoliang Dai, Peizhao Zhang, Yanghan Wang. Fei Sun, Yiming Wu, Yuandong Tian, Peter Vajda, Yangqing Jia, and Kurt Keutzer. FBNet: Hardware-awand efficient ConwNet design via differentiable neural architecture search. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10734-10742, 2019.
Bichen Wu, Chenfeng Xu, Xiaoliang Dai, Alvin Wan, Peizhao Zhang, Zhicheng Yan, Masayoshi Tomizuka, Joseph Gonzalez, Kurt Keutzer, and Peter Vajda. Visual transformers: Token-based image representation and processing for computer vision. arXiv preprint wrXiv:2006.0367 2020.
Lemeng Wu, Xingchao Liu, and Qiang Liu. Centroid transformers: Learning abstract with atten- tion. arXiv preprint arXiv:2102.08600 2021.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems (NeurlPS), 2020.
Hengshuang Zhao, Xiaojuan Qi, Xiaoyong Shen, Jianping Shi, and Jiaya Jia. ICNet for real-time semantic segmentation on high-resolution images. In European Conference on Computer Vision (ECCV), pp. 405-420, 2018.
Sixiao Zheng, Jiachen Lu, Hengshuang Zhao, Xiatian Zhu, Zekun Luo, Yabiao Wang, Yanwei Fu, Jianfeng Feng, Tao Xiang, Philip HS Torr, et al. Rethinking semantic segmentation from a sequence-to-scquence perspective with transformers. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 6881 6881-6890, 2021.
13

---

Under review as a conference paper at ICLR 2022
A
WEIGHTED CLUSTERING ALGORITHMS
Weighted K-Means minimizes the following objective w.r.t. F (fi,.... fk] C RR::
((F Σ min will fyl1?
(13)
F:F tee
The extension of the K- K-Means algorithm to the weighted case iterates the following steps:
a(1) + arg min llf. - fill
vi E f1.2....,NN
(14)
j
fy Σ [a(i) = =lwuff
Vj (1,2, (1, Kj.,
(15)
-i=1(((i) = 5|wu
Weighted K-Medoids optimizes objective (13) under the medoid constraint f C F:
a(i) arg min [fr
fall
iA E 2 N),
(16)
n(j) + mig Σ llf. - f
Vj E 21,2.... K}.
(17)
d (x((=)) P=a("))-
fz fn(f)-
(18)
B CLUSTERING INITIALIZATION ABLATIONS
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 ini- tialization, which randomly selects tokens as initial cluster centers. As shown in Figure 6, Token Pooling is robust to the initialization methods.
82
80
accuracy (%)
78
76
K-Means
K-Means, rand cluster init
74
K-Medoids
K-Medoids, rand cluster
72
2
3
Gflops
Figure 6: Default initialization vs. random initialization. Our Token Pooling is robust to initialization of clustering algorithms. The default initialization is top-K w.r.t. significance score, see Algorithm 1.
C TRAINING DETAILS & HYPERPARAMETERS
A Token Pooling layer has one parameter-tie 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 average pooling layers, the sizes of the kernels and strides play the same role as K. We use the method proposed by Goyal al. (2020) to automatically select Ks of all Token Pooling layers (see below for details). All other hyperparameters inherited from DeiTT (Touvron et al., 2021).
14

---

Under review as a conference paper at ICLR 2022
To fairly compare PoWER-BER BERT and the baseline methods with the proposed Token Pooling, all methods (except convolution downsampling) use the same target number of tokens for downsam- pling layers after each transformer block. Specifically, we run the second stage of the PoWER-BERS BERT training (Goyal et al., 2020) for 30 epochs with various values of the token-selection weight parame- ter À producing family of models. The single parameter À controls the overall efficiency of model, and the numbers of retained tokens of all Token Pooling layers are selected automatically based on the choice of A. Each of the resulted models has different number of retained tokens each of its L transformer blocks: K = [Kl..... Kg). 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 Dei'T 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 similar number of training steps.
82
80
aecuracy (5e)
78
76
DeiT (Ti S)
74
DeiT-Distil (Ti and S)
DeiT-2x (Ti + S)
72
1
2
Gflops
Figure 7: This figure shows the results of the pretrained DeiT models provided by Touvron et al. (2021) (DeiT) and the DeiT models trained with our protocol (DeiT-2x). Our training protocal uses the same hyper-parameters provided by Touvron et al. (2021), but after the model trained, we finetune the model using the same also show DeiT- Distil results (cited from (Touvron al., 2021)), which use knowledge distillation.
(i.e.. restart the learning rate schedule). We
D CONVOLUTION DOWNSAMPLING
As mentioned in the main paper, enumerate combinations of layers 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 DeîT-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 layers at depths 1 < l2 < lg. we restrict 12 - 1 L3 l2. With these constraints, we enumerate all possible downsampling configurations. Each of the combinations produces a model with a different ie., the best accuracy these models achieve at a given flop. See Figure 8.
ceppuutationcccuraaa trade-off, and we report the Pareto front,
Note that the convolution downsampling layer is equivalent to the patch-merging layer used by Liu 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 generalized version of the commonly used average pooling layer. Average pooling can be implemented by convolving the feature map with constant kernel of size n n containing f and subsample (stride) every 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.
NEW
15

---

Under review as a conference paper at ICLR 2022
80
78
accuracy (%
76
74
(WK-Medoids)
DeiT (Ti
Convolution (best)
Convolution
1.5
2
2.5
3
3.5
4.5
Gflops
Figure 8: Results of convolution downsampling
E RESULTS WITHOUT FINETUNING
Since Token Pooling minimizes the reconstruction error due to token downsampling, in this section, evaluate the performance when insert Token Pooling layers into 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 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-Meamg 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 cluster center and carry the count when we compute -attat--atention the next transformer blocks. This operation preserves the attention weights, and since the models are not trained after inserting Token Pooling layers, it pre- serves the most accuracy. In practice, when the models are trained with Token Pooling layers, the carry operation is not important.
80
60
нсcuracy (%)
DeiT-S
K Means, no finetuning
WK-Means, no finetuning
40
WK- KeManns, carry, finetuning
K-Medoids, no finetuning
WK- MK-Medoids, no finetuning
WK-Medoide earry, finetuning
20
2.5
3
3.5 -100. Flops
4
4.5
Figure 9: The figure shows the performance results when we insert Token Pooling layers into pretrained DeiT-S without finetuning the model (i.e., skip step 3 of the training protocol in Section 5).
16

---

Under review as a conference paper at ICLR 2022
82
82
80
80
78
78
accuracy (%
aceuracy (%
76
74
Ours (WK-Medoids)
72
72
Ours (WK-Medoide normalized
PoWER-BEE -BERT
70
DeiT (Ti S), normalized
70
POWER-BERT, normalized
POWER-BERT, normalized
DeiT (Ti
68
68
Ours (WK-Medoids), normalized
DeiT (Ti S), normalized
66
66
1
Gflops
Gflops
(a)
(b) Learned for each layer
Figure 10: Results of models using normalized key and query vectors with (a) a = and (b) learned
(6). The base model architecture DeïT-S.
F NORMALIZED KEY AND QUERY VECTORS
Our analysis of softmax-attention in $3.3 assumes the query and the key vectors are normalized to have constant e² 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, verify this observation by training various DeiT, POWER-BERT, and Token Pooling with normalized keys and queries. We found that the scalar a in the softmax attention layer (6) can affect the performance of transformer with normalized keys and queries. As can be seen in Figure 10a, setting the scalar = slightly deteriorates the performance of the model. Instead of using fixed a, we let the model learn the a for each layer. As can be seen in Figure 10b, learning a enables the resulting models to achieve similar cost-accuracy trade-off as the standard (unnormalized) models. With or without the normalization and the learnable a, the proposed Token Pooling significantly improves the cost-accuracy trade-off of DeiT and outperforms PoWER-BERT.
G CLUSTERING ALGORITHM ABLATIONS
Tables 2-5 detail the results of POWER-BERT and Token Pooling on the DeiT architectures that tested (DeiTT-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-BERS via varying token sparsity and feature dimensionality of DeiT.
Apart from the standard K-Means and K- 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 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 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.
17

---

Under review as a conference paper at ICLR 2022
78
80
76
74
accuracy (56)
accuracy (96
76
72
70
74
68
Ours (K- (K-Medoids)
Ours (WK-Medoids)
72
Ours (K-Medoids)
(WK-Medoids)
66
POWER BERT
PoWER- BERT
DeiT (Ti + <252)
70
DeiT' (Ti - c318)
64
1.5
1.5
2.5
Gflops
Gflops
(a) DeiT-e252 the base architecture
(b) DeiT-e318 as the base architecture
Figure 11: This figure compares the cost-aceuracy curves of Token Pooling with PoWER-BER
using (a) D25 and (b) De1T-e318 as the base architectures.
DefT' (Ti + S)
PoWER-BERS
Token Pooling (K- (K-Medoids)
Token Pnoling (WK-Medoids)
Gflops
Accuracy (%)
Gflops
Accuracy (%)
Gflops
Accuracy (%)
Gflops
Accuracy (%)
0.46
65.0
0.48
68.5
0.48
69.4
0.55
68.0
0.57
70.9
0.57
71.7
0.70
71.7
0.73
73.6
0.73
74.0
0.89
72.0
0.93
74.6
0.93
75.0
0.96
74.8
1.00
75.9
1.00
76.4
1.05
75.4
1.11
76.4
1.11
76.6
1.25
73,9
1.19
76.6
1.25
77.0
1.25
77.2
1.38
77.1
1,46
77.4
1,46
77.6
1.46
77.4
1.52
78.4
1.52
78.9
1.62
78.4
1.68
78,9
1.68
79.2
2.07
77.7
1.81
78.8
1.88
79.3
1.88
79,4
2.11
79.3
2.13
79,9
2.13
80.1
2.27
79.4
2.35
80.1
2.35
$0.4
2.52
80.0
2.61
80.5
2.61
80.6
3.21
80.0
2.93
80.6
3.04
80.7
3.04
80.7
4.60
81.2
4.28
81.2
4.44
81.4
4.44
81.2
Table 2: Best cost-accuracy trade-off achieved by PoWER- BERT and the proposed Token Pooling via varying sparsity level and feature dimensionality.
18

---

Under review as a conference paper at ICLR 2022
clustering method Weighting
ImageNet Accuracy
GFlops
Sparsity level 0: K
196, 196, 195, 194, 189, 180, 173, 173, 173, 173, 173, 173]
PoWER-BERS
N/A
81.2
4.3
81.3
4.7 (+0.4)
K-Means
81.1
4.7 (40.4)
81.4
4.4 (+0.1)
K-Medoids
81.2
4.4 (+0.1)
Sparsity level 1: 1166,
195,
193, 188, 169,
140, 121, 110, 73, 38, 7,
POWER-BERT N/A
80.6
2.9
80.7
3.3 (+0.4)
K-Means
80.8
3.3 (+0.4)
80.7
3.0 (+0.1)
K-Medoids
80.7
3.0 (+0.1)
Sparsity level K 196, 195, 190, 177, 141, 108, 84, 69, 35, 18,3,
POWER-BERT
N/A
80.0
2.5
80.5
2.8 (+0.3)
K-Means
80.5
2.8 (+0.3)
80.5
2.6 (+0.1)
K-Medoids
80.6
2.6 (+0.1)
Sparsity level 3: K 196,, 194, 187, 163, 118, 85, 58,
47, 20, 12,
POWER-BERT
N/A
79.4
2.3
80.1
2.5 (+0.2)
K-Means
80.4
2.5 (+0.2)
80.1
2.3 (+0.08)
K-Medoids
80.4
2.3 (+0.08)
Sparsity level K 1996, 193,
179, 142, 97, 64, 46, 34, 13, 9, 1,
POWER-BERT
N/A
78.8
2.1
79.8
2.3 (+0.2)
K-Means
80.0
2.3 (+0.2)
79.9
2.1 (+0.07)
K-Medoids
80.1
(+0.07)
Sparsity level 5: K [194, 183,
142, 89, 41, 20, 10, 7, 0,
PoWER-BERT N/A
76.2
1.5
77.6
1.7 (+0.2)
K-Means
78.1
1.7 (+0.2)
77.8
1.6 (+0.05)
K-Medoids
78.1
1.6 (+0.05)
Sparsity level K 186, 162,
102, 56, 13, 4, 0,0000
POWER-BERY
N/A
73.3
1.2
75.0
1.4 (+0.2)
K-Means
75.6
1.4 (+0.2)
75.4
1.3 (+0.04)
K-Medoids
75.7
1.3 (+0.04)
Sparsity level 7: K 1662,
129,
66, 33,
0, 0,0,
PoWER-BERT N/A
69.6
1.0
72.4
1.1 (+0.1)
K-Means
73.0
1.1 (+0.1)
72.3
1.0 (+0.03)
K-Medoids
73.0
1.0 (+0.03)
Table 3: Results of applying Token Pooling and PoWER-BERE on DeîT-S model. The models are grouped by deseribed in Appendix C. The integer list denotes the maximal number tokens re- tained after each transformer block. These numbers do not take into account the classification token, which always retained. Thus, "0" means that only the classification token remains. Additional flops (denoted by the parentheses) are due to clustering.
19

---

Under review as a conference paper at ICLR 2022
clustering method
Weighting
ImageNet Accuracy
GFlops
Sparsity level 0: K 1196, 196, 196, 194, 192, 184, 181, 181, 170, 170, 170, 170]
POWER-BERS N/A
80.1
3.0
80.1
3.1 (+0.1)
K-Medoids
80.0
3.1 (+0.1)
Sparsity level 1: K [196, 195, 195, 190, 171, 147, 132, 122, 66, 43, 15, 0]
PoWER-BERE N/A
79.3
2.1
79.6
2.2 (+0.1)
K-Medoids
79.7
2.2 (+0.1)
Sparsity level K [196, 195, 193, 180, 143, 109, 92,
77, 40, 21, 4,
PoWER-BERT
N/A
78.8
1.8
79.3
1.9 (+0.09)
K-Medoids
79.4
1.9 (+0.09)
Sparsity level [196, 195, 190, 165, 119, 84, 65, 50, 26, 14, 3, 0)
POWER-BERS BERT
N/A
78.3
1.6
78.9
1.7 (+0.07)
K-Medoids
79.2
1.7 (+0.07)
Sparsity level 4: [196, 194, 184, 149, 97, 65, 47, 33, 12,9,2,0)
POWER-BERT
N/A
77.4
1.5
78 4
1.5 (+0.06)
K-Medoids
78.9
1.5 (+0.06)
Sparsity level 5: K [196, 186, 153, 93, 40, 15, 12, 9, 0,0,0,0)
POWER-BERT N/A
74.5
1.1
76.2
1.1 (+0.05)
K-Medoids
76.8
1.1 (+0.05)
Sparsity level 6: K 193., 173, 109, 52, 16, 4, 4, 0,
POWER-BERT
N/A
72.0
0.89
74.6
0.93 (+0.04)
K-Medoids
75.0
0.93 (+0.04)
Sparsity level 7: K [183, 145, 80, 33, 5, 1, 2, 0, 0,
POWER-BERT N/A
69.4
0.75
72.1
0.78 (+0.03)
K-Medoids
72.6
0.78 (+0.03)
Table 4: Results of applying Token Pooling and PoWER-BER on DeiT-e318 model. The mod- els are grouped by described in Appendix C. The integer list denotes maximal number of tokens retained after each transformer block. These numbers do not take into account the classifi- cation token, which is always retained. Thus, "0" means that only the classification token remains. Additional flops (denoted by the parentheses) due to clustering.
20

---

Under review as a conference paper at ICLR 2022
clustering method
Weighting
ImageNet Accuracy
GFlops
Sparsity level 0: K 1196, 196, 196, 195, 193, 189, 177, 177, 177, 177, 177, 177]
POWER-BERS N/A
78.0
2.0
77.7
2.1 (+0.1)
K-Medoids
77.9
2.1 (+0.1)
Sparsity level 1: K [196, 196, 195, 189, 176, 161, 123, 122, 76, 48, 17, 0]
PoWER-BERE N/A
77.1
1.4
77.4
1.5 (+0.07)
K-Medoids
77.6
1.5 (+0.07)
Sparsity level K [196, 196, 193, 179, 154, 123, 88, 39, 22, 5, 0].
PoWER-BERT
N/A
76.6
1.2
77.0
1.3 (+0.06)
K-Medoids
77.2
1.3 (+0.06)
Sparsity level [196, 195, 188, 167, 131, 96, 63, 52, 13, 11, 2, OJ.
POWER-BERS BERT
N/A
75.4
1.1
76.4
1.1 (+0.05)
K-Medoids
76.6
1.1 (+0.05)
Sparsity level 4: [196, 194, 181, 147, 111, 75, 48, 36, 5, 7,
POWER-BERT
N/A
74.8
0.96
75.9
1.0 (+0.04)
K-Medoids
76.4
1.0 (+0.04)
Sparsity level 5: K [196, 187, 129, 88, 54, 20, 10, 10, 0, 1,0, 0]
POWER-BERT N/A
71.7
0.70
73.6
0.73 (+0.03)
K-Medoids
74.0
0.73 (+0.03)
Sparsity level 6: K 194,, 163, 83, 44, 22, 6, 1, 0,0,0,0)
POWER-BERT
N/A
68.0
0.55
70.9
0.57 (+0.02)
K-Medoids
71.7
0.57 (+0.02)
Sparsity level 7: K [188, 122, 58, 28, 12, 0, 2, 0, 0,
POWER-BERT N/A
65.0
0.46
68.5
0.48 (+0.02)
K-Medoids
69.4
0.48 (+0.02)
Table 5: Results of applying Token Pooling and POWER-BEIT on DeiT-e252 model. The mod- els are grouped by described in Appendix The integer list denotes maximal number of tokens retained after each transformer block. These numbers do not take into account the classifi- cation token, which is always retained. Thus, "0" means that only the classification token remains. Additional flops (denoted by the parentheses) due to clustering.
21

---

Under review as a conference paper at ICLR 2022
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 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 reference on how our models actually perform given our implementation and the type of GPU used. Note that our PyTorch imple- mentation of the clustering algorithms can be significantly improved, for example, via implementing as CUDA kernel. To determine the throughput, we run the model inference several times with dif- ferent 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 1791 fps), and of DeiT-Ti by 14% (3220 vs. 2834 fps).
82
DefT-S
DeiT (various dimensions)
0.6
Proposed Token Pooling
80
25%
DeiT-e318
accuracy (%)
78
DeiT-e252
22%
76
74
DeiT-Ti 14%
72
1.5
2.5
3.5
throughput (103 fps)
Figure 12: Throughput of our PyTorch implementation. Note that our implementation has not been optimized for the throughput. The throughput is measured in frames (images) per second (fps). These numbers are measured on a single Nvidia VI00 GPU. As can be seen, despite our non- optimized code, Token Pooling significantly improves both the flops (see Figure 1b) and throughput of DeiT models.
22


Reviews:
[{"review_id": "nrX897ro8E8", "reviewer": "Reviewer_RU8U", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n+ Solid experimental results\n\nCompared with a well-optimized DeiT baseline, token pooling significantly improves accuracy and reduces computation. Ablations in the appendix on different clustering, convolution downsampling, models are well written. \n\n+ Well-structured analysis\n\nThe analysis in Table 1 clearly show the portion of different operators in vision transformer and figure 2 shows the limitation of score-based token downsampling, which motivates the authors to reduce redundancies with a novel token pooling method.\n\n+ Impact on community\n\nVision transformers have been deployed on multiple vision tasks and are being a mainstream solution. Efficient vision transformers are important to be developed and the token pooling proposed in this paper can be used for other vision transformers without major modifications. It can replace convolution/patch-based downsampling and contribute to the community.\n\nWeaknesses:\n\n- Attention as a low-pass filter\n\nI appreciate the analysis in section 3.3 on attention vs low-pass filter, but I don't get the relation between this observation and the token pooling methods. In the experiment at Appendix F, it seems token pooling is not affected by normalized or unnormalized Q/K vectors.", "summary_of_the_paper": "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.", "summary_of_the_review": "Based on my main review above, I think this paper proposes a solid token pooling method for vision transformers that can be potentially applied to other vision transformers based on tokens. I recommend an acceptance to this paper. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "QSn90hr_KGZ", "reviewer": "Reviewer_jHZj", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n1. The authors analyzed the computational cost of vision-transformer components in detail. These pie charts (in Table 1) visually illustrate the computational bottlenecks of vision transformers are fully-connected layers.\n2. The authors demonstrated that softmax-attention is a low-pass filter under mild assumptions, thus attention layers will generate redundant representations.\n3. This paper is well written and easy to understand. The authors provide a clear explanation and motivation behind their token pooling method. As far as I'm concerned, it has an impressive feature in that it adopts clustering algorithms (such as K-Means and K-Medoids) to reduce the number of tokens.\n\nWeaknesses\n1. Despite claims that the proposed token pooling is an effective operator that can benefit many architectures, but in this paper, the authors only examined based on the DeiT architecture. If token pooling is a general method for vision transformers, the paper should apply it to at least the current SOTA vision transformers.\n2. As mentioned in the introduction, max pooling and average pooling are widely-used downsampling methods. From my experience, average pooling is a competitive method to reduce the number of tokens. Why not do an ablation study with max/average pooling?\n3. Personally, I would appreciate experiments with more vision transformers. If the authors proved that the proposed token pooling could be applicable to other vision transformers besides DeiT, it would strengthen its value. \n4. It would be better to provide more results and analysis of time consumption and inference speed, including the clustering process.\n\n", "summary_of_the_paper": "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.", "summary_of_the_review": "In summary, this work presents a new token downsampling method for efficient vision transformers. The paper is well written and properly structured, but the experiment is only performed on DeiT and ImageNet-1k benchmark, and no time analysis is provided, which is not up to par.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "JLNxkaNljN4", "reviewer": "Reviewer_E53C", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n+ The submission is technically sound. In Section 3.3, the analysis on attention as a low-pass filter is interesting and makes sense. The K-Means based token down-sampling method is straightforward.\n+ The submission is clearly written and well-organized.\n\nWeakness:\n- The K-Means clustering method for efficient transformers has been introduced in [1], which narrows the contribution of the paper. \n- According to Table 2, 3 and 4, the proposed framework is sensitive to the selection of hyperparameters. It seems that the method is tricky and heavily relies on manual tuning.\n- The method is only evaluated on the DeiT. Can it apply to other vision transformers, e.g., PVT and Swin?\n- Reducing the number of queries is unfriendly to apply to downstream tasks.\n- It would be better to report the actual throughput and compare the proposed method with the previous dynamic vision transformers, such as [2] and [3].\n\n[1] Roy A, Saffar M, Vaswani A, et al. Efficient content-based sparse attention with routing transformers[J]. Transactions of the Association for Computational Linguistics, 2021, 9: 53-68.\n\n[2] Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsieh. DynamicViT: Efﬁcient vision transformers with dynamic token sparsiﬁcation. ArXiv, abs/2106.02034, 2021.\n\n[3] Wang Y, Huang R, Song S, et al. Not All Images are Worth 16x16 Words: Dynamic Vision Transformers with Adaptive Sequence Length[J]. arXiv preprint arXiv:2105.15075, 2021.\n", "summary_of_the_paper": "The paper proposes a token pooling approach based on the K-Means algorithm to improve the efficiency of vision transformers. The method is evaluated on image classification, where it reduces the computational complexity by half while maintaining similar performance.", "summary_of_the_review": "The paper proposes a new token pooling method to achieve efficient vision transformers, which is interesting and well-motivated. But the method heavily relies on manual tuning and the generalization is limited. I will increase the rate if the concerns are well addressed.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This submission receives mixed ratings initially. Two reviewers lean negatively while one reviewer is positive. The raised issues include 
whether the proposed method can be adapted to other vision transformers, the design choice of pooling strategy, the computational time cost, the similarity to an existing work, and the influence of the proposed method on downstream tasks. In the rebuttal, the authors have addressed several issues such as pooling strategy analysis and time consumption.

There are still some issues not completely solved. The proposed method introduces K-mean clustering on tokens between different layers. The K-mean clustering is prevalent and the weighted clustering does not make the technical contribution sufficient. Also as a general token pooling operation, the proposed method shall be integrated into various types of vision transformers (e.g., vanilla ViT [a], ConViT [b]), rather than one single DeiT. Besides, the downstream tasks in DeiT are not conducted in the proposed method. 

Overall, the AC feels the proposed method, although interesting, requires a major revision that addresses existing issues.  The authors are suggested to further improve the current submission and welcome to submit for the next venue. 

[a]. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. Dosovitskiy et al. ICLR 2021.

[b]. ConViT: Improving Vision Transformers with Soft Convolutional Inductive Biases. Ascoli et al. ICML 2021.
Decision: Reject