Token Pooling in Vision Transformers

Despite the recent success in many applications , the high computational requirements of vision transformers limit their use in resource-constrained settings . While many existing methods improve the quadratic complexity of attention , in most vision transformers , self-attention is not the major computation bottleneck , e.g. , more than 80 % of the computation is spent on fully-connected layers . To improve the computational complexity of all layers , we propose a novel token downsampling method , called Token Pooling , efficiently exploiting redundancies in the images and intermediate token representations . We show that , under mild assumptions , softmax-attention acts as a high-dimensional low-pass ( smoothing ) filter . Thus , its output contains redundancy that can be pruned to achieve a better trade-off between the computational cost and accuracy . Our new technique accurately approximates a set of tokens by minimizing the reconstruction error caused by downsampling . We solve this optimization problem via cost-efficient clustering . We rigorously analyze and compare to prior downsampling methods . Our experiments show that Token Pooling significantly improves the cost-accuracy trade-off over the state-of-the-art downsampling . Token Pooling is simple and effective and can benefit many architectures with global attention . Applied to DeiT , NEW it achieves the same ImageNet top-1 accuracy using 42 % fewer computations . 1 INTRODUCTION . Vision transformers ( Dosovitskiy et al. , 2020 ; Touvron et al. , 2021 ; Liu et al. , 2021 ; Heo et al. , 2021 ; Zheng et al. , 2021 ) have demonstrated state-of-the-art results in many vision applications , from image classification to segmentation . However , the high computational cost limits their use in resource-restricted , real-time , or low-powered applications . While most prior work in Natural Language Processing ( NLP ) improve the time-complexity of attention ( Tay et al. , 2020b ; Ilharco et al. , 2020 ) , in vision transformers the main computation bottleneck is the fully-connected layers , as we show in §3.1 . The computational complexity of these layers is determined by the number of tokens and their feature dimensionality . While reducing the dimensionality improves computational cost , it sacrifices model capacity and often significantly deteriorates the accuracy of the model . On the other hand , since images often contain mostly smooth surfaces with sparsely located edges and corners , they contain similar ( and thus redundant ) features . Moreover , we show that , under mild assumptions , softmax-attention is equivalent to low-pass filtering of tokens and thereby produces tokens with similar features , as empirically observed by Goyal et al . ( 2020 ) and Rao et al . ( 2021 ) . This redundancy in representations suggests that we can reduce the number of tokens , i.e. , downsampling , without a significant impact to the accuracy , achieving a better cost-accuracy trade-off than reducing feature dimensionality alone . Downsampling is widely used in Convolutional Neural Network ( CNN ) architectures to improve computational efficiency , among other purposes . Given a grid of pixels or features , downsampling gradually reduces the grid dimensions via combining neighboring vertices on the grid . The prevailing max/average pooling and sub-sampling are examples of ( spatially uniform ) grid-downsampling that only uses locations on the grid to decide which vertices to combine . Such methods do not efficiently address non-uniformly distributed redundancy in images and features ( Recasens et al. , 2018 ; Marin et al. , 2019 ) . Unlike CNNs that require grid preservation , transformers allow a wider range of nonuniform data-aware downsampling layers , where a better operator can be designed . We propose Token Pooling , a novel nonuniform data-aware downsampling operator for transformers efficiently exploiting redundancy in features . See the illustration and performance metric in Fig- ures 1a & 1b . Motivated by nonuniform sampling and image compression ( Marvasti , 2012 ; Unat et al. , 2009 ; Belfor et al. , 1994 ; Rabbani , 2002 ) , we formulate token downsampling as an optimization problem that minimizes the reconstruction error caused by downsampling . We show that clustering algorithms , K-Means and K-Medoids , efficiently solve this problem , see illustration in Figure 1c . To the best of our knowledge , we are the first to use this formulation and simple clustering analysis for token donwsampling in transformers . We compare with various prior downsampling techniques , including grid-downsampling ( Pan et al. , 2021 ; Liu et al. , 2021 ; Wang et al. , 2021 ) and token pruning ( Goyal et al. , 2020 ; Rao et al. , 2021 ) . The proposed Token Pooling outperforms existing methods and provides the best trade-off between computational cost and classification accuracy . Contributions . The paper makes the following contributions : • We conduct an extensive study of prior downsampling techniques for visual transformers by comparing their computation-accuracy trade-offs . • We analyze the computational cost of vision-transformer components and the limitations of the prior score-based downsampling methods . We also show that attention layers behave like low-pass filtering and thus produce redundant tokens . • Motivated by the redundancy in images and features , we propose a novel token downsampling technique , Token Pooling , for transformers with global attention via error minimization and NEW achieve a significant improvement in the computation-accuracy trade-off . 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.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 is originally designed for NLP by Vaswani et al . ( 2017 ) and further popularized by Radford et al . ( 2018 ) and Devlin et al . ( 2019 ) . In a high level , a vision transformer is a composition ofL transformer blocks that take a set of input tokens and return another set of output tokens . In vision , input tokens are features representing individual nonoverlapping image patches . To perform classification , a classification token is inserted to estimate the probabilities of individual classes . To achieve the state-of-the-art , ViT ( Dosovitskiy et al. , 2020 ) used pretraining on JFT-300M , a proprietary dataset much larger than standard ImageNet1k ( Deng et al. , 2009 ) . Recently , DeiT ( Touvron et al. , 2021 ) achieved state-of-the-art results with advanced training on ImageNet1k only . Let the set of tokens at depth l be F l = { f l0 , . . . , f lN } where f li ∈ RM is the feature values of the i-th token . A typical transformer block φ at depth l processes F l by a Multi-head Self-Attention ( MSA ) layer and a point-wise Multi-Layer Perceptron ( MLP ) . Let matrix F ∈ RN×M be a row-wise concatenation of tokens F l. Then1 , φ ( F ) = MLP ( MSA ( F ) ) such that ( 1 ) MSA ( F ) = [ O1 , O2 , . . . , OH ] W O ( 2 ) where H is the number of heads , matrix WO ∈ RM×M is a learnable parameter of the block , [ , ] is column-wise concatenation and Oh ∈ RN×d is the output of h-th attention head for d =M/H : Oh = AhVh such that Ah = softmax ( QhK > h /√ d ) ∈ RN×N . ( 3 ) Keys Kh , queries Qh and values Vh are linear projections of the input tokens ( QKV projections ) : Qh = FW Q h , Kh = FW K h , Vh = FW V h ( 4 ) where WQh ∈ RM×d , WKh ∈ RM×d , W Vh ∈ RM×d are learnable linear transformations . Note , the number of tokens is not affected by the transformer blocks , i.e. , |F l+1| = |F l| . 2.2 EFFICIENT TRANSFORMERS . Similar to many machine learning models , the efficiency of transformers can be improved via metaparameter search ( Howard et al. , 2017 ; Tan & Le , 2019 ) , automated neural architecture search ( Elsken et al. , 2019 ; Tan et al. , 2019 ; Wu et al. , 2019 ) , manipulating the input size and resolution of feature maps ( Paszke et al. , 2016 ; Howard et al. , 2017 ; Zhao et al. , 2018 ) , pruning ( LeCun et al. , 1990 ) , quantization ( Jacob et al. , 2018 ) , and sparsification ( Gale et al. , 2019 ) , etc . For example , Dosovitskiy et al . ( 2020 ) and Touvron et al . ( 2021 ) obtain a family of ViT and DeiT models , respectively , by varying the input resolution , the number of heads H , and the feature dimensionality M . Each of the models operates with a different computational requirement and accuracy . In the following , we review techniques developed for transformers . 2.2.1 EFFICIENT SELF-ATTENTION . The softmax-attention layer ( 3 ) has a quadratic time complexity w.r.t . the number of tokens , i.e. , O ( N2 ) . In many NLP applications where every token represents a word or a character , N can be large , making attention a computation bottleneck ( Dai et al. , 2019 ; Rae et al. , 2019 ) . While many works improve the time complexity of attention layers , as we will see in §3.1 , they are not the bottleneck in most current vision transformers . The time complexity of an attention layer can be reduced by restricting the attention field-of-view and thus imposing sparsity on Ah in ( 3 ) . This can be achieved using the spatial relationship between tokens in the image/text domain ( Parmar et al. , 2018 ; Ramachandran et al. , 2019 ; Qiu et al. , 2020 ; Beltagy et al. , 2020 ; Child et al. , 2019 ; Zaheer et al. , 2020 ) or based on token values using localitysensitive hashing , sorting , compression , etc . ( Kitaev et al. , 2020 ; Vyas et al. , 2020 ; Tay et al. , 2020a ; Liu et al. , 2018 ; Wang et al. , 2020 ; Tay et al. , 2021 ) . Prior works have also proposed attention mechanisms with lower time complexity , e.g. , O ( N ) or O ( N logN ) ( Katharopoulos et al. , 2020 ; Peng et al. , 2021 ; Choromanski et al. , 2021 ; Tay et al. , 2021 ) . Roy et al . ( 2021 ) cluster queries and NEW keys to sparsify attention matrices to speed-up the attention , but they do not downsample the tokens . Note that the goal of these methods is to reduce the time complexity of the attention layer—the number of tokens remains the same across the transformer blocks . In contrast , our method reduces the number of tokens after attention has been computed . Thereby , we can utilize these methods to further improve the overall efficiency of transformers . Recently , Wu et al . ( 2020 ) proposed a new attention-based layer that learns a small number of query vectors to extract information from the input feature map . Similarly , Wu et al . ( 2021 ) replace selfattention with a new recurrent layer that outputs a smaller number of tokens . In comparison , our method directly minimizes the token reconstruction error due to token downsampling . Also , our layer has no learnable parameters and can be easily incorporated into existing vision transformers .

Summary: 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.