A UTO -S CALING V ISION T RANSFORMERS WITHOUT T RAINING
A BSTRACT
This work targets automated designing and scaling of Vision Transformers (ViTs).
The motivation comes from two pain spots: 1) the lack of efﬁcient and principled methods for designing and scaling ViTs; 2) the tremendous computational cost of training ViT that is much heavier than its convolution counterpart.
To tackle these issues, we propose As-ViT , an auto-scaling framework for ViTs without training, which automatically discovers and scales up ViTs in an efﬁcient and principled manner.
Speciﬁcally, we ﬁrst design a “seed” ViT topology by leveraging a trainingfree search process.
This extremely fast search is fulﬁlled by a comprehensive study of ViT’s network complexity, yielding a strong Kendall-tau correlation with ground-truth accuracies.
Second, starting from the “seed” topology, we automate the scaling rule for ViTs by growing widths/depths to different ViT layers.
This results in a series of architectures with different numbers of parameters in a single run.
Finally, based on the observation that ViTs can tolerate coarse tokenization in early training stages, we propose a progressive tokenization strategy to train ViTs faster and cheaper.
As a uniﬁed framework, As-ViT achieves strong performance on classiﬁcation (83.5% top1 on ImageNet-1k) and detection (52.7% mAP on
COCO) without any manual crafting nor scaling of ViT architectures: the end-toend model design and scaling process cost only 12 hours on one V100 GPU .
Our code is available at https://github.com/VITA-Group/AsViT .
I NTRODUCTION
Transformer (Vaswani et al., 2017), a family of architectures based on the self-attention mechanism, is notable for modeling long-range dependencies in the data.
The success of transformers has evolved from natural language processing to computer vision.
Recently, Vision Transformer (ViT)
(Dosovitskiy et al., 2020), a transformer architecture consisting of self-attention encoder blocks, has been proposed to achieve competitive performance to convolution neural networks (CNNs) (Simonyan &
Zisserman, 2014; He et al., 2016) on ImageNet (Deng et al., 2009).
However, it remains elusive on how to effectively design, scale-up, and train ViTs, with three important gaps awaiting.
First, Dosovitskiy et al.
(2020) directly hard-split the 2D image into a series of local patches, and learn the representation with a pre-deﬁned number of attention heads and channel expansion ratios.
These ad-hoc “tokenization” and embedding mainly inherit from language tasks (Vaswani et al., 2017) but are not customized for vision, which calls for more ﬂexible and principled designs .
Second, the learning behaviors of ViT, including (loss of) feature diversity (Zhou et al., 2021), receptive ﬁelds (Raghu et al., 2021) and augmentations (Touvron et al., 2020; Jiang et al., 2021), differ vastly from CNNs.
Beneﬁting from self-attention, ViT can capture global information even with shallow layers, yet its performance is quickly plateaued as going deeper.
Strong augmentations are also vital to avoid ViTs from overﬁtting.
These observations indicate that ViT architectures may require uniquely customized scaling-up laws to learn a more meaningful representation hierarchy.
Third, training ViTs is both data and computation-consuming.
To achieve state-of-the-art performance, ViT requires up to 300 million images and thousands of TPU-days.
Although recent works attempt to enhance ViT’s data and resource efﬁciency (Touvron et al., 2020; Hassani et al., 2021; Pan et al., 2021; Chen et al., 2021d), the heavy computation cost (e g., quadratic with respect to the number of tokens) is still overwhelming, compared with training CNNs
We point out that the above gaps are inherently connected by the core architecture problem: how to design and scale-up ViTs?
Different from the convolutional layer that directly digests raw pixels, ViTs embed coarse-level local patches as input tokens.
Shall we then divide an image into non-overlapping tokens of smaller size, or larger but overlapped tokens?
The former could embed more visual details in each token but ignores spatial coherency, while the latter sacriﬁces the local details but may beneﬁt more spatial correlations among tokens.
A further question is on ViT’s depth/width trade-off: shall we prefer a wider and shallower ViT, or a narrower but deeper one?
A similar dilemma also persists for ViT training: reducing the number of tokens would effectively speed up the ViT training, but meanwhile might sacriﬁce the training performance if sticking to coarse tokens from end to end.
In this work, we aim to reform the discovery of novel ViT architectures.
Our framework, called
As-ViT (Auto-scaling ViT), allows for extremely fast, efﬁcient, and principled ViT design and scaling.
In short, As-ViT ﬁrst ﬁnds a promising “seed” topology for ViT of small depths and widths, then progressively “grow” it into different sizes (number of parameters) to meet different needs.
Speciﬁcally, our “seed” ViT topology is discovered from a search space relaxed from recent manual
ViT designs.
To compare different topologies, we automate this process by a training-free architecture search approach and the measurement of ViT’s complexity, which are extremely fast and efﬁcient.
This training-free search is supported by our comprehensive study of various network complexity metrics, where we ﬁnd the expected length distortion has the best trade-off between time costs and Kendall-tau correlations.
Our “seed” ViT topology is then progressively scaled up from a small network to a large one, generating a series of ViT variants in a single run.
Each step, the increases of depth and width are automatically and efﬁciently balanced by comparing network complexities.
Furthermore, to address the data-hungry and heavy computation costs of ViTs, we make our ViT tokens elastic, and propose a progressive re-tokenization method for efﬁcient ViT training.
We summarize our contributions as below:
We for the ﬁrst time automate both the backbone design and scaling of ViTs.
A “seed” ViT topology is ﬁrst discovered (in only seven V100 GPU-hours), and then its depths and widths are grown with a principled scaling rule in a single run (ﬁve more V100 GPU-hours).
To estimate ViT’s performance at initialization without any training, we conduct the ﬁrst comprehensive study of ViT’s network complexity measurements.
We empirically ﬁnd the expected length distortion has the best trade-off between the computation costs and its Kendall-tau correlations with ViT’s ground-truth accuracy.
During training, we propose a progressive re-tokenization scheme via the change of dilation and stride, which demonstrates to be a highly efﬁcient ViT training strategy that saves up to 56.2% training FLOPs and 41.1% training time, while preserving a competitive accuracy.
Our As-ViT achieves strong performance on classiﬁcation (83.5% top-1 on ImageNet-1k) and detection (52.7% mAP on COCO).
W HY WE NEED AUTOMATED DESIGN AND SCALING PRINCIPLE FOR V I T? Background and recent development of ViT
To transform a 2D image into a sequence,
ViT (Dosovitskiy et al., 2020) splits each image into 14 × 14 or 16 × 16 patches and embeds them into a ﬁxed number of tokens; then following practice of the transformer for language modeling, ViT applies self-attention to learn reweighting masks as relationship modeling for tokens, and leverages FFN (Feed-Forward Network) layers to learn feature embeddings.
To better facilitate the visual representation learning, recently works try to train deeper ViTs (Touvron et al., 2021;
Zhou et al., 2021), incorporate convolutions (Wu et al., 2021; d’Ascoli et al., 2021; Yuan et al., 2021a), and design multi-scale feature extractions (Chen et al., 2021b; Zhang et al., 2021; Wang et al., 2021).
Why manual design and scaling may be suboptimal?
As the ViT architecture is still in its infant stage, there is no principle in its design and scaling.
Early designs incorporate large token sizes, constant sequence length, and hidden size (Dosovitskiy et al., 2020; Touvron et al., 2020), and recent trends include small patches, spatial reduction, and channel doubling (Zhou et al., 2021; Liu et al., 2021).
They all achieve comparably good performance, leaving the optimal choices unclear.
Moreover, different learning behaviors of transformers from CNNs make the scaling law of ViTs highly unclear.
Recent works (Zhou et al., 2021) demonstrated that attention maps of ViTs gradually become similar in deeper layers, leading to identical feature maps and saturated performance.
ViT also generates more uniform representations across layers, enabling early aggregation of global context (Raghu et al., 2021).
This is contradictory to CNNs as deeper layers help the learning of visual global information (Chen et al., 2018).
These observations all indicate that previously studied scaling laws (depth/width allocations) for CNNs (Tan & Le, 2019) may not be appropriate to ViTs.
What principle do we want?
We aim to automatically design and scale-up ViTs, being principled and avoiding manual efforts and potential biases.
We also want to answer two questions: 1) Does ViT have any preference in its topology (patch sizes, expansion ratios, number of attention heads, etc.)?
2) Does ViT necessarily follow the same scaling rule of CNNs?
A UTO - DESIGN & SCALING OF V I T S WITH NETWORK COMPLEXITY
[Figure]
To accelerate in ViT designing and avoid tedious manual efforts, we target efﬁcient, automated, and principled search and scaling of ViTs.
Speciﬁcally, we have two problems to solve: 1) with minimal or even zero training cost (Section 3.2), how to efﬁciently ﬁnd the optimal ViT architecture topology (Section 3.3)?
2) how to scale-up depths and widths of the ViT topology to meet different needs of model sizes (Section 3.4)?
E XPANDED T OPOLOGY S PACE FOR V I T S
Before designing and scaling, we ﬁrst brieﬂy introduce our expanded topology search space for our As-ViT ( blue italics in Figure 1).
We ﬁrst embed the input image to patches of a 14 -scale resolution, and adopt a stage-wise spatial reduction and channel doubling strategy.
This is for the convenience of dense prediction tasks like detection that require multi-scale features.
Table 1 summarizes details of our topology space [Table] and will be explained below.
Elastic kernels.
Instead of generating nonoverlapped image patches, we propose to search for the kernel size.
This will enable patches to be overlapped with their neighbors, introducing more spatial correlations among tokens.
Moreover, each time we downsample the spatial resolution, we also introduce overlaps when re-embedding local tokens (implemented by either a linear or a convolutional layer)
[Table]
[Table]
Elastic attention splits.
Splitting the attention into local windows is an important design to reduce the computation cost of self-attention without sacriﬁcing much performance (Zaheer et al., 2020; Liu et al., 2021).
Instead of using a ﬁxed number of splits, we propose to search for elastic attention splits for each stage 2 .
Note that we try to make our design general and do not use shifted windows (Liu et al., 2021).
More search dimensions.
ViT (Dosovitskiy et al., 2020) by default leveraged an FFN layer with 4 × expanded hidden dimension for each attention block.
To enable a more ﬂexible design of ViT architectures, for each stage we further search over the FFN expansion ratio.
We also search for the ﬁnal number of heads for the self-attention module.
3. A SSESSING V I T COMPLEXITY AT INITIALIZATION VIA MANIFOLD PROPAGATION
Training ViTs is slow: hence an architecture search guided by evaluating trained models’ accuracies will be dauntingly expensive here.
We note a recent surge of training-free neural architecture search methods for ReLU-based CNNs, leveraging local linear maps (Mellor et al., 2020), gradient sensitivity (Abdelfattah et al., 2021), number of linear regions (Chen et al., 2021e;f), or network topology Bhardwaj et al.
However, ViTs are equipped with more complex non-linear functions: self-attention, softmax, and GeLU.
Therefore, we need to measure their learning capacity in a more general way.
In this way, we consider measuring the complexity of manifold propagation through ViT, to estimate how complex functions can be approximated by ViTs. Intuitively, a complex network can propagate a simple input into a complex manifold at its output layer, thus likely to possess a strong learning capacity.
In our work, we study the manifold complexity of mapping a simple circle input through the ViT: h ( θ ) = √ N (cid:2) u 0 cos( θ ) + u 1 sin( θ ) (cid:3) .
Here, N is the dimension of ViT’s input (e.g. N = 3 × 224 × 224 for ImageNet images), u 0 and u 1 form an orthonormal basis for a 2-dimensional subspace of R N in which the circle lives.
We further deﬁne
ViT network as N , its input-output Jacobian v ( θ ) =
∂ θ N ( h ( θ )) at the input θ , and a ( θ ) =
∂ θ v ( θ ) .
We will calculate expected complexities over a certain number of θ s uniformly sampled from [0 , 2 π ) .
In our work, we study three different types of manifold complexities:
Curvature is deﬁned as the reciprocal of the radius of the osculating circle on the ViT’s output manifold.
Intuitively, a larger curvature indicates that N ( θ ) changes fast at a certain θ .
According to Riemannian geometry (Lee,
; Poole et al., 2016), the curvature can be explicitly calculated as
[Equation]
=
[Equation]
Length Distortion in Euclidean space is deﬁned as L E = length
[Equation] N
[Equation] = [Equation] v
[Equation] θ ) [Equation] dθ .
It measures when the network takes a unit-length curve as input, what is the length of the output curve.
Since the ground-truth function we want to estimate (using N ) is usually very complex, one may also expect that networks with better performance should also generate longer outputs.
The problem of L E is that, longer outputs not necessarily translate to complex outputs.
A simple example: even an appropriately initialized linear network could grow a straight line into a long output (i.e. a large norm of input-output Jacobian).
Therefore, one could instead use Length
Distortion taking curvature into consideration to measure how quickly the normalized Jacobian ˆ v ( θ )
= v ( θ ) / (cid:112) v ( θ ) · v ( θ ) changes with respect to θ , deﬁned as L
[Equation]
[Table]
[Table]
In our study, we aim to compare the potential of using these three complexity metrics to guide the ViT architecture selection.
As the core of neural architecture search is to rank the performance of different architectures, we measure the Kendall-tau correlations ( τ ) between these metrics and models’ ground-truth accuracies.
We randomly sampled 87 ViT topologies from Table 1 (with L 1 =
L 2 =
L 3 =
L 4 = 1 , C = 32 ), fully train them on ImageNet-1k for 300 epochs (following the same training recipe of DeiT (Touvron et al., 2020)), and also measure their κ, L E , L Eκ at initialization.
As shown in Figure 2, we can clearly see that both κ and L E exhibit high
Kendall-tau correlations.
κ has a negative correlation, which may indicate that changes of output manifold on the tangent direction are more important to ViT training, instead of on the perpendicular direction.
Meanwhile, κ costs too much computation time due to second derivatives.
We decide to choose L E as our complexity measure for highly fast ViT topology search and scaling.
3. L E AS REWARD FOR SEARCHING V I T T OPOLOGIES
We now propose our training-free search based on L E (Algorithm 1).
Most NAS (neural architecture search) methods evaluate the accuracies or loss values of single-path or super networks as proxy inference.
This training-based search will suffer from more computation costs when applied to ViTs.
Instead of training ViTs, for each architecture we sample, we calculate L E and treat it as the reward to guide the search process.
In addition to L E , we also include the NTK condition number κ Θ = λ max λ min to indicate the trainability of ViTs (Chen et al., 2021e; Xiao et al., 2019; Yang, 2020; Hron et al., 2020).
λ max and λ min are the largest and smallest eigenvalue of NTK matrix Θ .
[Table]
[Table]
We use reinforcement learning (RL) for search.
The RL policy is formulated as a joint categorical distribution, and is updated by policy gradient (Williams, 1992).
We update our policy forsteps, which is observed enough for the policy to converge (entropy drops from 15.3 to 5.7).
The search process is extremely fast: only seven GPU-hours (V100) on ImageNet-1k, thanks to the simple calculation of L E that bypasses the ViT training.
To address the different magnitude of L E and κ Θ , we normalize them by their relative value ranges (line 5 in Algorithm 1).
We summarize the
ViT topology statistics from our search in Table 3.
We can see that L E and κ Θ highly prefer: (1) tokens with overlaps ( K 1 ∼ Kare all larger than strides), and (2) larger FFN expansion ratios in deeper layers ( E 1 < E 2 < E 3 < E 4 ).
No clear preference of L E and κ Θ are found on attention splits and number of heads.
3. A UTOMATIC AND PRINCIPLED SCALING OF V I T S
After obtaining an optimal topology, another question is: how to balance the network depth and width?
Currently, there is no such rule of thumb for ViT scaling.
Recent works try to scale-up or grow convolutional networks of different sizes to meet various resource constraints (Liu et al., 2019a;
Tan & Le, 2019).
However, to automatically ﬁnd a principled scaling rule, training ViTs will cost enormous computation costs.
It is also possible to search different ViT variants (as in Section 3.3), but that requires multiple runs.
Instead, “scaling-up” is a more natural way to generate multiple model variants in one experiment.
We are therefore motivated to scale-up our searched basic “seed”
ViT to a larger model in an efﬁcient training-free and principled manner.
We depict our auto-scaling method in Algorithm 2.
The starting-point architecture has one attention block for each stage, and an initial hidden dimension C = 32 .
In each iteration, we greedily ﬁnd the optimal depth and width to scale-up next.
For depth, we try to ﬁnd out which stage to deepen (i.e., add one attention block to which stage); for width, we try to discover the best expansion ratio (i.e., widen the channel number to what extent).
The rule to choose how to scale-up is by comparing the propagation complexity among a set of scaling choices.
For example, in the case of four backbone stages (Table 1) and four expansion ratio choices ( [0 .
× , 0 .
1 × , 0 . 15 × , 0 . 2 × ] ), we have 4 × 4 =scaling choices in total for each step.
We calculate L E and κ Θ after applying each choice, and the one with the best L E / κ Θ trade-off (minimal sum of rankings by L E and κ Θ ) will be selected to scale-up with.
The scaling stops when a certain limit of parameter number is reached.
In our work, we stop the scaling process once the number of parameters reaches 100 million, and the scaling only takes ﬁve GPU hours (V100) on ImageNet-1k. [Figure]
The scaling trajectory is visualized in Figure 3.
By comparing our automated scaling against random scaling, we ﬁnd our scaling principle prefers to sacriﬁce the depths to win more widths, keeping a shallower but wider network.
Our scaling is more similar to the rule developed by Zhai et al.
In contrast, ResNet and Swin Transformer (Liu et al., 2021) choose to be narrower and deeper.
1 Input:
ViT topology criterion parameters) P , t
, channel expansion ratio choices C
{ 1 . 05 × , 1 . 1 × , 1 . 15 × , 1 . 2 ×} (to increase the width by 5%, 10%, 15%, or 20%), depth choices D
{ (+1 , 0 , 0 , 0) , (0 , +1 , 0 , 0) , (0 , 0 , +1 , 0) , (0 , 0 , 0 , +1) } (to add one more layer to one of the four stages in Table 1).
E FFICIENT V I T TRAINING VIA PROGRESSIVE ELASTIC RE - TOKENIZATION
[Figure]
Recent works (Jia et al., 2018;
Zhou et al., 2019; Fu et al., 2020) show that one can use mixed or progressive precision to achieve an efﬁcient training purpose.
The rationale behind this strategy is that, there exist some “short-cuts” on the network’s loss landscape that can be manually created to bypass perhaps less important gradient descent steps, especially during early training phases.
As in ViT, both self-attention and FFN have quadratic computation costs to the number of tokens.
It is therefore natural to ask: do we need full-resolution tokens during the whole training process?
We provide an afﬁrming answer by proposing a progressive elastic re-tokenization training strategy.
To update the number of tokens during training without affecting the shape of weights in linear projections, we adopt different sampling granularities in the ﬁrst linear projection layer.
Taking the ﬁrst projection kernel K 1 = 4 with stride = 4 as an example: during training we gradually change the ( stride , dilation ) pair 3 of the ﬁrst projection kernel to ( 16 , 5 ), ( 8 , 2 ), and ( 4 , 1 ), keeping the shape of weights and the architecture unchanged.
[Table]
This re-tokenization strategy stimulates curriculum learning for ViTs: when the training begins, we introduce coarse sampling to significantly reduce the number of tokens.
In other words, our As-ViT quickly learns coarse information from images in early training stages at extremely low computation cost (only 13 .
FLOPs of full-resolution training).
Towards the late phase of training, we progressively switch to ﬁne-grained sampling, restore the full token resolution, and maintain the competitive accuracy.
As shown in Figure 4, when the ViT is trained with coarse sampling in early training phases, it can still obtain high accuracy while requiring extremely low computation cost.
The transition between different sampling granularity introduces a jump in performance, and eventually the network restores its competitive ﬁnal performance.
E XPERIMENTS 5.
[Equation]
[Table]
We show our searched As-ViT topology in
Table 4.
This architecture facilitates strong overlaps among tokens during both the ﬁrst projection (“tokenization”) step and three re-embedding steps.
FFN expansion ratios are ﬁrst narrow then become wider in deeper layers.
A small number of attention splits are leveraged for better aggregation of global information.
The seed topology is automatically scaledup, and three As-ViT variants of comparable sizes with previous works will be benchmarked.
Our scaling rule prefers shallower and wider networks, and layers are more balanced among different resolution stages.
5. I MAGE C LASSIFICATION
Settings.
We benchmark our As-ViT on
ImageNet-1k (Deng et al., 2009).
We use
Tensorﬂow and Keras for training implementations and conduct all training on TPUs.
We set the default image size as 224 × 224 , and use AdamW (Loshchilov & Hutter, 2017) as the optimizer with cosine learning rate decay (Loshchilov & Hutter, 2016).
A batch size of 1024, an initial learning rate of 0.001, and a weight decay of 0.05 are adopted.
Table 5 demonstrates comparisons of our
As-ViT to other models.
Compared to the previous both Transformer-based and CNNbased architectures, As-ViT achieves stateof-the-art performance with a comparable number of parameters and FLOPs.
More importantly, our As-ViT framework achieves competitive or stronger performance than concurrent NAS works for ViTs with much more search efﬁciency.
As-ViTs are designed with highly reduced human or NAS efforts.
All our three As-ViT variants are generated in only 12 GPU hours (on a single V100 GPU).
In contrast, BoneNAS (Li et al., 2021) requires 10 GPU days to search a single architecture.
For each variant of ViT-ResNAS (Liao et al., 2021), the super-network training takes 16.7 ∼ 21 hours, followed by another 5.5 ∼ 6 hours of evolutionary search.
[Table]
Efﬁcient
Training.
We leverage the progressive elastic retokenization strategy proposed in
Section 4 to reduce both FLOPs and training time for large ViT models.
As illustrated in Figure 4, we progressively apply 4 × and 2 × reductions on the number of tokens during training by changing both the dilation and the stride of the ﬁrst linear projection layer.
We tune the epochs allocated to each token reduction stage and show the results in Table 6.
Standard training takes 42.8 TPU days, whereas our efﬁcient training could save up to 56.2% training FLOPs and 41.1% training TPU days, still achieving a strong accuracy.
[Table]
Disentangled Contributions from Topology and Scaling.
To better verify the contribution from our searched topology and scaling rule, we conduct more ablation studies (Table 7).
First, we directly train the searched topology before scaling.
Our searched seed topology is better than the best from 87 random topologies in Figure 2.
Second, we compare our complexity-based scaling rule with “random scaling + As-ViT topology”.
At different scales, our automated scaling is also better than random scaling.
5. O BJECT D ETECTION ON COCO
Settings Beyond image classiﬁcation, we further evaluate our designed As-ViT on the detection task.
Object detection is conducted on COCO 2017 that contains 118,000 training and 5000 validation images.
We adopt the popular Cascade Mask R-CNN as the object detection framework for our
As-ViT. We use an input size of 1024 × 1024 , AdamW optimizer (initial learning rate of 0.001), weight decay of 0. [Table] and a batch size of 256.
Efﬁciently pretrained ImageNet-1K checkpoint (82.9% in Table 6) is leveraged as the initialization.
We compare our As-ViT to standard CNNs (ResNext, SpineNet (Du et al., 2020)) and previous
Transformer networks (Swin (Liu et al., 2021)).
The comparisons are conducted by changing only the backbones with other settings unchanged.
In Table 8 we can see that our As-ViT can also capture multi-scale features and achieve state-of-the-art detection performance, although being designed on
ImageNet and its complexity is measured for classiﬁcation
[Table] [Table]
R ELATED WORKS 6. V ISION T RANSFORMER
Transformers (Vaswani et al., 2017) leverage the self-attention to extract global correlation, and become the dominant models for natural language processing (NLP)
(Devlin et al., 2018;
Radford et al., 2018; Brown et al., 2020; Liu et al., 2019b).
Recent works explored transformers to vision problems: image classiﬁcation (Dosovitskiy et al., 2020), object detection (Carion et al., 2020; Zhu et al., 2020; Zheng et al., 2020; Dai et al., 2020; Sun et al., 2020), segmentation (Chen et al., 2020;
Wang et al., 2020), etc.
The Vision Transformer (ViT)
(Dosovitskiy et al., 2020) designed a pure transformer architecture and achieved SOTA performance on image classiﬁcation.
However, ViT heavily relies on large-scale datasets (ImageNet-21k (Deng et al., 2009), JFT-300M (Sun et al., 2017)) for pretraining, requiring huge computation resources.
DeiT (Touvron et al., 2020) proposed
Knowledge Distillation (KD) (Hinton et al., 2015; Yuan et al., 2020) via a special KD token to improve both performance and training efﬁciency.
In contrast, our proposed As-ViT introduces more ﬂexible tokenization, attention splitting, and FFN expansion strategies, with automated discovery.
6. N EURAL A RCHITECTURE D ESIGN AND S CALE
Manual design of network architectures heavily relies on human prior, which is difﬁcult to scale-up.
Recent works leverage AutoML to ﬁnd optimal combinations of operators/topology in a given search space (Zoph & Le, 2016; Real et al., 2019; Liu et al., 2018; Dong & Yang, 2019).
However, the searched model are small due to the ﬁxed and hand-crafted search space, far from being scaled-up to modern networks.
For example, models from NASNet space (Zoph et al., 2018) only have ∼ 5M parameters, much smaller than real-world ones (20 to over 100M).
One main reason for not being scalable is because NAS is a computation-consuming task, typically costing 1 ∼ 2 GPU days to search even small architectures.
Meanwhile, many works try to grow a “seed” architecture to different variants.
EfﬁcientNet (Tan & Le, 2019) manually designed a scaling rule for width and depth.
Give a template backbone with ﬁxed depth, Liu et al.
(2019a) grow the width by gradient descent.
For
ViT, we for the ﬁrst time bring both architecture design and scaling together in one framework.
To overcome the computation-consuming problem in the training of transformers, we directly use the complexity of manifold propagation as a surrogate measure towards a training-free search and scale.
6. E FFICIENT T RAINING
A number of methods have been developed to accelerate the training of deep neural networks, including mixed precision (Jia et al., 2018), distributed optimization (Cho et al., 2017), large-batch training (Goyal et al., 2017; Akiba et al., 2017; You et al., 2018), etc.
Jia et al.
combined distributed training with a mixed-precision framework.
Wang et al.
(2019) proposed to save deep
CNN training energy cost via stochastic mini-batch dropping and selective layer update.
In our work, customized progressive tokenization via the changing of stride/dilation can effectively reduce the number of tokens during ViT training, thus largely saving the training cost.
C ONCLUSIONS
To automate the principled design of vision transformers without tedious human efforts, we propose
As-ViT, a uniﬁed framework that searches and scales ViTs without any training.
Compared with hand-crafted ViT architecture, our As-ViT leverages more token overlaps, increased FFN expansion ratios, and is wider and shallower.
Our As-ViT achieves state-of-the-art accuracies on both ImageNet1K classiﬁcation and COCO detection, which veriﬁes the strong performance of our framework.
Moreover, with progressive tokenization, we can train heavy ViT models with largely reduced training FLOPs and time.
We hope our methodology could encourage the efﬁcient design and training of ViTs for both the transformer and the NAS communities.
A CKNOWLEDGEMENT
Z.W. is in part supported by the NSF AI Institute for Foundations of Machine Learning (IFML) and a
Google TensorFlow Model Garden Award.
A I MPLEMENTATIONS
Training-free Topology Search and Scaling.
We calculate L E by uniformly sampling 10 θ s in [0 , 2 π ) .
For one architecture, the calculation of L E is repeated ﬁve times with different (random) network initializations, and L E is set as their mean.
Image Classiﬁcation.
We use 20 epochs of linear warm-up, a batch size of 1,024, an initial learning rate of 0.001, and a weight decay of 0.05.
Augmentations including stochastic depth (Huang et al., 2016), Mixup (Zhang et al., 2017), Cutmix (Yun et al., 2019), RandAug (Cubuk et al., 2020),
Exponential Moving Average (EMA) are also applied.
Object Detection.
Our training adopts a batch size of 256 for 36 epochs, with also stochastic depth.
We do not use any stronger techniques like HTC (Chen et al., 2019), multi-scale testing, sotf-NMS (Bodla et al., 2017), etc.
B C ONVERGENCE OF T RAINING - FREE S EARCH
To demonstrate the convergence of the policy learned by our RL search, we show the entropy during learning the policy in Figure 5.
We can see that a training ofsteps is enough for the policy to converge to low entropy (high conﬁdence).