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: Initial variants of ViTs borrow their architecture from the seminal ‘attention is all you need’ work. Some recent methods amend the default transformer architecture to incorporate convolutions in the design (e.g., CCT: Compact Convolutional Transformers), or a hierarchical architecture with feature pyramids (e.g., Swin, PyramidViT etc) to make ViTs suitable for dense prediction tasks (detection and segmentation).  Instead of relying on hand-designed architectures, the paper proposes to automatically search for a ViT architecture, that is both efficient & accurate.