TNASP: A Transformer-based NAS Predictor with a Self-evolution Framework

1 Introduction . Neural Architecture Search ( NAS ) aims to automatically find out superb architectures in a predefined search space . The NAS models have outperformed human-designed models in many domains [ 37 , 18 , 13 , 3 , 15 ] . However , traditional NAS methods like Reinforcement learning and Evolutionary learning require enormous computation resources , i.e. , hundreds or thousands of GPU hours , to train sub-models to obtain their performance estimation . The intensive computation prohibits NAS models from deployments to real applications . Although differentiable NAS methods [ 26 ] have less search time than traditional methods , they usually suffer from performance collapse due to several problems such as optimization gap [ 6 ] , discretization discrepancy and the unfair advantage [ 9 ] , trapping into sharp minimum and the dominant eigenvalue [ 49 , 5 ] . To reduce the search cost in NAS , predictor-based NAS methods [ 30 , 23 , 43 , 29 , 32 , 7 , 46 , 4 , 1 ] use performance predictors to predict the accuracy of architectures quickly instead of training all architectures to get the accuracy . Simple training-free predictors have been shown promising in some applications , however , those performances are usually not good enough in practice . Many 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . works focus on how to design effective training-based predictors [ 43 , 29 , 32 , 7 , 46 ] . A training-based predictor usually consists of an encoder module and a regressor module , and only needs to learn from a few architecture-accuracy sampled pairs , thus leading to a fast learning procedure . After the performance predictor learns the relationship between the architecture and its corresponding accuracy , it is able to predict the performance of other unseen architectures in the search space , which greatly accelerates the search process of NAS . One of the key components of a performance predictor is to encode discrete architectures into continuous feature representations . Neural predictor [ 43 ] and CTNAS [ 7 ] applied GCN [ 20 ] to capture the feature representation of input model structures . Both SemiNAS [ 29 ] and GATES [ 32 ] learned an embedding matrix for the candidate operations in the search space , and represented architectures as a combination of different embeddings . ReNAS [ 46 ] calculated the types matrix , flops matrix , and parameters matrix , and then concatenated them together to form a feature tensor to represent a specific architecture . Unlike previous methods , we propose a Transformer-based NAS performance Predictor ( TNASP ) and use the linear transformation output over the Laplacian matrix of the model structure graph to be the positional encoding . There are several advantages of the Transformer that can be used to train a good performance predictor . First , the self-attention module can help explore better feature representations from the graph structure data . Second , the multi-head mechanism can further help encode the different subspace information at different positions from the graph structure data , as also claimed by the original Transformer paper . Third , the Laplacian matrix based positional encoding method also fits well to find topology position information on the graph . In summary , we demonstrate that Transformer is an effective method to extract feature representation from discrete architecture graphs , and also has superb generalization abilities for processing unseen data , as shown in our experiments . Since the predictors are trained on a small size proxy dataset but the test dataset is much larger , training-based NAS predictors usually have poor generalization abilities . The powerful Transformer encoder can mitigate this problem to a certain degree , due to its good ability to encode topology information . To further improve our predictor , we introduce a self-evolution framework , which makes full use of temporal information to guide the training . The framework iteratively involves each evaluation score of the previously predicted results on a validation dataset into a gradient based optimization iteration as constraints , to push the predictions close to ground truths gradually . We demonstrate the framework can make our predictor have better generalization , and also achieve better performance than previous predictors . Our proposed framework is fruitful and practical in several scenarios , for example , all of the machine learning competitions . In competition , the test set that the host provided to users is exactly the validation role in our setting . And the temporal information , such as the feedback of each submission , is able to gradually improve original training iteratively with our framework . Generally speaking , our contributions can be summarized as follows : • We propose a Transformer-based NAS performance predictor ( TNASP ) to better encode the spatial topology information , utilizing the multi-head self-attention mechanism to map the discrete architectures to a meaningful feature representation and applying the linear transformation of the Laplacian matrix as the positional encoding . • By leveraging each evaluation score information in history as constraints in training , and applying a gradient based optimization method to iteratively solve the constrained optimization problem , we introduce a generic self-evolution framework to further improve the performance of the proposed predictor , making full use of temporal information . • Our proposed method surpasses the previous state-of-the-art methods under the same proxy training dataset and achieves state-of-the-art results on NAS-Bench-101 [ 48 ] , NAS-Bench201 [ 14 ] , and DARTS [ 26 ] search space . 2 Related Work . Due to the enormous search cost of traditional NAS methods such as reinforcement learning based methods [ 52 , 2 ] and evolutionary methods [ 34 , 27 ] , network performance predictor based NAS methods have become a pretty active topic recently . Most works only require a few architectureaccuracy data pairs to train a predictor and estimate the performance of unseen architectures , which can be categorized as training-based network performance predictors . Other works propose to calculate some kinds of metrics over network structure to represent network performance without training , which are denoted as training-free network performance predictors . Training-based network performance predictors Training-based predictors follow a unified paradigm to learn the correlation from network architectures and their corresponding accuracy . As it ’ s hard to learn useful features directly from the discrete network architectures , various methods have been explored to map the discrete representation into a continuous latent space and can be roughly divided into sequence-based schemes and graph-based methods . Sequence-based schemes denote each architecture by a discrete sequence with fixed length and use MLP [ 25 , 44 , 46 ] , embedding matrix [ 11 , 32 ] , Auto-Encoder [ 30 , 51 , 29 ] or GBDT [ 28 ] to convert the sequence into continuous representation . However , graph-based methods treat the architecture as a graph and use the graphformat data i.e . adjacency matrix and node features as inputs . Diverse graph processing methods such as GHN [ 50 ] , GCN [ 22 , 43 , 36 , 7 ] , GIN [ 47 ] and WL-Kernel [ 35 ] have been tried . Similar to these methods , our Transformer-based predictor also belongs to the training-based network performance predictors category . Differently , we propose a novel encoding scheme by utilizing the Laplacian matrix as the positional encoding and leveraging the multi-head self-attention mechanism from the Transformer [ 41 ] to encode network architectures to excavate more representative features . Training-free network performance predictors Recently , several works proposed to predict the network performance by designing some metrics , without training a model to perform prediction . [ 31 ] scored networks at initialization by computing correspondence between binary codes in the whole mini-batch . TE-NAS [ 4 ] analyzed trainability and expressivity to evaluate each network performance . Zero-Cost NAS [ 1 ] compared six proxies to predict network performance , such as grasp [ 42 ] , fisher [ 39 , 40 ] , synflow [ 38 ] , and so on , but failed to rank architectures in NAS-Bench-101 [ 48 ] . Zen-NAS [ 24 ] measured the expressivity of a network by expected Gaussian complexity to represent its performance . Although these training-free network performance predictors are efficient and have achieved satisfying results on several datasets , their robustness can not be guaranteed and the performance fluctuated dramatically among different tasks . Compared to these methods , our training-based predictor is more time-consuming , but its performance is dramatically better . 3 Methods . In this section , we first briefly review the common paradigm for the training-based NAS predictors as preliminary . Following the same paradigm , we then introduce the details of our Transformer-based predictor . Furthermore , we introduce a general framework to utilize the information in evaluation histories over a validation dataset as constraints in a gradient based optimization to advance the predictor by a self-evolution framework . 3.1 Training-based network performance predictors . Previous works [ 30 , 43 , 46 , 7 ] proposed to first apply an encoder fE to encode the discrete architectures into continuous feature representations , which can be formulated as : e = fE ( A , κ ) ( 1 ) where A ∈ RN×N denotes the adjacency matrix and stands for the directed acyclic connections between nodes , N denotes the number of the nodes . κ ∈ RN×F stands for the feature matrix and represents the characteristics of the nodes , F denotes the output dimension of the embedding extractor . For NAS predictors , the adjacency matrix A shows the topology information of an architecture and κ usually indicates the representation of operations for nodes or edges . Among previous works , the encoder fE can be a GCN or LSTM or simply an embedding matrix , and the embedding vector e can be interpreted as a latent representation of a specific architecture . After encoding the discrete architecture into a continuous representation using an encoder , it is easier and more accurate to estimate the network accuracy ŷ by a simple regressor fR once the embedding vector e meaningfully represents the architecture in the latent space . ŷ = fR ( e ) ( 2 ) Generally , we apply the MSE loss to supervise the training of the encoder and the regressor , and we use θ to denote the parameters of both encoder fE and regressor fR . The optimal parameters θ ? can be optimized by stochastic gradient descent method : θ ? = arg min θ 1 n n∑ i=1 ( ŷi − yi ) 2 ( 3 ) where n represents the total number of the training data , yi denotes the ground truth of the i-th sample . 3.2 Transformer-based predictor . Recently , the Transformer architecture [ 41 ] has attracted much attention and outperformed many state-of-the-art models due to its powerful representation ability , which motivates us to borrow the decent encoder structure to be our predictor backbone . As shown in Fig.1 , we first get the operation feature e1 ∈ RN×M by transforming the operation vector κ with an embedding matrix E ∈ RF×M . e1 = E ( κ ) ( 4 ) To characterize the positional information , we have tried various encoding methods , and find that the Laplacian matrix is an effective scheme and has the potential to better represent topology information of graph structure data as it contains both the connectivity of a graph and the centrality of each node . Thus we propose to encode the positional information by utilizing the Laplacian matrix ( L ∈ RN×N ) , computed from the adjacency matrix ( A ∈ RN×N ) and the degree matrix ( D ∈ RN×N ) : L = D −A ( 5 ) We did not choose the normalized Laplacian matrix , because we find it does not perform well in our experiments . Then we use a linear layer to map the the Laplacian matrix ( L ) to a continuous feature vector e2 , whose dimension is same as the embedding vector e1 : e2 = MLP ( L ) ( 6 ) Finally , we get the the continuous representation from the Transformer encoder with multi-head self-attention module : e = ETransformer ( e1 + e2 ) ( 7 ) As the continuous representation e contains richer information about the architecture ( including the operations and the connections ) in the latent space , it ’ s easier to learn the relationship between this representation and its corresponding accuracy . Therefore , we only choose a simple regressor , specifically 2 Multi-Layer Perceptions ( MLP ) , to estimate the final accuracy . The experiments demonstrate the promising performance of our whole designs .

Summary: This work leverages a transformer-based predictor network to approximate the accuracies of unseen architectures from a search space, for the purpose of NAS. The authors further proposed a ADMM-based method to accelerate the convergence of the predictor training. Experiments are conducted on NAS benchmarks and DARTS spaces.