Auto-scaling Vision Transformers without Training

1 INTRODUCTION . 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-defined 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 flexible and principled designs . Second , the learning behaviors of ViT , including ( loss of ) feature diversity ( Zhou et al. , 2021 ) , receptive fields ( Raghu et al. , 2021 ) and augmentations ( Touvron et al. , 2020 ; Jiang et al. , 2021 ) , differ vastly from CNNs . Benefiting 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 overfitting . 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 efficiency ( 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 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 sacrifices the local details but may benefit 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 sacrifice 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 , efficient , and principled ViT design and scaling . In short , As-ViT first finds 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 . Specifically , 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 efficient . This training-free search is supported by our comprehensive study of various network complexity metrics , where we find 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 efficiently 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 efficient ViT training . We summarize our contributions as below : 1 . We for the first time automate both the backbone design and scaling of ViTs . A “ seed ” ViT topology is first discovered ( in only seven V100 GPU-hours ) , and then its depths and widths are grown with a principled scaling rule in a single run ( five more V100 GPU-hours ) . 2 . To estimate ViT ’ s performance at initialization without any training , we conduct the first comprehensive study of ViT ’ s network complexity measurements . We empirically find the expected length distortion has the best trade-off between the computation costs and its Kendall-tau correlations with ViT ’ s ground-truth accuracy . 3 . During training , we propose a progressive re-tokenization scheme via the change of dilation and stride , which demonstrates to be a highly efficient ViT training strategy that saves up to 56.2 % training FLOPs and 41.1 % training time , while preserving a competitive accuracy . 4 . Our As-ViT achieves strong performance on classification ( 83.5 % top-1 on ImageNet-1k ) and detection ( 52.7 % mAP on COCO ) . 2 WHY WE NEED AUTOMATED DESIGN AND SCALING PRINCIPLE FOR VIT ? . Background and recent development of ViT1 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 fixed 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 1We generally use the term “ ViT ” to indicate deep networks of self-attention blocks for vision problems . We always include a clear citation when we specifically discuss the ViTs proposed by Dosovitskiy et al . ( 2020 ) . 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 ? 3 AUTO-DESIGN & SCALING OF VITS WITH NETWORK COMPLEXITY . To accelerate in ViT designing and avoid tedious manual efforts , we target efficient , automated , and principled search and scaling of ViTs . Specifically , we have two problems to solve : 1 ) with zero training cost ( Section 3.2 ) , how to efficiently find 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 ) ? 3.1 EXPANDED TOPOLOGY SPACE FOR VITS . Before designing and scaling , we first briefly introduce our expanded topology search space for our As-ViT ( blue italics in Figure 1 ) . We first embed the input image into 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 , and will be explained below . Elastic kernels . Instead of generating non- overlapped 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 ) . Elastic attention splits . Splitting the attention into local windows is an important design to reduce the computation cost of self-attention without sacrificing much performance ( Zaheer et al. , 2020 ; Liu et al. , 2021 ) . Instead of using a fixed number of splits , we propose to search for elastic attention splits for each stage2 . 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 flexible design of ViT architectures , for each stage we further search over the FFN expansion ratio . We also search for the final number of heads for the self-attention module . 2Due to spatial reduction , the 4th stage may already reach a resolution at 7× 7 on ImageNet , and we set its splitting as 1 . 3.2 ASSESSING VIT COMPLEXITY AT INITIALIZATION VIA MANIFOLD PROPAGATION . Training ViTs is slow : hence an architecture search guided by evaluating trained models ’ accuracies will be dauntingly expensive . 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. , 2021 ) . 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 our work , 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 [ u0 cos ( θ ) + u1 sin ( θ ) ] . Here , N is the dimension of ViT ’ s input ( e.g . N = 3× 224× 224 for ImageNet images ) , u0 and u1 form an orthonormal basis for a 2-dimensional subspace ofRN in which the circle lives . We further define the 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 : 1 . Curvature can be defined 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 , 2006 ; Poole et al. , 2016 ) , the curvature can be explicitly calculated as κ = ∫ ( v ( θ ) · v ( θ ) ) −3/2 √ ( v ( θ ) · v ( θ ) ) ( a ( θ ) · a ( θ ) ) − ( v ( θ ) · a ( θ ) ) 2dθ . 2 . Length Distortion in Euclidean space is defined as LE = length ( N ( θ ) ) length ( θ ) = ∫ √ ‖v ( θ ) ‖2dθ . 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 ( usingN ) is usually very complex , one may also expect that networks with better performance should also generate longer outputs . 3 . The problem of LE is that , stretched 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 ( θ ) / √ v ( θ ) · v ( θ ) changes with respect to θ , defined as LEκ = ∫ √ ‖∂θv̂ ( θ ) ‖2dθ . 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 L1 = L2 = L3 = L4 = 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 κ , LE , LEκ at initialization . As shown in Figure 2 , we can clearly see that both κ and LE 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 LE as our complexity measure for highly fast ViT topology search and scaling .

Summary: In this paper, the author proposes auto-scaling ViT, which is to search seed ViT topology based on number of kernels, attention splits, expansion ratio, depth and width jointly. The searched ViT is trained with a progressive re-tokenization scheme that saves ~40% training time and preserves the accuracy with less parameters. Experiments on ImageNet and MSCOCO show that AS-ViT can reach strong performances on classification and detection.