NASV I T: N EURAL A RCHITECTURE S EARCH FOR E FFICIENT V ISION T RANSFORMERS WITH G RADIENT C ONFLICT - AWARE S UPERNET T RAINING
Anonymous authors
Paper under double-blind review
A BSTRACT
Designing accurate and efﬁcient vision transformers (ViTs) is an important but challenging task.
Supernet-based one-shot neural architecture search (NAS) enables fast architecture optimization and has achieved state-of-the-art results on convolutional neural networks (CNNs).
However, directly applying the supernetbased 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 conﬂict issue: the gradients of different sub-networks conﬂict 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 simpliﬁed data augmentation and regularization training recipe.
The proposed techniques signiﬁcantly 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. 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.
I NTRODUCTION
Transformers have recently been applied to various vision tasks, including image classiﬁcation (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) beneﬁt from high model capacity, large receptive ﬁeld, 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% topaccuracy 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 signiﬁcant 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 ﬁelds beneﬁcial for small models?
To answer the question above, in this work, we target at developing a family of efﬁcient ViTs with
FLOPs ranging from 200M to 800M .
A natural approach is to leverage NAS, which has achieved state-of-the-art (SOTA) accuracy-efﬁciency 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 ﬁrst 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 ﬁnd 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 conﬂict 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 conﬂicting gradients, we propose three different techniques to improve the supernet training.
Firstly, instead of directly adding the gradients from different sub-networks together, we ﬁnd it beneﬁcial to prioritize the training of the sub-networks over the supernet, as our main purpose is to build efﬁcient sub-networks.
We achieve this with a projection gradient algorithm which removes the component of the supernet gradient that is conﬂict with the sub-network gradient.
Secondly, to alleviate the gradient conﬂicts 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 conﬂicts.
Thirdly, we propose to use a weak data augmentation scheme and reduce the regularization in training to decrease the optimization difﬁculty and hence reduce gradient conﬂicts.
Our proposed techniques signiﬁcantly alleviate the gradient conﬂict 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., EfﬁcientNet (Tan & Le, 2019).
We refer readers to appendix for more discussions about related works.
NAS FOR EFFICIENT TRANSFORMERS
Our goal is to design efﬁcient 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 efﬁcient models leveraging a hybrid architecture of convolutions and transformers.
In LeViT, the convolutions are introduced to handle high resolution inputs thanks to their efﬁciency 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
[Figure]
[Table]
al 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 ﬁrst 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 ﬁnd accurate ViTs We ﬁrst closely follow the previous best practices in AlphaNet (Wang et al., 2021a) for the supernet training.
We train the supernet forepochs 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.
Speciﬁcally, 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
EfﬁcientNet-B5 (Tan & Le, 2019) with 83.3% top-1 accuracy on ImageNet as the teacher to train
our
ViT supernet unless otherwise speciﬁed.
We plot the training curves of the smallest sub-network and the largest sub-network in Figure 2.
We ﬁnd both the smallest sub-network and the largest sub-network from our search space converge poorly compared to the CNN baseline.
Speciﬁcally, the validation accuracy of both the smallest and the largest sub-network is saturated at around the 250-th epoch, and the ﬁnal accuracy is much worse
I m a g e N e t T opv a l acc
ViT
CNN (AlphaNet)
I m a g e N e t T opv a l acc   360 79.
ViT
CNN (
[Figure]
Training Epoch Training Epoch
[Figure]
[Table] [Table]
[Table]
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 speciﬁc training recipes and summarize our ﬁndings 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 conﬂicts?
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 conﬂicting gradients due to the heterogeneous and complex structures of networks, and the conﬂict 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 conﬂict gradients and tend to update model parameters in opposite directions.
To quantitatively examine the gradient conﬂict issue, we go through the entire ImageNet training set and calculate the percentage of negative cosine similarity between the gradients of
[Table]
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 ofﬁcial code 2 .
As shown in Table 4, our ViT supernet suffers from more severe gradient conﬂicts compared to the CNN baseline.
According to existing works in multi-task learning, large gradient conﬂict ratios may result in signiﬁcant accuracy drop even for binary classiﬁcation problems (see Figure 3 in Du et al.
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.
GRADIENT CONFLICT AWARE SUPERNET TRAINING
[Figure]
Activated components are the neurons selected in the forward path for one sub-network.
‘c1’ and ‘c2’ represent the number of channels activated in a selfattention layer and MLP, respectively.
We propose to improve the ViT supernet training by addressing the gradient conﬂict issue between the supernet and the sub-networks from three aspects: 1) manually resolving the gradient conﬂict 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 to provide easier training signals.
Gradient projection to prioritize sub-networks update
Our ﬁrst idea is to focus on training the sub-networks whenever gradients from the supernet and the gradients from the sub-networks conﬂicted with each other.
As we are interested in the sub-networks from the 200M to 800M FLOPs range, we propose to prioritize the optimization of the sub-networks over the supernet when gradient conﬂicts are observed.
Let ∇ (cid:96) sup and ∇ (cid:96) sub denote the gradients of the supernet and the sub-networks, respectively.
To prioritize sub-networks training, we always project ∇ (cid:96) sup to the normal vector of ∇ (cid:96) sub to avoid gradient conﬂicts when the cosine similarity between ∇ (cid:96) sup and ∇ (cid:96) sub is negative.
The overall accumulated gradient at each training iteration with projection can be written as follows,
[Equation]
otherwise .
Note that cos( ∇ (cid:96) sub , proj( ∇ (cid:96) sup ))
=
if cos( ∇ (cid:96) sup , ∇ (cid:96) sub ) < 0 , which ensures the gradient cosine similarity is always non-negative.
In sandwich sampling, since several sub-networks are sampled in each iteration, (cid:96) sub is computed as the summation of the gradients from all sub-networks.
Note similar ideas of gradient projection have also been explored in multi-tasks learning, see e.g., Yu et al.
(2020b)
; Du et al.
; Real et al.
; Dery et al.
While the gradient projection in Eqn.
(1) eliminates the gradient conﬂicts, it may lead to slow convergence as the resulting gradients are biased.
Hence, we also propose the follow two techniques to reduce the gradient conﬂicts from a search space design and training strategy reﬁnement perspective.
Switchable scaling layer Motivated by Slimmable NN (Yu et al., 2018)
, we introduce additional switchable scaling layers to allow sub-networks with different layer widths and depths to re-scale their features in a privatized way.
Speciﬁcally, for each transformer layer, a switchable scaling layer is introduced at the output of the self-attention (SA) and the MLP, respectively, as shown in
Figure 3.
Assume x
[ c,d ] ∈ R c is a input feature of a scaling layer, with c the feature dimension (i.e. the number of selected channels in the forward path) and d the index of this layer in a transformer search block.
The scaling layer transforms x [ c,d ] as w [ c,d ] (cid:12) x
[ c,d ] .
Here w [ c,d ] ∈ R c are learnable parameters and (cid:12) denotes element-wise multiplication.
For each transformer block (see Table 1), each different conﬁguration 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 a small value (e.g. 10 − 4 ) 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 ﬂexibility.
Reduced data augmentation and regularization
Furthermore, we observe that the supernet and the sub-networks are more likely to conﬂict with each other in the presence of stronger data augmentations 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 1, and remove the regularization, e.g. DropConnect (Wan et al., 2013), dropout and weight decay, from the training; see Table 5 for a comparison.
[Table]
E XPERIMENTS
[Figure]
Note that Mobile-Former (Chen et al., 2021c) and
MNasNet (Tan et al., 2019) are trained without additional teacher models.
We ﬁrst retrain our ViT supernet with our proposed gradient conﬂict reduction techniques on
ImageNet (Deng et al., 2009); we then conduct an evolutionary search on a subset of the ImageNet training dataset to search the accuracy vs.
FLOPs Pareto following (Wang et al., 2021b).
We refer the reader to Appendix C.1 for more details.
Note all the models are directly sampled from the supernet without retraining or ﬁnetuning.
We call our discovered as NASViT models and compare with state-of-the-art efﬁcient CNNs and ViTs, including FBNetV3 (Dai et al., 2020),
AlphaNet (Wang et al., 2021a), LeViT (Graham et al., 2021) and Segformer (Xie et al., 2021), on both image classiﬁcation (e.g.
, ImageNet) and semantic segmentation benchmarks (e
g., Cityscapes and ADE20K).
4. I MAGE N ET
We compare our NASViT models with state-of-the-art NAS-based CNNs, including AlphaNet (Wang et al., 2021a) and FBNetV3 (Dai et al., 2020), and recently-proposed efﬁcient ViTs, e.g., LeViTs.
Settings Note that our ViT supernet is trained with a pretrained Efﬁcient-B5 teacher model (83.3% top-1) model.
For fair comparison, we retrain AlphaNet with the same teacher.
For FBNet-V3 and
[Table]
The “*” indicates that the ViTs are trained without external teacher models.
LeViT models, these models already use teachers with better performance than Efﬁcient-B5 for training, and therefore we directly report their results following their papers.
Speciﬁcally, FBNet-V use a RegNetY-32G with 84.5% top-
and LeViT use a RegNetY-16G with 83.6% top-1 as the teacher model, respectively.
Results We summarize our results in both Table 6 and Figure 4.
Our discovered NASViT models outperform all evaluated CNN and ViT baselines.
Our models are the ﬁrst models with transformers blocks that outperform state-of-the-art efﬁcient CNNs with similar FLOPs on ImageNet.
For example, with < 600M FLOPs, our NASViT-A4 achieves 81.4% top-1 accuracy on ImageNet.
As a reference point, a ResNet-50 model (4G FLOPs) achieves 81.5% top-1 accuracy by distilling from a BiT (Kolesnikov et al., 2020) teacher (87.5% top-1 accuracy) with 1200 epochs of training (Beyer et al., 2021).
[Figure]
4. S EMANTIC S EGMENTATION
We evaluate the transfer learning performance of our discovered NASViT models by ﬁne-tuning them on downstream semantic segmentation tasks.
In particular, we ﬁne-tune NASViT-A1 to
NASViT-A4 as backbones and we show that our NASViT models yield the best segmentation performance compared to the results from the previous efﬁcient CNN backbones, e.g., AlphaNet and MobileNetV3 (Howard et al., 2019), as well as the recently proposed transformer-based Segformer (Xie et al., 2021).
Settings
We evaluate on two benchmark datasets, Cityscapes (Cordts et al., 2016) and
ADE20K
(Zhou et al., 2017).
To handle large input resolutions efﬁciently, for each transformer block, we set the number of windows to be the input feature map size divided by 7 instead of using our 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, to achieve a better accuracy efﬁciency trade-off.
Results As shown in Figure 5, our models yield the best FLOPs vs. mIoU trade-offs.
For example, our model achieves 76.1% mIoU and 41.4% mIOU with less than 30G FLOPs on the Cityscapes and the ADE20K validation set, respectively.
4. A BLATION S TUDIES ON GRADIENT CONFLICT AWARE TRAINING
We conduct ablation studies on ImageNet to have a better understanding on the effectiveness of our proposed methods.
We mainly study 1) how our techniques can mitigate the gradient conﬂict issues and improve the performance 2) whether CNN supernets can also beneﬁt from our techniques.
All the comparisons in this section are conducted on ImageNet.
On the effectiveness of our gradient conﬂict aware training techniques As demonstrated in
Table 7, both weak data augmentation and regularization and switchable scaling layer and can signiﬁcantly reduce the gradient conﬂict ratios and in the meantime, improve the top-1 accuracy of both the smallest sub-network 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 ∇ (cid:96) sub .
As demonstrated in the last column of Table 7, this training strategy (denoted by Prioritize (sup) ) leads to a slightly improved supernet while resulting in less competitive performance on the smallest sub-network.
[Table]
Improving CNN-based supernets We verify the generalizability of our three techniques to the
CNN supernets.
In this setting, we applied all three techniques together to improve CNN-based supernets.
We show in Table 8 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/ KL) , we follow ours studies in Table 4 and compute the gradient conﬂict ratio for AlphaNet (w/
KL) at epoch 1st, 90th, 180th, 270th, and 360th, and the corresponding gradient conﬂicts ratio is 25%, 18%, 24%, 28% and 31%, respectively.
The gradient conﬂict issue is more severe for AlphaNet (w/ α -KL) compared with AlphaNet trained with α -divergence based KD ( AlphaNet (w/ α -div) ).
Our ﬁndings indicate that our techniques are not restricted to the ViT supernet training and might be beneﬁcial for all supernets in which a large ratio of gradient conﬂicts presents.
4. A BLATION STUDIES ON SEARCH SPACE
In this part, we provide additional ablation studies to 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.
[Table]
Global attention vs. local attention In our search space (Table 1), we mainly use the global self-attention for the best representation learning capacity.
However, with the more computationally efﬁcient local and linear self-attention schemes, we would be able to use a slightly bigger model under similar FLOPs constraints with a sacriﬁce of global context modeling.
To test this trade-off, we train the aforementioned model from scratch with different types of self-attention strategies.
Speciﬁcally, 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 et al., 2021), VOLO (Yuan et al., 2021b) and LSH (Kitaev et al., 2020).
We uniformly scale the width of the transformer blocks to ensure all models have similar compute FLOPs.
In Table 9, we 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.
[Table]
The placement of transformer blocks Our ViT supernet has a convolution stem with 3 downsampling strides.
We further test the optimal choice of where to switch to transformer blocks.
We uniformly scale the width of transformer layers of our baseline model to ensure similar FLOPs for different architecture designs.
As shown in Table 10, we notice that our current design, a convolution stem with 3 strides, yields the best performance.
[Table] [Table]
Head Dimension
In previous ViT works (e.g Touvron et al., 2020; 2021;
Zhou et al., 2021; Wu et al., 2021; Liu et al., 2021), the feature dimension of each self-attention head is usually set to be 32, 64, or larger.
In this work, as shown in Table 11, we found a smaller head feature dimension (e.g., 8) yields better performance.
C ONCLUSION
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 to inferior
NAS performance.
We ﬁx this gradient conﬂict issue by introducing a 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 to provide easier training signals.
With our improved ViT supernet training techniques, our method ﬁnds a family of efﬁcient models, called
NASViT, that establishes a new state-of-the-art top-1 accuracy vs. FLOPs trade-offs on ImageNet.
Our NASViT models are the ﬁrst ViT variants that outperform prior-art efﬁcient CNNs on the mobile
FLOPs regime.
A R ELATED W ORK
ViTs ViT (Dosovitskiy et al., 2020) and its follow-ups (Wu et al., 2021; Liu et al., 2021;
Zhou et al., 2021; Touvron et al., 2021) have been demonstrated to be an alternative choice to CNNs for challenging vision tasks, especially for image classiﬁcation.
In (Dosovitskiy et al., 2020), a 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. Liu et al., 2021;
Dong et al., 2021; Yuan et al., 2021b; Wu et al., 2021).
For smaller data sizes, researchers successfully achieve good performance using ViTs without extra data.
For example, VOLO (Yuan et al., 2021b) achieves 87.3% top-1 accuracy on ImageNet without extra data.
For smaller model sizes, however, ViTs have not achieved comparable results to efﬁcient CNNS smaller than 1G FLOPs, even if additional CNN layers are introduced (Graham et al., 2021).
Dynamic
ViTs (Rao et al., 2021; Chen et al., 2021b) propose to dynamically ﬁlter the tokens to reduce the computation cost, and the efﬁciency is not comparable to efﬁcient CNNs.
LeViT (Graham et al., 2021), Xiao et al.
(2021) and PiT (Heo et al., 2021) processes the high resolution inputs with early convolution layers or spatial-aware layers and also adopt more efﬁcient self-attention designs.
Mobileformer (Chen et al., 2021c) proposes a two-branch neural network: one is efﬁcient 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 efﬁcient neural architecture design.
It often targets at searching for the best model in a search space under given efﬁciency-related constraints.
Earlier NAS solutions often build on reinforcement learning (e.g. Zoph & Le, 2016; Zoph et al., 2018;
Howard et al., 2019) and evolutionary algorithms (e.g. Real et al., 2019; 2017; Wan et al., 2020).
More
NAS practices have made the search more efﬁcient 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 of training all candidate networks from scratch and accelerates the NAS process signiﬁcantly.
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 signiﬁcantly improves the supernet performance (e.g. Yu et al., 2020a;
Wang et al., 2021b;a).
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; Li et al.
(2020) proposes to distill the feature level knowledge from an additional teacher model to improve the NAS performance.
NAS for ViT Most recently, several related works, e.g., AutoFormer (Chen et al., 2021a), and
ViTAS (Su et al., 2021), have been proposed to search for ViTs.
AutoFormer is the ﬁrst paper that leverages NAS for ViT optimization.
A comprehensive search space for the ViT architecture is proposed and a weight-entanglement training strategy is developed to improve the NAS efﬁciency.
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 1G FLOPs.
Gradient cosine similarity in multi-mask learning A line of prior approaches have observed that the difﬁcult 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 to as a regularization or an indicator.
Real et al.
(2019) adds a regularization term to force the cosine similarity between two different losses to be larger than zero.
Du et al.
and Dery et al.
(2021) propose to use gradient cosine similarity to identify whether auxiliary tasks can beneﬁt the main task.
In Du et al.
(2018), once the cosine similarity is negative (gradient conﬂict), 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-networks, which have similar intuition as auxiliary losses.
In continual learning, many works uses orthogonal gradient descent to restrict the direction of gradient updates of new tasks in order to avoid catastrophic forgetting (e.g. Farajtabar et al., 2020; Bennani et al., 2020; Saha et al., 2021).
B S EARCH SPACE
[Figure]
B. E FFICIENT TRANSFORMER BUILDING LAYER
In this section, we give a 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. efﬁciency, convergence, lone-term dependency, etc.).
Motivated recent works, e.g.,
LeVit, SWIN-tranformer and VOLO (Graham et al., 2021; Liu et al., 2021; Yuan et al., 2021b), we develop a transformer layer for the purpose of efﬁciency and effectiveness in vision tasks.
A demonstration of our transformer layer is shown in Figure 6.
To enhance the learning capacity of our
ViT models, we incorporate talking head (Shazeer et al., 2020) layers and depth-wise convolution layer in the self-attention module.
Additionally, following LeViT (Graham et al., 2021), 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 efﬁciency, we reduce the MLP expansion ratio to { 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 or additional depth-wise convolution layers (Dong et al., 2021) to enhance the local information.
We introduce two additional depth-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-wise convolution layer in the MLP layer and another depth-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 design which expand the dimension of V by an expansion ratio 4.
researchers have explored how many layers (Lu et al., 2019) should we use for MLP in a self-attention block.
We follow the strategy developed by (Lu et al., 2019) which adds one more MLP layer for each self-attention block, but reduce the MLP expansion ratio to { 1, 2 } for efﬁciency.
As displayed in Figure 6, we place an additional MLP layer after the ﬁrst MLP layer.
Normalization Layers and activation Many recent works apply additional batch normalization layers, layer normalization layers or activation functions to the network.
Taking the computation cost of layer normalization layers into consideration, we do not introduce any new normalization layers to the basic self-attention layers.
Talking-head attention and number of heads Most of the existing ViTs set the dimension of each head to be 24/32 (e.g Touvron et al., 2020; 2021;
Zhou et al., 2021; Wu et al., 2021; Liu et al., 2021).
However, for a model with few channels, a 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 to improve the capacity of different heads.
Talking-head attention (Shazeer et al., 2020) introduce two additional linear transformation between all the heads, one is before softmax and another is after softmax.
Classiﬁcation head Due to the use of depth-wise convolution layers and down-sampling, we remove the classiﬁcation token for simplicity.
While LeViT 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 onelayer fully-connected layer head with MobilenetV3 (Howard et al., 2019) head so as to reducing computation cost.
Scaling Factor To train very deep transformer models, Touvron et al.
(2021) introduces additional learnable channel-wise scaling factors initialized with 10 − 4 into the models.
The channel-wise scaling factors are introduced to the output of each MLP and multi-head attention (MHA) layer in the model.
Notice that many of sub-networks in our search space are very deep, we introduce architecture-dependent switchable scaling factors into the supernet (see Section 3).
C I MPLEMENTATION DETAILS ON TRAINING AND SEARCHING C. T RAINING AND S EARCH A LGORITHM
Training Consider a supernet with trainable parameter θ and the candidate sub-networks set A .
The goal of training a supernet is to learn model parameter θ target at optimizing all the sub-networks in A and simultaneously achieving good accuracy.
Let s ∼ A , p ( x ; θ ) and q ( x ; θ s ) denote the output probability of the supernet and the sub-network s , we have the loss
L = (cid:96) ( θ ) + k (cid:88) i = (cid:96) KD
( θ s i , θ detach ) , where (cid:96)
KD
( θ s i , θ detach ) =
D (cid:18) p ( x ; θ detach ) ||
q ( x ; θ s i ) (cid:19) .
Here, (cid:96) ( θ ) denotes the loss of the supernet, θ s i represents the parameters for a sampled sub-network s i , and D ( p || q ) is a divergence that measures the difference between p and q .
The θ detach denotes a copy of θ whose gradient is stopped during back-propagation.
batch data
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 k random sub-networks from the search space and train them with KD by using the supernet network as the teacher model
while
Search After training, a random forest based neural predictor is trained to ﬁt the map from the architecture hyper-parameters to the model performance (e.g., accuracy).
A 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., 2021a) 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) we apply crossover and random mutation (see Zhou et al.
(2011) for more details about evolutionary algorithms) on the best performing 128 sub-networks.
We ﬁx both the crossover size and mutation size to be 128, yielding 256 new sub-networks.
We then evaluate the performance of these sub-networks; 3) We repeat the second step 40 times.
The total number of sub-networks thus evaluated is around 10K.
1: Input: a pretrained supernet with ﬁxed weights 2: Randomly sample 1024 sub-networks and evaluate their performance on a withhold training set (which is not 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 accuracy given the sub-network architecture hyperparameters as the input.
4: Run evolution algorithm to search the Pareto of sub-networks.
The sub-network accuracy is given by the random forest based predictor.
C. A BLATION S TUDIES
Latency-aware Searching Note that the networks in the paper were optimized for the best FLOPs vs. 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, to achieve the best latency vs. accuracy trade-off, we re-search three NASViTs (B0/B1/B2) that form better latency vs. accuracy trade-off compared to the results from NASViTs (A0/A1/A2).
For latency comparison, we evaluate the latency of NASViTs-A0/A1/A2 and AlphaNet-A0/A2/A on Intel(R)
Xeon CPUs with a batch size of 1.
See the result in the table below.
[Table]
D A RCHITECTURE VISUALIZATION OF NASV I T MODELS
NASViT-A
NASViT-A NASViT-A
NASViT-A
Conv c:c:c:c:d:d:d:d:ks:ks:ks:ks:s:s:s:s:
MBConvc:c:c:c:d:d:d:d:ks:ks:ks:ks:e:e:e:e:s:s:s:s:
MBConvc:c:c:c:d:d:d:d:ks:ks:ks:ks:e:e:e:e:s:s:s:s: MBConvc:c:c:c:d:d:d:d:ks:ks:ks:ks:e: e:e:e: s:s:s:s:
Transformerc:c:c:c:d:d:d:d:k:k:k:k:e:e:e:e:s:s:s:s: Transformerc:c:c:c:d:d:d:d:e:e:e:e:s:s:s:s:
Transformerc:c:c:c:d:d:d:d:e:e:e:e:s:s:s:s: Transformerc:c:c:c:d:d:d:d:e:e:e:e:s:s:s:s:
MBPool c:c:c:c:
Resolution