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