NASViT: Neural Architecture Search for Efficient Vision Transformers with Gradient Conflict aware Supernet Training

1 INTRODUCTION . Transformers have recently been applied to various vision tasks , including image classification ( Liu et al. , 2021 ; Dong et al. , 2021 ; Bao et al. , 2021 ) , object detection ( Carion et al. , 2020 ; Zhu et al. , 2020 ) , semantic segmentation ( Xie et al. , 2021 ; Cheng et al. , 2021 ) , video understanding ( Bertasius et al. , 2021 ; Fan et al. , 2021 ) , etc . Vision transformers ( ViTs ) benefit from high model capacity , large receptive field , and grouping effect , etc ( Dosovitskiy et al. , 2020 ) , and demonstrate superior performance compared to convolutional neural networks ( CNNs ) especially with the scaling of the model size and training data size . For example , CoAtNet ( Dai et al. , 2021 ) achieves 90.88 % top-1 accuracy on Imagenet by scaling the model to 2586G FLOPs and pre-training the model on JFT-3B dataset ( Sun et al. , 2017 ) . Though promising in the high computation budget regime , the performance of ViTs is still inferior to that of the CNN counterparts on small- or medium-sized architectures , especially compared to CNN architectures that are highly optimized by neural architecture search ( NAS ) , e.g. , AlphaNet ( Wang et al. , 2021a ) , FBNetV3 ( Dai et al. , 2020 ) , etc . For example , the initial DeiT-Tiny ( Touvron et al. , 2020 ) only achieves 72.2 % top-1 accuracy with 1.2G FLOPs . The recently proposed LeViT ( Graham et al. , 2021 ) makes significant progress to achieve 76.6 % top-1 accuracy with 305M FLOPs with convolution/transformer hybrid architectures and a 3x longer training schedule . In contrast , AlphaNet ( Wang et al. , 2021a ) achieves 77.8 % top-1 accuracy with only 203M FLOPs . The large accuracy gap illustrated above raises a natural question : are transformer blocks that build large and dynamic receptive fields beneficial for small models ? To answer the question above , in this work , we target at developing a family of efficient ViTs with FLOPs ranging from 200M to 800M . A natural approach is to leverage NAS , which has achieved state-of-the-art ( SOTA ) accuracy-efficiency trade-off for CNNs ( Wang et al. , 2021a ; Dai et al. , 2020 ; Cai et al. , 2019 ) . The recently proposed supernet-based NAS , e.g. , BigNAS ( Yu et al. , 2020a ) and AlphaNet ( Wang et al. , 2021a ) , builds a weight-sharing graph including all the sub-networks in the architecture search space . A sandwich sampling rule with inplace knowledge distillation ( KD ) ( Yu et al. , 2018 ) is leveraged to simultaneously optimize the supernet and sub-networks for each mini-batch , which stabilizes the training and improves the training convergence . To leverage the supernet-based NAS , we first modify the LeViT model to build the architecture search space for ViTs and then jointly optimize the model architectures and parameters following AlphaNet . However , we find that directly applying AlphaNet achieves poor performance on the ViT search space , even worse compared to training single ViTs . To understand the root cause of the poor performance , we examine the supernet training procedure and observe that the gradients of the supernet and the different sub-networks conflict with each other during the sandwich sampling , which makes the training loss saturates much more quickly for ViTs , thus leading to slow convergence . To alleviate the issue of conflicting gradients , we propose three different techniques to improve the supernet training . Firstly , instead of directly adding the gradients from different sub-networks together , we find it beneficial to prioritize the training of the sub-networks over the supernet , as our main purpose is to build efficient sub-networks . We achieve this with a projection gradient algorithm which removes the component of the supernet gradient that is conflict with the sub-network gradient . Secondly , to alleviate the gradient conflicts among different sub-networks , we propose to augment each transformer layer with switchable channel-wise scaling layers . The weights of different scaling layers are not shared among different transformer blocks to reduce gradient conflicts . Thirdly , we propose to use a weak data augmentation scheme and reduce the regularization in training to decrease the optimization difficulty and hence reduce gradient conflicts . Our proposed techniques significantly alleviate the gradient conflict issue and empirically improve the convergence of supernet training . Compared to the baseline supernet training algorithm in AlphaNet , we can improve the top-1 accuracy to 78.2 % for the small model with 205M FLOPs and achieve 81.8 % for the large model with 757M FLOPs . Meanwhile , the resulting model family , NASViT , outperforms all the SOTA CNN and ViT models across a wide range of computation constraints . NASViT also demonstrates good performance on downstream tasks . When transferring to semantic segmentation tasks , NASViT backbones outperform previous CNN and ViT backbones on both Cityscape and ADE20K datasets , achieving 73.2 % and 37.9 % mIoU with 5G FLOPs , respectively . Related Works Recently , researchers have used supernet-based NAS to optimize the architecture for transformers . For example , HAT ( Han et al. , 2021 ) uses supernet for hardware-aware transformer optimization . HAT mainly focuses on NLP tasks and features a design space with heterogeneous transformer layers . AutoFormer ( Chen et al. , 2021a ) and ViTAS ( Su et al. , 2021 ) leverages supernetbased NAS to optimize the ViT architecture . By searching the width , depth , K/Q/V dimension , MLP ratio , etc , better accuracy is achieved compared to the baseline DeiT models ( Chen et al. , 2021a ) . However , these works focus on large ViT models with more than 1G FLOPs and their accuracy is still inferior to the CNN backbones with similar compute , e.g. , EfficientNet ( Tan & Le , 2019 ) . We refer readers to appendix for more discussions about related works . 2 NAS FOR EFFICIENT TRANSFORMERS . Our goal is to design efficient small- and medium-sized ViTs in the FLOPs regime from 200M to 800M . We build our search space inspired by the recently proposed LeViT ( Graham et al. , 2021 ) . LeViT is a family of efficient models leveraging a hybrid architecture of convolutions and transformers . In LeViT , the convolutions are introduced to handle high resolution inputs thanks to their efficiency from local computation while the transformers are leveraged for lower resolution features to extract global information . We closely follow LeViT to build our baseline search space ; see Figure 1 for an overview . Search Space We summarize the detailed search dimensions of our search space in Table 1 . For each CNN block , we directly follow the design in AlphaNet ( Wang et al. , 2021a ; b ) and search for the optimal channel widths , block depths , expansion ratios and kernel sizes ; for each transformer block , we search for the best number of windows , hidden feature dimensions ( denoted as Width in Table 1 ) 1 , depths and MLP expansion ratios . Compared to CNN blocks , one special search dimension 1Hidden feature dimension equals the number of heads times the feature dimension of each head . In our search space , we fix the head dimension to be 8 , and only searching for the number of heads . for transformer blocks is the number of windows k. When the number of windows k is greater than 1 , we follow Swin transformer ( Liu et al. , 2021 ) and partition the input tokens into k groups . We then compute the self-attention weights for each group separately to reduce computational cost . Standard global self-attention is a special case of k = 1 . In this work , we only search the number of windows for the first transformer block , as the input resolutions to the other transformer blocks are already small after 4 times of down-sampling . Similar to the search range of AlphaNet , the smallest sub-network in our search space has 190M FLOPs and the largest sub-network has FLOPs of 1,881M . we refer the reader to Appendix B for more description of our search space . Naive supernet-based NAS fails to find accurate ViTs We first closely follow the previous best practices in AlphaNet ( Wang et al. , 2021a ) for the supernet training . We train the supernet for 360 epochs on ImageNet ( Deng et al. , 2009 ) . At each training step , we adopt the sandwich sampling rule ( Yu et al. , 2018 ) and sample four sub-networks : the smallest sub-network , the supernet ( a.k.a . the largest sub-network ) , and two random sub-networks . All small sub-networks are supervised by the supernet with α-divergence-based KD ; see Algorithm 1 in Appendix C.1 for an overview of the supernet training procedure . Additionally , as our candidate networks contain transformer blocks , we further incorporate the best training recipe from LeViT ( Graham et al. , 2021 ) by replacing the SGD optimizer with Adam ( Kingma & Ba , 2014 ) and leveraging an external pre-trained teacher model for the best accuracy . Specifically , we use the pre-trained teacher to supervise the supernet and still constrain all other small sub-networks to learn from the supernet . In this work , we always use an EfficientNet-B5 ( Tan & Le , 2019 ) with 83.3 % top-1 accuracy on ImageNet as the teacher to train our ViT supernet unless otherwise specified . We plot the training curves of the smallest sub-network and the largest sub-network in Figure 2 . We find both the smallest sub-network and the largest sub-network from our search space converge poorly compared to the CNN baseline . Specifically , the validation accuracy of both the smallest and the largest sub-network is saturated at around the 250-th epoch , and the final accuracy is much worse than the CNN baselines . To understand the inferior model performance , we investigate the potential issues of our ViT supernet training from the following three directions . Investigation 1 : Is our search space designed badly ? We seek to understand if the performance gap is caused by a bad search space design . To verify , we randomly pick four sub-networks from the search space with computation cost ranging from 190M to 591M FLOPs . Then , we train these networks from scratch with the same data augmentation and regularization . As we can see from Table 2 , the sub-networks trained from scratch outperform the sub-networks sampled from the supernet . Note that from previous works ( e.g . Yu et al. , 2020a ) , supernet often learns more accurate sub-networks compared to the training from scratch performance , by taking advantage of inplace knowledge KD and weight-sharing . Our observations in Table 2 indicate that the poor performance does not come from the search space but from the interference with the training of the supernet . Investigation 2 : Are the training settings suitable for ViTs ? Our default training settings from AlphaNet are originally optimized for CNNs only . Compared with AlphaNet , recent ViT methods , e.g. , DeiT and LeViT , suggest to use stronger data augmentation schemes ( e.g. , a combination of CutMix ( Yun et al. , 2019 ) , Mixup ( Zhang et al. , 2017 ) , randaugment ( Cubuk et al. , 2020 ) , random erasing ( Zhong et al. , 2020 ) , and repeated augmentation ) and stronger regularization ( e.g. , large weight decay , large drop path probability ) for training . We evaluate the effectiveness of these ViT specific training recipes and summarize our findings in Table 3 . As we can see from Table 3 , DeiT- or LeViT-based training recipe produces even worse accuracy compared to the results from AlphaNet-based training . Investigation 3 : Saturated supernet training due to gradient conflicts ? Compared to the standard single network training , a major difference of supernet training is that multiple networks are sampled and trained at each step . We hypothesize that the training loss from the supernet and that from the sub-networks may yield conflicting gradients due to the heterogeneous and complex structures of networks , and the conflict gradients may consequently lead to slow convergence and undesirable performance . To verify this hypothesis , we compute the cosine similarity between the gradients from the supernet and the averaged gradients from the sub-networks . A negative cosine similarity indicates the supernet and sub-networks produce conflict gradients and tend to update model parameters in opposite directions . To quantitatively examine the gradient conflict issue , we go through the entire ImageNet training set and calculate the percentage of negative cosine similarity between the gradients of supernet and sub-networks among all training images at a per layer granularity . The gradients are computed under the same data augmentation and regularization as the supernet training stage . For AlphaNet , we train the model using its official code 2 . As shown in Table 4 , our ViT supernet suffers from more severe gradient conflicts compared to the CNN baseline . According to existing works in multi-task learning , large gradient conflict ratios may result in significant accuracy drop even for binary classification problems ( see Figure 3 in Du et al . ( 2018 ) and Figure 4 ( b ) in Yu et al . ( 2020b ) ) . We hypothesize that the inferior performance of our ViT supernet is mainly caused by the large percentage of disagreements between the supernet gradients and the subnetworks gradients .

Summary: The paper proposes a NAS algorithm for vision transformers. The search space contains both CNN blocks and transformer blocks. The authors solve the gradient conflict dilemma in this process and find out  a good architecture. Experiments on classification and segmentation show its effectiveness over other vison transformers like swin, cvt, etc.