Not All Images are Worth 16x16 Words: Dynamic Transformers for Efficient Image Recognition

Vision Transformers ( ViT ) have achieved remarkable success in large-scale image recognition . They split every 2D image into a fixed number of patches , each of which is treated as a token . Generally , representing an image with more tokens would lead to higher prediction accuracy , while it also results in drastically increased computational cost . To achieve a decent trade-off between accuracy and speed , the number of tokens is empirically set to 16x16 or 14x14 . In this paper , we argue that every image has its own characteristics , and ideally the token number should be conditioned on each individual input . In fact , we have observed that there exist a considerable number of “ easy ” images which can be accurately predicted with a mere number of 4x4 tokens , while only a small fraction of “ hard ” ones need a finer representation . Inspired by this phenomenon , we propose a Dynamic Transformer to automatically configure a proper number of tokens for each input image . This is achieved by cascading multiple Transformers with increasing numbers of tokens , which are sequentially activated in an adaptive fashion at test time , i.e. , the inference is terminated once a sufficiently confident prediction is produced . We further design efficient feature reuse and relationship reuse mechanisms across different components of the Dynamic Transformer to reduce redundant computations . Extensive empirical results on ImageNet , CIFAR-10 , and CIFAR-100 demonstrate that our method significantly outperforms the competitive baselines in terms of both theoretical computational efficiency and practical inference speed . Code and pre-trained models ( based on PyTorch and MindSpore ) are available at https : //github.com/ blackfeather-wang/Dynamic-Vision-Transformer and https : //github . com/blackfeather-wang/Dynamic-Vision-Transformer-MindSpore . 1 Introduction . Transformers , the dominant self-attention-based models in natural language processing ( NLP ) [ 10 , 40 , 3 ] , have been successfully adapted to image recognition problems [ 11 , 55 , 38 , 17 ] recently . In particular , vision Transformers achieve state-of-the-art performance on the large scale ImageNet benchmark [ 9 ] , while exhibit excellent scalability with the further growing dataset size ( e.g. , on JFT300M [ 11 ] ) . These models split each image into a fixed number of patches and embed them into 1D tokens as inputs . Typically , representing the data using more tokens contributes to higher prediction accuracy , but leads to intensive computational cost , which grows quadratically with respect to the ∗Equal contribution . †Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . token number in self-attention blocks . For a proper trade-off between efficiency and effectiveness , existing works empirically adopt 14x14 or 16x16 tokens [ 11 , 55 ] . In this paper , we argue that it may not be optimal to treat all samples with the same number of tokens . In fact , there exist considerable variations among different images ( e.g. , contents , scales of objects , backgrounds , etc. ) . Therefore , the number of representative tokens should ideally be configured specifically for each input . This issue is critical for the computational efficiency of the models . For example , we train a T2T-ViT-12 [ 55 ] with varying token numbers , and report the corresponding accuracy and FLOPs in Table 1 . One can observe that adopting the officially recommended 14x14 tokens only correctly recognizes ∼15.9 % ( 76.7 % v.s . 60.8 % ) more test samples compared to that of using 4x4 tokens , while increases the computational cost by 8.5x ( 1.78G v.s . 0.21G ) . In other words , computational resources are wasted on applying the unnecessary 14x14 tokens to many “ easy ” images for which 4x4 tokens are sufficient . relationships to learn more accurate attention maps . Illustrative examples of our method are given in Figure 1 . Notably , DVT is designed as a general framework . Most of the state-of-the-art image recognition Transformers , such as ViT [ 11 ] , DeiT [ 38 ] , and T2T-ViT [ 55 ] , can be straightforwardly deployed as its backbones for higher efficiency . Our method is also appealing in its flexibility . The computational cost of DVT is able to be adjusted online by simply adapting the early-termination criterion . This characteristic makes DVT suitable for the cases where the available computational resources fluctuate dynamically or a minimal power consumption is required to achieve a given performance . Both situations are ubiquitous in real-world applications ( e.g. , searching engines and mobile apps ) . The performance of DVT is evaluated on ImageNet [ 9 ] and CIFAR [ 26 ] with T2T-ViT [ 55 ] and DeiT [ 38 ] . Experimental results show that DVT significantly improves the efficiency of the backbones . For examples , DVT reduces the computational cost of T2T-ViT by 1.6-3.6x without sacrificing accuracy . The real inference speed on a NVIDIA 2080Ti GPU is consistent with our theoretical results . 2 Related Work . Vision Transformers . Inspired by the success of Transformers on NLP tasks [ 10 , 40 , 3 , 43 ] , vision Transformers ( ViT ) have recently been developed for image recognition [ 11 ] . Although ViT by itself is not comparable with state-of-the-art convolutional networks ( CNN ) on the standard ImageNet benchmark , it attains excellent results when pre-trained on the larger JFT-300M dataset . DeiT [ 38 ] studies the training strategy of ViT and proposes a knowledge distilling-based approach , surpassing the performance of ResNet [ 19 ] . Some following works such as T2T-ViT [ 55 ] , TNT [ 17 ] , CaiT [ 39 ] , DeepViT [ 62 ] , CPVT [ 6 ] , LocalViT [ 28 ] and CrossViT [ 5 ] focus on improving the architecture design of ViT . Another line of research proposes to integrate the inductive bias of CNN into Transformers [ 49 , 8 , 54 , 16 ] . There are also attempts to adapt ViT for other vision tasks ( e.g. , object detection , semantic segmentation , etc . ) [ 30 , 44 , 12 , 20 , 57 , 60 , 14 ] . The most majority of these concurrent works represent each image with a fix number of tokens . To the best of our knowledge , we are the first to consider configuring token numbers conditioned on the inputs . Efficient deep networks . Computational efficiency plays a critical role in real-world scenarios , where the executed computation translates into power consumption , carbon emission or latency . A number of works have been done on reducing the computational cost of CNNs [ 22 , 33 , 21 , 53 , 59 , 31 , 35 ] . However , designing efficient vision Transformers is still an under-explored topic . T2T-ViT [ 55 ] proposes a light-weighted tokens-to-token module and obtains a competitive accuracy-parameter trade-off compared to MobileNetV2 [ 33 ] . LeViT [ 16 ] accelerates the inference of Transformer models by involving convolutional layers . Swin Transformer [ 30 ] introduces an efficient shifted window-based approach in multi-stage vision Transformers . Compared to these models with fixed computational graphs , the proposed DVT framework improves the efficiency by adaptively changing the architecture of the network on a per-sample basis . Dynamic models . Designing dynamic architectures is an effective approach for efficient deep learning [ 18 ] . In the context of recognition tasks , MSDNet and its variants [ 23 , 52 , 27 ] develop a multi-classifier CNN architecture to perform early exiting for easy samples . Another type of dynamic CNNs skips redundant layers [ 41 , 45 , 51 ] or channels [ 29 ] conditioned on the inputs . Besides , the spatial adaptive paradigm [ 15 , 4 , 47 , 42 , 46 ] has been proposed for efficient image and video recognition . Although these works are related to DVT on the spirit of adaptive computation , they are developed based on CNN , while DVT is tailored for vision Transformers . Notably , our work may seem similar to the multi-exit networks [ 23 , 52 , 27 , 47 ] from the lens of early-exiting . However , DVT differentiates itself from these existing works in several important aspects . For example , a major efficiency gain in both MSDNet [ 23 ] and RANet [ 52 ] comes from the adaptive depth , namely processing “ easier ” samples using fewer layers within a single network . In comparison , DVT cascades multiple Transformers to infer several entire networks at test time . Besides , DVT does not leverage the multi-scale architecture or dense connection in MSDNet [ 23 ] . Compared with RANet [ 52 ] , DVT introduces the adaptive token number and the relationship reuse mechanism tailored for ViT . In addition , IMTA [ 27 ] studies the training techniques of multi-exit models , which are actually complementary to DVT . Another difference is that all these networks [ 23 , 52 , 27 , 47 ] adopt pure convolutional models . 3 Dynamic Vision Transformer . Vision Transformers [ 11 , 17 , 38 , 55 ] split each 2D image into 1D tokens , while model their long range interaction with the self-attention mechanism [ 40 ] . As aforementioned , to correctly recognize some “ hard ” images and achieve high accuracy , the number of tokens usually needs to be large , leading to the quadratically grown computational cost . However , “ easier ” images that make up the bulk of the datasets typically require far fewer tokens and much less costs ( as shown in Table 1 ) . Inspired by this observation , we propose a Dynamic Vision Transformers ( DVT ) , aiming to improve the computational efficiency of Transformers via adaptively reducing the number of representative tokens for each input . In specific , we propose to deploy multiple Transformers trained with increasing number of tokens , such that one can sequentially activate them for each test image until obtaining a convincing prediction ( e.g. , with sufficient confidence ) . The computation is allocated unevenly across different samples for improving the overall efficiency . It is worth noting that , if all the Transformers are learned separately , the computation performed by upstream models will simply be abandoned once a downstream Transformer is activated , resulting in considerable inefficiency . To alleviate this problem , we introduce the efficient feature and relationship reuse mechanisms . 3.1 Overview . Inference . We start by describing the inference procedure of DVT , which is shown in Figure 2 . For each test sample , we first coarsely represent it using a small number of 1D token embeddings . This can be achieved by either straightforwardly flattening the split image patches [ 11 , 17 ] or leveraging techniques like the tokens-to-token module [ 55 ] . We infer a vision Transformer with these few tokens to obtain a quick prediction . This process enjoys high efficiency since the computational cost of Transformers grows quadratically with respect to token number . Then the prediction will be evaluated with certain criterion to determine whether it is reliable enough to be retrieved immediately . In this paper , early-termination is performed when the model is sufficiently confident ( details in Section 3.3 ) . Once the prediction fails to meet the termination criterion , the original input image will be split into more tokens for more accurate but computationally more expensive inference . Note that , here the dimension of each token embedding remains unchanged , while the number of tokens increases , enabling more fine-grained representation . An additional Transformer with the same architecture as the previous one but different parameters will be activated . By design , this stage trades off computation for higher accuracy on some “ difficult ” test samples . To improve the efficiency , the new model can reuse the previously learned features and relationships , which will be introduced in Section 3.2 . Similarly , after obtaining a new prediction , the termination criterion will be applied , and the above procedure will proceed until the sample exits or the final Transformer has been inferred . Training . For training , we simply train DVT to produce correct predictions at all exits ( i.e. , each with the corresponding number of tokens ) . Formally , the optimization objective is minimize 1 |Dtrain| ∑ ( x , y ) ∈Dtrain [ ∑ i LCE ( pi , y ) ] , ( 1 ) where ( x , y ) denote a sample in the training set Dtrain and its corresponding label . We adopt the standard cross-entropy loss function LCE ( · ) , while pi denotes the softmax prediction probability output by the ith exit . We find that such a simple training objective works well in practice . Transformer backbone . DVT is proposed as a general and flexible framework . It can be built on top of most existing vision Transformers like ViT [ 11 ] , DeiT [ 38 ] and T2T-ViT [ 55 ] to improve their efficiency . The architecture of Transformers simply follows the implementation of these backbones .

Summary: This paper introduces an efficient inference strategy for vision transformers (ViTs). The motivation is that there are some “easy” images which can be classified using a few tokens. Based on the motivation, the paper proposes a cascading inference framework which do not proceed to next stage if the prediction confidence is enough high (i.e., early exit  strategy). For further efficiency of the cascading scheme, the previous stage’s features are fused with the next stage features. Experimental results show that the proposed method outperforms the baseline methods both in inference speed and accuracy.