Original: T OKEN P OOLING IN V ISUAL T RANSFORMERS
Revised: T OKEN P OOLING IN V ISION T RANSFORMERS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Despite the recent success in many applications, the high computational requirements of vision transformers limit their use in resource-constrained settings.
Revised: Despite the recent success in many applications, the high computational requirements of vision transformers limit their use in resource-constrained settings.

Original: While many existing methods improve the quadratic complexity of attention, in most vision transformers, self-attention is not the major computation bottleneck, e .
Revised: While many existing methods improve the quadratic complexity of attention, in most vision transformers, self-attention is not the major computation bottleneck, e .

Original: g ., more than 80% of the computation is spent on fully-connected layers.
Revised: g ., more than 80% of the computation is spent on fully-connected layers.

Original: To improve the computational complexity of all layers, we propose a novel token downsampling method, called Token Pooling , efﬁciently exploiting redundancies in the images and intermediate token representations.
Revised: To improve the computational complexity of all layers, we propose a novel token downsampling method, called Token Pooling , efﬁciently exploiting redundancies in the images and intermediate token representations.

Original: We show that, under mild assumptions, softmax-attention acts as a high-dimensional low-pass (smoothing) ﬁlter.
Revised: We show that, under mild assumptions, softmax-attention acts as a high-dimensional low-pass (smoothing) ﬁlter.

Original: Thus, its output contains redundancy that can be pruned to achieve a better trade-off between the computational cost and accuracy.
Revised: Thus, its output contains redundancy that can be pruned to achieve a better trade-off between the computational cost and accuracy.

Original: Our new technique accurately approximates a set of tokens by minimizing the reconstruction error caused by downsampling.
Revised: Our new technique accurately approximates a set of tokens by minimizing the reconstruction error caused by downsampling.

Original: We solve this optimization problem via cost-efﬁcient clustering.
Revised: We solve this optimization problem via cost-efﬁcient clustering.

Original: We rigorously analyze and compare to prior downsampling methods.
Revised: We rigorously analyze and compare to prior downsampling methods.

Original: Our experiments show that Token Pooling signiﬁcantly improves the cost-accuracy trade-off over the state-of-the-art downsampling.
Revised: Our experiments show that Token Pooling signiﬁcantly improves the cost-accuracy trade-off over the state-of-the-art downsampling.

Original: Token Pooling isa simple and effective operator that can beneﬁt many architectures.
Revised: Token Pooling is simple and effective and can beneﬁt many architectures with global attention.

Original: Applied to DeiT,  it achieves the same ImageNet top-1 accuracy using 42% fewer computations.
Revised: Applied to DeiT, NEW it achieves the same ImageNet top-1 accuracy using 42% fewer computations.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Vision transformers (Dosovitskiy et al., 2020; Touvron et al., 2021;
Revised: Vision transformers (Dosovitskiy et al., 2020; Touvron et al., 2021;

Original: 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 classiﬁcation to segmentation.
Revised: 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 classiﬁcation to segmentation.

Original: However, the high computational cost limits their use in resource-restricted, real-time, or low-powered applications.
Revised: However, the high computational cost limits their use in resource-restricted, real-time, or low-powered applications.

Original: While most prior work in Natural
Revised: While most prior work in Natural

Original: 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.
Revised: 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.

Original: The computational complexity of these layers is determined by the number of tokens and their feature dimensionality.
Revised: The computational complexity of these layers is determined by the number of tokens and their feature dimensionality.

Original: While reducing the dimensionality improves computational cost, it sacriﬁces model capacity and often signiﬁcantly deteriorates the accuracy of the model.
Revised: While reducing the dimensionality improves computational cost, it sacriﬁces model capacity and often signiﬁcantly deteriorates the accuracy of the model.

Original: On the other hand, since images often contain mostly smooth surfaces with sparsely located edges and corners, they contain similar (and thus redundant) features.
Revised: On the other hand, since images often contain mostly smooth surfaces with sparsely located edges and corners, they contain similar (and thus redundant) features.

Original: Moreover, we show that, under mild assumptions, softmax-attention is equivalent to low-pass ﬁltering of tokens and thereby produces tokens with similar features, as empirically observed by Goyal et al.
Revised: Moreover, we show that, under mild assumptions, softmax-attention is equivalent to low-pass ﬁltering of tokens and thereby produces tokens with similar features, as empirically observed by Goyal et al.

Original: and Rao et al.
Revised: and Rao et al.

Original: This redundancy in representations suggests that we can reduce the number of tokens, i .
Revised: This redundancy in representations suggests that we can reduce the number of tokens, i .

Original: e ., downsampling, without a signiﬁcant impact to the accuracy, achieving a better cost-accuracy trade-off than reducing feature dimensionality alone.
Revised: e ., downsampling, without a signiﬁcant impact to the accuracy, achieving a better cost-accuracy trade-off than reducing feature dimensionality alone.

Original: Downsampling is widely used in Convolutional Neural Network (CNN) architectures to improve computational efﬁciency, among other purposes.
Revised: Downsampling is widely used in Convolutional Neural Network (CNN) architectures to improve computational efﬁciency, among other purposes.

Original: Given a grid of pixels or features, downsampling gradually reduces the grid dimensions via combining neighboring vertices on the grid.
Revised: Given a grid of pixels or features, downsampling gradually reduces the grid dimensions via combining neighboring vertices on the grid.

Original: 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.
Revised: 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.

Original: Such methods do not efﬁciently address non-uniformly distributed redundancy in images and features (Recasens et al., 2018;
Revised: Such methods do not efﬁciently address non-uniformly distributed redundancy in images and features (Recasens et al., 2018;

Original: Marin et al., 2019).
Revised: Marin et al., 2019).

Original: Unlike CNNs that require grid preservation, transformers allow a wider range of nonuniform data-aware downsampling layers, where a better operator can be designed.
Revised: Unlike CNNs that require grid preservation, transformers allow a wider range of nonuniform data-aware downsampling layers, where a better operator can be designed.

Original: We propose Token Pooling , a novel nonuniform data-aware downsampling operator for transformers efﬁciently exploiting redundancy in features.
Revised: We propose Token Pooling , a novel nonuniform data-aware downsampling operator for transformers efﬁciently exploiting redundancy in features.

Original: See the illustration and performance metric in Fig-
Revised: See the illustration and performance metric in Fig-

Original: [Figure]
Revised: [Figure]

Original: ures 1a & 1b.
Revised: ures 1a & 1b.

Original: 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.
Revised: 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.

Original: We show that clustering algorithms, K-Means and K-Medoids, efﬁciently solve this problem, see illustration in
Revised: We show that clustering algorithms, K-Means and K-Medoids, efﬁciently solve this problem, see illustration in

Original: Figure 1c.
Revised: Figure 1c.

Original: To the best of our knowledge, we are the ﬁrst to use this formulation and simple clustering analysis for token donwsampling in transformers.
Revised: To the best of our knowledge, we are the ﬁrst to use this formulation and simple clustering analysis for token donwsampling in transformers.

Original: Wealso compare with various prior downsampling techniques, including grid-downsampling (Pan et al.,  2021) and token pruning (Goyal et al., 2020; Rao et al., 2021).
Revised: 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).

Original: Our results show that the proposed Token Pooling outperforms existing methods and provides the best trade-off between computational cost and classiﬁcation accuracy.
Revised: The proposed Token Pooling outperforms existing methods and provides the best trade-off between computational cost and classiﬁcation accuracy.

Original: Contributions.
Revised: Contributions.

Original: The paper makes the following contributions:
Revised: The paper makes the following contributions:

Original: • We conduct an extensive study of prior downsampling techniques for visual transformers by comparing their computation-accuracy trade-offs.
Revised: • We conduct an extensive study of prior downsampling techniques for visual transformers by comparing their computation-accuracy trade-offs.

Original: • We analyze the computational cost of vision-transformer components and the limitations of the prior score-based downsampling methods.
Revised: • We analyze the computational cost of vision-transformer components and the limitations of the prior score-based downsampling methods.

Original: We also show that attention layers behave like lowpass ﬁltering and thus produce redundant tokens.
Revised: We also show that attention layers behave like low-pass ﬁltering and thus produce redundant tokens.

Original: • Motivated by the redundancy in images and features, we propose a novel token downsampling technique, Token Pooling, for transformers  via error minimization and  achieve a signiﬁcant improvement in the computation-accuracy trade-off.
Revised: • 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 signiﬁcant improvement in the computation-accuracy trade-off.

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: In this section, we introduce vision transformers, and review existing methods that improve the efﬁciency of transformers including existing token-downsampling methods.
Revised: In this section, we introduce vision transformers, and review existing methods that improve the efﬁciency of transformers including existing token-downsampling methods.

Original: 2. V ISION TRANSFORMERS
Revised: 2. V ISION TRANSFORMERS

Original: Vision transformers (Dosovitskiy et al., 2020; Touvron et al., 2021; Heo et al., 2021;
Revised: Vision transformers (Dosovitskiy et al., 2020; Touvron et al., 2021; Heo et al., 2021;

Original: Liu et al., 2021;
Revised: Liu et al., 2021;

Original: Pan et al., 2021) utilize the transformer architecture that is originally designed for NLP by Vaswani et al.
Revised: Pan et al., 2021) utilize the transformer architecture that is originally designed for NLP by Vaswani et al.

Original: (2017) and further popularized by Radford et al.
Revised: (2017) and further popularized by Radford et al.

Original: (2018) and Devlin et al.
Revised: (2018) and Devlin et al.

Original: In a high level, a vision transformer is a composition of L transformer blocks that take a set of input tokens and return another set of output tokens.
Revised: In a high level, a vision transformer is a composition of L transformer blocks that take a set of input tokens and return another set of output tokens.

Original: In vision, input tokens are features representing individual nonoverlapping image patches.
Revised: In vision, input tokens are features representing individual nonoverlapping image patches.

Original: To perform classiﬁcation, a classiﬁcation token is inserted to estimate the probabilities of individual classes.
Revised: To perform classiﬁcation, a classiﬁcation token is inserted to estimate the probabilities of individual classes.

Original: 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).
Revised: 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).

Original: Recently, DeiT (Touvron et al., 2021) achieved state-of-the-art results with advanced training on ImageNet1k only.
Revised: Recently, DeiT (Touvron et al., 2021) achieved state-of-the-art results with advanced training on ImageNet1k only.

Original: Let the set of tokens at depth l be F l = { f l 0 , . .
Revised: Let the set of tokens at depth l be F l = { f l 0 , . .

Original: , f lN } where f li ∈ R M is the feature values of the i -th token.
Revised: , f lN } where f li ∈ R M is the feature values of the i -th token.

Original: 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).
Revised: 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).

Original: Let matrix F ∈ R N ×
Revised: Let matrix F ∈ R N ×

Original: M be a row-wise concatenation of tokens F l .
Revised: M be a row-wise concatenation of tokens F l .

Original: Then 1 ,
Revised: Then 1 ,

Original: [Equation]
Revised: [Equation]

Original: such that
Revised: such that

Original: [Equation]
Revised: [Equation]

Original: where H is the number of heads, matrix
Revised: where H is the number of heads, matrix

Original: W O ∈ R M × M is a learnable parameter of the block, [ , ] is column-wise concatenation and O h ∈ R N × d is the output of h -th attention head for d = M/H :
Revised: W O ∈ R M × M is a learnable parameter of the block, [ , ] is column-wise concatenation and O h ∈ R N × d is the output of h -th attention head for d = M/H :

Original: O [Equation]
Revised: O h =

Original:  such that
Revised: A h V h such that

Original: [Equation]
Revised: [Equation]

Original: Keys K h , queries Q h and values V h are linear projections of the input tokens (QKV projections):
Revised: Keys K h , queries Q h and values V h are linear projections of the input tokens (QKV projections):

Original: [Equation]
Revised: [Equation]

Original: where W Qh ∈ R [Equation] [Equation] Kh 
Revised: where W Qh ∈ R M × d , W Kh ∈ R M × d ,

Original: [Equation] R M × d , W Vh ∈R M ×
Revised: W Vh ∈

Original:  d are learnable linear transformations.
Revised: R M × d are learnable linear transformations.

Original: Note, the number of tokens is not affected by the transformer blocks, i .
Revised: Note, the number of tokens is not affected by the transformer blocks, i .

Original: e .,
Revised: e .,

Original: |F l +1 | =
Revised: |F l +1 | =

Original: |F l | .
Revised: |F l | .

Original: 2. E FFICIENT TRANSFORMERS
Revised: 2. E FFICIENT TRANSFORMERS

Original: Similar to many machine learning models, the efﬁciency 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 sparsiﬁcation (Gale et al., 2019), etc .
Revised: Similar to many machine learning models, the efﬁciency 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 sparsiﬁcation (Gale et al., 2019), etc .

Original: For example, Dosovitskiy et al.
Revised: For example, Dosovitskiy et al.

Original: (2020) and Touvron et al.
Revised: (2020) and Touvron et al.

Original: (2021) obtain a family of ViT and DeiT models, respectively, by varying the input resolution, the number of heads H , and the feature dimensionality
Revised: (2021) obtain a family of ViT and DeiT models, respectively, by varying the input resolution, the number of heads H , and the feature dimensionality

Original: M .
Revised: M .

Original: Each of the models operates with a different computational requirement and accuracy.
Revised: Each of the models operates with a different computational requirement and accuracy.

Original: In the following, we review techniques developed for transformers.
Revised: In the following, we review techniques developed for transformers.

Original: 2.2. E FFICIENT SELF - ATTENTION
Revised: 2.2. E FFICIENT SELF - ATTENTION

Original: The softmax-attention layer (3) has a quadratic time complexity w.r.t.
Revised: The softmax-attention layer (3) has a quadratic time complexity w.r.t.

Original: the number of tokens, i .
Revised: the number of tokens, i .

Original: e .,
Revised: e .,

Original: O ( N 2 ) .
Revised: O ( N 2 ) .

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: The time complexity of an attention layer can be reduced by restricting the attention ﬁeld-of-view and thus imposing sparsity on A h in (3).
Revised: The time complexity of an attention layer can be reduced by restricting the attention ﬁeld-of-view and thus imposing sparsity on A h in (3).

Original: This can be achieved using the spatial relationship between tokens in the image/text domain (Parmar et al., 2018;
Revised: This can be achieved using the spatial relationship between tokens in the image/text domain (Parmar et al., 2018;

Original: Ramachandran et al., 2019; Qiu et al., 2020;
Revised: Ramachandran et al., 2019; Qiu et al., 2020;

Original: Beltagy et al., 2020; Child et al., 2019; Zaheer et al., 2020) or based on token values using localitysensitive hashing, sorting, compression, etc .
Revised: Beltagy et al., 2020; Child et al., 2019; Zaheer et al., 2020) or based on token values using localitysensitive hashing, sorting, compression, etc .

Original: (Kitaev et al., 2020;Roy et al., 2021;
Revised: (Kitaev et al., 2020;

Original: Vyas et al., 2020; Tay et al., 2020a; Liu et al., 2018; Wang et al., 2020; Tay et al., 2021).
Revised: Vyas et al., 2020; Tay et al., 2020a; Liu et al., 2018; Wang et al., 2020; Tay et al., 2021).

Original: Prior works have also proposed attention mechanisms with lower time complexity, e .
Revised: Prior works have also proposed attention mechanisms with lower time complexity, e .

Original: g .
Revised: g .

Original: , O ( N ) or O ( N log N ) (Katharopoulos et al., 2020;
Revised: , O ( N ) or O ( N log N ) (Katharopoulos et al., 2020;

Original: Peng et al., 2021; Choromanski et al., 2021; Tay et al., 2021). 
Revised: Peng et al., 2021; Choromanski et al., 2021; Tay et al., 2021). Roy et al. (2021) cluster queries and

Revised: NEW keys to sparsify attention matrices to speed-up the attention, but they do not downsample the tokens.

Original: 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.
Revised: 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.

Original: In contrast, our method reduces the number of tokens after attention has been computed.
Revised: In contrast, our method reduces the number of tokens after attention has been computed.

Original: Thereby, we can utilize these methods to further improve the overall efﬁciency of transformers.
Revised: Thereby, we can utilize these methods to further improve the overall efﬁciency of transformers.

Original: Recently, Wu et al.
Revised: Recently, Wu et al.

Original: (2020) proposed a new attention-based layer that learns a small number of query vectors to extract information from the input feature map.
Revised: (2020) proposed a new attention-based layer that learns a small number of query vectors to extract information from the input feature map.

Original: Similarly, Wu et al.
Revised: Similarly, Wu et al.

Original: (2021) replace selfattention with a new recurrent layer that outputs a smaller number of tokens.
Revised: (2021) replace selfattention with a new recurrent layer that outputs a smaller number of tokens.

Original: In comparison, our method directly minimizes the token reconstruction error due to token downsampling.
Revised: In comparison, our method directly minimizes the token reconstruction error due to token downsampling.

Original: Also, our layer has no learnable parameters and can be easily incorporated into existing vision transformers.
Revised: Also, our layer has no learnable parameters and can be easily incorporated into existing vision transformers.

Original: 2.2. D OWNSAMPLING METHODS FOR TRANSFORMERS
Revised: 2.2. D OWNSAMPLING METHODS FOR TRANSFORMERS

Original: Grid-downsampling.
Revised: Grid-downsampling.

Original: The input tokens of the ﬁrst vision transformer block are computed from image patches.
Revised: The input tokens of the ﬁrst vision transformer block are computed from image patches.

Original: 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.
Revised: 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.

Original: The regular grid structure allows typical downsampling methods, such as max/mean pooling, uniform sub-sampling.
Revised: The regular grid structure allows typical downsampling methods, such as max/mean pooling, uniform sub-sampling.

Original: For example, Liu et al.
Revised: For example, Liu et al.

Original: (2021); Heo et al.
Revised: (2021), Heo et al.

Revised: (2021), and Wang et al.

Original: (2021) use convolutions with stride to downsample the feature maps formed by the tokens.
Revised: (2021) use convolutions with stride to downsample the feature maps formed by the tokens.

Original: Also, these works increase the feature dimensionality with depth, similarly to CNNs.

Original: Score-based token downsampling.
Revised: Score-based token downsampling.

Original: In the area of NLP, Goyal et al.
Revised: In the area of NLP, Goyal et al.

Original: (2020) introduce the idea of dropping tokens.
Revised: (2020) introduce the idea of dropping tokens.

Original: Their approach, called PoWER-BERT, is based on a measure of signiﬁcance score , which is deﬁned as the total attention given to a token from all other tokens.
Revised: Their approach, called PoWER-BERT, is based on a measure of signiﬁcance score , which is deﬁned as the total attention given to a token from all other tokens.

Original: Speciﬁcally, the signiﬁcance scores of all tokens in the l -th transformer block, s l ∈ R N , is computed by 
Revised: Speciﬁcally, the signiﬁcance scores of all tokens in the l -th transformer block, s l ∈ R N , is computed by s

Original: [Equation]
Revised: [Equation]

Original: where A lh is the attention weights of head h deﬁned in (3).
Revised: where A lh is the attention weights of head h deﬁned in (3).

Original: They only pass K l tokens with the highest scores in s l to the next transformer block.
Revised: They only pass K l tokens with the highest scores in s l to the next transformer block.

Original: The pruning is performed on all blocks.
Revised: The pruning is performed on all blocks.

Original: PoWER-BERT is trained using a three-stage process.
Revised: PoWER-BERT is trained using a three-stage process.

Original: First, given a base architecture, they pretrain a model without pruning.
Revised: First, given a base architecture, they pretrain a model without pruning.

Original: In the second stage, a soft-selection layer is inserted after each transformer block, and the model is ﬁnetuned for a small number of epochs.
Revised: In the second stage, a soft-selection layer is inserted after each transformer block, and the model is ﬁnetuned for a small number of epochs.

Original: Once learned, the number of tokens to keep, K l , for each layer is computed from the soft-selection layers.
Revised: Once learned, the number of tokens to keep, K l , for each layer is computed from the soft-selection layers.

Original: Last, the model is ﬁnetuned again with the tokens pruned using the K l from the second stage.
Revised: Last, the model is ﬁnetuned again with the tokens pruned using the K l from the second stage.

Original: See Goyal et al.
Revised: See Goyal et al.

Original: (2020) for details.
Revised: (2020) for details.

Original: Recently, Rao et al.
Revised: Recently, Rao et al.

Original: (2021) proposed Dynamic-ViT that also uses scores to prune tokens.
Revised: (2021) proposed Dynamic-ViT that also uses scores to prune tokens.

Original: Unlike
Revised: Unlike

Original: PoWER-BERT, which computes signiﬁcance scores from attention weights, Rao et al.
Revised: PoWER-BERT, which computes signiﬁcance scores from attention weights, Rao et al.

Original: use a dedicated sub-network with learned parameters.
Revised: use a dedicated sub-network with learned parameters.

Original: The method requires knowledge distillation, GumbelSoftmax, and straight-through estimators on top of the DeiT training.
Revised: The method requires knowledge distillation, GumbelSoftmax, and straight-through estimators on top of the DeiT training.

Original: We will analyze the limitations of score-based methods in § 3.2.
Revised: We will analyze the limitations of score-based methods in §3.2.

Revised: Note that while downsampling

Revised: NEW tokens can complicate some downstream tasks, e .

Revised: g ., semantic segmentation and pose estimation, many techniques Wu et al.

Revised: ; Marin et al.

Revised: (2019) have been proposed to deal with the problem.

Original:  A NALYSIS
Revised:  A NALYSIS

Original: This section addresses three questions.
Revised: This section addresses three questions.

Original: First, it identiﬁes the computational bottleneck of vision transformers.
Revised: First, it identiﬁes the computational bottleneck of vision transformers.

Original: Second, it discusses the limitations of score-based downsampling.
Revised: Second, it discusses the limitations of score-based downsampling.

Original: Third, it analyzes how the softmax-attention affects the redundancy in tokens.
Revised: Third, it analyzes how the softmax-attention affects the redundancy in tokens.

Original: 3. C OMPUTATION ANALYSIS OF VISION TRANSFORMERS
Revised: 3. C OMPUTATION ANALYSIS OF VISION TRANSFORMERS

Original: We analyze the time complexity and computational costs (measured in ﬂops) of commonly used vision transformers, namely ViT and DeiT.
Revised: We analyze the time complexity and computational costs (measured in ﬂops) of commonly used vision transformers, namely ViT and DeiT.

Original: We breakdown the computation into four categories:
Revised: We breakdown the computation into four categories:

Original: Layer

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: In the ﬁgure, the x-axis represents the token values (in one dimension), and the y-axis represents their scores.
Revised: In the ﬁgure, the x-axis represents the token values (in one dimension), and the y-axis represents their scores.

Original: Suppose four tokens are to be selected.
Revised: Suppose four tokens are to be selected.

Original: Score-based methods select tokens with higher scores.
Revised: Score-based methods select tokens with higher scores.

Original: Since the scoring function is continuous, all tokens in the left lobe will be selected, resulting in redundancy and information loss in the right lobe
Revised: Since the scoring function is continuous, all tokens in the left lobe will be selected, resulting in redundancy and information loss in the right lobe

Original: The proposed method ﬁrst forms four clusters to approximate the set of tokens, then selects the cluster centers.
Revised: The proposed method ﬁrst forms four clusters to approximate the set of tokens, then selects the cluster centers.

Original: Thus, the output tokens are a more accurate representation of the original token set than the score-based methods.
Revised: Thus, the output tokens are a more accurate representation of the original token set than the score-based methods.

Original: softmax-attention (3), QKV projections (4), O projection (2) and MLP (1).
Revised: softmax-attention (3), QKV projections (4), O projection (2) and MLP (1).

Original: 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.
Revised: 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.

Original: In comparison, softmax-attention only takes less than 15%.
Revised: In comparison, softmax-attention only takes less than 15%.

Original: 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.
Revised: 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.

Original: This decomposition reveals that the QKV and O projections spend most of the computations of the multi-head self-attention.
Revised: This decomposition reveals that the QKV and O projections spend most of the computations of the multi-head self-attention.

Original: 3. L IMITATIONS OF SCORE - BASED TOKEN DOWNSAMPLING
Revised: 3. L IMITATIONS OF SCORE - BASED TOKEN DOWNSAMPLING

Original: Existing score-based token downsampling methods like PoWER-BERT and Dynamic-ViT utilize scoring functions to determine the tokens to keep (or prune).
Revised: Existing score-based token downsampling methods like PoWER-BERT and Dynamic-ViT utilize scoring functions to determine the tokens to keep (or prune).

Original: They keep tokens with the top-K scores and discard the rest.
Revised: They keep tokens with the top-K scores and discard the rest.

Original: Since these scoring functions are continuous with limited Lipschitz constants, tokens that are close in the feature space will be assigned similar scores.
Revised: Since these scoring functions are continuous with limited Lipschitz constants, tokens that are close in the feature space will be assigned similar scores.

Original: Therefore, these similar tokens will likely either be all kept or discarded, as illustrated in Figure 2a.
Revised: Therefore, these similar tokens will likely either be all kept or discarded, as illustrated in Figure 2a.

Original: 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.
Revised: 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.

Original: 3. A TTENTION AS A LOW - PASS FILTER
Revised: 3. A TTENTION AS A LOW - PASS FILTER

Original: Given a query vector q , a set of key vectors K = { k 1 , . . .
Revised: Given a query vector q , a set of key vectors K = { k 1 , . . .

Original: , k N } , the corresponding value vectors
Revised: , k N } , the corresponding value vectors

Original: V = { v 1 , . .
Revised: V = { v 1 , . .

Original: . , v N } and a scalar α > 0 , softmax-attention computes the output via o ( q ) =
Revised: . , v N } and a scalar α > 0 , softmax-attention computes the output via o ( q ) =

Original: z ( q )
Revised: z ( q )

Original: N (cid:88) i =
Revised: N (cid:88) i =

Original: exp( α q · k i ) v i where z ( q ) =
Revised: exp( α q · k i ) v i where z ( q ) =

Original: N (cid:88) i =
Revised: N (cid:88) i =

Original: exp( α q · k i ) .
Revised: exp( α q · k i ) .

Original: Note that we write o ( q ) to indicate that the output vector o is a function of the query q .
Revised: Note that we write o ( q ) to indicate that the output vector o is a function of the query q .

Original: If the query vector and all key vectors are normalized to have a ﬁxed (cid:96) 2 norm, we can rewrite (6) as o ( q ) =
Revised: If the query vector and all key vectors are normalized to have a ﬁxed (cid:96) 2 norm, we can rewrite (6) as o ( q ) =

Original: [Equation]
Revised: [Equation]

Original: q)
Revised: q

Original: [Equation]
Revised: [Equation]

Original: (cid:19)
Revised: (cid:19)

Original: [Equation]
Revised: [Equation]

Original: (7) where ∗ represents high-dimensional convolution, z (cid:48)
Revised: (7) where ∗ represents high-dimensional convolution, z (cid:48)

Original: ( q ) =
Revised: ( q ) =

Original: (cid:80) i exp (cid:0) − α 2 (cid:107) q − k i (cid:107) 2 (cid:1)
Revised: (cid:80) i exp (cid:0) − α 2 (cid:107) q − k i (cid:107) 2 (cid:1)

Original: =
Revised: =

Original: G (cid:0) q ; 1 α (cid:1) ∗
Revised: G (cid:0) q ; 1 α (cid:1) ∗

Original: S ( q ; K , 1) is the normalization scalar function, G (cid:0) q ; σ 2 (cid:1) =
Revised: S ( q ; K , 1) is the normalization scalar function, G (cid:0) q ; σ 2 (cid:1) =

Original: exp (cid:0) −(cid:107) q (cid:107) 2 (cid:14) 2 σ 2 (cid:1) is an isometric Gaussian kernel, and S ( q ; K , V )
Revised: exp (cid:0) −(cid:107) q (cid:107) 2 (cid:14) 2 σ 2 (cid:1) is an isometric Gaussian kernel, and S ( q ; K , V )

Original: = (cid:80) Ni =1 δ ( q − k i )
Revised: = (cid:80) Ni =1 δ ( q − k i )

Original: v i is a high-dimensional sparse signal, which is composed of N delta functions located at k i with value v i .
Revised: v i is a high-dimensional sparse signal, which is composed of N delta functions located at k i with value v i .

Original: According to (7), given query vectors q 1 , . . .
Revised: According to (7), given query vectors q 1 , . . .

Original: , q N , there are two conceptual steps to compute softmax-attention:
Revised: , q N , there are two conceptual steps to compute softmax-attention:

Original: ﬁlter S ( q ; K , V ) with a Gaussian kernel to get o ( q ) , and 2.
Revised: ﬁlter S ( q ; K , V ) with a Gaussian kernel to get o ( q ) , and 2.

Original: sample o ( q ) at coordinates q 1 , . .
Revised: sample o ( q ) at coordinates q 1 , . .

Original: , q N to get the output vectors o 1 , . .
Revised: , q N to get the output vectors o 1 , . .

Original: o N .
Revised: o N .

Original: Since Gaussian ﬁltering is low-pass, o ( q ) is a smooth signal.
Revised: Since Gaussian ﬁltering is low-pass, o ( q ) is a smooth signal.

Original: Therefore, the output tokens of the attention layer, i .
Revised: Therefore, the output tokens of the attention layer, i .

Original: e ., discrete samples of o ( q ) , contain similar feature values.
Revised: e ., discrete samples of o ( q ) , contain similar feature values.

Original: 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.
Revised: 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.

Original: 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 signiﬁcantly affecting the performance of a transformer, as shown by Kitaev et al.
Revised: 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 signiﬁcantly affecting the performance of a transformer, as shown by Kitaev et al.

Original: It has also been empirically observed by Goyal et al.
Revised: It has also been empirically observed by Goyal et al.

Original: and Rao et al.
Revised: and Rao et al.

Original: (2021) that even without the normalization, transformers produce tokens with similar values.
Revised: (2021) that even without the normalization, transformers produce tokens with similar values.

Original: To demonstrate this, in all our experiments, we use standard multi-head attention without normalizing keys and queries.
Revised: To demonstrate this, in all our experiments, we use standard multi-head attention without normalizing keys and queries.

Original: We conduct an ablation study with normalized keys and queries in Appendix F.
Revised: We conduct an ablation study with normalized keys and queries in Appendix F.

Original:  T OKEN POOLING
Revised:  T OKEN POOLING

Original: Pruning tokens inevitably loses information.
Revised: Pruning tokens inevitably loses information.

Original: In this section, we formulate a new token downsampling principle enabling strategical tokens selection that preserves the most information.
Revised: In this section, we formulate a new token downsampling principle enabling strategical tokens selection that preserves the most information.

Original: Based on this principle, we formulate and discuss several Token Pooling algorithms.
Revised: Based on this principle, we formulate and discuss several Token Pooling algorithms.

Original: Given a set of output tokens F = { f 1 , . . .
Revised: Given a set of output tokens F = { f 1 , . . .

Original: , f N } of a transformer block, our goal is to ﬁnd a smaller set of tokens
Revised: , f N } of a transformer block, our goal is to ﬁnd a smaller set of tokens

Original: ˆ F = { ˆ f 1 , . .
Revised: ˆ F = { ˆ f 1 , . .

Original: , ˆ f K } that after upsampling minimizes the reconstruction error of F .
Revised: , ˆ f K } that after upsampling minimizes the reconstruction error of F .

Original: Speciﬁcally, the reconstruction of u ( f i ; ˆ F ) of token f
Revised: Speciﬁcally, the reconstruction of u ( f i ; ˆ F ) of token f

Original: i is computed via interpolating the tokens in ˆ F .
Revised: i is computed via interpolating the tokens in ˆ F .

Original: The reconstruction error is then deﬁned as
Revised: The reconstruction error is then deﬁned as

Original: [Equation]
Revised: [Equation]

Original: To simplify our formulation and reduce computation, we use nearest-neighbor interpolation as u .
Revised: To simplify our formulation and reduce computation, we use nearest-neighbor interpolation as u .

Original: As a consequence, the reconstruction error (8) becomes
Revised: As a consequence, the reconstruction error (8) becomes

Original: [Equation]
Revised: [Equation]

Original: which is the asymmetric Chamfer divergence between the two sets (Barrow et al., 1977; Mechrez et al., 2018).
Revised: which is the asymmetric Chamfer divergence between F and ˆ F (Barrow et al., 1977; Mechrez et al., 2018).

Original: The loss (9) can be minimized by the K-Means algorithm, i .
Revised: The loss (9) can be minimized by the K-Means algorithm, i .

Original: e ., clustering the tokens in F into K clusters.
Revised: e ., clustering the tokens in F into K clusters.

Original: The proposed Token Pooling layer is deﬁned in Algorithm 1.
Revised: The proposed Token Pooling layer is deﬁned in Algorithm 1.

Original: It downsamples input tokens via clustering the tokens and returns the cluster centers (the average of the tokens in a cluster).
Revised: It downsamples input tokens via clustering the tokens and returns the cluster centers (the average of the tokens in a cluster).

Original: As we have shown above, this operation directly minimizes the reconstruction error (9) caused by the downsampling.
Revised: As we have shown above, this operation directly minimizes the reconstruction error (9) caused by the downsampling.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: K-Means.
Revised: K-Means.

Original: We use the K-Means algorithm to minimize (9) via the following iterations:
Revised: We use the K-Means algorithm to minimize (9) via the following iterations:

Original: [Equation]
Revised: [Equation]

Original: where [ ] is the Iverson bracket and a is the cluster assignment function.
Revised: where [ ] is the Iverson bracket and a is the cluster assignment function.

Original: The overall algorithm complexity is O ( TKNM ) where T is the number of iterations.
Revised: The overall algorithm complexity is O ( TKNM ) where T is the number of iterations.

Original: The vast majority of the computation is spent on the repetitive evaluation of the distances between tokens and centroids in step (10).
Revised: The vast majority of the computation is spent on the repetitive evaluation of the distances between tokens and centroids in step (10).

Original: K-Medoids.
Revised: K-Medoids.

Original: We can use the more efﬁcient K-Medoids algorithm by replacing step (11) with:
Revised: We can use the more efﬁcient K-Medoids algorithm by replacing step (11) with:

Original: [Equation]
Revised: [Equation]

Original: and ˆ f j ← fn
Revised: and ˆ [Equation] ← f

Original: [Equation]
Revised: [Equation]

Original: j

Original: [Equation]
Revised: [Equation]

Original: These steps minimize objective (9) under the medoid constraint : ˆ F ⊆ F .
Revised: These steps minimize objective (9) under the medoid constraint : ˆ F ⊆ F .

Original: The advantage of the K-Medoids algorithm is its time complexity O
Revised: The advantage of the K-Medoids algorithm is its time complexity O

Original: ( TKN + N 2 M ) , which is substantially lower in practice as we only compute the distances between tokens once.
Revised: ( TKN + N 2 M ) , which is substantially lower in practice as we only compute the distances between tokens once.

Original: In our experience, it requires less than 5 iterations to converge.
Revised: In our experience, it requires less than 5 iterations to converge.

Original: Apart from the distance matrix computation, the cost of the K-Medoids algorithm is negligible when compared with the cost of the other layers.
Revised: Apart from the distance matrix computation, the cost of the K-Medoids algorithm is negligible when compared with the cost of the other layers.

Original: Weighted clustering.
Revised: Weighted clustering.

Original: Reconstruction error (9) treats every token equally; however, each token contributes differently to the ﬁnal output of an attention layer.
Revised: Reconstruction error (9) treats every token equally; however, each token contributes differently to the ﬁnal output of an attention layer.

Original: Thus, we also consider a weighted reconstruction error: (cid:96) ( F , ˆ F 
Revised: Thus, we also consider a weighted reconstruction error: (cid:96) ( F , ˆ F ; w ) =

Original: [Equation] w ) = [Equation]
Revised: (cid:80) f

Original: i ∈ F min ˆ f j ∈
Revised: i ∈ F min ˆ f j ∈

Original: ˆ F w i (cid:107) f i − ˆ f j (cid:107) 2 where w =
Revised: ˆ F w i (cid:107) f i − ˆ f j (cid:107) 2 where w =

Original: [ w 1 , . .
Revised: [ w 1 , . .

Original: . , w N ] are the positive weights corresponding to the individual tokens in F .
Revised: . , w N ] are the positive weights corresponding to the individual tokens in F .

Original: Appendix A details the clustering algorithms for the weighted case.
Revised: Appendix A details the clustering algorithms for the weighted case.

Original: A good choice of the weights is the signiﬁcance scores (5), i .
Revised: A good choice of the weights is the signiﬁcance scores (5), i .

Original: e ., w i
Revised: e ., w i

Original: =
Revised: =

Original: s i .
Revised: s i .

Original: The signiﬁcance score identiﬁes the tokens that inﬂuence the current transformer block most and thus should be approximated more precisely.
Revised: The signiﬁcance score identiﬁes the tokens that inﬂuence the current transformer block most and thus should be approximated more precisely.

Revised: Note, when applying Token Pooling, it is important to keep special tokens like the classiﬁcation

Revised: NEW token.

Revised: If spatial locations on an image are need, e .

Revised: g .

Revised: , 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.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Our implementation is based on DeiT (Touvron et al., 2021) where we insert downsampling layers after each transformer block.
Revised: Our implementation is based on DeiT (Touvron et al., 2021) where we insert downsampling layers after each transformer block.

Original: To evaluate the pure effect of downsampling, we keep all metaparameters of DeiT, including the feature dimensionality, network depth, learning rate schedule, etc .
Revised: To evaluate the pure effect of downsampling, we keep all metaparameters of DeiT, including the feature dimensionality, network depth, learning rate schedule, etc .

Original: We also do not use knowledge distillation during training.

Original: We use ImageNet1k classiﬁcation benchmark (Russakovsky et al., 2015).

Original: Our cost and performance metrics are ﬂops and top-1 accuracy.

Original: [Figure]
Revised: [Figure]

Revised: NEW

Revised: Methods.

Revised: We evaluate the following token downsampling methods:

Original: Convolution downsampling.
Revised: Convolution downsampling.

Original: We implement uniform grid-downsampling via convolution with stride 2,
Revised: We implement uniform grid-downsampling via convolution with stride 2,

Original: i
Revised: i

Original: e ., the tokens corresponding to the adjacent image patches are concatenated and mapped to R M .
Revised: e ., the tokens corresponding to the adjacent image patches are concatenated and mapped to R M .

Original: This design is used in Liu et al.
Revised: This design is used in Liu et al.

Original: (2021); Heo et al.
Revised: (2021); Heo et al.

Original: See details in Appendix D. 2.
Revised: See details in Appendix D. 2.

Original: PoWER-BERT.
Revised: PoWER-BERT.

Original: We implement PoWER-BERT on DeiT, following (Goyal et al., 2020).
Revised: We implement PoWER-BERT on DeiT, following Goyal et al.

Original: Random selection , a simple baseline randomly selecting K tokens without replacement.
Revised: Random selection , a simple baseline randomly selecting K tokens without replacement.

Original: Importance selection chooses K tokens by drawing from a categorical distribution without replacement with probabilities proportional to the signiﬁcance score (5) of each token.
Revised: Importance selection chooses K tokens by drawing from a categorical distribution without replacement with probabilities proportional to the signiﬁcance score (5) of each token.

Original: Token Pooling use K-Means or K-Medoids algorithms, or their weighted versions, WK-Means or WK-Medoids, respectively.
Revised: Token Pooling use K-Means or K-Medoids algorithms, or their weighted versions, WK-Means or WK-Medoids, respectively.

Original: The weights are the signiﬁcance scores (5).
Revised: The weights are the signiﬁcance scores (5).

Original: Selection of K =
Revised: Selection of K =

Original: [ K 1 , . .
Revised: [ K 1 , . .

Original: . , K L ] .
Revised: . , K L ] .

Original: 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.
Revised: 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.

Original: Appendix C details the selection of K .
Revised: Appendix C details the selection of K .

Original: Training protocol.
Revised: Training protocol.

Original: (1) A base DeiT model ( e .
Revised: (1) A base DeiT model ( e .

Original: g ., DeiT-S) is trained using the original training (Touvron et al., 2021).
Revised: g ., DeiT-S) is trained using the original training (Touvron et al., 2021).

Original: (2) We then ﬁnetune the base DeiT model using the second stage of PoWERBERT’s training and acquire K .
Revised: (2) We then ﬁnetune the base DeiT model using the second stage of PoWERBERT’s training and acquire K .

Original: (3) We further ﬁnetune the downsampling methods using K .
Revised: (3) We further ﬁnetune the downsampling methods using K .

Original: (4) We also ﬁnetune the base DeiT model.
Revised: (4) We also ﬁnetune the base DeiT model.

Original: 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 .
Revised: 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 .

Original: 5. M AIN RESULTS
Revised: 5. M AIN RESULTS

Original: 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 .
Revised: 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 .

Original: While convolution with stride is better than adjusting M at the low-compute regime, it fails in high-compute regimes.
Revised: While convolution with stride is better than adjusting M at the low-compute regime, it fails in high-compute regimes.

Original: Importance selection improves upon random selection but is still outperformed by PoWER-BERT.
Revised: Importance selection improves upon random selection but is still outperformed by PoWER-BERT.

Original: Our
Revised: Our

Original: Token Pooling (with weighted K-Medoids) achieves the best trade-off in all regimes.
Revised: Token Pooling (with weighted K-Medoids) achieves the best trade-off in all regimes.

Original: Next, we apply Token Pooling to DeiT models with different M (DeiT-Ti, DeiT-e252, DeiT-e318, and DeiT-S).
Revised: Next, we apply Token Pooling to DeiT models with different M (DeiT-Ti, DeiT-e252, DeiT-e318, and DeiT-S).

Original: Figure 4 shows trade-off curves for each DeiT model.
Revised: Figure 4 shows trade-off curves for each DeiT model.

Original: Token Pooling enables each of the models to achieve a better computation-accuracy trade-off than simply varying the feature dimensionality M .
Revised: Token Pooling enables each of the models to achieve a better computation-accuracy trade-off than simply varying the feature dimensionality M .

Original: For each computational budget, we ﬁnd the combination of M and K that gives the highest accuracy.
Revised: For each computational budget, we ﬁnd the combination of M and K that gives the highest accuracy.

Original: The best balance is achieved by applying Token Pooling and selecting the best
Revised: The best balance is achieved by applying Token Pooling and selecting the best

Revised: [Figure]

Revised: [Figure]

Original: M .
Revised: M .

Original: We do the same for PoWER-BERT.
Revised: We do the same for PoWER-BERT.

Original: Figure 3b shows the results of the proposed Token Pooling[Figure] [Figure] and state-of-the-art methods.
Revised: Figure 3b shows the results of the proposed Token Pooling and state-of-the-art methods.

Original: We cite the results of Rao et al.
Revised: We cite the results of Rao et al.

Original: (2021), d’Ascoli et al.
Revised: (2021), d’Ascoli et al.

Original: (2021), and Pan et al.
Revised: (2021), and Pan et al.

Original: Our Token Pooling achieves the best accuracy across all evaluated compute regimes.
Revised: Our Token Pooling achieves the best accuracy across all evaluated compute regimes.

Original: 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.
Revised: 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.

Original: As can be seen, utilizing Token
Revised: As can be seen, utilizing Token

Original: Pooling, we signiﬁcantly improve the computation costs of DeiT-Ti by 42% and improve the accuracy
Revised: Pooling, we signiﬁcantly improve the computation costs of DeiT-Ti by 42% and improve the topaccuracy by 3.3 points at the same ﬂops.

Revised: Similar beneﬁts can be seen on DeiT-e252 and DeiT-e318.

Original: 5. A BLATION STUDIES
Revised: 5. A BLATION STUDIES

Original: Figure 5a compares methods utilizing signiﬁcance scores.
Revised: Figure 5a compares methods utilizing signiﬁcance scores.

Original: As can be seen, using importance selection improves upon the simple random selection.
Revised: As can be seen, using importance selection improves upon the simple random selection.

Original: By minimizing the reconstruction error (9), our method achieves better cost-accuracy trade-off.
Revised: By minimizing the reconstruction error (9), our method achieves better cost-accuracy trade-off.

Original: Figure 5b evaluates Token Pooling with different clustering algorithms.
Revised: Figure 5b evaluates Token Pooling with different clustering algorithms.

Original: Weighted versions outperform regular versions of K-Means and K-Medoids.
Revised: Weighted versions outperform regular versions of K-Means and K-Medoids.

Original: Due to the higher time complexity, K-Means is outperformed by K-Medoids (the curves are shifted toward the right).
Revised: Due to the higher time complexity, K-Means is outperformed by K-Medoids (the curves are shifted toward the right).

Original: See the metrics and ﬂops used by clustering in table format in Appendix G. Still, all Token Pooling variants outperform the baseline, demonstrating the effectiveness of our method.
Revised: See the metrics and ﬂops used by clustering in table format in Appendix G. All Token Pooling variants outperform the baseline, demonstrating the effectiveness of our method.

Original: More ablation studies are in the appendices, including the effect of cluster-center initialization (Appendix B),results of convolution downsampling (Appendix D), results of Token Pooling before ﬁnetuning (Appendix E), results with normalized keys and queries of softmax attention (Appendix F), and detailed information ( K , ﬂops, accuracy, clustering cost) of our models (Appendix G).
Revised: More ablation studies are in the appendices: clustering initialization (Appendix B), convolution downsampling (D), Token Pooling without ﬁnetuning (E), results with normalized keys and queries (F), and detailed information ( K , ﬂops, accuracy, clustering cost) of our models (G).

Original:  C ONCLUSIONS
Revised:  C ONCLUSIONS

Original: This paper provides two insights of vision transformers: ﬁrst, their computational bottleneck is the fully-connected layers, and second, attention layers generate redundant representations due to the connection to Gaussian ﬁltering.
Revised: This paper provides two insights of vision transformers: ﬁrst, their computational bottleneck is the fully-connected layers, and second, attention layers generate redundant representations due to the connection to Gaussian ﬁltering.

Original: Token Pooling, our novel nonuniform data-aware downsampling operator, utilizes these insights and signiﬁcantly improves the computation-accuracy balance of DeiT, compared to existing downsampling techniques.
Revised: Token Pooling, our novel nonuniform data-aware downsampling operator, utilizes these insights and signiﬁcantly improves the computation-accuracy balance of DeiT, compared to existing downsampling techniques.

Original: We believe that Token Pooling is a timely development and can be used with other techniques ( e .
Revised: We believe that Token Pooling is a timely development and can be used with other techniques ( e .

Original: g ., architecture search, quantization) to spur the design of efﬁcient vision transformers.
Revised: g ., architecture search, quantization) to spur the design of efﬁcient vision transformers.

Original: A W EIGHTED CLUSTERING ALGORITHMS
Revised: A W EIGHTED CLUSTERING ALGORITHMS

Original: Weighted K-Means minimizes the following objective
Revised: Weighted K-Means minimizes the following objective

Original: [Equation]
Revised: [Equation]

Original: The extension of the K-Means algorithm to the weighted case iterates the following steps:
Revised: The extension of the K-Means algorithm to the weighted case iterates the following steps:

Original: [Equation]
Revised: [Equation]

Original: Weighted K-Medoids optimizes objective (13) under the medoid constraint ˆ F ⊂ F :
Revised: Weighted K-Medoids optimizes objective (13) under the medoid constraint ˆ F ⊂ F :

Original: [Equation]
Revised: [Equation]

Original: B C LUSTERING INITIALIZATION ABLATIONS
Revised: B C LUSTERING INITIALIZATION ABLATIONS

Original: We examine the effect of the cluster center initialization.
Revised: We examine the effect of the cluster center initialization.

Original: We compare our default initialization, which uses the tokens with top-K signiﬁcance scores as initial cluster centers, with random initialization, which randomly selects tokens as initial cluster centers.
Revised: We compare our default initialization, which uses the tokens with top-K signiﬁcance scores as initial cluster centers, with random initialization, which randomly selects tokens as initial cluster centers.

Original: As shown in Figure 6, Token
Revised: As shown in Figure 6, Token

Original: Pooling is robust to the initialization methods.
Revised: Pooling is robust to the initialization methods.

Original: [Figure]
Revised: [Figure]

Original: C T RAINING DETAILS 
Revised: C T RAINING DETAILS & HYPERPARAMETERS

Revised: A Token Pooling layer has one parameter—the number of the output tokens K (or equivalently the

Revised: NEW downsampling ratio)—which is the same as downsampling layers like max or average pooling.

Revised: For max or average pooling layers, the sizes of the kernels and strides play the same role as K .

Original:  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.
Revised: We use the method proposed by Goyal et al. (2020) to automatically select K s 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.

Original: Speciﬁcally, we run the second stage of the PoWERBERT training  for 30 epochs with various values of the token-selection weight parameter  producing a family of models.
Revised: Speciﬁcally, 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.

Revised: The single parameter λ controls the overall efﬁciency of a model, and the numbers of retained tokens of all Token Pooling layers are selected automatically based on the choice of λ .

Original: Each of the  models has a different number of retained tokens at each of its L transformer blocks: K =
Revised: Each of the resulted models has a different number of retained tokens at each of its L transformer blocks: K =

Original: ( K 1 , . .
Revised: ( K 1 , . .

Original: . , K L ) .
Revised: . , K L ) .

Original: Appendix G lists all combinations of  K .
Revised: Appendix G lists all combinations of automatically determined K .

Original: We then ﬁnetune the base DeiT model (from the ﬁrst stage training of PoWER-BERT) with the selected K using the third (last) stage of  PoWER-BERT, random selection, importance selection or our Token Pooling using the same K .
Revised: We then ﬁnetune 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.

Original: We ﬁnd that the DeiT models provided by Touvron et al.
Revised: We ﬁnd that the DeiT models provided by Touvron et al.

Original: (2021) are under-ﬁt, and their accuracy improves with additional training, see Figure 7.
Revised: (2021) are under-ﬁt, and their accuracy improves with additional training, see Figure 7.

Original: After the standard DeiT training, we restart the training.
Revised: After the standard DeiT training, we restart the training.

Original: This ensures that downsampling models and DeiT with almost the same amount of training.
Revised: This ensures that downsampling models and DeiT with a similar number of training steps.

Original: [Figure]
Revised: [Figure]

Revised: DeiT-Distil results (cited from (Touvron et al., 2021)), which use knowledge distillation.

Original: D C ONVOLUTION DOWNSAMPLING
Revised: D C ONVOLUTION DOWNSAMPLING

Original: [Figure]

Original: As mentioned in the main paper, we enumerate combinations of layers to insert the convolution downsampling layer.
Revised: As mentioned in the main paper, we enumerate combinations of layers to insert the convolution downsampling layer.

Original: We use 2x2 convolution with stride 2 like in Liu et al.
Revised: We use 2x2 convolution with stride 2 like in Liu et al.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: When using 3 layers at depths l 1 < l 2 < l 3 , we restrict l 2 − l 1 = l 3 − l 2 .
Revised: When using 3 layers at depths l 1 < l 2 < l 3 , we restrict l 2 − l 1 = l 3 − l 2 .

Original: With these constraints, we enumerate all possible downsampling conﬁgurations.
Revised: With these constraints, we enumerate all possible downsampling conﬁgurations.

Original: Each of the combinations produces a model with a different computation-accuracy trade-off, and we report the Pareto front, i .
Revised: Each of the combinations produces a model with a different computation-accuracy trade-off, and we report the Pareto front, i .

Original: e ., the best accuracy these models achieve at a given ﬂop.
Revised: e ., the best accuracy these models achieve at a given ﬂop.

Original: See Figure 8.
Revised: See Figure 8.

Revised: Note that the convolution downsampling layer is equivalent to the patch-merging layer used by Liu

Original: E R ESULTS WITHOUT FINETUNING
Revised: 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).

Revised: It is also a generalized version of the commonly used average pooling layer.

Revised: Average pooling can be implemented by convolving the feature map with a constant kernel of size n × n containing 1 n 2 and subsample (stride) every n pixels.

Revised: By learning its kernel, the convolution downsampling layer is able to ﬁnd a kernel suitable for the downsampling and the current task than average pooling.

Revised: [Figure]

Original: 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 ﬁnetuning.
Revised: 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 ﬁnetuning.

Original: 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,
Revised: 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,

Original: Token Pooling preserves information that enables the model to retain accuracy during token downsampling.
Revised: Token Pooling preserves information that enables the model to retain accuracy during token downsampling.

Original: In Figure 9, we also show the results when we replace tokens to their cluster centers (WK-Means, carry, no ﬁnetuning and WK-Medoids, carry, no ﬁnetuning).
Revised: In Figure 9, we also show the results when we replace tokens to their cluster centers (WK-Means, carry, no ﬁnetuning and WK-Medoids, carry, no ﬁnetuning).

Original: Speciﬁcally, in addition to outputting
Revised: Speciﬁcally, in addition to outputting

Original: 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.
Revised: 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.

Original: This operation preserves the attention weights, and since the models are not trained after inserting Token Pooling layers, it preserves the most accuracy.
Revised: This operation preserves the attention weights, and since the models are not trained after inserting Token Pooling layers, it preserves the most accuracy.

Original: In practice, when the models are trained with Token Pooling layers, the carry operation is not important.
Revised: In practice, when the models are trained with Token Pooling layers, the carry operation is not important.

Original: [Figure]
Revised: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]
Revised: [Figure]

Original: F N ORMALIZED KEY AND QUERY VECTORS
Revised: F N ORMALIZED KEY AND QUERY VECTORS

Original: Our analysis of softmax-attention in § 3.3 assumes the query and the key vectors are normalized to have a constant (cid:96) 2 norm.
Revised: Our analysis of softmax-attention in §3.3 assumes the query and the key vectors are normalized to have a constant (cid:96) 2 norm.

Original: It has been observed by Kitaev et al.
Revised: It has been observed by Kitaev et al.

Original: (2020) that normalizing key and query vectors does not change the performance of a transformer.
Revised: (2020) that normalizing key and query vectors does not change the performance of a transformer.

Original: Thus, for all experiments, we train models without the norm normalization.
Revised: Thus, for all experiments, we train models without the norm normalization.

Original: In this section, we verify this observation by training various DeiT,
Revised: In this section, we verify this observation by training various DeiT,

Original: PoWER-BERT, and Token Pooling with normalized keys and queries.
Revised: PoWER-BERT, and Token Pooling with normalized keys and queries.

Original: We found that the scalar α in the softmax attention layer (6) can affect the performance of a transformer with normalized keys and queries.
Revised: We found that the scalar α in the softmax attention layer (6) can affect the performance of a transformer with normalized keys and queries.

Original: As can be seen in Figure 10a, setting the scalar α = 1 slightly deteriorates the performance of the model.
Revised: As can be seen in Figure 10a, setting the scalar α = 1 slightly deteriorates the performance of the model.

Original: Instead of using a ﬁxed α , we let the model learn the α for each layer.
Revised: Instead of using a ﬁxed α , we let the model learn the α for each layer.

Original: As can be seen in Figure 10b, learning α enables the resulting models to achieve similar cost-accuracy trade-off as[Figure] [Figure] the standard (unnormalized) models.
Revised: As can be seen in Figure 10b, learning α enables the resulting models to achieve similar cost-accuracy trade-off as the standard (unnormalized) models.

Original: With or without the normalization and the learnable α , the proposed Token Pooling signiﬁcantly improves the cost-accuracy trade-off of DeiT and outperforms
Revised: With or without the normalization and the learnable α , the proposed Token Pooling signiﬁcantly improves the cost-accuracy trade-off of DeiT and outperforms

Original: PoWER-BERT.
Revised: PoWER-BERT.

Original: G C LUSTERING ALGORITHM ABLATIONS
Revised: G C LUSTERING ALGORITHM ABLATIONS

Original: 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).
Revised: 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).

Original: Figure 11 shows the results.
Revised: Figure 11 shows the ablation results with different backbone models.

Original: Table 5 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.
Revised: 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.

Original: Apart from the standard K-Means and K-Medoids, other clustering approaches could be used.
Revised: Apart from the standard K-Means and K-Medoids, other clustering approaches could be used.

Original: Many methods are not suitable due to efﬁciency constraints.
Revised: Many methods are not suitable due to efﬁciency constraints.

Original: For example, normalized cut (Shi & Malik, 2000) uses expensive spectral methods.
Revised: For example, normalized cut (Shi & Malik, 2000) uses expensive spectral methods.

Original: One prerequisite of using K-Means and K-Medoids is the number of clusters, K .
Revised: One prerequisite of using K-Means and K-Medoids is the number of clusters, K .

Original: While selecting K for each of the layers may be tedious and difﬁcult, one can choose K via computational budgets and heuristics.
Revised: While selecting K for each of the layers may be tedious and difﬁcult, one can choose K via computational budgets and heuristics.

Original: In this work, we use the automatic search procedure proposed by Goyal et al.
Revised: In this work, we use the automatic search procedure proposed by Goyal et al.

Original: (2020) to determine K , see Appendix C.
Revised: (2020) to determine K , see Appendix C.

Original: Since both K-Means and K-Medoids require specifying the number of clusters K in advance, one may consider using methods automatically determining K .
Revised: Since both K-Means and K-Medoids require specifying the number of clusters K in advance, one may consider using methods automatically determining K .

Original: Nevertheless, such methods typically have other parameters, which are less interpretable than K .
Revised: Nevertheless, such methods typically have other parameters, which are less interpretable than K .

Original: For example, mean-shift (Cheng, 1995) or quick-shift (Vedaldi & Soatto, 2008) require specifying the kernel size.
Revised: For example, mean-shift (Cheng, 1995) or quick-shift (Vedaldi & Soatto, 2008) require specifying the kernel size.

Original: From our experience, determining these parameters is challenging.
Revised: From our experience, determining these parameters is challenging.

Original: Also, since the number of clusters (and hence the number of output tokens) is determined on the ﬂy during inference, the computational requirement can ﬂuctuate, making deployment of these models difﬁcult.
Revised: Also, since the number of clusters (and hence the number of output tokens) is determined on the ﬂy during inference, the computational requirement can ﬂuctuate, making deployment of these models difﬁcult.

Original: [Table]
Revised: [Figure]

Original: (+0.03) [Table] [Table] [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: Gflops
Revised: Pooling

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure] Gflops
Revised: Pooling (WK

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure] [Figure]
Revised: Medoids

Original: [Figure]
Revised: [Figure]

Original: PoWER
Revised: Gﬂops

Original: [Figure]
Revised: [Figure]

Original: BERT

Original: Token Pooling (K-Medoids)

Original: Token Pooling (WK-Medoids)

Original: Gﬂops Accuracy (%)
Revised: Gﬂops Accuracy (%)

Original: Gﬂops Accuracy (%)

Original: Gﬂops
Revised: Gﬂops

Original: Accuracy (%)
Revised: Accuracy (%)

Original: Gﬂops
Revised: Gﬂops

Original: Accuracy (%) - - 0.
Revised: Accuracy (%) - - 0.

Original: [Table] [Table]
Revised: [Table] [Table]

Revised: (+0.03)

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: NEW Throughput reﬂects 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 speciﬁc hardware, implementation quality, and the workload and the state ( e .

Original: [Table]
Revised: g ., temperature) of the machines.

Original: [Table] [Table] [Table]
Revised: Since it is usually difﬁcult 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 .

Original: [Table]
Revised: e ., ﬂops) that is known to be highly correlated with the energy consumption on device.

Original: [Table] [Table] [Table]
Revised: 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.

Original: [Table] [Table]
Revised: Note that our PyTorch implementation of the clustering algorithms can be signiﬁcantly improved, for example, via implementing as a CUDA kernel.

Original: [Table] [Table] [Table]
Revised: 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.

Original: [Table] [Table] [Table]
Revised: As can be seen, with the same accuracy, using Token Pooling increases the throughput of DeiT-e 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).