# Searching for Efficient Multi-Stage
Vision Transformers

Yi-Lun Liao, Sertac Karaman, Vivienne Sze
Massachusetts Institute of Technology
{ylliao, sertac<EMAIL_ADDRESS>
###### Abstract

Vision Transformer (ViT) demonstrates that Transformer for natural language
processing can be applied to computer vision tasks and result in comparable
performance to convolutional neural networks (CNN), which have been studied
and adopted in computer vision for years. This naturally raises the question
of how the performance of ViT can be advanced with design techniques of CNN.
To this end, we propose to incorporate two techniques and present ViT-ResNAS,
an efficient multi-stage ViT architecture designed with neural architecture
search (NAS). First, we propose residual spatial reduction to decrease
sequence lengths for deeper layers and utilize a multi-stage architecture.
When reducing lengths, we add skip connections to improve performance and
stabilize training deeper networks. Second, we propose weight-sharing NAS with
multi-architectural sampling. We enlarge a network and utilize its sub-
networks to define a search space. A super-network covering all sub-networks
is then trained for fast evaluation of their performance. To efficiently train
the super-network, we propose to sample and train multiple sub-networks with
one forward-backward pass. After that, evolutionary search is performed to
discover high-performance network architectures. Experiments on ImageNet
demonstrate that ViT-ResNAS achieves better accuracy-MACs and accuracy-
throughput trade-offs than the original DeiT and other strong baselines of
ViT. Code is available at https://github.com/yilunliao/vit-search.

##  1 Introduction

Self-attention and Transformers [43], which originated from natural language
processing (NLP), have been widely adopted in computer vision (CV) tasks,
including image classification [2, 20, 30, 35, 49, 62], object detection [6,
35, 63], and semantic segmentation [46, 47]. Many works utilize hybrid
architectures and incorporate self-attention mechanisms into convolutional
neural networks (CNN) to model long-range dependence and improve the
performance of networks. On the other hand, Vision Transformer (ViT) [10]
demonstrates that a pure transformer without convolution can achieve
impressive performance on image classification when trained on large datasets
like JFT-300M [38]. Additionally, DeiT [41] shows that ViT can outperform CNN
when trained on ImageNet [33] with stronger regularization. It is appealing to
have powerful Transformers for CV tasks since it enables using the same type
of neural architecture for applications in both CV and NLP domains.

A parallel line of research is to design efficient neural networks with neural
architecture search (NAS) [4, 5, 13, 19, 36, 37, 39, 45, 51, 54, 56, 57, 64,
65]. Pioneering works use reinforcement learning to design efficient CNN
architectures. They sample many networks in a pre-defined search space and
train them from scratch for a few epochs to approximate their performance,
which requires expensive computation. To accelerate the process, weight-
sharing NAS has become popular. Instead of training individual networks in a
search space, weight-sharing NAS trains a super-network whose weights are
shared across all networks in the search space. Once the super-network is
trained, we can directly use its weights to approximate the performance of
different networks in the search space. These methods successfully result in
CNN architectures outperforming manually designed ones.

While CNN architectures have been studied and adopted in CV for years and
optimized with NAS, recently ViT demonstrates superior performance over CNN in
some scenarios. Despite its promising performance, ViT adopts the same
architecture as Transformer for NLP [43]. This naturally leads to the question
of how the performance of ViT can be further advanced with design techniques
of CNN. Therefore, in this work, we propose to incorporate two design
techniques of CNN, which are spatial reduction and neural architecture search,
and present ViT-ResNAS, an efficient multi-stage ViT architecture with
residual spatial reduction and designed with NAS.

First, we propose residual spatial reduction to decrease sequence lengths and
increase embedding sizes for deeper transformer blocks. As illustrated in Fig.
1, we transform the original single-stage architecture into a multi-stage one,
with each stage having the same sequence length and embedding size. This
architecture follows that of CNN, where the resolution of feature maps
decreases and the channel size increases for deeper layers. Additionally, we
add skip connections when reducing sequence lengths, which can further improve
performance and stabilize training deeper networks. ViT with residual spatial
reduction is named ViT-Res. Second, we propose weight-sharing neural
architecture search with multi-architectural sampling to improve the
architecture of ViT-Res as shown in Fig. 2. We enlarge ViT-Res network by
increasing its depth and width. Its sub-networks are utilized to define a
search space. Then, a super-network covering all sub-networks is trained to
directly evaluate their performance. For each training iteration and given a
batch of examples, we sample and train multiple sub-networks with one forward-
backward pass to efficiently train the super-network. Once the super-network
is trained, evolutionary search [31, 32] is applied to discover high-
performance ViT-ResNAS networks. Experiments on ImageNet [33] demonstrate the
effectiveness of our proposed ViT-ResNAS. Compared to the original DeiT [41],
ViT-ResNAS-Tiny achieves 8.6% higher accuracy than DeiT-Ti with slightly
higher multiply-accumulate operations (MACs), and ViT-ResNAS-Small achieves
similar accuracy to DeiT-B while having $6.3\times$ less MACs and $3.7\times$
higher throughput. Additionally, ViT-ResNAS achieves better accuracy-MACs and
accuracy-throughput trade-offs than other strong baselines of ViT such as PVT
[48] and PiT [17].

Our main contributions are as follows: (1) We propose residual spatial
reduction to improve the efficiency of ViT. (2) We propose weight-sharing NAS
with multi-architectural sampling to improve ViT with residual spatial
reduction (ViT-Res). (3) Experiments on ImageNet demonstrate that our ViT-
ResNAS achieves comparable accuracy-MAC trade-offs to previous works.

Figure 1: Architecture of ViT-Res. We propose residual spatial reduction
(light orange) to reduce sequence length and increase embedding size for
deeper blocks, which divides the network into several stages. Each stage has
the same sequence length and embedding size and consists of several
transformer blocks. All stages except the first one start with residual
spatial reduction blocks. Figure 2: Algorithm flow of NAS. First, we train a
ViT-Res super-network with multi-architectural sampling. The performance of
sub-networks can be directly evaluated using the super-network’s trained
weights without further training. Then, we perform evolutionary search to find
high-performance sub-networks. Finally, the best sub-network becomes our
searched network and is trained from scratch to convergence.

##  2 Method

We first review the architecture of Vision Transformer (ViT) [10]. Then, we
discuss residual spatial reduction and weight-sharing NAS with multi-
architectural sampling to improve its architecture. Other extra techniques
that help improve performance are presented as well.

###  2.1 Background on Vision Transformer

Main components are tokenization, position embedding, multi-head self-
attention (MHSA), feed-forward network (FFN), and layer normalization (LN).
MHSA and FFN form a transformer block.

#### Tokenization.

The input to ViT is a fixed-size image and is split into patches of pixels.
Each patch is transformed into a vector of dimension $d_{embed}$ called patch
embedding or patch token with a linear layer. Thus, an image is viewed as a
sequence of patch embeddings. To perform classification, a learnable vector
called classification token is appended to the sequence. After being processed
with transformer blocks, the feature captured by the classification token is
used to predict classes.

#### Position Embedding.

To preserve relative positions of image patches, position embeddings in the
form of pre-defined or learnable vectors are added to patch embeddings. Then,
joint patch and position embeddings are processed with transformer blocks as
shown on the left of Fig. 1.

#### MHSA.

The attention function in MHSA operates on a sequence of embeddings
$X\in\mathbb{R}^{N\times d_{embed}}$, with $N$ being sequence length and
$d_{embed}$ being embedding dimension. It first generates three vectors, key,
query and value, for each embedding with linear transformations, $W^{K}$,
$W^{Q}$ and $W^{V}$, and packs them into matrices $K$, $Q$, and $V$
$\in\mathbb{R}^{N\times d}$. Then, for each query corresponding to an
embedding, we calculate its inner products with all $N$ keys, divide each
product by $\sqrt{d}$, and apply softmax function over all products to obtain
$N$ normalized weights. The output of the attention function for the embedding
is the weighted sum of $N$ value vectors. The process is conducted for all
embeddings and the attention function results in an output as follows:

$\displaystyle\text{Attn}(X,W^{K},W^{Q},W^{V})=\text{softmax}(\frac{QK^{T}}{\sqrt{d}})V\text{,
where }K=XW^{K},Q=XW^{Q},V=XW^{V}$ (1)

One MHSA contains $h$ parallel attention functions or attention heads. The $h$
different outputs are concatenated and projected with a linear transformation
to produce the output of MHSA.

#### FFN.

It consists of two linear layers separated by GeLU [16]. The first layer
expands the dimension of each embedding to hidden size $d_{hidden}$ and the
second projects the dimension back to $d_{embed}$.

#### LN.

Layer normalization [1] normalizes each embedding $x$ in the sequence $X$
separately as follows:

$\text{LN}(x)=\frac{x-\mu}{\sigma}\circ\gamma+\beta$ (2)

$\mu$ and $\sigma$ are the mean and standard deviation of $x$ and the
calculation is batch-independent.

###  2.2 Residual Spatial Reduction

ViT maintains the same sequence length throughout the network, which is
similar to maintaining the same resolution of feature maps. In contrast, CNNs
decrease the resolution of feature maps and increase the channel size for
deeper layers. The former is based on the fact that there is spatial
redundancy in feature maps, and the latter is to assign more channels to high-
level features in deeper layers and to maintain similar computation when
resolution is decreased. The design technique has been widely adopted in high-
performance CNNs [15, 29], which motivates whether it can be introduced to
improve the efficiency of ViT as well. To this end, we propose residual
spatial reduction. As illustrated in Fig. 3, since there is spatial
relationship in patch embeddings, we reshape the 1D sequence into a 2D feature
map and then apply layer normalization [1] and strided convolution (“Norm” and
“Conv” in Fig. 3) to reduce the sequence length of patch embeddings and
increase embedding dimension. Since the sequence length is changed, we update
the relative position information by adding new position embeddings. To
maintain the same channel size of all embeddings, we apply layer normalization
and a linear layer (“Norm” and “Linear” in Fig. 3) to the classification
token. These constitute the residual branch.

Figure 3: Structure of residual spatial reduction. We use strided convolution
to reduce sequence length of patch embeddings. A skip connection as shown in
the lower branch is added to stabilize training and improve performance.

Although introducing only the residual branch can significantly improve the
accuracy-MAC trade-offs, training deeper networks can be unstable.
Specifically, under the training setting of DeiT-distill [41], using the
residual branch significantly improves the accuracy of DeiT-Tiny from 74.5% to
79.6% with little increase in MACs. However, training loss can become “NaN”
when training deeper networks like our ViT-Res super-networks. To remedy this,
we introduce an extra skip connection [15] without any learnable operations as
motivated by the residual structure of transformer blocks [43, 53]. We use 2D
average pooling to reduce sequence length (“Average pooling” in Fig. 3) and
concatenate embeddings with zero tensors (“Zero pad” in Fig. 3) to increase
embedding dimension. The structure results in the main branch and helps to
stabilize training and improve performance.

Figure 4: Searchable dimensions in each stage of ViT-Res super-network.

The residual and main branches form residual spatial reduction. We insert 2
residual spatial reduction blocks into ViT with 12 blocks and divide the
architecture evenly into 3 stages. Following the design rule of ResNet [15],
we double embedding dimension when halving spatial resolution or reducing
sequence length of patch embeddings by $4\times$. The resulting ViT
architecture is called ViT-Res as illustrated in Fig. 1. Since the spatial
resolution is to be reduced by $4\times$ and input images are of size
$224\times 224$, the patch size in tokenization before the first stage is set
to $14$, resulting in spatial resolution $16$ ($=\frac{224}{14}$) and
$16\times 16$ patch embeddings in the first stage. After passing through two
residual spatial reduction blocks, the sequence length of patch embeddings is
reduced to $4\times 4$. Please refer to Appendix A for further details.

###  2.3 Weight-Sharing NAS with Multi-Architectural Sampling

Another design technique used in CNNs is neural architecture search (NAS). Due
to its efficiency, we adopt weight-sharing NAS to improve the archtiecture of
ViT-Res in terms of designing better architectural parameters like numbers of
attention heads and transformer blocks as shown in Fig. 4.

#### Algorithm Overview.

We enlarge ViT-Res network by increasing its depth and width. Sub-networks
with smaller depths and widths define a search space. Then, a super-network
covering all sub-networks is trained to directly evaluate their performance.
For each training iteration and given a batch of examples, we propose to
sample multiple different sub-networks with one forward-backward pass to
efficiently train the super-network. After that, evolutionary search is
applied to discover architectures of high-performance sub-networks satisfying
some resource constraints like MACs. Finally, the best sub-network evaluated
becomes the searched network and is trained from scratch.

#### Search Space.

We construct a large network by uniformly increasing depth and width of all
stages of ViT-Res, and sub-networks contained in the large network define a
search space. As shown in Fig. 4, for each stage, we search embedding
dimension and the number of transformer blocks. For different blocks, we
search different numbers of attention heads $h$ in MHSA and different hidden
sizes $d_{hidden}$ in FFN. The range of each searchable dimension is pre-
defined. During super-network training and evolutionary search, sub-networks
of different configurations are sampled for training and evaluation. Details
of search space are presented in Appendix B.

Figure 5: Multi-architectural sampling for super-network training. Left: one
sub-network is sampled and trained with one forward-backward pass given a
batch of examples. Right: multi-architectural sampling samples $N_{a}$ ($=3$)
different sub-networks with one forward-backward pass.

#### Multi-Architectural Sampling for Super-Network Training.

To evaluate the performance of sub-networks, their weights have to be
optimized to reflect their relative quality. Therefore, we train a super-
network whose architecture is the same as the largest network defining the
search space and whose weights are shared across sub-networks. During super-
network training, we sample and train different sub-networks for different
training iterations. Generally, the more sub-networks are sampled, the more
accurate the relative quality of sub-networks indicated by the trained super-
network can be. Previous works on NAS for CNNs [4, 13, 57] sample and train a
single sub-network with one forward-backward pass for each training iteration.
Given a fixed amount of iterations, this, however, results in room for
improvement when training ViT-Res super-networks as we can sample multiple
sub-networks with one forward-backward pass for each training iteration as
illustrated in Fig. 5.

Unlike batch normalization [22] in CNNs, layer normalization (LN) [1] in ViT
avoids normalizing along batch dimension, which enables sampling different
sub-networks with one forward-backward pass without mixing their statistics.
Specifically, for each training iteration, we sample $N_{a}$ sub-networks and
divide a batch into $N_{a}$ sub-batches. Each sub-network is trained with its
corresponding sub-batch. This can be achieved efficiently with one forward-
backward pass and channel masking (ordered dropout) [44, 55, 57]. As shown in
Fig. 6 (a), different masks are generated for different feature maps
corresponding to different examples in order to zero out different channels
and simulate sampling different sub-networks. The shapes of feature maps
remain the same, which maintains regular batch computation and therefore
enables a single forward-backward pass.

Additionally, we re-calibrate the statistics in LN when sampling multiple sub-
networks in order to prevent discrepancy between super-network training and
standard training. As shown in Fig. 6 (b) and (c), LN can incorrectly
normalize over a larger channel dimension when sampling networks with smaller
channel sizes. This is because channel masking only changes the values of
feature maps not shapes. To avoid the problem, we propose masked layer
normalization (MLN), which re-calibrates the statistics of LN when its input
is masked. Instead of only looking at the shape of input tensors, MLN
calculates the ratio of the number of masked channels to the total number of
channels and uses that ratio to re-calibrate statistics as in Fig. 6 (d). With
MLN, the statistics become the same as we train sub-networks separately as in
Fig. 6 (b). Eventually, with channel masking and MLN, we can sample multiple
sub-networks with one forward-backward pass, which improves sample efficiency
when training ViT-Res super-networks and therefore the performance of searched
networks. Other details of super-network training can be found in Appendix C.

Figure 6: Channel masking and masked layer normalization for multi-
architectural sampling. (a) We simulate sampling different sub-networks by
masking different channels. (b) The channel dimension LN normalizes over when
training a sub-network alone. (c) The channel dimension LN incorrectly
normalizes over when sampling a sub-network during super-network training. (d)
MLN re-calibrates statistics by considering the number of masked channels.

#### Evolutionary Search.

Throughout the search process, we only sample sub-networks satisfying pre-
defined resource constraints (e.g., MACs) and evaluate their performance
(e.g., accuracy) with trained super-network weights. Evolutionary search
maintains a population of networks and refines top-performing ones for many
iterations. We start with an initial population of $P_{0}$ randomly sampled
sub-networks. At every iteration, $N_{parent}$ sub-networks with the highest
performance in the population serve as parent networks that generate
$N_{child}$ new sub-networks through mutation and crossover. For mutation, we
randomly select one sub-network from parent networks and modify every
architectural parameter with a pre-defined probability $p_{mutate}$. For
crossover, we choose two random sub-networks from parent networks and randomly
fuse every architectural parameter. Mutation and crossover generates the same
amount of new sub-networks, and they are added to the population. The process
is repeated for $T_{search}$ times, and the best sub-network in the population
becomes the searched network and is trained from scratch to convergence.

###  2.4 Extra Techniques

We present other techniques in training and architecture that help improve
performance below.

#### Token Labeling with CutMix and Mixup.

We incorporate token labeling [11, 23] and propose to improve it with Mixup
[60]. Token labeling provides labels for all patches in an input image. This
enables training ViT to predict class labels of all patch tokens (patch
embeddings) in addition to predicting class label of an input image with
classification token and can improve training Transformers [8, 11, 23]. Token
labeling generates an input image, its class label, and patch-wise class
labels by patch-wise CutMix [59] described as follows. Given two images
$x_{1},x_{2}$ and their class labels $y_{1},y_{2}$, we divide each image into
$K$ patches, flatten the patches, and associate each patch with its original
image-level class label. In this way, we have sequences of flattened patches
$X_{1}=[x_{1}^{1},...,x_{1}^{K}]$ and $X_{2}=[x_{2}^{1},...,x_{2}^{K}]$ and
sequences of patch-wise labels $Y_{1}=[y_{1},...,y_{1}]$ and
$Y_{2}=[y_{2},...,y_{2}]$. Then, we randomly generate a binary mask $M$ to
combine the images and associated patch-wise labels: $X=X_{1}\odot
M+X_{2}\odot(1-M)$ and $Y=Y_{1}\odot M+Y_{2}\odot(1-M)$. We restructure the 1D
sequence $X$ into a 2D image $x$. The class label $y$ associated with $x$ is
obtained by combining the original class labels: $y=\lambda
y_{1}+(1-\lambda)y_{2}$ where $\lambda$ is the average of elements in the
binary mask $M$. For ViT-Res, we choose $K=4\times 4$, which is equal to the
sequence length of patch embeddings in the last stage.

Additionally, we find that applying Mixup [60] along with token labeling can
improve performance, which is contrary to previous results [11, 23]. We
surmise that they first apply Mixup and then perform patch-wise CutMix could
lead to noisy training images and labels. In contrast, we improve token
labeling with Mixup through switching between patch-wise CutMix and Mixup.
Specifically, we choose either patch-wise CutMix or Mixup to generate an
image, its label and patch-wise labels. When the former is selected, we follow
the practice as mentioned above. When the latter is chosen, the image and its
label are generated in the same way as Mixup, and this image-level label is
assigned to all patches to produce the patch-wise labels.

#### Convolution before Tokenization.

Following previous works [10, 12, 23], we add some convolutional layers before
tokenization. Specifically, we add three extra $3\times 3$ convolutional
layers, each with $C$ output channels. The stride of the first layer is set to
$2$, and others are $1$. A residual skip connection is added between outputs
of the first and the third layers. The computation of these convolutional
layers is relatively low. Since the spatial resolution is reduced by
$2\times$, the patch size in tokenization is set to $7$ ($=\frac{14}{2}$).
Further details can be found in Appendix A and B.

##  3 Experiments

In this section, we first describe the experiment setup. Then, we conduct
experiments to study the effectiveness of our proposed methods. Finally, the
comparison with related works is presented.

###  3.1 Experiment Setup

The dataset used is ImageNet [33]. Our implementation is based on timm library
[50] and that of DeiT [41]. Most of the training settings follow those in DeiT
[41] except that we do not use repeated augmentation [18]. We train models
with $8$ GPUs for $300$ epochs, and batch size is $1024$ and input resolution
is $224\times 224$. As for token labeling, we adopt the same loss function as
previous works [11, 23], and the associated loss is directly added to the
original classification loss without any scaling. Unless otherwise stated, we
incorporate token labeling with patch-wise CutMix and Mixup to train our
networks and include several convolutional layers before tokenization. For
ablation study presented in Section 3.2, when token labeling is used, drop
path [21] rate is increased from $0.1$ to $0.2$.

For neural architecture search, we increase the depth and width of our ViT-Res
network to build search spaces and search the architectures of ViT-ResNAS
networks of three different sizes, which we name Tiny, Small, and Medium. For
ViT-ResNAS-Tiny, we enlarge ViT-Res-Tiny to build a super-network and train
both the super-network and the searched network with drop path rate $0.2$. For
ViT-ResNAS-Small and Medium, we further enlarge ViT-Res-Tiny super-network and
share the same search space. The super-network is trained with drop path rate
$0.3$. The drop path rates for Small and Medium are $0.3$ and $0.4$,
respectively. The details of the search space are presented in Table B in
appendix. A super-network is trained for $120$ epochs, with other settings
being the same as mentioned above. We sample $N_{a}$ sub-networks with one
forward-backward pass during super-network training, experiment with different
values of $N_{a}$ and empirically set $N_{a}$ to $16$. For evolutionary
search, the resource constraint is MAC. We set search iteration
$T_{search}=20$, the number of parent networks $N_{parent}=75$, the initial
population size $P_{0}=500$, the number of new sub-networks for each iteration
$N_{child}=150$ and mutation probability $p_{mutate}=0.3$.

In addition, following DeiT [41], we fine-tune networks at larger resolutions
to obtain higher capacity. A network is fine-tuned for $30$ epochs, with batch
size $512$, learning rate $5\times 10^{-6}$, weight decay $10^{-8}$ and drop
path rate $0.75$.

Methods | Top-1 Acc.
---|---
Spatial reduction | 78.5
Residual spatial reduction | 78.8 (+0.3)
\+ Token labeling with CutMix | 79.6 (+0.8)
\+ Token labeling with CutMix & Mixup | 80.1 (+0.5)

Table 1: Additive study of improving multi-stage network with residual
connections and token labeling. We start with DeiT-Tiny with spatial reduction
and progressively introduce residual connections and token labeling. Adding
residual connections and combining token labeling with Mixup in the proposed
manner improve accuracy without increasing MACs.

###  3.2 Ablation Study

#### Multi-Stage Network with Residual Connection and Token Labeling.

We study how the performance of vanilla multi-stage networks can be enhanced
with the proposed residual connections and improved token labeling training.
We build such a network by starting with DeiT-Tiny network, adding three
convolutional layers before tokenization and inserting two spatial reduction
blocks (i.e., only the residual branch of residual spatial reduction). The
results are shown in Table. 1. Without token labeling, introducing only two
residual connections (i.e., main branch of residual spatial reduction) can
improve accuracy from $78.5\%$ to $78.8\%$ with negligible overhead. When
token labeling is used, incorporating Mixup in the proposed manner can further
improve the accuracy by $0.5\%$. With residual spatial reduction and token
labeling, we can improve the accuracy by $1.6\%$ without increasing MACs, and
our ViT-Res-Tiny achieves $80.1\%$ top-1 accuracy with $1.8$G MACs.

We also study the effectiveness of residual connections in the proposed
residual spatial reduction under another training setting. When using the
training setting of DeiT-distill [41], which includes repeated augmentation
and distilling knowledge of CNN, residual connections can improve the accuracy
from $79.6\%$ to $80.1\%$. Moreover, when we train deeper networks such as our
ViT-Res-Tiny super-network, without residual connections, the training can be
unstable with “NaN” training loss.

Number of sampled sub-networks per forward-backward pass ($N_{a}$) | Single-arch. | Multi-arch.
---|---|---
1 | 8 | 16 | 32
Top-1 accuracy (%) | 80.5 |  | 80.6
---
(+0.1)
| 80.8
---
(+0.3)
| 80.6
---
(+0.1)

Table 2: Effect of numbers of sampled sub-networks with one forward-backward
pass ($N_{a}$) on the performance of searched networks. The type of resource
constraint is MAC, and all searched networks have MAC around $1.8$G.
Empirically, $N_{a}=16$ results in the best searched network.

#### Weight-Sharing NAS with Multi-Architectural Sampling.

With the proposed residual spatial reduction and token labeling, we study how
weight-sharing NAS with multi-architectural sampling can further improve the
performance of ViT-Res. During super-network training, we sample and train
$N_{a}$ sub-networks with one forward-backward pass. Given the same amount of
training iterations and training examples, the value of $N_{a}$ controls the
trade-offs between sample efficiency (i.e., how many sub-networks are sampled)
and the quality of training each sub-network (i.e., how many examples are used
to train it). Therefore, instead of arbitrarily choosing a large value, we
experiment with different $N_{a}$. We use our ViT-Res-Tiny super-network to
study the effect of different $N_{a}$ on the performance of searched networks
and search for networks with MACs around $1.8$G. The results are presented in
Table. 2. Compared to sampling one sub-network for each training iteration
(“Single-arch.” in Table 2), sampling multiple sub-networks leads to better
searched networks. Among different values, $N_{a}=16$ empirically results in
the best searched network. Please note that with NAS, the top-1 accuracy is
increased from $80.1\%$ to $80.5\%$ and that with the proposed multi-
architectural sampling, the accuracy is further increased from $80.5\%$ to
$80.8\%$. Based on the results, we choose $N_{a}=16$ to design larger ViT-
ResNAS networks.

###  3.3 Comparison with Related Works

Method | Model Size (M) | MACs (G) |  | Top-1
---
Accuracy (%)
| Throughput
---
(images/second)
DeiT-Ti [41] | 5 | 1.3 | 72.2 | 1968
T2T-ViT-12 [58] | 7 | 2.2 | 76.5 | 1192
PiT-XS [17] | 11 | 1.4 | 78.1 | 1647
PVT-Tiny [48] | 13 | 1.9 | 75.1 | 1133
ViL-Tiny [61] | 7 | 1.3 | 76.3 | 754
ViT-Res-Tiny | 43 | 1.8 | 80.1 | 1807
ViT-ResNAS-Tiny | 41 | 1.8 | 80.8 | 1579
DeiT-Small [41] | 22 | 4.6 | 79.9 | 846
T2T-ViT-14 [58] | 22 | 5.2 | 81.5 | 682
PiT-S [17] | 24 | 2.9 | 80.9 | 986
PVT-Small [48] | 25 | 3.8 | 79.8 | 628
PVT-Medium [48] | 44 | 6.7 | 81.2 | 407
TNT-S [14] | 24 | 5.2 | 81.5 | 353
Swin-T [26] | 29 | 4.5 | 81.3 | 600
Twins-PCPVT-S [7] | 24 | 3.7 | 81.2 | 622
ViT-ResNAS-Small | 65 | 2.8 | 81.7 | 1084
DeiT-Base [41] | 86 | 17.6 | 81.8 | 290
T2T-ViT-19 [58] | 39 | 8.9 | 81.9 | 428
PiT-B [17] | 74 | 12.5 | 82.0 | 316
PVT-Large [48] | 61 | 9.8 | 81.7 | 284
ViL-Small [61] | 25 | 4.9 | 82.0 | 310
CvT-13 [52] | 20 | 4.5 | 81.6 | 587
ViT-ResNAS-Medium | 97 | 4.5 | 82.4 | 751
Swin-S [26] | 50 | 8.7 | 83.0 | 351
Twins-PCPVT-B [7] | 44 | 6.4 | 82.7 | 403
CvT-21 [52] | 32 | 7.1 | 82.5 | 379
ViT-ResNAS-Medium $\uparrow 280$ | 97 | 7.1 | 83.1 | 467
ViL-Medium-Wide [61] | 40 | 11.3 | 82.9 | 177
Twins-PCPVT-L [7] | 61 | 9.5 | 83.1 | 282
CaiT-S36 [42] | 68 | 13.9 | 83.3 | 191
ViT-ResNAS-Medium $\uparrow 336$ | 97 | 10.6 | 83.5 | 292
Swin-B [26] | 88 | 15.4 | 83.3 | 243
CaiT-XS24 $\uparrow 384$ [42] | 27 | 19.3 | 83.8 | 57
ViT-ResNAS-Medium $\uparrow 392$ | 97 | 15.2 | 83.8 | 194

Table 3: Comparison with related works on ViT. “$\uparrow R$” denotes that the
model is first trained at resolution $224$ and then fine-tuned at resolution
$R$. Other models are trained at resolution $224$. Throughput is measured on
one Titan RTX GPU with batch size $128$.

We design our ViT-ResNAS-Tiny, Small and Medium networks with NAS and MACs as
our search metric. Following DeiT [41], we fine-tune our ViT-ResNAS-Medium at
larger resolutions to obtain models with higher capacity. Since the patch size
before the first stage is $14$ and there are $2$ residual spatial reduction in
the network, the spatial resolution is reduced by $56$ $(=14\times 2\times 2)$
times in the last stage, and therefore we can only increase the input
resolution by multiples of $56$. We also report the performance of fine-tuning
ViT-ResNAS-Medium at resolutions $280$, $336$ and $392$. When comparing
different models, following the implementation of Swin Transformer [26], we
measure the inference throughput on a single Titan RTX GPU with batch size 128
as well.

Table 3 summarizes the comparison with previous works on ViT. Please note that
why our models have more parameters is that we further reduce the sequence
length and increase the channel size of DeiT models. Even though we have more
parameters, our ViT-ResNAS networks consistently achieves better accuracy-MACs
trade-offs as well as accuracy-throughput trade-offs.

Compared to the original single-stage DeiT [41], ViT-ResNAS-Tiny achieves
$8.6\%$ higher accuracy than DeiT-Ti while having only $0.5$G higher MACs and
slightly lower throughput and has $0.9\%$ higher accuracy than DeiT-Small with
$2.5\times$ less MACs and $1.9\times$ higher throughput. ViT-ResNAS-Small
obtains the similar level of accuracy to DeiT-Base while having $6.3\times$
lower MACs and $3.7\times$ higher throughput. Compared to other works on
multi-stage architectures like PVT [48] and PiT [17], ViT-ResNAS consistently
has better accuracy-MACs and accuracy-throughput trade-offs. The computation
saving becomes more apparent as accuracy becomes higher. For example, in
comparison to PiT-S [17], ViT-ResNAS-Tiny achieves similar accuracy while
having $1.6\times$ less MACs and $1.6\times$ higher throughput. When compared
with PiT-B [17], ViT-ResNAS-Medium achieves $0.4\%$ higher accuracy with
$2.8\times$ less MACs and $2.4\times$ higher throughput. This suggests the
effectiveness of NAS to scale up models.

Additionally, for lower MACs around 2.0G, ViT-ResNAS-Tiny achieves better
accuracy-MACs trade-offs than ViL-Tiny. However, when MACs are around $5.0$G,
ViT-ResNAS-Medium is on par with ViL-Small in terms of accuracy-MACs trade-
offs. This probably suggests that utilizing efficient attention mechanisms [7,
26, 61] with convolution-like locality to process high-resolution features is
necessary for models in higher MAC regimes to generalize better. Nevertheless,
those methods are orthogonal to our approaches, and the proposed weight-
sharing NAS with multi-architectural sampling could further improve their
accuracy-MACs trade-offs as well.

##  4 Related Works

#### Vision Transformers.

Vision Transformer (ViT) [10] demonstrates that a pure transformer without
convolution can perform well on image classification when trained on large
datasets like JFT-300M [38]. To make it data-efficient, DeiT [41] uses strong
regularization and adds a distillation token to its architecture for knowledge
distillation, and demonstrates comparable performance when trained on ImageNet
[33] only. Subsequent works improve the performance of ViT on ImageNet through
either better training [11, 23] or architectures [9, 14, 17, 26, 27, 42, 48,
52, 58, 61]. For example, they bring locality into network architectures by
using convolutions [48, 52] or efficient local attention [7, 14, 26, 61] or
similarly adopt multi-stage architectures [17, 27, 48, 52, 61]. Our proposed
methods are complementary to these works. Residual spatial reduction can
derive a more efficient multi-stage architectures, and weight-sharing NAS with
multi-architectural sampling can be applied to further improve performance
since they use layer normalization [1] as well.

#### Neural Architecture Search.

There have been increasingly interest in designing efficient architectures
with neural architecture search (NAS) [3, 4, 5, 13, 19, 24, 31, 32, 34, 36,
37, 39, 40, 45, 51, 54, 55, 56, 57, 64, 65]. Among different methods, weight-
sharing NAS [3, 4, 5, 13, 24, 25, 28, 36, 37, 45, 51, 56, 57] has become
popular due to efficiency. They train one over-parametrized super-network
whose weights are shared across all networks in a search space to conduct
architecture search, which saves computation cost significantly. However, many
of these works focus on CNN, which has been researched for years, and have
well-designed search space. In contrast, our proposed NAS with multi-
architectural sampling focuses on multi-stage ViT, which is much less studied,
and we utilize its batch-independent property to further improve the
performance of searched networks.

##  5 Conclusion

We incorporate two design techniques of CNN, which are spatial reduction and
NAS, into ViT and present ViT-ResNAS, an efficient multi-stage ViT designed
with NAS. The proposed residual spatial reduction enhances accuracy-MACs
trade-offs significantly and the residual connections can improve performance
and stabilize training. Weight-sharing NAS with multi-architectural sampling
trains super-networks more efficiently and results in better searched
networks. Experiments on ImageNet demonstrates the effectiveness of our
methods and the efficiency of our ViT-ResNAS networks.

## References

  * [1] Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton.  Layer normalization.  arXiv preprint arXiv:1607.06450, 2016.
  * [2] Irwan Bello, Barret Zoph, Ashish Vaswani, Jonathon Shlens, and Quoc V. Le.  Attention augmented convolutional networks.  In International Conference on Computer Vision (ICCV), 2019.
  * [3] Gabriel Bender, Hanxiao Liu, Bo Chen, Grace Chu, Shuyang Cheng, Pieter-Jan Kindermans, and Quoc V. Le.  Can weight sharing outperform random architecture search? an investigation with tunas.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
  * [4] Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han.  Once for all: Train one network and specialize it for efficient deployment.  In International Conference on Learning Representations (ICLR), 2020\.
  * [5] Han Cai, Ligeng Zhu, and Song Han.  ProxylessNAS: Direct neural architecture search on target task and hardware.  In International Conference on Learning Representations (ICLR), 2019\.
  * [6] 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 (ECCV), 2020.
  * [7] Xiangxiang Chu, Zhi Tian, Yuqing Wang, Bo Zhang, Haibing Ren, Xiaolin Wei, Huaxia Xia, and Chunhua Shen.  Twins: Revisiting the design of spatial attention in vision transformers.  arXiv preprint arXiv:2104.13840, 2021.
  * [8] Kevin Clark, Minh-Thang Luong, Quoc V. Le, and Christopher D. Manning.  ELECTRA: Pre-training text encoders as discriminators rather than generators.  In International Conference on Learning Representations (ICLR), 2020\.
  * [9] Stéphane d’Ascoli, Hugo Touvron, Matthew Leavitt, Ari Morcos, Giulio Biroli, and Levent Sagun.  Convit: Improving vision transformers with soft convolutional inductive biases.  arXiv preprint arXiv:2103.10697, 2021.
  * [10] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  In International Conference on Learning Representations (ICLR), 2021\.
  * [11] Chengyue Gong, Dilin Wang, Meng Li, Vikas Chandra, and Qiang Liu.  Improve vision transformers training by suppressing over-smoothing.  arXiv preprint arXiv:2104.12753, 2021.
  * [12] Benjamin 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 arXiv:22104.01136, 2021.
  * [13] Zichao Guo, Xiangyu Zhang, Haoyuan Mu, Wen Heng, Zechun Liu, Yichen Wei, and Jian Sun.  Single path one-shot neural architecture search with uniform sampling.  In European Conference on Computer Vision (ECCV), 2020.
  * [14] Kai Han, An Xiao, Enhua Wu, Jianyuan Guo, Chunjing Xu, and Yunhe Wang.  Transformer in transformer.  arXiv preprint arXiv:2103.00112, 2021.
  * [15] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
  * [16] Dan Hendrycks and Kevin Gimpel.  Gaussian error linear units (gelus).  arXiv preprint arXiv:1606.08415, 2016.
  * [17] Byeongho Heo, Sangdoo Yun, Dongyoon Han, Sanghyuk Chun, Junsuk Choe, and Seong Joon Oh.  Rethinking spatial dimensions of vision transformers.  arXiv preprint arXiv:2103.16302, 2021.
  * [18] Elad Hoffer, Tal Ben-Nun, Itay Hubara, Niv Giladi, Torsten Hoefler, and Daniel Soudry.  Augment your batch: Improving generalization through instance repetition.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
  * [19] Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, Quoc V. Le, and Hartwig Adam.  Searching for mobilenetv3.  In International Conference on Computer Vision (ICCV), 2019.
  * [20] Jie Hu, Li Shen, Samuel Albanie, Gang Sun, and Enhua Wu.  Squeeze-and-excitation networks.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 42(8):2011–2023, 2020.
  * [21] Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q. Weinberger.  Deep networks with stochastic depth.  In European Conference on Computer Vision (ECCV), 2016.
  * [22] Sergey Ioffe and Christian Szegedy.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In International Conference on Machine Learning (ICML), 2015.
  * [23] Zihang Jiang, Qibin Hou, Li Yuan, Daquan Zhou, Xiaojie Jin, Anran Wang, and Jiashi Feng.  Token labeling: Training a 85.5% top-1 accuracy vision transformer with 56m parameters on imagenet.  arXiv preprint arXiv:2104.10858, 2021.
  * [24] Changlin Li, Tao Tang, Guangrun Wang, Jiefeng Peng, Bing Wang, Xiaodan Liang, and Xiaojun Chang.  Bossnas: Exploring hybrid cnn-transformers with block-wisely self-supervised neural architecture search.  arXiv preprint arXiv:2103.12424, 2021.
  * [25] Hanxiao Liu, Karen Simonyan, and Yiming Yang.  Darts: Differentiable architecture search.  In International Conference on Learning Representations (ICLR), 2019\.
  * [26] 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.14030, 2021.
  * [27] Zizheng Pan, Bohan Zhuang, Jing Liu, Haoyu He, and Jianfei Cai.  Scalable visual transformers with hierarchical pooling.  arXiv preprint arXiv:2103.10619, 2021.
  * [28] Hieu Pham, Melody Guan, Barret Zoph, Quoc Le, and Jeff Dean.  Efficient neural architecture search via parameters sharing.  In International Conference on Machine Learning (ICML), 2018.
  * [29] Ilija Radosavovic, Raj Prateek Kosaraju, Ross B. Girshick, Kaiming He, and Piotr Dollár.  Designing network design spaces.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
  * [30] Prajit Ramachandran, Niki Parmar, Ashish Vaswani, Irwan Bello, Anselm Levskaya, and Jonathon Shlens.  Stand-alone self-attention in vision models.  arXiv preprint arXiv:1906.05909, 2019.
  * [31] Esteban Real, Alok Aggarwal, Yanping Huang, and Quoc V. Le.  Regularized evolution for image classifier architecture search.  In AAAI Conference on Artificial Intelligence (AAAI), 2019.
  * [32] Esteban Real, Sherry Moore, Andrew Selle, Saurabh Saxena, Yutaka Leon Suematsu, Jie Tan, Quoc V. Le, and Alexey Kurakin.  Large-scale evolution of image classifiers.  In International Conference on Machine Learning (ICML), 2017.
  * [33] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei.  ImageNet Large Scale Visual Recognition Challenge.  International Journal of Computer Vision (IJCV), 115(3):211–252, 2015.
  * [34] David So, Quoc Le, and Chen Liang.  The evolved transformer.  In International Conference on Machine Learning (ICML), 2019.
  * [35] Aravind Srinivas, Tsung-Yi Lin, Niki Parmar, Jonathon Shlens, Pieter Abbeel, and Ashish Vaswani.  Bottleneck transformers for visual recognition.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
  * [36] Dimitrios Stamoulis, Ruizhou Ding, Di Wang, Dimitrios Lymberopoulos, Bodhi Priyantha, Jie Liu, and Diana Marculescu.  Single-path mobile automl: Efficient convnet design and nas hyperparameter optimization.  arXiv preprint arXiv:1907.00959, 2019.
  * [37] 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 4 hours.  arXiv preprint arXiv:1904.02877, 2019.
  * [38] Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta.  Revisiting unreasonable effectiveness of data in deep learning era.  In International Conference on Computer Vision (ICCV), 2017.
  * [39] Mingxing Tan, Bo Chen, Ruoming Pang, Vijay Vasudevan, and Quoc V. Le.  Mnasnet: Platform-aware neural architecture search for mobile.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019.
  * [40] Haotian Tang, Zhijian Liu, Shengyu Zhao, Yujun Lin, Ji Lin, Hanrui Wang, and Song Han.  Searching efficient 3d architectures with sparse point-voxel convolution.  In European Conference on Computer Vision (ECCV), 2020.
  * [41] 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.12877, 2020.
  * [42] Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou.  Going deeper with image transformers.  arXiv preprint arXiv:2103.17239, 2021.
  * [43] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In Conference on Neural Information Processing (NeurIPS), 2017.
  * [44] Alvin Wan, Xiaoliang Dai, Peizhao Zhang, Zijian He, Yuandong Tian, Saining Xie, Bichen Wu, Matthew Yu, Tao Xu, Kan Chen, Peter Vajda, and Joseph E. Gonzalez.  Fbnetv2: Differentiable neural architecture search for spatial and channel dimensions.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
  * [45] Hanrui Wang, Zhanghao Wu, Zhijian Liu, Han Cai, Ligeng Zhu, Chuang Gan, and Song Han.  Hat: Hardware-aware transformers for efficient natural language processing.  In Annual Conference of the Association for Computational Linguistics (ACL), 2020.
  * [46] Huiyu Wang, Yukun Zhu, Hartwig Adam, Alan Yuille, and Liang-Chieh Chen.  Max-deeplab: End-to-end panoptic segmentation with mask transformers.  arXiv preprint arXiv:2012.00759, 2020.
  * [47] Huiyu Wang, Yukun Zhu, Bradley Green, Hartwig Adam, Alan Yuille, and Liang-Chieh Chen.  Axial-deeplab: Stand-alone axial-attention for panoptic segmentation.  In European Conference on Computer Vision (ECCV), 2020.
  * [48] Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao.  Pyramid vision transformer: A versatile backbone for dense prediction without convolutions.  In International Conference on Computer Vision (ICCV), 2021.
  * [49] Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He.  Non-local neural networks.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2018.
  * [50] Ross Wightman.  Pytorch image models.  https://github.com/rwightman/pytorch-image-models, 2019.
  * [51] Bichen Wu, Xiaoliang Dai, Peizhao Zhang, Yanghan Wang, Fei Sun, Yiming Wu, Yuandong Tian, Peter Vajda, Yangqing Jia, and Kurt Keutzer.  Fbnet: Hardware-aware efficient convnet design via differentiable neural architecture search.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019.
  * [52] 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.15808, 2021.
  * [53] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Yan Liu.  On layer normalization in the transformer architecture.  In International Conference on Machine Learning (ICML), 2020.
  * [54] Tien-Ju Yang, Andrew Howard, Bo Chen, Xiao Zhang, Alec Go, Mark Sandler, Vivienne Sze, and Hartwig Adam.  Netadapt: Platform-aware neural network adaptation for mobile applications.  In European Conference on Computer Vision (ECCV), 2018.
  * [55] Tien-Ju Yang, Yi-Lun Liao, and Vivienne Sze.  Netadaptv2: Efficient neural architecture search with fast super-network training and architecture optimization.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
  * [56] Jiahui Yu and Thomas Huang.  Autoslim: Towards one-shot architecture search for channel numbers.  arxiv preprint arxiv:1903.11728, 2019.
  * [57] Jiahui Yu, Pengchong Jin, Hanxiao Liu, Gabriel Bender, Pieter-Jan Kindermans, Mingxing Tan, Thomas Huang, Xiaodan Song, Ruoming Pang, and Quoc Le.  Bignas: Scaling up neural architecture search with big single-stage models.  In European Conference on Computer Vision (ECCV), 2020.
  * [58] Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Francis EH Tay, Jiashi Feng, and Shuicheng Yan.  Tokens-to-token vit: Training vision transformers from scratch on imagenet.  arXiv preprint arXiv:2101.11986, 2021.
  * [59] Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo.  Cutmix: Regularization strategy to train strong classifiers with localizable features.  In International Conference on Computer Vision (ICCV), 2019.
  * [60] Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz.  mixup: Beyond empirical risk minimization.  In International Conference on Learning Representations (ICLR), 2018\.
  * [61] Pengchuan Zhang, Xiyang Dai, Jianwei Yang, Bin Xiao, Lu Yuan, Lei Zhang, and Jianfeng Gao.  Multi-scale vision longformer: A new vision transformer for high-resolution image encoding.  arXiv preprint arXiv:2103.15358, 2021.
  * [62] Hengshuang Zhao, Jiaya Jia, and Vladlen Koltun.  Exploring self-attention for image recognition.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
  * [63] Xizhou Zhu, Weijie Su, Lewei Lu, Bin Li, Xiaogang Wang, and Jifeng Dai.  Deformable detr: Deformable transformers for end-to-end object detection.  In International Conference on Learning Representations (ICLR), 2021\.
  * [64] Barret Zoph and Quoc V. Le.  Neural architecture search with reinforcement learning.  In International Conference on Learning Representations (ICLR), 2017\.
  * [65] Barret Zoph, V. Vasudevan, Jonathon Shlens, and Quoc V. Le.  Learning transferable architectures for scalable image recognition.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2018.

## Appendix

##  Appendix A Network Architecture of ViT-Res

We build our ViT-Res-Tiny network by introducing two modifications to DeiT-
Tiny [41]. First, residual spatial reduction (RSR) is applied to evenly divide
the single-stage architecture into a multi-stage one as described in Sec. 2.2.
For “Conv” in Fig. 3, we use $3\times 3$ convolutions to downsample patch
embeddings. Second, we add three convolutional layers as mentioned in Sec. 2.4
and use $C=24$. Table 4 summarizes the archtiecture of ViT-Res-Tiny.

|  | Output Sequence
---
Length
ViT-Res-Tiny
Tokenization | $16\times 16+1$ | conv-3_24_2
$\left[\begin{tabular}[c]{@{}c@{}}conv-3\\_24\\_1\\\
conv-3\\_24\\_1\end{tabular}\right]$
conv-7_192_7
Stage 1 | $16\times 16+1$ | $\left[\begin{tabular}[c]{@{}c@{}}MHSA-64\\_3\\\ FFN-768\end{tabular}\right]\times 4$
Stage 2 | $8\times 8+1$ | RSR-384
$8\times 8+1$ | $\left[\begin{tabular}[c]{@{}c@{}}MHSA-64\\_6\\\ FFN-1536\end{tabular}\right]\times 4$
Stage 3 | $4\times 4+1$ | RSR-768
$4\times 4+1$ | $\left[\begin{tabular}[c]{@{}c@{}}MHSA-64\\_12\\\ FFN-3072\end{tabular}\right]\times 4$

Table 4: Architecture of ViT-Res-Tiny. “conv-$k$_$C$_$s$” stands for a
$k\times k$ convolutional layer with output channel $C$ and stride $s$.
“MHSA-$d$_$h$” is MHSA with head dimension $d$ and $h$ attention heads.
“FFN-$d_{hidden}$” is FFN with hidden size $d_{hidden}$. “RSR-$d_{embed}$” is
residual spatial reduction with output embedding dimension $d_{embed}$. Note
that the embedding dimension in Stage 1 is determined by the last
convolutional layer in tokenization.

##  Appendix B Search Space

We enlarge ViT-Res-Tiny to build ViT-Res-Tiny super-network. We increase
numbers of attention heads $h$ and decrease head dimensions $d$ for the first
two stage so that we have more choices of attention heads $h$. We do not
search configurations of convolutional layers in tokenization except that we
search the number of output channels in the last layer, which determines the
embedding dimension of the first stage. For the search space for ViT-ResNAS-
Tiny, each stage has three pairs of transformer blocks. The first block in
each pair always remains while the second one is skippable and can be removed.
Therefore, each stage can have three to six transformer blocks. Additionally,
we further enlarge the ViT-Res-Tiny super-network to construct ViT-Res-Small
and Medium super-network and search space by increasing width and adding one
extra block for each stage. Table 5 summarizes the search spaces for ViT-
ResNAS networks.

|  | Output
---
Length
| ViT-ResNAS-Tiny
---
Search Space
Token. | $d_{1}\in\\{256,224,192,176,160\\}$
$16\times 16+1$ | conv-3_24_2
$\left[\begin{tabular}[c]{@{}c@{}}conv-3\\_24\\_1\\\
conv-3\\_24\\_1\end{tabular}\right]$
conv-7_$d_{1}$_1_7
Stage 1 |  |  $h_{1}\in\\{6,5,4,3\\}$,
---
$f_{1}\in\\{768,704,640,576,512,448,384\\}$
$16\times 16+1$ | $\left[\begin{tabular}[c]{@{}c@{}}$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_$h_{1}$\\\ FFN-$f_{1}$\\\ \end{tabular}\right)$\\\ \\\ $\left(\begin{tabular}[c]{@{}c@{}}(skippable)\\\ MHSA-32\\_$h_{1}$\\\ FFN-$f_{1}$\\\ \end{tabular}\right)$\\\ \end{tabular}\right]\times 3$
Stage 2 |  |  $d_{2}\in\\{512,448,384,352,320\\}$,
---
$h_{2}\in\\{12,10,8,6\\}$,
$f_{2}\in\\{1536,1408,1280,1152,1024,$
$896,768\\}$
$8\times 8+1$ | RSR-$d_{2}$
$8\times 8+1$ | $\left[\begin{tabular}[c]{@{}c@{}}$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_$h_{2}$\\\ FFN-$f_{2}$\\\ \end{tabular}\right)$\\\ \\\ $\left(\begin{tabular}[c]{@{}c@{}}(skippable)\\\ MHSA-48\\_$h_{2}$\\\ FFN-$f_{2}$\\\ \end{tabular}\right)$\\\ \end{tabular}\right]\times 3$
Stage 3 |  |  $d_{3}\in\\{1024,896,768,704,640\\}$,
---
$h_{3}\in\\{12,10,8,6\\}$,
$f_{3}\in\\{3072,2816,2560,2304,2048,$
$1792,1536\\}$
$4\times 4+1$ | RSR-$d_{3}$
$4\times 4+1$ | $\left[\begin{tabular}[c]{@{}c@{}}$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_$h_{3}$\\\ FFN-$f_{3}$\\\ \end{tabular}\right)$\\\ \\\ $\left(\begin{tabular}[c]{@{}c@{}}(skippable)\\\ MHSA-64\\_$h_{3}$\\\ FFN-$f_{3}$\\\ \end{tabular}\right)$\\\ \end{tabular}\right]\times 3$

|  | Output
---
Length
| ViT-ResNAS-Small and
---
Medium Search Space
Token. | $d_{1}\in\\{320,280,240,220,200\\}$
$16\times 16+1$ | conv-3_24_2
$\left[\begin{tabular}[c]{@{}c@{}}conv-3\\_24\\_1\\\
conv-3\\_24\\_1\end{tabular}\right]$
conv-7_$d_{1}$_1_7
Stage 1 |  |  $h_{1}\in\\{8,7,6,5\\}$,
---
$f_{1}\in\\{960,880,800,720,640,560,480\\}$
$16\times 16+1$ |  | $\left[\begin{tabular}[c]{@{}c@{}}$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_$h_{1}$\\\ FFN-$f_{1}$\\\ \end{tabular}\right)$\\\ \\\ $\left(\begin{tabular}[c]{@{}c@{}}(skippable)\\\ MHSA-32\\_$h_{1}$\\\ FFN-$f_{1}$\\\ \end{tabular}\right)$\\\ \end{tabular}\right]\times 3$
---
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_$h_{1}$\\\ FFN-$f_{1}$\\\
\end{tabular}\right)$
Stage 2 |  |  $d_{2}\in\\{640,560,480,440,400\\}$,
---
$h_{2}\in\\{16,14,12,10\\}$,
$f_{2}\in\\{1920,1760,1600,1440,1280,$
$1120,960\\}$
$8\times 8+1$ | RSR-$d_{2}$
$8\times 8+1$ |  | $\left[\begin{tabular}[c]{@{}c@{}}$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_$h_{2}$\\\ FFN-$f_{2}$\\\ \end{tabular}\right)$\\\ \\\ $\left(\begin{tabular}[c]{@{}c@{}}(skippable)\\\ MHSA-48\\_$h_{2}$\\\ FFN-$f_{2}$\\\ \end{tabular}\right)$\\\ \end{tabular}\right]\times 3$
---
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_$h_{2}$\\\ FFN-$f_{2}$\\\
\end{tabular}\right)$
Stage 3 |  |  $d_{3}\in\\{1280,1120,960,880,800\\}$,
---
$h_{3}\in\\{16,14,12,10\\}$,
$f_{3}\in\\{3840,3520,3200,2880,2560,$
$2240,1920\\}$
$4\times 4+1$ | RSR-$d_{3}$
$4\times 4+1$ |  | $\left[\begin{tabular}[c]{@{}c@{}}$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_$h_{3}$\\\ FFN-$f_{3}$\\\ \end{tabular}\right)$\\\ \\\ $\left(\begin{tabular}[c]{@{}c@{}}(skippable)\\\ MHSA-64\\_$h_{3}$\\\ FFN-$f_{3}$\\\ \end{tabular}\right)$\\\ \end{tabular}\right]\times 3$
---
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_$h_{3}$\\\ FFN-$f_{3}$\\\
\end{tabular}\right)$

Table 5: ViT-ResNAS search space. Left: search space for ViT-ResNAS-Tiny.
Right: search space for ViT-ResNAS-Small and Medium. For each stage $i$, we
search embedding dimension $d_{embed}$ and numbers of transformer blocks. For
each transformer block, we search the number of attention heads $h$ in MHSA
and hidden size $d_{hidden}$ in FFN. Different blocks can have different
values of $h$ and $d_{hidden}$. The first rows in tokenization and each stage
define the range of each searchable dimension, with $d_{i}$, $h_{i}$, and
$f_{i}$ corresponding to $d_{embed}$, $h$, and $d_{hidden}$ in stage $i$,
respectively. Transformer blocks with “skippable” can be removed during super-
network training and evolutionary search, which supports different numbers of
blocks in searched networks.

##  Appendix C Additional Training Details

#### Super-Network Training.

We divide the original training set into sub-train and sub-validation sets.
The sub-validation set contains $25$K images, with 25 images for each class.
The rest of images form the sub-train set. We train super-networks on the sub-
train set, and during evolutionary search, we evaluate the accuracy of sub-
networks on the sub-validation set.

Additionally, during super-network training, we warm up training different
widths and depths (filter warmup [3, 55] or progressive shrinking [4, 40]).
Specifically, at the beginning, we only sample and train the largest sub-
network in a search space. As the super-network training proceeds, we
gradually sample and train sub-networks with smaller widths and depths. After
25% of total training epochs, sub-networks with any width and depth can be
sampled and trained.

#### Training Cost.

We report the time for training networks and performing NAS when 8 V100 (16GB)
GPUs are used. Training ViT-Res-Tiny takes about 32 hours. For ViT-ResNAS-
Tiny, super-network training takes about 16.7 hours, evolutionary search takes
5.5 hours, and it takes about 34.5 hours to train the searched ViT-ResNAS-
Tiny. For ViT-ResNAS-Small and Medium, the cost of training the shared super-
network is 21 hours. For ViT-ResNAS-Small, the evolutionary search takes 6
hours and training the searched network takes 41.6 hours. For ViT-ResNAS-
Medium, it takes 6 hours and 45 hours to perform evolutionary search and
training the searched network, respectively. Note that we report the training
time just for completeness and that the time can vary across different
platforms.

##  Appendix D Searched ViT-ResNAS Architecture

We report the searched architectures of ViT-ResNAS-Tiny, Small and Medium in
Table 6, 7, and 8, respectively. Compared to ViT-Res-Tiny, ViT-ResNAS-Tiny has
smaller embedding dimensions but more transformer blocks. From ViT-ResNAS-Tiny
to Small, stage 1 has one extra block, and embedding dimensions are increased
uniformly. As for ViT-ResNAS-Medium, the first two stages have more blocks
than the last stage. These suggest that having more blocks in earlier stages
is more important to scale up networks.

##  Appendix E Limitation

We discuss some limitations of our approaches. First, the performance of
searched networks designed with NAS relies on manually designed search spaces.
We build our search space by uniformly increasing widths and depths of ViT-
Res-Tiny, which is designed with some simple heuristics and without tuning.
Compared to search spaces of CNN (e.g., MobileNet), our search space is less
studied and less optimized. Optimizing search spaces such as setting a better
range for each architectural parameter could potentially result in additional
performance gain. Second, our searched networks use the same type of attention
mechanism as ViT, which could have large memory consumption when processing
high-resolution feature maps (long sequences). Incorporating more efficient
attention mechanisms [7, 26, 61] into search spaces could solve the issue and
result in networks with better performance.

|  | Output
---
Length
ViT-ResNAS-Tiny
Tokenization | $16\times 16+1$ | conv-3_24_2
$\left[\begin{tabular}[c]{@{}c@{}}conv-3\\_24\\_1\\\
conv-3\\_24\\_1\end{tabular}\right]$
conv-7_176_7
Stage 1 | $16\times 16+1$ |  | $\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_3, FFN-704\end{tabular}\right)$
---
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_3, FFN-576\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_3, FFN-640\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_4, FFN-576\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_4, FFN-704\end{tabular}\right)$
Stage 2 | $8\times 8+1$ | RSR-352
$8\times 8+1$ |  | $\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_10, FFN-1408\end{tabular}\right)$
---
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_8, FFN-1408\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_8, FFN-1280\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_8, FFN-1408\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_10, FFN-1280\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_10, FFN-1024\end{tabular}\right)$
Stage 3 | $4\times 4+1$ | RSR-704
$4\times 4+1$ |  | $\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_10, FFN-2560\end{tabular}\right)$
---
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_10, FFN-1792\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_10, FFN-2816\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_8, FFN-2816\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_8, FFN-2560\end{tabular}\right)$

Table 6: Architecture of ViT-ResNAS-Tiny.

|  | Output
---
Length
ViT-ResNAS-Small
Tokenization | $16\times 16+1$ | conv-3_24_2
$\left[\begin{tabular}[c]{@{}c@{}}conv-3\\_24\\_1\\\
conv-3\\_24\\_1\end{tabular}\right]$
conv-7_220_7
Stage 1 | $16\times 16+1$ |  | $\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_5, FFN-880\end{tabular}\right)$
---
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_5, FFN-880\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_7, FFN-800\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_5, FFN-720\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_5, FFN-720\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_5, FFN-720\end{tabular}\right)$
Stage 2 | $8\times 8+1$ | RSR-440
$8\times 8+1$ |  | $\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_10, FFN-1760\end{tabular}\right)$
---
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_10, FFN-1440\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_10, FFN-1920\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_10, FFN-1600\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_12, FFN-1600\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_12, FFN-1440\end{tabular}\right)$
Stage 3 | $4\times 4+1$ | RSR-880
$4\times 4+1$ |  | $\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_16, FFN-3200\end{tabular}\right)$
---
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_12, FFN-3200\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_16, FFN-2880\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_12, FFN-2240\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_14, FFN-2560\end{tabular}\right)$

Table 7: Architecture of ViT-ResNAS-Small.

|  | Output
---
Length
ViT-ResNAS-Medium
Tokenization | $16\times 16+1$ | conv-3_24_2
$\left[\begin{tabular}[c]{@{}c@{}}conv-3\\_24\\_1\\\
conv-3\\_24\\_1\end{tabular}\right]$
conv-7_240_7
Stage 1 | $16\times 16+1$ |  | $\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_7, FFN-960\end{tabular}\right)$
---
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_6, FFN-960\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_7, FFN-800\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_8, FFN-960\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_7, FFN-880\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_8, FFN-880\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-32\\_6, FFN-800\end{tabular}\right)$
Stage 2 | $8\times 8+1$ | RSR-640
$8\times 8+1$ |  | $\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_10, FFN-1120\end{tabular}\right)$
---
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_14, FFN-1760\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_14, FFN-1920\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_16, FFN-1760\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_14, FFN-1440\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_16, FFN-1760\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-48\\_16, FFN-1920\end{tabular}\right)$
Stage 3 | $4\times 4+1$ | RSR-880
$4\times 4+1$ |  | $\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_16, FFN-3200\end{tabular}\right)$
---
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_10, FFN-3840\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_16, FFN-3840\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_12, FFN-3200\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_16, FFN-3520\end{tabular}\right)$
$\left(\begin{tabular}[c]{@{}c@{}}MHSA-64\\_14, FFN-3520\end{tabular}\right)$

Table 8: Architecture of ViT-ResNAS-Medium.