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 applies NAS to the vision transformer architecture. Besides searching the architectures, it also automatically adjusts the search space guided by the E-T error that combines the error distribution and top model errors to evaluate the quality of different choices in a search dimension. By fitting a linear function to approximate the E-T error, the new search space is updated through a simple formula. They used the supernet to efficiently compute the E-T errors, and use evolutionary search to search for the best algorithm in a given search space. The discovered architecture is shown to be competitive to the state-of-the-art models, and also generalize to other vision tasks.