An Investigation on Hardware-Aware Vision Transformer Scaling

1 INTRODUCTION . Transformer ( Vaswani et al. , 2017 ) , which was initially proposed for natural language processing ( NLP ) and is a type of deep neural networks ( DNNs ) mainly based on the self-attention mechanism , has achieved significant breakthroughs in NLP tasks . Thanks to its strong representation capabilities , many works have developed ways to apply Transformer to computer vision ( CV ) tasks , such as image classification ( Dosovitskiy et al. , 2021 ) , object detection ( Carion et al. , 2020 ) , semantic segmentation ( Wu et al. , 2020 ) , and video classification ( Bertasius et al. , 2021 ) . Among them , Vision Transformer ( ViT ) ( Dosovitskiy et al. , 2021 ) stands out and demonstrates that a pure Transformer applied directly to sequences of image patches can perform very well on image classification tasks , e.g. , achieving a comparable ImageNet ( Deng et al. , 2009 ) top-1 accuracy as ResNet ( He et al. , 2016 ) . Motivated by ViT ’ s promising performance , a fast growing number of works follow it to explore pure Transformer architectures in order to push forward its accuracy-efficiency trade-off and deployment into real-world applications ( Touvron et al. , 2020 ; Heo et al. , 2021 ; Graham et al. , 2021 ; Liu et al. , 2021 ; Wu et al. , 2021 ) , achieving an even better performance than EfficientNetV1 ( Tan & Le , 2019 ) , a widely used efficient convolutional neural network ( CNN ) . The success of recent ViT works suggests that the model architecture is critical to ViT ’ s achievable performance . Therefore , in this work we explore ViT architectures from a new perspective , aiming to analyze and understand ViT ’ s architecture design space and real hardware-cost across different devices . Despite the recent excitement towards ViT models and the success of model scaling for CNNs , little effort has been made into exploring ViT ’ s model scaling strategies or hardware-cost . Note that directly applying the scaling strategies for CNNs ( Tan & Le , 2019 ; Sun et al. , 2017 ) or Transformer on NLP tasks ( Kaplan et al. , 2020 ; Henighan et al. , 2020 ) will lead to sub-optimality , as discussed in Section 3.2 . Furthermore , scaling strategies targeting one device/task might not be transferable to another device/task . Interestingly , we find that simply scaled ViT models can achieve comparable or even better accuracy-efficiency trade-off than dedicatedly designed ViT variants , as shown in Figure 1 . Motivated by this , we further explore the transferability of our scaling strategies ( 1 ) across different real hardware devices and ( 2 ) to different ViT variants and tasks . In particular , we make the following contributions : • We are the first to show that simply scaled vanilla ViT models can achieve comparable or even better accuracy-efficiency trade-off as compared to dedicatedly designed ViT variants ( Touvron et al. , 2021 ; Yuan et al. , 2021 ; Zhou et al. , 2021 ; Chen et al. , 2021 ; Wu et al. , 2021 ; Liu et al. , 2021 ; Chu et al. , 2021 ; Touvron et al. , 2020 ; Heo et al. , 2021 ) , as illustrated in Figure 1 . Specifically , as compared to DeiT-Tiny , our scaled model achieves a ↑ 1.9 % higher ImageNet top-1 accuracy under the same FLOPs and a ↑ 3.7 % better ImageNet top-1 accuracy under the same latency on an NVIDIA Edge GPU TX2 . • We study the transferability of the scaled ViT models across different devices and show that the transferability effectiveness depends on the underlying devices and deployment tools . For example , scaling strategies targeting FLOPs or the throughput on V100 GPU ( NVIDIA LLC . ) can be transferred to the Pixel3 ( Google LLC. , a ) device with little or even no performance loss , but those targeting the latency on TX2 ( NVIDIA Inc. , c ) may not be transferred to other devices due to the obvious performance loss . Additionally , we provide ViT models ’ cost breakdown and rank correlation between their hardware-cost on different devices for better understanding it . • We show that our scaling strategies can also be effectively transferred to different ViT variants and recognition tasks to further boost the achieved accuracy-efficiency trade-off , e.g. , achieving a ↑ 2.1 % higher accuracy under a similar FLOPs when being transferred to the PiT model and ↑ 0.7 % higher average precision under a similar inference throughput when being transferred to an object detection task . 2 RELATED WORKS . Vision Transformers . Transformer was first proposed for machine translation ( Vaswani et al. , 2017 ) . Motivated by its state-of-the-art ( SOTA ) performance in NLP tasks , there has been a growing interest in applying the Transformer/self-attention mechanism to CV tasks , e.g. , by proposing novel attention mechanisms for CNNs ( Hu et al. , 2018 ; Li et al. , 2019 ; Zhang et al. , 2020 ) , fusing Transformer and CNN designs within the same model ( Bello et al. , 2019 ; Carion et al. , 2020 ; Wu et al. , 2020 ) , or designing pure Transformer models ( Dosovitskiy et al. , 2021 ; Chen et al. , 2020 ) . Among them , ViT ( Dosovitskiy et al. , 2021 ) has achieved SOTA performance by directly applying the Transformer architecture for NLP tasks to the input raw image patches of vision tasks . Nevertheless , ViT ’ s powerful performance largely depends on its pre-training on JFT-300M ( Sun et al. , 2017 ) ( a giant private labelled dataset ) . As such , DeiT ( Touvron et al. , 2020 ) further develops an improved training recipe ( i.e. , the setting of optimization hyper-parameters ) , including a distillation setup and stronger data augmentation and regularization , to achieve comparable performance while removing the necessity of the costly pre-training . In order to build more efficient ViT models , ( Chen et al. , 2021 ) leverages multiple branches to extract and fuse features at different scales ; ( Heo et al. , 2021 ; Graham et al. , 2021 ; Liu et al. , 2021 ; Fan et al. , 2021 ; Wang et al. , 2021 ) apply a pyramid-like architecture commonly used in CNNs to ViT ; and ( Graham et al. , 2021 ; Liu et al. , 2021 ; Wu et al. , 2021 ) propose more efficient attention mechanisms or feature projection blocks . Model scaling . Prior works have explored scaling CNNs/NLP-Transformer ( i.e. , Transformer in NLP tasks ) to boost its accuracy or lower its computational resource requirements , e.g. , ResNet can be scaled along its depth dimension ( He et al. , 2016 ) and MobileNets can be scaled along its width ( i.e. , the number of channels ) and input resolution dimensions ( Sandler et al. , 2018 ; Howard et al. , 2019 ) . Notably , EfficientNet further points out that it is critical to scale CNNs in a compound manner ( i.e. , simultaneously scaling the model width , depth , and input resolution ) and does so to achieve SOTA accuracy-efficiency trade-off ( Tan & Le , 2019 ) . Nevertheless , as ( Bello et al. , 2021 ) demonstrates , the scaling strategies obtained from a specific model ( e.g. , EfficientNet-B0 ) can result in a sub-optimal accuracy-efficiency trade-off for another model ; motivated by this observation , they develop a more general scaling strategies extracted from grid search experiments based on the chosen training recipe rather than a specific model , achieving an improved trade-off . In addition to scaling the model architecture , ( Kaplan et al. , 2020 ; Henighan et al. , 2020 ) show that scaling up the dataset size and the number of computations used for training can also help to achieve a smaller cross-entropy loss for Transformer in NLP tasks . Recently , Zhou et al . ( 2021 ) demonstrates that the accuracy of ViT will decrease when it is scaled up along only the depth dimension ( i.e. , number of layers ) , and proposes Re-attention to resolve it . Nevertheless , none of the prior works has targeted scaling strategies for ViT with multiple scaling factors or study its real-hardware efficiency across different platforms featuring diverse computational and storage capabilities . Additionally , it is not clear whether their insights on scaling CNNs can be directly applied to ViT because of their different scaling factor definitions , e.g. , while the number of channels represents the width in CNNs , the number of heads and embedding dimensions can both represent the width in ViT . As such , scaling strategies dedicated to ViT are highly desirable and our scaling strategies can provide unique insights to inspire more innovations towards efficient ViT models . Although there is some model scaling strategy explorations in ( Dosovitskiy et al. , 2021 ; Zhai et al. , 2021 ) , our work distinguishes with them in providing more discoveries and insights . Specially , we focus more on the accuracy vs. efficiency trade-off when scaling ViTs instead of merely the accuracy ; and ( 2 ) provide additional analysis on the transferability of the extracted scaling strategies across different devices , ViT variants , and tasks . 3 SCALING VIT : HOW AND WHY DO WE SCALE VIT ? . In this section , we first analyze the scaling factors of ViT , then study the effectiveness of prior scaling strategies , which are dedicated to CNNs or Transformers , on ViT , and finally present our iterative greedy search approach to scale ViT . 3.1 SCALING FACTORS IN VIT As analyzed in ( Kaplan et al. , 2020 ) , the scaling factors in Transformers include the number of layers ( d ) , the number of heads ( h ) , the embedding dimension for each head ( e ) , and the linear projection ratio ( r ) . ViT , which directly adopts the Transformer architecture for NLP tasks and splits the raw images into patches to serve as the Transformer input , adds additional scaling factors , including image resolution ( I ) and patch size ( p ) . Figure 2 illustrates and summarizes our considered scaling factors for ViT . 3.2 PREVIOUS SCALING STRATEGIES FAIL ON VIT . CNN and ViT scaling factors do not match . Scaling strategies dedicated to CNNs ( Tan & Le , 2019 ; Sun et al. , 2017 ; Feichtenhofer , 2020 ) mostly come with CNN-specific scaling factor definitions ( e.g. , the number of channels in convolution layers represents the model width ) , which can not be directly transferred to ViT . For example , doubling ( 2× ) the width in CNNs can be achieved via various combinations of the number of heads ( h ) and embedding dimension for each head ( e ) in ViT . Furthermore , there are extra scaling factors for ViT , e.g. , the linear projection ratio ( r ) and the patch size ( p ) , which do not exist in the scaling factors for CNNs but are important for ViT as shown in Appendix D , thus directly transferring the scaling strategies from CNNs to ViTs can lead to ambiguity and sub-optimal performance . Transformer scaling strategies for NLP is sub-optimal on ViT . ( Kaplan et al. , 2020 ) noted that for NLP , model performance ( i.e. , accuracy or training loss ) depends “ strongly on the model scale ( i.e. , the number of parameters ) , but weakly on the model shape ” . However , when scaling ViT along the factors sum- marized in Figure 2 , our observations suggest that this is not true for ViT . As shown in Figure 3 , when performing an extensive search on top of DeiT-Small ( Touvron et al. , 2020 ) following ( Sun et al. , 2017 ) , we observe that a model ’ s shape has a great impact on the performance . Specifically , if we change the aspect ratio , i.e. , the ratio between the embedding dimension ( e× h ) and the number of layers ( d ) , while keeping the model parameters to be the same , the accuracy drifts as much as 18.61 % . This set of experiments motivates exploring scaling strategies dedicated to ViT .

Summary: The paper studies different hyper-parameters in the ViT design space (e.g., the number of layers, the dimensions per head, the number of head, the MLP ratio, the input image size). It starts with a very small, basic model, and then scales it up based on an iterative, greedy search. Each step it will select the best model (in terms of speed-accuracy trade-off) as the starting point of the next step. The "speed" aspect is firstly measured by FLOPs, and then later replaced by empirical measurements like latency on some devices. Overall, the resulting architectures show better speed-accuracy trade-offs than previous DEiT architectures, and shows consistent gains when transferred to other architectures (PiT) and tasks (e.g., object detection).