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: This paper proposes a model scaling strategy to change the model size. With a greedy search strategy, the author proposes a hardware-aware scaling method. Compared with the existing ViT variants, the proposed method achieves better performance.