EViT: Expediting Vision Transformers via Token Reorganizations

1 INTRODUCTION . Computer vision research has evolved into Transformers since ViTs ( Dosovitskiy et al. , 2021 ) . Equipped with global self-attention , ViTs have shown impressive capability upon local convolution ( i.e. , CNNs ) on prevalent visual recognition scenarios , including image classification ( Dosovitskiy et al. , 2021 ; Touvron et al. , 2021a ; Jiang et al. , 2021 ; Graham et al. , 2021 ) , object detection ( Carion et al. , 2020 ) , and semantic segmentation ( Xie et al. , 2021 ; Liu et al. , 2021 ; Wang et al. , 2021a ; c ) , with both supervised and unsupervised ( self-supervised ) training ( Pan et al. , 2021 ; Ge et al. , 2021 ) configurations . Based on the main spirit of ViTs ( i.e. , MHSA ) , there are wide investigations ( Liu et al. , 2021 ; Chu et al. , 2021 ; Wang et al. , 2021a ) to explore the network structure of ViT models for continuous recognition performance improvement . Along with the development of ViT models , the computation burden is becoming an issue . The global self-attention between image tokens and long-range dependency make the model converge slow compared to CNNs . As illustrated in ( Dosovitskiy et al. , 2021 ) , training a ViT from scratch ∗Major work done during an internship at Tencent AI Lab . †Corresponding author . typically requires larger datasets ( e.g. , ImageNet-21k ( Deng et al. , 2009 ) and JFT-300M ( Sun et al. , 2017 ) ) than those of CNNs ( e.g. , CIFAR-10/100 ( Krizhevsky et al. , 2009 ) and ImageNet-1k ) . Also , using more training iterations is a necessity for network convergence ( Dosovitskiy et al. , 2021 ; Touvron et al. , 2021a ; Wang et al. , 2022 ) . Without such large scale training , the ViT models are not fully exploited and perform inferior on visual recognition scenarios . These issues motivate us to expedite ViTs for practice usage . The model acceleration of ViTs is important to reduce computational complexity . However , there are few studies focused on ViT acceleration . This is because the significant model difference between CNNs and ViTs prevents CNN model acceleration ( e.g. , pruning and distillation ) from applying on ViTs . Nevertheless , we analyze ViT from another perspective . We observe that not all image tokens in ViTs contribute positively to the final predictions . Fig . 1 shows some examples where part of the input image tokens are randomly dropped out . In Fig . 1b , removing image tokens related to the visual content of the corresponding category makes the ViT predict incorrectly . In comparison , removing unrelated image tokens does not affect ViT predictions , as shown in Fig . 1a . On the other hand , we notice that ViTs divide images into non-overlapping tokens and perform self-attention ( Vaswani et al. , 2017 ) computation on these tokens . A notable characteristic of self-attention is that it can process a varying number of tokens . These observations motivate us to reorganize image tokens for ViT model accelerations . In this work , we propose a token reorganization method to identify and fuse image tokens . Given all the image tokens as input , we compute token attentiveness between these tokens and the class token for identification . Then , we preserve the attentive image tokens and fuse the inattentive tokens into one token to allow the gradient back-propagate through the inattentive tokens for better attentive token identification . In this way , we gradually reduce the number of image tokens as the network goes deeper to decrease computation cost . Also , the capacity of the ViT backbone can be flexibly controlled via the identification process where no additional parameters are introduced . We adopt our token reorganization method on representative ViT models ( i.e. , DeiT ( Touvron et al. , 2021a ) and LV-ViT ( Jiang et al. , 2021 ) ) for ImageNet classification evaluation . The experimental results show our advantages . For instance , with the same amount of input image tokens , our method speeds up the DeiT-S model by 50 % , while only sacrificing 0.3 % recognition accuracy on ImageNet classification . On the other hand , we extend our methods to boost the ViT model recognition performance under the same computational cost . By increasing the input image resolution , our method facilitates Vision Transformers in taking more image tokens to achieve higher classification accuracy . Numerically , we improve the ImageNet classification accuracy of the DeiT-S model by 1 % under the same computational cost . Moreover , by using an oracle ViT to guide the token reorganization process , our method can increase the accuracy of the original DeiT-S from 79.8 % to 80.7 % while reducing its computation cost by 36 % under the multiply-accumulate computation ( MAC ) metric . 2 RELATED WORK . 2.1 VISION TRANSFORMERS . Transformers ( Vaswani et al. , 2017 ) have drawn much attention to computer vision recently due to its strong capability of modeling long-range relation . A few attempts have been made to add self- attention layers or Transformers on top of CNNs in image classification ( Hu et al. , 2019 ) , object detection ( Carion et al. , 2020 ) , segmentation ( Wang et al. , 2021c ) , image retrieval ( Lu et al. , 2019 ) and even video understanding ( Sun et al. , 2019 ; Girdhar et al. , 2019 ) . Vision Transformer ( ViT ) ( Dosovitskiy et al. , 2021 ) first introduced a set of pure Transformer backbones for image classification and its follow-ups modify the ViT architecture for not only better visual recognition ( Touvron et al. , 2021a ; Yuan et al. , 2021 ; Zhou et al. , 2021 ) but many other high-level vision tasks , such as object detection ( Carion et al. , 2020 ; Zhu et al. , 2020 ; Liu et al. , 2021 ) , semantic segmentation ( Wang et al. , 2021a ; Xie et al. , 2021 ; Chu et al. , 2021 ) , and video understanding ( Bertasius et al. , 2021 ; Fan et al. , 2021 ) . Vision Transformers have shown its strong potential as an alternative to the previously dominant CNNs . 2.2 MODEL ACCELERATION . Neural networks are typically overparameterized ( Allen-Zhu et al. , 2019 ) , which results in significant redundancy in computation in deep learning models . To deploy deep neural networks on mobile devices , we must reduce the storage and computational overhead of the networks . Many adaptive computation methods are explored ( Bengio et al. , 2015 ; Wang et al. , 2018 ; Graves , 2016 ; Hu et al. , 2020 ; Wang et al. , 2020b ; Han et al. , 2021b ) to alleviate the computation burden . Parameter pruning ( Srinivas & Babu , 2015 ; Han et al. , 2015 ; Chen et al. , 2015b ) reduces redundant parameters which are not sensitive to the final performance . Some other methods leverage knowledge distillation ( Hinton et al. , 2015 ; Romero et al. , 2014 ; Luo et al. , 2016 ; Chen et al. , 2015a ) to obtain a small and compact model with distilled knowledge of a larger one . These model acceleration strategies are limited to convolutional neural networks . There are also some attempts to accelerate the computation of the Transformer model , including proposing more efficient attention mechanisms ( Wang et al. , 2020a ; Kitaev et al. , 2020 ; Choromanski et al. , 2020 ) and the compressed Transformer structures ( Liu et al. , 2021 ; Heo et al. , 2021 ; Wang et al. , 2021a ) . These methods mainly focus on reducing the complexity of the network architecture through artificially designed modules . Another approach to ViT acceleration is reducing the number of tokens involved in the inference of ViTs . Notably , Wang et al . ( 2021b ) proposed a method to dynamically determine the number of patches to divide on an image . The ViT will stop inference for an input image if it has sufficient confidence in the prediction of the intermediate outputs . Remarkably , Ryoo et al . ( 2021 ) proposed TokenLearner to expedite ViTs , where a relatively small amount of tokens are learned by aggregating the entire feature map weighted by a dynamic attention map conditioned on the feature map . This can be seen as a sophisticated method for tokenizing the input images . Different from TokenLearner , our work focuses on the progressive selection of informative tokens during training . Another related work is DynamicViT ( Rao et al. , 2021 ) , which introduces a method to reduce token for a fully trained ViT , where an extra learnable neural network is added to ViT to select a subset of tokens . Our work provides a novel perspective for reducing the computational overhead of inference by proposing a token reorganization method to progressively reduce and reorganize image tokens . Unlike DynamicViT , our method does not need a fully trained ViT to help the training and brings no additional parameters into ViT . 3 TOKEN REORGANIZATIONS . Our method EViT is built upon ViT ( Dosovitskiy et al. , 2021 ) and its variants for visual recognition . We first review ViT and then present how to incorporate our method into the ViT training procedure . Each component of EViT , including attentive token identification and inattentive token fusion , will be elaborated . Furthermore , we analyze the effectiveness of our method by visualizing the attentive tokens at different layers and discuss training on higher resolution images with EViT . 3.1 VIT OVERVIEW . Vision Transformers ( ViTs ) are first introduced by Dosovitskiy et al . ( 2021 ) into visual recognition . They perform tokenization by dividing an input image into patches and projecting each patch to a token embedding . An extra class token [ CLS ] is added to the set of image tokens and is responsible for aggregating global image information and final classification . All of the tokens are added by a learnable vector ( i.e. , positional encoding ) and fed into the sequentially-stacked Transformer encoders consisting of a multi-head self-attention ( MHSA ) layer and a feed-forward network ( FFN ) . In MHSA , the tokens are linearly mapped and further packed into three matrices , namely Q , K , and V . The attention operation is conducted as follows . Attention ( Q , K , V ) = Softmax ( QK > √ d ) V . ( 1 ) where d is the length of the query vector . The result of Softmax ( QK > / √ d ) is a square matrix which is called the attention map . The first row of attention map represents the attention from [ CLS ] to all tokens and will be used to determine the attentiveness ( importance ) of each token ( detailed in the next subsection ) . The output tokens of MHSA are sent to FFN , consisting of two fully connected layers with a GELU activation layer ( Hendrycks & Gimpel , 2016 ) in between . At the final Transformer encoder layer , the [ CLS ] token is extracted and utilized for object category prediction . More details of Transformers can be found in Vaswani et al . ( 2017 ) . 3.2 ATTENTIVE TOKEN IDENTIFICATION . Let n denote the number of input tokens to a ViT encoder . In the last encoder of ViT , the [ CLS ] token is taken out for classification . The interactions between [ CLS ] and other tokens are performed via the attention mechanism ( Vaswani et al. , 2017 ) in the ViT encoders : xclass = Softmax ( qclass ·K > √ d ) V = a · V . ( 2 ) where qclass , K , and V denote the query vector of [ CLS ] , the key matrix , and the value matrix , respectively , in an attention head . In other words , the output of the [ CLS ] token xclass is a linear combination of the value vectors V = [ v1 , v2 , . . . , vn ] > , with the combination coefficients ( denoted by a in Eq . 2 ) being the attention values from [ CLS ] with respect to all tokens . Since vi comes from the i-th token , the attention value ai ( i.e. , the i-th entry in a ) determines how much information of the i-th token is fused into the output of [ CLS ] ( i.e. , xclass ) through the linear combination . It is thus natural to assume that the attention value ai indicates the importance of the i-th token . Moreover , Caron et al . ( 2021 ) also showed that the [ CLS ] token in ViTs pays more attention ( i.e. , having a larger attention value ) to class-specific tokens than to the tokens on the non-object regions . To this end , we propose to use the attentiveness of the [ CLS ] token with respect to other tokens to identify the most important tokens . Based on these arguments , a simple method to reduce computation in ViT is to remove the tokens with the smallest attention values . However , we find that directly removing those tokens severely deteriorates the classification accuracy , as shown in Table 1 . Therefore , we propose to incorporate image token reorganization during the ViT training process . In multi-head self-attention layer , there are multiple heads performing the computation of Eq . 1 in parallel . Thus , there are multiple [ CLS ] attention vectors a ( h ) , h = [ 1 , . . . , H ] , with H being the total number of attention heads ( Vaswani et al. , 2017 ) . We compute the average attentiveness value of all heads by ā = ∑H h=1 a ( h ) /H . As shown in Figure 2 , we identify and preserve the tokens corresponding to the k largest ( top-k ) elements in ā ( k is a hyperparameter ) , which we call the attentive tokens , and further fuse the other tokens ( which we call the inattentive tokens ) into a new token . The fusion of tokens is detailed in the following paragraph . We define the token keeping rate as κ = k/n .

Summary: This paper aims to expediting vision transformers by reducing the number of tokens. The main contribution is  the attentive token identification, which is based on calculating the attentiveness of the class token with respect to each image token. Experiments on DeiT and LV-ViT show that the proposed approach is able to successfully reduce the computational cost of the baseline models with nearly no performance drop.