NASViT: Neural Architecture Search for Efficient Vision Transformers with Gradient Conflict aware Supernet Training
Authors: Chengyue Gong
Dilin Wang
Meng Li
Xinlei Chen
Zhicheng Yan
Yuandong Tian
qiang liu
Vikas Chandra
Venue: ICLR 2022 Poster (2022)
Keywords: vision transformer
gradient conflict
neural architecture search

Abstract: Designing accurate and efficient vision transformers (ViTs) is a highly important but challenging task. Supernet-based one-shot neural architecture search (NAS) enables fast architecture optimization and has achieved state-of-the-art (SOTA) results on convolutional neural networks (CNNs). However, directly applying the supernet-based NAS to optimize ViTs leads to poor performance - even worse compared to training single ViTs. In this work, we observe that the poor performance is due to a gradient conflict issue: the gradients of different sub-networks conflict with that of the supernet more severely in ViTs than CNNs, which leads to early saturation in training and inferior convergence. To alleviate this issue, we propose a series of techniques, including a gradient projection algorithm, a switchable layer scaling design, and a simplified data augmentation and regularization training recipe. The proposed techniques significantly improve the convergence and the performance of all sub-networks. Our discovered hybrid ViT model family, dubbed NASViT, achieves top-1 accuracy from 78.2% to 81.8% on ImageNet from 200M to 800M FLOPs, and outperforms all the prior art CNNs and ViTs, including AlphaNet and LeViT, etc. When transferred to semantic segmentation tasks, NASViTs also outperform previous backbones on both Cityscape and ADE20K datasets, achieving 73.2% and 37.9% mIoU with only 5G FLOPs, respectively. Code is available at
https://github.com/facebookresearch/NASViT.


Published as a conference paper at ICLR 2022
NASVIT: NEURAL ARCHITECTURE SEARCH FOR EF- FICIENT VISION TRANSFORMERS WITH GRADIENT CONFLICT-AWARE SUPERNET TRAINING
Chengyue Gong?; Dilin Wang', Meng Li', Xinlei Chen', Zhicheng Yan',
Yuandong Tian', Qiang Liu", Vikas Chandra'
Meta Reality Labs University of Texas at Austin [wdilin, meng.li, xinleic, zyan3, yuandong, vchandra} @fb.com, {cygong, Iqiang) } @cs. utexas.edu
ABSTRACT
Designing accurate and efficient vision transformers (ViTs) is an important but challenging task. Supernet-based one-shot neural architecture search (NAS) en- ables fast architecture optimization and has achieved state-of-the-art results on convolutional neural networks (CNNs)). However, directly applying the supernet- based NAS to optimize VITs leads to poor performance even worse compared to training single ViTs. In this work. we observe that the poor performance is due a gradient conflict issue: the gradients of different sub-networks conflict with that of the supernet more severely in ViT' than CNNs, which leads to early saturation training and inferior convergence. To alleviate this issue, propose series of techniques, including gradient projection algorithm, a switchable layer scaling design, and simplified data augmentation and regularization training recipe. The proposed techniques significantly improve the convergence and the performance of all sub- nub-networks. Our discovered hybrid ViT model family, dubbed NASViT, achieves top-1 accuracy from 78.2% to 81.8% on ImageNet from 200M to 800M FLOOPs, and outperforms all the prior CNNs and ViTs, including AlphaNet and ViT. When transferred to semantic segmentation tasks, NASViTs also out- perform previous backbones on both Cityscape and ADE20K datasets, achieving 73.2% and 37, 77% mloU with only 5G FLOPs, respectively. Code is available at http:///g..hatccccooeeeeeereeeee NASSVT
INTRODUCTION
Transformers have recently been applied to various vision tasks, including image classification (Liu
(tt al. | Dong al. 2021). object detection (Carion et al.) 2020 Zhu et al. 2020), semantic
segmentation (Xie et al. 2021:Cheng Cheng et al. 2021). video understanding (Bertasius et al. 2021 Fan
let al. 202T), etc. Vision transformers (ViTs) benefit from high model capacity, large receptive field, and grouping effect, etc Doooovissky 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., scaling the model to 2586G FLOPS and pre-training the model JFT-3B dataset (Sun et al. (007).
2021) achieves 90.88% top-1 accuracy on Imagenet by
Though promising in the high computation budget regime, the performance of ViTs is still inferior to that of the CNN counterparts on small- medium-sized architectures, especially compared CNN architectures that are highly optimized by neural architecture search (NAS), e.g., AlphaNet (Wang 2021a). FBNetV3 (Dai et al. 2020), etc. For example, the initial DeiT-Tiny (Touvron al. only achieves 72.2% top-1 accuracy with 1.2G FLOPs. The recently proposed LeViT (Graham 202T makes significant progress to achieve 76.6% top-1 accuracy with 305M FLOPs with convolution/ransffrmee hybrid architectures and 3x longer training schedule. In contrast, AlphaNet (Wang et al.) (2211a) achieves 77.8% top-1 accuracy with only 203M FLOPS. The large accuracy gap illustrated above raises natural question: are transformer blocks that build large and dynamic receptive fields beneficial for small models?
al.
2020
al.
WWrk done during intersship Meta Reality Labs

---

Published as a conference paper at ICLR 2022
To answer the question above, in this work, we target at developing family of efficient ViTs with FLOPs ranging from 200M to 800M A natural approach is to leverage NAS, which has achieved state-of-the-nnn art (SOTA) aceuracy-efficieney trade-off for CNNs (Wang et al., 2021a Dai et al. 2020
Cai et al.
2019p. The recently proposed supernet-based NAS, e.g., BigNAS Yu et al., 2020a) and
AlphaNet Wang et al. 2021a), builds 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 ViT's. 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, find it beneficial to prioritize the training of the neb-networks over the supemet, as 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. 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 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. (02T) uses supernet for hardware-aware transformer optimization. HAT mainly focuses on NLP tasks and features design space with heterogencous transformer layers. AutoFormer (Chen et al. 202 and ViTAS (Su et al.) 2021) leverages supernet- based 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 2021a). However, these works focus on large ViT models with more than IG 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 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 search space; see Figure T for an overview.
Search Space We summarize the detailed search dimensions of our search space in Table I For each CNN block, we directly follow the design in AlphaNet (Wang et al.) 2221aa b) and search for the optimal channel widths, block depths, expansion ratios and kernel sizes; for each transformer
2

---

Published as a conference paper at ICLR 2022
Dynamic CNN blocks
Dynamic transformer blocks
Dynamic Inputs
Conv
MBConv
MBConw
Transformer
Transformer
Softmax
Figure 1: An illustration of VIT search space. MBConv refers inverted residual blocks (Sandler et al. 2018). All CNN and transformer blocks contain a stack of dynamic layers with searchable architecture configurations. Additionally, we also search for the input resolutions.
Block
Width
Depth
Kermel size Expansion ratio SE Stride Number of Windows
Conv
16. 24
3
2
MBConv-1
[16, 24
(1.2)
(3.5)
N
MBConv-2
(24,
(3.4,5)
[3,5]
(,
N
MBConv-3
[32, 40)
(3.4.566
[3.5] (4.5.6)
Y
2
Transformer-4
(64, 72)
(3,4.5.6)
(1,2)
2
Transformer-5
[112. 120. 128)
53.4.5.6.7.8
[1.2]
2
Transformer-6
(160. 168, 176, 184)
(3,4,566,788)
[1.2]
1
Tramsformer-7
(208. 216, 224)
(3.4.5.6)
[1.2)
2
MBPool
(1792, 1984)
Input resolution
192. 224, 256, 288)
Table 1: An illustration of our search space. MBConv refers the inverted residual block (Sandler et al. 2018). MBPool denotes the efficient last stage (Howard et al. (201)). SE represents the squeeze and excite layer (Hu al. 2018). Transformer stands for the transformer blocks (Vaswani al.. 2017). For MBConv blocks, the expansion ratio refers to the expansion ratio of the depth-wise convolution layer. For transformer layers, it refers to the MLP expansion ratio. For each transformer block, we use x 3 depth-wise wise convolution with stride 2 for down-sampling and the down-sampling layer is placed as the first layer for that block.
block, we search for the best number of windows, hidden feature dimensions (denoted as Width in Table depths and MLP expansion ratios. Compared CNN blocks, one special search dimension for transformer blocks is the number of windows k. When the number of windows & is greater than 1, we follow Swin transformer (Liu et al. 2021) and partition the input tokens into 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- samp-ing.inn.. Similar to the search range of AlphaNet, the smallest netwoe in our search space has 190M FLOPs and the largest sub-network has FLOPs of 1,88 M.. we refer the reader 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 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 a-divergence-basee based KD; see Algorithm in Appendix C.1 for an overview of the supernet training procedure. Additionally, as our candidate networks contain transformer blocks, we further incorporate best training recipe from LeViT (Graham et al. 202T) by replacing the SGD optimizer with Adam (Kingma & Ba 2014) and leveraging 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 netbnnsssors to learn from the supernet. In this work, always EfficientNet-B5 (Tan & Le 2019) with 83.3% top-1 accuracy on ImageNet as the teacher train our VIT supernet unless otherwise specified.
We plot the training curves of the smallest sub-network and the largest sub-network Figure 2 We find both the smallest sub-netwont and the largest sub-ne network from our search space converge poorly compared to the CNN baseline. Specifically, the validation accuracy of both the smallest and the
Hidden feature dimension equals the number of heads times the feature dimension of each head. In search space, fix the head dimension he 8, and only searching for the number of heads.
3

---

Published as a conference paper at ICLR 2022
ImageNet Top-I val ace
78
ImageNer Top-I val ace
82.0
75
80.5
VIT
VIT
CNN (AlphaNet)
CNN (AlphaNet)
72
79.0
79,00
100
200
300
360
200
300
360
Training Epoch
Training Epoch
(a) The smallest sub-network
(b) The largest sub-network (the supernet)
Figure 2: (a-b) show the training curves of the smallest sub-network and the largest sub-network (i.e., the supernet), respectively, Note that AlphaNet is trained without external teacher models.
FLOPs (M)
190
208
309
591
AlphaNet
DeiT
LeViT
Scratch
77.2
77.5
79.1
80.4
Smallest Largest
77.0
76.6
76.8
Supernet
76.4
76.6
78.5
80.6
82.4
82.2
82.2
ImageNet top-1 accuracy from sub- networks trained from scratch vs. results from sub-networks sampled from the supernet.
Table 2:
Table 3: ImageNet Top-1 accuracy from the smallest and the largest sub-network by using different training recipes.
largest sub-network saturated at around the 250-th epoch, and the final accuracy 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 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 Le VIT, suggest to use stronger data augmentation schemes (e.g., a combination of CutMix Yun et al. 2019), Mixup (Zhang et al. (0177), 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 As we can see from Table 3 DeiT- or ViT--badeed training recipe produces even worse accuracy compared to the results from AlphaNet-based training.
Investigation 3: Saturated supernet training due to gradient conflicts? Compared the stan- dard single network training, a major difference of supernet training is that multiple networks are sampled and trained each step. We hypothesize that the training loss from the supernet and that from the aub-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, go through the entire ImageNet
4

---

Published as a conference paper at ICLR 2022
Epoch
1st
90th
180th
270th
360th
AlphaNet VIT
27%
20%
21%
24%
28%
36%
27%
27%
32%
34%
Table 4: An estimation of negative cosine similarity ratio (gradient conflict ratio) between the supernet gradient and the averaged gradient of the sub-networks.
training and calculate the percentage of negative cosine similarity between the gradients of supernet and sub-networks among all training images at per layer granularity. The gradients are computed under the same data augmentation and regularization AlphaNet, we train the model using its official code As shown in Table 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 Du et al. (2018) and Figure 4(b) Yu et al. (2020b)). We hypothesize that the inferior performance of our ViT supernet is mainly caused by the large percentage disagreements between the supernet gradients and the subnetworks gradients.
the supernet training stage. For
GRADIENT CONFLICT AWARE SUPERNET TRAINING
activeted layer transformer
We propose to improve the ViT supernet training by address- ing the gradient conflict issue between the supernet and the sub-networks from three aspects: I) manually resolving the gra- dient conflict by projecting the supernet gradients to the normal vector of the sub-networks gradients; 2) introducing switchable scaling layers to the search space to give more optimization freedom for sub-networks; 3) reducing data augmentation and regularization provide easier training signals.
Output
Scaling
Charmel-wise
Switthhbbe
Scaling
Channel-wize
Gradient projection to prioritize sub-networks update Our first idea is focus on training the sub-networks when- ever gradients from the supernet and the gradients from the sub-networks conflicted with each other. As we are interested in the sub-networks from the 200M to 800M FLOOS range, we propose to prioritize the optimization of the sub-networks the supernet when gradient confficts observed. Let Vlsup and Vlsab denote the gradients of the supernet and the sub-networks, respectively. To prioritize sub-networks train- ing, we always project Vlsup to the normal vector of V avoid gradient conflicts when the cosine similarity between Vlsup and Veeu is negative. The overall accumulated gradi- ent each training iteration with projection can be written as follows.
activated
Input
Inactivated
Figure 3: A basic transformer layer with scaling. Activated compo- nents are the neurons selected in the forward path for one sub-network 'cl' and 'c2' represent the num- ber of channels activated in self- attention layer and MLP, respec- tively.
9 = Vesub + proj (Vesp) with
V Veup
if cos( Veapp Vluub) 2 0,
proj(Vlsup)
=
V Vlau
VL Vlaub [V
Vland
otherwise.
0 if cos (osups Vlaub) < O, which ensures the gradient cosine similarity is non-negative. In sandwich sampling, since several sub-networks are sampled each iteration, Lxub computed as the summation the gradients from all sub-networks. Note similar ideas of gradient projection have also been explored in multi-task learning, e.g., Yu et al. (2020b); Du et al. (2018); Real et al. (2019); Dery et al. (202T).
Note that cos( Vlyub,
While the gradient projection in Eqn. I climinates the gradient conflicts, it may lead to slow convergence as the resulting gradients are biased. Hence, we also propose the follow two techniques to reduce the gradient conflicts from search space design and training strategy refinement perspective.
httpps:/ htttesss//atttoo..cc///aeeeeeee wttp:h/ccw..ccccccoocreeeeeee/ee htt.://t////attooccc/oooooeoeeea
5

---

Published as a conference paper at ICLR 2022
Switchable scaling layer Motivated by Slimmable NN (Yu et al. 2018), introduce additional switchable scaling layers to allow sub-networks with different layer widths and depths to re-scale their features in privatized way. Specifically, for each transformer layer, switchable scaling layer is introduced at the output of the self-attention (SA) and the MIP, respectively, as shown in Figure 3 Assume [f[,d]] RC is input feature of scaling layer, with the feature dimension (i.e. the number of selected channels in the forward path) and d the index of this layer in transformer search block. The scaling layer transforms Tie,d| as W|e,d| O d]c,d]- Here Wic,d| E R are learnable parameters and @ denotes element-wise multiplication. For each transformer block (see Table ( each different configuration of (c. d will specify a set of independent switchable scaling layers. Following CaiT (Touvron et al. 2021). we initialize all scaling factors w to small value (e.g. 10 10-") for fast convergence and stable training. Intuitively, the switchable scaling layers effectively increase the model capacity of sub-networks and give the sub-networks more optimization flexibility.
Reduced data augmentation and regularization Furthermore, we observe that the supernet and the sub-networks are more likely to conflict with each other in the presence of stronger data aug- mentations and stronger regularization, e.g., large weight decay, large DropConnect (Wan et al. 2013). Hence, we simplify the AlphaNet training recipe and use a weaker data augmentation scheme RandAugment (Cubuk et al. 2019) with both the number of augmentation transformations and the magnitude set to I, and remove the regularization, e.g. DropConnect (Wan et al. 2013), dropout and weight decay, from the training: see Table 5 for comparison.
Method AlphaNet Ours
Data augmentation
Weight decay
DropConnect
Dropout
AutoAugment
10
0.2
0.2
RandAugment (n 1,m
0
0
0
Table 5: An illustration of our simplified training settings, where n is the number of augmentation transformations and m the number of magnitudes in RandAuument. typical setting of RandAug- ment is n=2 and m=9 for training single network; see Cubuk et al. (2020): Liu et al. (2021).
EXPERIMENTS
Top-1 validation accuracy
81.5
We first retrain our ViT supernet with our pro- posed gradient conflict reduction techniques on ImageNet (Deng et al. 2009); we then conduct an evolutionary search on subset of the Ima- geNet training dataset to search the accuracy vs. FL.OPs Pareto following (Wang et al. 2021b). We refer the reader to Appendix C for more details. Note all the models are directly sampled from the supernet without retraining or finetun- ing.
79.5
Ours (NASVIT T
AlphaNet T
FBNetV3
AlphaNet
77.5
Mobillr-Former
LeVIT MNasNet
200
400
600
We call our discovered as NASVIT models and compare with of-te-of-thh-are efficient CNNs and ViTs, including VB (Dai et al. 2020). AlphaNet (Wang et al. 2021a). ViT (Graham et al.) 2021) and Segformer (Xie et al.) (021) on both image classification (c.g., ImageNet) and semantic segmentation benchmarks (e.g.. Cityscapes and ADE20K).
M FLOPPs
Figure 4: Comparison with prior-art CNNs and ViTs on ImageNet. Here "+7"' indicates meth- ods that are trained with external teacher models. Note that Mobile-Fonnnn
(Chen et al.] 2021c) and MNasNet (Tan et al. 2019) are trained without additional teacher models.
4.1 IMAGENET
We compare our NASVIT models with sfate-of-the-art -ased CNNs, including AlphaNet (Wang et al. 2021a) and FBNetV3 (Dai al. 2020). and propossd efficient ViTs, e.g., ViTs.
Settings Note that our VIT supernet is trained with a pretrained Efficient-BB B5 teacher model (83.3% 1op-1) model. For fair comparison, retrain AlphaNet with the same teacher. For FBNet-V3 and

---

Published as a conference paper at ICLR 2022
Group
Method
FLOPPs
Top-1 accuracy (%)
AlphaNet-A0
203
77.9
200-300 (M)
NASViT-AO (ours)
208
78.2
LeViT (Graham al.
2021
300
76.6
TASViT-AT(oun)
309
79.7
300-400 (M)
AlphaNet-A2
317
FBNetV3 (Dai ct al.
2020
357
79.6
LeViT
406
78.6
400-500 (M)
NASViT-A2 (ours)
421
80.5
AlphaNet-A4
444
80.4
NASViT-A3 (ours)
528
81.0
FBNetV3
557
80.8
500-600 (M)
NASVIT-A4 (ours)
591
81.4
AlphaNet
596
81.1
LeViT
658
80.0
600 1000 (M)
NASVIT-A5 (ours)
757
81.8
FBNetV3
762
81.5
AutoFormer* (Chen
2021a
1,300
74.7
PYT-XS (Hco cf al. 2021
1,400
79.1
ViTAS-D* (Su et 2021
1,600
76.2
NASVIT (supernel) (ours)
1,881
82.9
1000 (M)
CVT-13-NAS* (Wu et al.
2021
4,100
82.2
Swin-Tiny" (Liu
2021
4,500
81.3
CVT-13* (Wu
4,500
81.6
T2T-ViT-14 Yuan el al. 2021a
5,200
81.5
Die Zhou et al. 2021
6,200
82.3
Table Comparison with prior efficient CNNs and ViTs on ImageNet. The reported AlphaNet models are trained with an external teacher model. The indicates that the ViTs are trained without external teacher models.
LeViT models, these models already use teachers with better performance than Efficient-B5 for training, and therefore directly report their results following their papers. Specifically, V3 use a RegNerY-320G with 84.5% top-1 and LeViT use RegNetY-16G with 83.6% top-1 as the teacher model, respectively.
Our discovered NASViT models outperform all evaluated CNN and VIT baselines. Our models are the first models with transformers blocks that outperform state-of-the oftt--ofthe-arttt efficient CNNs with similar FLOPs on ImageNet. For example, with < 600M FL.OPs, our NASViT-A4 achieves 81 81.4% top-1 accuracy on ImageNet. As a reference point, a ResNet-50 model (4G FI.OPs) achieves 81.5% top-1 accuracy by distilling from BiT (Kolesnikov et 2020) teacher (87.5% top-1 accuracy) with 1200 epochs of training (Beyer et
Results We summarize our results in both Table and Figure 4
2021
76
41
mloU
mIoU
39
Ours
74
Segformer
AlphaNet
37
MobileNetv3
72
(a) FLOPs (Cityscapes)
(b) FLOPs (ADE20K)
Figure 5: Results of our method and baselines semantic segmentation, (a-b) show the results the Cityscapes and ADE20K validation set, respectively.
4.2 SEMANTIC SEGMENTATION
We evaluate the transfer learning performance of our discovered NASVIT models by fine-tuning them downstream semantic segmentation tasks. In particular, fine-tune NASViT-AI to NASViT-A4 as backbones and we show that our NASVIT models yield the best segmentation performance compared to the results from the previous efficient CNN backbones, e.g., AlphaNet
7

---

Published as a conference paper at ICLR 2022
and MobileNetV3 (Howard et al., 2019), well as the recently proposed transformer-based Seg- former (Xie et al. 2021).
two benchmark datasets, Cityscapes (Cordts et al. 2016 and ADE20K (Zhou et al. 2017). To handle large input resolutions efficiently, for each transformer block, set the number of windows to be the input feature map size divided by 7 instead of using searched settings on ImageNet. We use the recent proposed light-weight Segformer head (Xie et al.) 2021 as the decoder head for all the backbones, achieve a better accuracy efficiency trade-off.
Settings We evaluate
Results As shown in Figure 5 our models yield the best FLOPs vs. mloU trade-offs. For example, model achieves 76.1% mloU and 41.4% mIOU with less than 30G FLOPS on the Cityscapes and the ADE20K validation set, respectively.
4.3 ABLATION STUDIES ON GRADIENT CONFLICT AWARE TRAINING
We conduct ablation studies on ImageNet to have better understanding on the effectiveness of our proposed methods. We mainly study 1) how our techniques can mitigate the gradient conflict issues and improve the performance 2) whether CNN supernets can also benefit from our techniques. All the comparisons in this section are conducted on ImageNet.
On the effectiveness of our gradient confiit aware training techniques As demonstrated in Table 77, both weak data augmentation and regularization and switchable scaling layer and can significantly reduce the gradient conflict ratios and in the meantime, improve the top-1 accuracy of both the smallest sub-netwon and the supernet. By further applying gradient projection to prioritize the sub-networks update (denoted by Prioritize (sub)), the performance of both the smallest and largest sub-network is boosted by around 0.3% on top-1 accuracy.
Prioritizing the supernet update Instead of focusing on training the sub-networks, we retrain our ViT supernet and prioritize the supernet update by moving the proj(-) term in Eqn. (1) to V As demonstrated in the last column of Table 7 this training strategy (denoted by Prioritize (sup)) leads to slightly improved supernet while resulting in less competitive performance on the smallest sub-network.
Baseline
Weak DA & Reg
Switchable scaling
Prioritize (sub)
Prioritize
(sup)
Top-1 I (smallest) Top-1 (supernet)
76.6
77.4
77.6
78.1
77.9
82.2
82.5
82.6
82.9
83.0
Negative Cosine Similarity Ratio
34%
30%
29%
0%
0%
Table 7:
Ablation study results on ImageNet. We show the top-1 validation accuracy of the smallest and largest sub-network, and the negative cosine similarity ratio for each case. Note that switchable scaling layer is applied on top of Weak DA & Reg; and Prioritize (sub) is applied on top of both Weak DA reg and switchable scaling layer.
Improving CNN-based supernets We verify the generalizability of our three techniques to the CNN supernets. this setting, we applied all three techniques together to improve CNN-based supernets. We show in Table that our method is especially helpful for AlphaNet trained with KL based KD (denoted by AlphaNet (w/ KL)). To further understand the large improvements on AlphaNet (w/ KI.), we follow ours studies in Table Table compute the gradient conflict ratio for AlphaNet (w/ KL) at epoch 1st, 90th, 180th, 270th, and 360th, and the corresponding gradient conflicts ratio is 25%, 18%, 24%. 28% and 31%, respectively. The gradient conflict issuc more severe for AlphaNet (w/ a-KL) compared with AlphaNet trained with a-divergence based KD (AlphaNet (w/ --di Our findings indicate that techniques not restricted to the ViT supernet training and might be beneficial for all supernets in which a large ratio of gradient conflicts presents.
4.4 ABLATION STUDIES ON SEARCH SPACE
In this part, we provide additional ablation studies support some key design choices of our ViT search space. For all the studies, we use a 250M sub-network that is randomly sub-sampled from our search space for consistency.
8

---

Published as a conference paper at ICLR 2022
Architectures (M FLOP))
AO (203M)
(279M)
(2
(357MM))
A4(444M)
A5 (491M)
A6 (709M)
AlphaNet (w/ KL)
77.0
78.2
78.5
78.8
79.3
79.6
80.1
AlphaNet (w/ KL) + Ours
77.5
78.6
78.9
79.2
79.8
80.1
80.7
AlphaNet (w/ a-divv
77.8
78.9
79.2
79.4
80.0
80.3
80.8
AlphaNet (w/ a-div) + Ours
77.8
78.9
79.2
79.4
80.0
80.4
80.9
Table 8: Improving CNN-based supernets on ImageNet. AlphaNet (w/ KL) and AlphaNet (w/ a-div) denote AlphaNets trained with KL and a-divergence based knowledge distillation, respectively. AO A6 are the architectures reported in AlphaNet (Wang et al. 2021a). Note that the AlphaNet supernets here are trained without external teacher models.
Global attention vs. local attention In our search space (Table D. we mainly use the global self-attention for the best representation learning capacity. However, with the more computationally efficient local and linear self-attention schemes, we would be able to use slightly bigger model under similar FLOPs constraints with sacrifice of global context modeling. To test this trade-off, we train the aforementioned model from scratch with different types of self-attention strategies. Specifically, in addition to the global self-attention, we further evaluate a number of local and linear self-attention mechanisms, including Swin (Liu et al. 2021). CSwin (Dong al. 2021), VOLO (Yuan et al., 20221b) and LSS (Kitaev et al.) 2020). We uniformly scale the width of the transformer blocks to ensure all models have similar compute FL.OPs. In Table
show that the standard global attention achieves the best accuracy compared to other faster local and linear self-attention methods. Additionally, from our evolutionary search results, we also notice that the sub-networks with all standard global self-attention layers often yield the best accuracy vs. FLOPs trade-offs.
Global
Swin
CSwin
VOLO
LSH
78.5
78.0
77.9
78.1
78.0
Table 9: ImageNet top-1 accuracy with different types of self-attention mechanisms.
The placement of transformer blocks Our VIT supernet has a convolution stem with 3 down- sampling strides. We further test the optimal choice of where switch transformer blocks. We uniformly scale the width of transformer layers of our baseline model to ensure similar FLOPs for different architecture designs. As shown Table [0 we notice that our current design, a convolution stem with 3 strides, yields the best performance.
#strides
2
3
4
Head Dimension
8
16
32
Top-1 Accuracy
78.1
78.5
78.4
Top-1 Accuracy
78.5
78.3
78.2
Table 1I:
Ablation studies on the impact of head dimension.
Table 10:
Ablation studies where to switch to transformer blocks.
Head Dimension previous VIT works (e.g Touvron et al., 2020 2021 Zhou et al., 2021 Wu et al. [021]] [Lu et al.) (021). the feature dimension of each self atte-atttttonn head is usually set to be 32, 64, or larger. In this work, as shown in Table II we found smaller head feature dimension (e.g.. 8) yields better performance.
5 CONCLUSION
In this work, we identify one key issue of ViT supernet training that the supernet gradients and the sub-network gradients are likely to disagree with each other, and consequently leading inferior NAS performance. We fix this gradient conflict issue by introducing gradient projection method to prioritize the sub-networks update, designing switchable scaling layers to increase the model capacities of sub-networks and simplifying the training recipe provide easier training signals. With our improved ViT supernet training techniques, our method finds a family of efficient models, called NASVIT, that establishes new state-of of-the-art top-1 accuracy FLOPs trade-offt ImageNet. Our NASViT models are the first ViT variants that outperform prior-art efficient CNNs on the mobile FL FLOPs regime.

---

Published as a conference paper at ICLR 2022
Acknowledgements Chengyue Gong and Qiang Liu supported in part by CAREER- 1846421, SenSE-200 2032220,,,67, EAGER-2041322 and Office of Navy Research, and NSF Al Institute for Founda- tions of Machine Learning (1FML). We would like thank the anonymous reviewers and the area chair for their thoughtful comments and efforts towards improving our manuseript.
REFERENCES
Mehdi Abbana Bennani, Thang Doan, and Masashi Sugiyama. Generalisation guarantees for continual learning with orthogonal gradient descent. arXiv preprint arXiv:2006. / 1942, 2020.
Gedas Bertasius, Heng Wang, and Lorenzo Torresani. Is space-time attention all you need for video understanding? ICML 2021, 2021.
Lucas Beyer, Xiaohua Zhai, Amélie Royer, Larisa Markeeva, Rohan Anil, and Alexander Kolesnikov. Knowledge distillation: A good teacher is patient and consistent. arXiv preprint arXiv:2106.05237, 2021.
Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han. Ance-for-all Train one network and specialize for efficient deployment. arXiv preprint arXiv: 09791, 2019.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European Conference on Computer Vision, pp. 213-229, Springer, 2020.
Minghao Chen, Houwen Peng, Jianlong Fu, and Haibin Ling. Autoformer: Searching transformers for visual recognition. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 12270-12280, 2021a.
Tianlong Chen, Yu Cheng, Zhe Gan, Lu Yuan, Lei Zhang, and Zhangyang Wang. Chasing sparsity in vision transformers: An end-to-end exploration. arXiv preprint arXiv:2 04533, 2021b.
Yinpeng Chen, Xiyang Dai, Dongdong Chen, Mengehen Liu, Xiaoyi Dong, Lu Yuan, and Zicheng Liu. Mobile-former: Bridging mobilenet and transformer. arXiv preprint arXiv:2/08.0589. 2021c.
Bowen Cheng, Alexander G Schwing, and Alexander Kirillov. Per-pixel classification is not all you need for semantic segmentation. arXiv preprint arXiv:2 2021.
Marius Cordts, Mohamed Omran, Sebastian Ramos, Timo Rehfeld, Markus Enzweiler, Rodrigo Benenson, Uwe Franke, Stefan Roth, and Bernt Schiele. The cityscapes dataset for semantic urban scene understanding. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3213-3223, 2016.
Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le. Autoaugment: Learning augmentation strategies from data. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 113-123, 2019.
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with reduced search space. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops. pp. 702-703, 2020.
Xiaoliang Dai, Alvin Wan, Peizhao Zhang, Bichen Wu, Zijian He, Zhen Wei, Kan Chen, Yuandong Tian, Matthew Yu, Peter Vajda, et al. Fbnetv3: Joint architecture-rccpe search using neural acquisition function. arXiv e-prints, pp. arXiv-2006. 2020.
Zihang Dai, Hanxiao Liu, Quoc V Le, and Mingxing Tan. Coatnet: Marrying convolution and attention for all data sizes. arXiv preprint wxxv:22166.0888 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scall scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. 255. Teee, 2009.
Lucio M Dery, Yann Dauphin, and David Grangier. Auxiliary task update decomposition: The good, the had and the neutral. arXiv preprint arXiv:2 108. /346, 2021.
Xiaoyi Dong. Jianmin Bao, Dongdong Chen, Weiming Zhang, Nenghai Yu, Lu Yuan, Dong Chen, and Baining Guo. Cswin transforerr: A general vision transformer backbone with cross-shaped windows. arXiv preprint arXiv:2107.00652. 2021.
10

---

Published as a conference paper at ICLR 2022
Xuanyi Dong and Yi Yang. Nas-bench-201. Extending the scope of reproducible neural architecture search. arXiv preprint arXiv:2001.00326. 2020.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 1929, 2020.
Yunshu Du, Wojciech M Czarnecki, Siddhant M Jayakumar, Mehrdad Farajtabar, Razvan Pascanu, and Balaji Lakshminarayanan. Adapting auxiliary losses using gradient similarity. arXiv preprini arXiv: /812.0222, 2018.
Haoqi Fan, Bo Xiong. Karttikeya Mangalam, Yanghao Li. Zhicheng Yan, Jitendra Malik, and Christoph Feichtenhofer. Multiscale vision transformers. arXiv preprint arXiv:2 104. 1227, 2021.
Mehrdad Farajtabar, Navid Azizan, Alex Mott, and Ang Li. Orthogonal gradient descent for continual learning. In International Conference on Artificial Intelligence and Statistics, pp. 3762-3773.
PMLR, R, 2020.
Ben Graham, Alaaeldin El-Nouby, Hugo Touvron, Pierre Stock, Armand Joulin, Hervé Jégou, and Matthijs Douze. Levit: a vision transformer in convnet's clothing for faster inference. arXiv preprint urXiv:2104. 106,,,662000 2021.
Kai Han, An Xiao, Enhua Wu, Jianyuan Guo, Chunjing Xu, and Yunhe Wang. transformer. arXiv preprint arXiv:2103.001 /2, 2021.
Transformer in
Byeongho Heo, Sangdoo Yun, Dongyoon Han, Sanghyuk Chun, Junsuk Choe, and Seong Joon Oh. Rethinking spatial dimensions of vision transformers. arXiv preprint urXiv:2103. 16302, 2021.
Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, al. Searching for mobilenetv3. In Proceedings of the JEEE/CVF International Conference on Computer Vision, pp. 1314-1324, 2019.
Jie Hu, Li Shen, and Gang Sun. Squeeze-and-exciattion networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7132-7141, 2018.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980 2014.
Nikita Kitaev, Eukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv
preprint arXiv:2001.0445// 2020.
Alexander Kolesnikov, Lucas Beyer. Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning. In Computer ViCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part V 16, pp. 491-507. Springer, 2020.
Changlin Li, Jiefeng Peng, Liuchun Yuan, Guangrun Wang, Xiaodan Liang, Liang Lin, and Xiaojun Chang. Block wiselly supervised neural architecture search with knowledge distillation. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 1989-1998, 2020.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search, arXiv preprint arXiv:/806. 11880600955.. 2018.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint arXiv:2103. /4030, 2021.
Yiping Lu. Zhuohan Li, Di He, Zhiqing Sun, Bin Dong, Tao Qin, Liwei Wang, and Tie- Yan Liu. Understanding and improving transformer from multi-particle dynamic system point of view. arXiv preprint arXiv: 0e2///906007666 2019.
Houwen Peng, Hao Du, Hongyuan Yu, Qi Li, Jing Liao, and Jianlong Fu. Cream of the crop: Distilling prioritized paths for one-shot neural architecture search. arXiv preprint arXiv:2010, 15821, 2020.
Hieu Pham, Melody Guan, Barret Zoph, Quoe Le, and Jeff Dean. Efficient neural architecture search via parameter sharing. arXiv preprint arXiv: 08802000366 2018
Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsich. Dynamicvit: Efficient vision transformers with dynamic token sparsification. arXiv preprint 2Xiv:2/06.00000
11

---

Published as a conference paper at ICLR 2022
2021.
Esteban Real, Sherry Moore, Andrew Selle, Saurabh Saxena, Yutaka Leon Suematsu, Jie Tan, Quoe V Le. and Alexey Kurakin. Large-scale evolution of image classifiers. In International Conference on Machine Learning. pp. 2902-299 PMLR, 2017.
Esteban Real, Alok Aggarwal, Yanping Huang, and Quoe Le. Regularized evolution for image classifier architecture search. In Proceedings of the aaai conference artificial intelligence. volume 33, pp. 4780-4789, 2019.
Gobinda Saha, Isha Garg, and Kaushik Roy. Gradient projection memory for continual learning. arXiv preprint wwxiv:2103.09765 2021.
Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chiel Chen. Mo- bilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 4510-4520, 2018.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. 2018.
arXiv preprint arXiv:/803.002
Noam Shazeer, Zhenzhong Lan, Youlong Cheng, Nan Ding, and Le Hou. Talking- attention.
arXiv preprint ar arXiv:2003.02433333 2020.
Dimitrios Stamoulis, Ruizhou Ding, Di Wang, Dimitrios Lymberopoulos, Bodhi Priyantha, Jie Liu, and Diana Marculescu. Single-path nas: Designing hardware-efficient convnets in less than hours. arXiv preprint
2019.
Xiu Su, Shan You, Jiyang Xie, Mingkai Zheng, Fei Wang, Chen Qian, Changshui Zhang, Xiaogang Wang, and Chang Xu. Vision transformer architecture search. arXiv preprint arXiv:2106.13700, 2021.
Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. Revisiting unreasonable effectiveness of data in deep learning era. In Proceedings of the IEEE international conference on computer vision, pp. 843-852, 2017.
Mingxing Tan and Quoc Le. Efficientnet: Rethinking model scaling for convolutional neural networks In International Conference on Machine Learning, pp. 6105-61 14. PMLR, 2019.
Mingxing Tan, Bo Chen, Ruoming Pang, Vijay Vasudevan, Mark Sandler, Andrew Howard, and Quoc V Le. Mnasnet: Platform-aware neural architecture search for mobile. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2820-2828, 2019.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé
Jégou. Training data-efficient image transformers & distillation through attention. arXiv preprint arXiv:2012.12887 2020.
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. arXiv preprint arXiv:2103. 17239, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv preprint arXiv: 03762, 2017.
Alvin Wan, Xiaoliang Dai, Peizhao Zhang, Zijian He, Yuandong Tian, Saining Xie, Bichen Wu, Matthew Yu, Tao Xu, Kan Chen, et al. Fbnetv2: Differentiable neural architecture search for spatial and channel dimensions. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattrrn Recognition, pp. 12965-12974, 2020.
Li Wan, Matthew Zeiler, Sixin Zhang, Yann Le Cun, and Rob Fergus. Regularization of neural networks using dropconnect. In International conference on machine learning. pp. 1058-1066 PMLR, 2013.
Dilin Wang, Chengyue Gong, Meng Li, Qiang Liu, and Vikas Chandra. Alphanet: Improved training of supernet with aipha-divergence. arXiv preprint arXiv:2102.07954 2021a.
Dilin Wang, Meng Li, Chengyue Gong, and Vikas Chandra. Attentivenas: Improving neural architecture search via attentive sampling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1418-6427, 2021b.
Haiping Wu, Bin Xiao, Noel Codella, Mengchen Liu, Xiyang Dai, Lu Yuan, and Lei Zhang. Cvt: Introducing convolutions to vision transformers. arXiv preprint arXiv:2103. 2021.
12

---

Published as a conference paper at ICLR 2022
Tete Xiao, Mannat Singh, Eric Mintun, Trevor Darrell, Piotr Dollár, and Ross Girshick. Early convolutions help transformers see better. arXiv preprint arXiv:2106. 14881, 2021.
Enze Xie, Wenhai Wang, Zhiding Yu, Anima Anandkumar, Jose M Alvarez, and Ping Luo. Seg- former: Simple and efficient design for semantic segmentation with transformers. arXiv preprint arXiv:2105.15203 2021.
Jiahui Yu, Linjie Yang. Ning Xu, Jianchao Yang, and Thomas Huang. Slimmable neural networks.
arxiv preprint arXiv: 1812.08928 2018.
Jiahui Yu, Pengchong Jin, Hanxiao Liu, Gabriel Bender, Pieter-Jan Kindermans, Mingxing Tan, Thomas Huang, Xiaodan Song, Ruoming Pang, and Quoe Le. Bignas: Scaling up neural archi- tecture search with big single-stage models. European Conference on Computer Vision, pp- 702-717. Springer, 2020a.
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. arXiv preprint 2Xiv:20011066822 2020b.
Li Yuan, Yunpeng Chen, Tao Wang. Weihao Yu, Yujun Shi, Zihang Jiang, Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-token vit: Training vision transformers from scratch on imagenet. arXiv preprini arXiv:2 1986, 2021a.
Li Yuan, Qibin Hou, Zihang Jiang, Jiashi Feng, and Shuicheng Yan. Volo: Vision outlooker for visual recognition. arXiv preprini arXiv:2106. 12, 2021b.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 6023-6032, 2019.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint arXiv: 07.40.9112 2017.
Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yi Yang. Random erasing data aug- mentation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 13001-13008, 2020.
Aimin Zhou, Bo- Yang Qu, Hui Li, Shi-Zheng Zhao, Ponnuthurai Nagaratnam Suganthan, and Qingfu Zhang. Multiobjective evolutionary algorithms: A survey of the state of the art. Swarm and evolutionary computation, 1(1):32-49, 2011.
Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Scene parsing through ade20k dataset. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 633-641, 2017.
Daquan Zhou, Bingyi Kang, Xiaojie Jin, Linjie Yang. Xiaochen Lian, Qibin Hou, and Jiashi Feng. Deepvit: Towards deeper vision transformer. arXiv preprint arXiv:2103. 1886. 2021.
Xizhou Zhu, Weijie Su, Lewei Lu, Bin Li, Xiaogang Wang, and Jifeng Dai. Deformable detr: Deformable transformers for end-to-end object detection. arXiv preprint arXiv:2010.04159, 2020.
Barret Zoph and Quoc V Le. Neural architecture search with reinforcement learning. arXiv preprint arXiv: 1X1101111 2016.
Barret Zoph, Vijay Vasudevan, Jonathon Shlens, and Quoc Ve. Learning transferable architectures for scalable image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 8697-8710, 2018.
13

---

Published as a conference paper at ICLR 2022
A RELATED WORK
ViTs
ViT Dosovitskiy et al., 2020) and its follow-ups (Wu al. 2021 Liu et 2021 Zhou
et al. 2021: Touvron et al. 2021) have been demonstrated to be an alternatvee choice to CNNs for challenging vision tasks, especially for image classification. In (Dosovitskiy et al. 2020 large-scale ViT-Large model is trained on JFT-300 to obtain good performance. The follow-ups mainly focus on making the data size and model size smaller without loss of accuracy. A line of works introduce inductive bias or CNN layers to keep the good performance of ViTs while reducing the data and model sizes (e.g. [Luu et al. 2021 Dong et al., 2021 Yuan et al., 202 Wu et al. (021). For smaller data sizes, researchers successfully ly achieve good performance using ViT's without extra data. For example, VOLO (uuan et al. 2021b achieves 87.3% top-1 accuracy on ImageNet without extra data.
For smaller model sizes, however, ViT' have not achieved comparable results to efficient CNNS smaller than 1G FI.Oss even if additional CNN layers introduced (Graham al.) 2021). Dynamic ViT's (Rao et al. 2021 Chen et al.] 2021b) propose to dynamically filter the tokens to reduce the computation cost, and the efficiency is not comparable to efficient CNNs. ViViT (Graham al.)
2021). Xiao et al. 2 and PiT (Heo et al. (02T) processes the high resolution inputs with early convolution layers or spatial- aware layers and also adopt more efficient self-attention designs. Mobile- former (Chen et al. (021c) proposes two-branch neural network: one is efficient CNN layers and the other is transformer layers with a small number of tokens (e.g. 6, 8).
NAS NAS is a powerful tool for automating efficient neural architecture design. It often targets
searching for the hest model in a search space under given efficiency- constraints. Earlier NAS
solutions often build on reinforcement learning (e.g. Zoph & Le, 2016 Zoph et al., 2018 Howard
iet al. (019) and evolutionary algorithms (e.g. Real et al.) 2019 2017; Wan et al. 2020 More
NAS practices have made the search more efficient through weight-sharing and search architectures
with gradient-based methods (e.g. Liu et al., 2018 Pham et al. 2018 Stamoulis et al.) 2019). This helps alleviate the heavy computational burden training all candidate networks from scratch and accelerates the NAS process significantly. and researchers work on how to rank the model performance more accurately (Dong& & Yang 2020). Recently, training a large supernet without retraining candidate sub-networks with inplace KD is shown to be an effective mechanism that significantly improves the supernet performance (c.g. Yu et al. 2020a Wang et al. (2021baa) In addition to inplace KD, various of KD variants have also been investigated in the literature. For example, Peng et al. (2020) proposes to search a prioritized path as the teacher; [i et al. (2020) proposes to distill the feature level knowledge from additional teacher model to improve the NAS performance.
NAS for VIT Most recently, several related works, e.g., AutoFormer (Chen al.) 2021a). and ViTAS (Su et al. 2021), have been proposed to search for ViTs. AutoFormer is the first paper that leverages NAS for ViT optimization. A comprehensive search space for the ViT architecture is proposed and weight-entanglement training strategy is developed to improve the NAS efficiency. ViTAS leverages a similar supernet-based NAS method compared to AutoFormer and introduces private class token and self-attention maps to cater for the variance of distinct ViT architectures. Both work demonstrate promising accuracy improvement compared to the baseline DeiT models for large models with more than IG FLOPs.
Gradient cosine similarity in multi-mask learning A line of prior approaches have observed that the difficult training with multiple objectives can be improved by using the cosine similarity between gradients (e.g. Du et al. 2018: Yu et al.] 2020b] Real et al. 2019). The cosine similarity is used as a regularization or an indicator, Real et al. (2019) adds regularization term to force the cosine similarity between two different losses to be larger than zero. Du et al. (2188) and et al. (202I) propose to use gradient cosine similarity to identify whether auxiliary tasks can benefit the main task. In Du et al. (2018), once the cosine similarity is negative (gradient conffict), the weight of the auxiliary task is set to be zero. Yu et al. (2020b) is most related to our projection method, which projects the gradient of every loss to achieve orthogonal gradients. To avoid negative gradient cosine similarity, we project the gradient of the supernet to prioritize the training of sub-netwon networks, which have similar intuition as auxiliary losses. In continual learning, many works uses orthogonal
14

---

Published as a conference paper at ICLR 2022
gradient descent to restrict the direction of gradient updates of tasks in order avoid catastrophic forgetting (e.g. Farajtabar et al. 2020: Bennani et al. 2020 Saha et al. (2221)
B SEARCH SPACE
Proj
Activation
Softmax
Hoad
MLP
Talk ng Head
MLP
Multi-Head Attention
Input
Figure 6: A demonstration of our self-attention module. 'RPE', 'Dw Cony, 'Talking Head', 'Proj' and 'MLP' refer relative positional embedding, depth-wise convolutional layer, talking head attention, projection layer and MLP layer, respectively.
B.1 EFFICIENT TRANSFORMER BUILDING LAYER
In this section, we give detailed introduction about our transformer building layer with self-attention. In the literature, researchers have develop many variants of the standard self-attention with different focuses (e.g. efficiency, convergence, lone-term dependency, etc.). Motivated recent works, e.g..
and VOLO (Graaam et al.) 2021 Liu al., 2021] Yuan al.) 202Tb). we develop a transformer layer for the purpose of efficiency and effectiveness in vision tasks. A demonstration of our transformer layer shown in Figure
Le
To enhance the learning capacity of our
VIT models, incorporate talking head (Shazeer et al.) 2020) layers and depth-wiss westh-wise convolution layer in the self-attention module. Additionally, following LeViT (Graham et al., (021), we expand the dimension of V matrix by expansion ratio 4 and introduce activation function after the projection matrix. Following Swin Transformer (Liu et al.) 2021), we use relative positional embedding for the attention scores. For efficiency, we reduce the MLP expansion ratio {1, 2) and add one additional MLP layer to keep the model complexity following MacaronNet (Lu et al. 2019).
Positional information The positional embedding in transformer architectures is location-dependent trainable parameters. Recent works propose absolute positional embedding, relative positional embedding additional depth-wise convolution layers (Dong al. (021) enhance the local information. We introduce two additional depth-wiss wise convolution layers into a MHSA with relative positional embedding. For relative positional embedding, we directly follow the implementation in NLP (Shaw et al. 2018). For depth-wise convolution, we add one depth- depth-wise wise convolution layer in the MLP layer and another depth-wiss wise convolution layer after the linear transformations of V matrix.
Expansion Ratio In the self-attention design space, researchers have explored whether expanding the channels can have good performance. LeViT proposes to expand the dimension of V. We follow LeVIT's s design which expand the dimension of V by expansion ratio 4. researchers have explored how many layers (Lu et al. 2019) should we use for MLP a self-attention block. We follow the strategy developed by (Lu al. 2019) which adds but reduce the MLP expansion ratio to 1,, 2} for efficiency. As displayed in Figure 6 we place an additional MLP layer after the first MLP layer.
more MLP layer for each self-attention block,
Normalization Layers and activation Many recent works apply additional batch normalization layers, layer normalization layers activation functions the network. Taking the computation cost
15

---

Published as a conference paper at ICLR 2022
of layer normalization layers into consideration, we do not introduce any normalization layers to the basic self-attention layers.
Talking-heat attention and number of heads Most of the existing VITs set the dimension of each
head to be 24/32 g Touvron et al. 2020; 2021 Zhou et al. 202T] Wu et al. 2021) et al. 2021). However, for model with few channels, large head dimension leads to few number of heads. We set a smaller head dimension (e.g. 8, 16) to make the number of heads to be large, and further introduce the talking-head attention improve the capacity of different heads. Talking-head attention (Shazeer et al., (020)) introduce two additional linear transformation between all the heads, one is before softmax and another after softmax.
Classification head Due to the use of deett-wwisc convolution layers and down wwwn-sampling, we remove the classification token for simplicity. While ViT and DeiT use two heads for the teacher knowledge distillation and supervised labels, we use one head for all the training settings and replace the one- layer fully-connected layer head with MobilenetV3 (Howard et al., 2019) head so as to reducing computation cost.
To train very deep transformer models, Touvron al. 2021) introduces additional learnable channel-wise scaling factors initialized with 10-2 into the models. The channel-wise scaling factors are introduced to output of each MI.P and multi-head attention (MHA) layer in the model. Notice that many of sub-networks in our search space very deep. we introduce architecture-dependen switchable scaling factors into the supernet (see Section
Scaling Factor
C IMPLEMENTATION DETAILS ON TRAINING AND SEARCHING
C.1 TRAINING AND SEARCH ALGORITHM
Training Consider a supernet with trainable parameter 0 and the candidate sub-networks set A. The goal of training a supernet is to learn model parameter 0 target optimizing all the sub-networks in A and simultaneously achieving good accuracy. Let s ~ A, p(æ; 0) and q(r; 0,) denote the output probability of the supernet and the sub-network s, we have the loss
{(0) + Σ ZlKD(Bs, Odetach), where (OD(Ox,, Odetach) = DD p(r; (detach)) q(x;; (..)
Here, {()) denotes the loss of the supernet, 9s, represents the parameters for sampled sub-networt 8; and (Dp || q) is divergence that measures the difference between and The Odetach denotes copy of 0 whose gradient is stopped during propaaaaaaaan.nn
Algorithm 1 Algorithm: Supernet based NAS training
1: while not converged do
2: Sample a mini-batch data from dataset
3: Sample the supernet (i.e., the biggest sub-network) from the search space and train the supernet with with ground truth labels (or with KD from an external teacher model)
4: Sample & random sub-netwe networks from the search space and train them with KD by using the supernet network as the teacher model
5: end while
Search After training, a random forest based neural predictor is trained to fit the map from the architecture parametrtsss the model performance (e.g.. accuracy). number of sub-networks are sampled from the trained supernet to train the neural predictor.
We then follow the strategy in previous works (e.g. Cai et al. 2019 Wang et al. 202 la) to do evolutionary search: 1) we randomly sample 1024 sub-networks from the supernet and estimate their accuracy on a sub-sampled subset of the ImageNet training set, which is never used during the supernet training: 2) apply crossover and random mutation (see (hou et al. (201T) for more details about evolutionary algorithms) on the best performing 128 sub-networks. We fix both the crossover size and mutation size to be 128, yielding 256 new sub-networks. We then evaluate the performance of these bebneetwrrss; 3) We repeat the second step 40 times. The total number of sub-networks thus evaluated is around 10K.
16

---

Published as a conference paper at ICLR 2022
Algorithm 2 Algorithm: Supernet based NAS searching
Input: pretrained supernet with fixed weights
2: Randomly sample 1024 sub-networks and evaluate their performance on a withhold training set (which used during training).
3: Partition 1024 sub-networks into training and validation subset with equal size. Train a random forest regressor to predict sub-netowrk's s accuracy given the sub-network architecture hyper- parameters as the input.
4: Run evolution algorithm to search the Pareto of nub-networks. The sub-network accuracy is given by the random forest based predictor.
C.2 ABLATION STUDIES
Latency-aware Searching Note that the networks the paper were optimized for the best FLOPPs accuracy trade-off. However, it is expected that the networks that achieve the best FLOPs vs. accuracy trade-off don't necessary yield the best latency vs. accuracy trade-off in the same time. Therefore, achieve the best latency accuracy trade-off, that form better latency vs. accuracy trade-off compared to the results from NASViTs (A0/A1/A2). For latency comparison, Intel(R) Xeon CPUs with a batch size of 1. See the result in the table below. Here, we do not remove BN LN layers when estimating latency.
re-search three NASVITs (B0/B1/B2)
evaluate the latency of NASVITS-A0AALAA2 and AlphaNet-AO//A/// on
Model
Accuracy (%)
CPU latency (ms)
NASVIT-BO
78,2
21.0=0.4
NASVIT-AO
78.2
21.6==0.5
AlphaNet-AO
77.9
21.4=0.5
NASViT-BI
79.6
26.6:10.6
NASViT-Al
79.7
27.2+0.6
AlphaNet-A2
79.4
2..
NASViT-B2
80.6
29.1
NASVIT-A2
80.5
29.8
AlphaNet-A4 80.4
30.4
Table 12: NASVIT models searched with for better latency vs. accuracy trade-off.
17

---

Published as a conference paper at ICLR 2022
D ARCHITECTURE VISUALIZATION OF NASVIT MODELS
NASVIT-AI
NASViT-A2
NASVIT-A3
NASViT-A4
c: 16
e:
c: T6
e:
d: 1
d: 1 d: 1
d: 1
Conv
ks:
ks:
ks: 3
ks: 3
8: 2
s: 2:2
8: 2:2
2
c: 16
c: 16
c: 16
c: 16
d: 1
d: 1
d: 1
d: 1
MBConv-1
ks: 3
ks: 3
ks: 3
ks: 3
e: 3
e: 3
e:
S:
S:
S:
S:
c: 24
c: 24
c:
c:
d: 3
d: 3
d: 3
d: 3
MBConv-2
ks:
ks: 3
ks: 3
ks: 3
e: 4
e: 4
e:
8: 2
s:
8:
s:
c: 32
c: 32
c:
32
d: 3
d: 3
d: 3
d: 3
MBConv-3
ks:
ks: 3
ks: 3
ks: 3
c:4
6
e: 5
e:6
s: 2
s:
8:
c: 64
c: 64
64
d: 4
d: 4
d: 4
d: 4
Transformer-4
k:
k: 1
k: 1
k:
e: 1
e: 1
c:
s:
s:
c: 112
c: 112
112
120
d: 3
d: 3
d: 4
d: 3
Transformer-5
1
s: 2
s:
s:
s:
c: 160
c: 160
c: 160
160
d: 3
d: 5
d: 7
d: 6
Transformer-6
e:
e:
e:
8:
s:
c: 216
c: 208
c: 216
c: 216
d: 3
d: 4
d: 5
d: 6
Transformer-7
e: 1
e: 1
e:
e:
2
MBPool
c: 1792
c: 1792
c: 1984
1984
Resolution
192
224
256
288
Table 13: Here, denotes the number of output channels, 'd' denotes number of layers, 'ks' denotes kernel size, 'e' denotes expansion ratio, 'k' denotes number of windows, 's" denotes stride.
18


Reviews:
[{"review_id": "fV-QTazAZcJ", "reviewer": "Reviewer_oqFz", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "#Strength\n1. The goal of introducing NAS to ViTs sounds reasonable.\n2. The discovered gradient conflict issue looks interesting.\n3.The experimental results seems solid.\n\n#Weakness\n1. It seems that the authors only focus on the one-shot NAS. Could other types of NAS methods be applied to ViTs? From this perspective, the discovered gradient conflict issue seems not universal.\n\n2. The proposed three techniques, i.e., the gradient projection, the switchable layer scaling and the data augmentation, look like a “A+B+C” pattern. There seems to be no obvious connection among the three techniques. \n\n3. Why do the authors remove the component of the super-net gradient that is conflict with the sub-network gradient? This operation is not intuitive. What are the negative effects of this operation? \n\n4. The introducing of the weak data augmentation scheme is also sudden. It is highly suggested to give more intuitive explanations in Section 1.\n\n5. What is the motivation of the raised three investigations? \n\n\n", "summary_of_the_paper": "This paper aims at applying one-shot Neural Architecture Search (NAS) to Vision Transformers (ViTs). The authors claim that directly using existing CNN based NAS method to ViTs will lead to a gradient conflict issue. In order to tackle this issue, the authors propose three techniques, including a gradient projection, a switchable layer scaling, and a data augmentation. The experimental results demonstrate the effectiveness of the proposed method to some extent. ", "summary_of_the_review": "In general, the writing of this paper makes it difficult to follow the authors’ core thinking.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "aX0PPVfWwOW", "reviewer": "Reviewer_G3NT", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I really like the exposition of the paper, going through hypotheses and presenting experiments supporting or disproving the hypotheses. This was great to read. I also really appreciate the links to literature in different fields, speicifically multi-task learning regarding conflicting gradients.\n\n1. My main concern is that the paper is about finding efficient and well-performing models, however it does not include a single timing anywhere I could see. I would really appreciate if Table6 included a latency column, and there was a variant of Figure4 with latency as x-axis. Especially the small attention head size has me concerned, whether the good FLOPs will actually translate to good efficiency.\n2. What about gradient agreement between individual random subnetworks? Is it better, equal, or worse than the super network vs average f sub networks?", "summary_of_the_paper": "The paper investigates applying a SOTA NAS method (from AlphaNet) to a SOTA ViT search-space (from LeViT) and notices disappointing results. Following an investigation, conflicting gradients are identified as the root cause, alleviated via three components: gradient projection, learnable scaling factors, and heavily reduced augmentations.", "summary_of_the_review": "I find the paper generally of high quality, but the lack of any timings is concerning. If this can be fixed, it would be a pretty good paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "XQs0thbeWRG", "reviewer": "Reviewer_XjJo", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "+ The writing is easy to follow and read.\n+ The evaluation results are promising.\n-  The motivation of this work is clear, but the explanation of poor performance (gradient conflict issue) is somewhat  not convincible.\n-  The architecture in the search space consists of MBConv and transformer, it is better to compare with the existing works that also combine the CNN and transformers, like PiT, etc.\n", "summary_of_the_paper": "To solve the gradient conflict issue that affect the performance，this work propose a series of techniques, including a gradient projection algorithm, a switchable layer scaling design, and a simplified data augmentation and regularization training recipe to improve the convergence and the performance of all sub-networks.", "summary_of_the_review": "It is a good work with clear motivation and  promising performance, the reason of poor performance by directly applying the supernet-based NAS to optimize ViTs need to be discussed more.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "RuBytihEhTH", "reviewer": "Reviewer_Y6a2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1.\tThe motivation is clear, and the three proposed solutions are simple and generalizable.\n2.\tThe proposed NASViT achieves promising performance in classification and segmentation with low computational cost.\n\nWeakness:\n1.\tIt is an interesting phenomenon that ViT supernet training suffers from the gradient conflict issue. But I am curious about the performance of other supernet training methods that don’t utilize KD from supernet, like once-for-all (ICLR20), DNANAS (CVPR20) and Cream (NeurIPS20).\n2.\tI am not clear why the authors introduce the switchable scaling. I don’t see the relationship between it and the gradient conflict issue, also the improvement in Table 7 looks marginal. Also Figure 3 is not very clear, e.g., what does c1 and c2 means?\n3.\tThere is no experiment to support the statement “the supernet and the sub-networks are more likely to conflict with each other in the presence of stronger data augmentations and strong regularization.”. There is performance improvement shown in Table 7, but it is not clear whether this is caused by the conflict between supernet and sub-networks. And if so, I am wondering is there still an improvement for weak DA & Reg if the gradient projection is already applied.\n4.\tIn Table 8, it seems that AlphaNets trained with alpha-divergence doesn’t benefit from your method. Maybe you should show the gradient conflicts ratio for it to give some insights.\n5.    Since the gradient projection leads to slow convergence, the authors should provide the training time cost before and after using the proposed gradient projection, as the training time cost is also important.\n6.\tHow about the training time consuming after using the other two techniques.  Does it reduce the training time caused by gradient projection for its slow convergence? Detailed information should be presented.\n7.\tWhy do the authors use EfficientNet-B5 as the teacher? Why not B6，B7 or other networks? The reasons are not mentioned in this paper. It is confusing.\n8.    To be more general and convincing, though the evaluations of the smallest and the largest networks are provided, I think the results of the middle-size network in needed experiments are also critical. For example, in Fig 2, Tab 3, and so on.\n9.    The authors only adopt gradient projection on CNN-based supernet while conducting ablation studies. The generalizability of the other two techniques in CNN should be verified as well like Tab.7. \n\n\n", "summary_of_the_paper": "This work presents the gradient conflict issue in ViT training, i.e., the gradients of sub-networks conflict with that of the supernet, leads to inferior performance of ViT supernet training. The authors fix this issue by 1) a gradient project method to prioritize the sub-network update; 2) Use switchable layers to increase the model capacities of sub-networks; 3) Simplify the training recipe. The proposed NASViT shows the state-of-the-art top-1 accuracy v.s. FLOPs trade-offs on ImageNet.", "summary_of_the_review": "The motivation of this paper is not very convincing, need to provide more evidence to prove. For more details, pls see the weaknesses.\nAlso, there are some recent works on NAS+ViT missing in the related works, such as \"CATE: Computation-aware Neural Architecture Encoding with Transformers\", \"Vision Transformer Architecture Search\" and \"AutoFormer: Searching Transformers for Visual Recognition\". It is suggested to add them into comparison and discussion, especially for Autoformer which might be the first ViT+NAS work and has been published in ICCV21. I will re-consider the rate if all the concerns are well addressed or not. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "HHCe6STWbn", "reviewer": "Reviewer_GnU6", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths:\n\n1. The paper does not only transfer the NAS algorithms from CNN to transformer, but also observes the gradient conflict dilamma and resolves it partly.\n\n2. The searched architecture seems over-performs swin, cvt, and some other vision transformers.\n\n3. Besides gradient conflict, the authors also improve the performance by changing data augmentation.\n\n## Weaknesses & Questions:\n\n1. I am not sure what \"Group\" in Table 6 means. Does it mean the number of parameters? As I know, most vision transformers have parameters less than 100M.\n\n2. The improvement on segmentation tasks seems a bit marginal.\n\n3. It is recommended to explore how other data augmentation methods affect the NAS performance, e.g., AutoAugment, color jitter, cropping, etc. All these are common data augmentation methods for vision transformers.", "summary_of_the_paper": "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.", "summary_of_the_review": "Though there are still some minor flaws in the paper, I think it is a good one and is nearly ready for publication.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper tackles a very timely problem. 
Scores of 5,6,6,8 put it in the borderline region, but in the private discussion the more negative reviewer noted that they would also be OK with the paper being accepted. I therefore recommend acceptance. 
Going through the paper I missed any mention of available source code. I strongly recommend that the authors make code available; this would greatly increase the paper's impact.
Decision: Accept (Poster)