Searching the Search Space of Vision Transformer

1 Introduction . Vision transformer recently has drawn great attention in computer vision because of its high model capability and superior potentials in capturing long-range dependencies . Building on top of transformers [ 41 ] , modern state-of-the-art models , such as ViT [ 10 ] and DeiT [ 40 ] , are able to achieve competitive performance on image classification and downstream vision tasks [ 13 , 20 ] compared to CNN models . There are fast-growing efforts on manually designing more effective architectures to further unveil the power of vision transformer [ 27 , 53 , 15 ] . Neural architecture search ( NAS ) as a powerful technique for automating the design of networks has shown its superiority to manual design [ 58 , 39 ] . For NAS methods , search space is crucial because it determines the performance bounds of architectures during the search . It has been observed that the improvement of search space induced a favorable impact on the performance of many state-of-the-art models [ 2 , 16 , 50 ] . Researchers have devoted ample efforts to the design of CNN space [ 45 , 12 , 42 ] . However , limited attention is paid to Transformer counterparts , leading to the dilemma of search space design when finding more effective and efficient transformer architectures . In this paper , we propose to Search the Search Space ( S3 ) by automatically refining the main changeable dimensions of transformer architectures . In particular , we try to answer the following two key questions : 1 ) How to effectively and efficiently measure a specific search space ? 2 ) How to automatically change a defective search space to a good one without human prior knowledge ? For the first question , we propose a new measurement called E-T Error for search space quality estimation and utilize a once-for-all supernet for effective and efficient computing . Specifically , the “ E ” part in the E-T error is similar to the error empirical distribution function used in [ 35 ] focusing on ∗Equal contributions . Work done when Minghao , Kan , Bolin are interns of MSRA . †Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the overall quality , while the “ T ” part considers the quality of top-tier architectures in a search space . In contrast to [ 35 ] that trains each subnet for only a few epochs , we train a once-for-all supernet with AutoFormer [ 6 ] to obtain a large amount of well-trained subnets , providing a more reliable proxy for performance estimation . For the second question , as shown in Fig . 1 we decompose the search space by multiple dimensions , including depth , embedding dimension , MLP ratio , window size , number of heads and Q-K-V dimension , and progressively evolve each dimension to compose a better space . Particularly , we use linear parameterization to model the tendencies of different dimensions so as to guide the search . In addition , we analyze the search process of search space and provide some interesting observations and design guidelines : 1 ) The third stage is the most important one and increasing its number of blocks leads to performance improvement . 2 ) Shallow layers should use a small window size while deep layers should use a large window size . 3 ) MLP ratio is better to be progressively increased along with network depth . 4 ) Q-K-V dimension could be smaller than the embedding dimension without performance drop . We hope these observations will help both the manual design of vision transformer and the space design for automatic search . In summary , we make the following contributions : • We propose a new scheme , i.e . S3 , to automate the design of search space for vision transformer . We also present a new pipeline for vision transformer search with minimal human interactions . Moreover , we provide analysis and guidelines on vision transformer architecture design , which might be helpful for future architecture search and design . • The experiments on ImageNet verify the proposed automatic search space design method can improve the effectiveness of design space , thus boosting the performance of searched architectures . The discovered models achieve superior performance compared to the recent ViT [ 10 ] and Swin [ 27 ] transformer families under aligned settings . Moreover , the experiments on downstream vision and vision-language tasks , such as object detection , semantic segmentation and visual question answering , demonstrate the generality of the method . 2 Approach . 2.1 Problem Formulation . Most of existing NAS methods can be formulated as a constrained optimization problem as : α∗A = arg min α∈A L ( W ∗α ; Dval ) s.t . W ∗α = arg min Wα L ( Wα ; Dtrain ) , g ( α ) < c , ( 1 ) whereWα is the weights of the network architecture α , A is a predefined search space , L ( · ) is the loss function , Dtrain and Dval represent train and validation datasets respectively , g ( · ) denotes a function calculating resource consumption of a model while c denotes a given resource constraint . Ideally , the search space A is an infinite space Ω that contains all possible architectures . In practice , A is usually a small subspace of Ω due to limited computation resource . In this paper , we break the convention of searching in a fixed space by considering the architecture search together with the search space design . More concretely , we decouple the constrained optimization problem of Eq . ( 1 ) into three separate steps : Step 1 . Searching for an optimal search space under a specific constraint : A∗ = min A⊂Ω Q ( A ; Dval ) s.t . |A| ≤M , ( 2 ) where Q is a measurement of search space to be elaborated in Sec . 2.3 , and M represents the maximum search space size . Step 2 . Following one-shot NAS [ 12 , 34 ] , we encode the search space into a supernet and optimize its weights . This is formulated as solving the following problem : W ∗A∗ = arg min W Eα∈A∗ [ L ( Wα ; Dtrain ) ] , ( 3 ) where W is the weights of the supernet , Wα is the weights in W specified by the architecture α . Step 3 . After obtaining the well-trained supernet , we search for the optimal architecture via ranking the performance using the learned weights , which is formulated as : α∗A∗ = arg min α∈A∗ L ( W ∗α ; Dval ) , s.t . g ( α ) < c , ( 4 ) where W ∗α is the weights of architecture α inherited from W ∗ A∗ . 2.2 Basic Search Space . We set up the general architectures of vision transformer by following the settings in ViT [ 10 ] and Swin-T [ 27 ] as shown in Fig . 1 ( right ) . To be specific , given an input image , we uniformly split it into non-overlap patches . Those patches are linearly projected to vectors named patch embeddings . The embeddings are then fed into the transformer encoder , which performs most of the computation . At last , a fully connected layer is adopted for classification . The transformer encoder consists of four sequential stages with progressively reduced input resolutions . Each stage contains blocks with the same embedding dimension . Therefore , stage i has two search dimensions : the number of blocks di and embedding dimension hi . Each block contains a window-based multi-head self-attention ( WSA ) module and a feed-forward network ( FFN ) module . We do not force the blocks in one stage to be identical . Instead , block j in the ith stage have several search dimensions including window sizewij , number of heads n i j , MLP ratiomi , Q-K-V embedding dimension qij . 2.3 Searching the Search Space . Space Quality Evaluation . For a given space A , we propose to use E-T Error , denoted as Q ( A ) , for evaluating its quality . It is the average of two parts : expected error rate Qe ( A ) and top-tier error Qt ( A ) . For a given space A , the definition of Qe ( A ) is given as : Qe ( A ) = Eα∈A , g ( α ) < c [ eα ] , ( 5 ) where eα is the top-1 error rate of architecture α evaluated on ImageNet validation set , g ( · ) is the computation cost function and c is the computation resource constraint . It measures the overall quality of a search space . In practice , we use the average error rate of N randomly sampled architectures to approximate the expectation term . The top-tier error Qt ( A ) is the average error rate of top 50 candidates under resource constraints , representing the performance upper bound of the search space . Once-for-all Supernet . In constrast to RegNet [ 35 ] that trains hundreds of models with only a few epochs ( ∼10 ) and uses their errors as a performance proxy of the fully trained networks , we encode the search space into a supernet and optimize it similar to AutoFormer [ 6 ] . Under such training protocol , we are able to train an once-for-all transformer supernet that serves as a reliable and efficient performance indicator . Searching the Search Space . The search of search space mainly contains two iterative steps : 1 ) supernet optimization and 2 ) space evolution . Algorithm 1 Searching the Search Space Input : Infinite search space Ω , max iteration T , threshold τ for subspace evolved , the subspaces S ( 0 ) 1 , · · · , S ( 0 ) D , evolve steps for subspaces γ1 , · · · , γD Output : The most promising search space A∗ 1 : Initialize a search space A ( 0 ) from Ω with minimal prior human knowledge 2 : while t ∈ [ 0 , 1 , · · · , T ] do 3 : Optimize the weights W ( t ) A of supernet corresponding to the space A ( t ) 4 : Randomly sample N architectures from the well converged supernet 5 : for space S ( t ) i in S ( t ) 1 , · · · , S ( t ) D do 6 : Get the subspace { α|α ∈ A ( t ) ∩ αi = v ( t ) l } for each choice v ( t ) l ∈ S ( t ) i , where αi is the ith dimension value of the architecture α 7 : Calculate the E-T Error of each subspace according to Eq . ( 5 ) 8 : Evolve the subspace and get S ( t+1 ) Di according to Eq . ( 8 ) 9 : end for 10 : A∗ = A ( t+1 ) = S ( t+1 ) 1 × S ( t+1 ) 2 × · · · × S ( t+1 ) D 11 : end while Step 1 . For the search space A ( t ) of tth iteration , we encode it into a supernet . We adopt sandwich training [ 52 ] to optimize the weights of supernet , i.e. , sampling the largest , the smallest , and two middle models for each iteration and fusing their gradients to update the supernet . Notably , we do not use inplace-distill since it leads to collapse for transformer training . Step 2 . We first decompose the search space A ( t ) by the search dimensions , defined in Sec .2.2 , and four stages . Denote their corresponding subspaces as S ( t ) 1 , S ( t ) 2 , · · · , S ( t ) D , where D is the product of the number of stage and the number of search dimensions . The search space A ( t ) now can be expressed as the Cartesian product of all search spaces : A ( t ) = S ( t ) 1 × S ( t ) 2 × · · · × S ( t ) D , ( 6 ) Then the search space A ( t ) could be evolved to a better one A ( t+1 ) , with the update of S ( t ) i . Particularly , For a certain subspace S ( t ) i , we first get the subspace { α|α ∈ A ( t ) ∩αi = v ( t ) l } for each choice v ( t ) l ∈ S ( t ) i , where αi is the i th dimension value of the architecture α . Then we compute the E-T Errors of the subspaces . We hypothesize there is a strong correlation between the continuous choices and the E-T Error in most cases , so we fit a linear function to estimate its tendency for simplification : y = wx + b , ( 7 ) where x ∈ R is the choices in S ( t ) i , and y ∈ R is the corresponding E-T Error . w and b are both scalar parameters . Then the new subspace S ( t+1 ) i is given as : S ( t+1 ) i = { v ( t ) j − ⌊w τ ⌋ γi ∣∣∣ v ( t ) j ∈ S ( t ) i } , ( 8 ) where τ is the threshold for search space evolving , γi is the steps for search dimension Si . Note that we remove the choices if they become less than zero . The detailed algorithm is also shown in Alg . 1 .

Summary: This work propose to first search for the search space for transformers and then applies once-for-all NAS methods to search for architectures within the searched search space. In particular, the work proposes an E-T error metric to evaluate the quality of the search spaces. From the results of the whole search process, this work also summarizes some design guidelines for vision transformers. The experiments are extensive and they show the superiority of the proposed search method.