Contrastive Embeddings for Neural Architectures

The performance of algorithms for neural architecture search strongly depends on the parametrization of the search space . We use contrastive learning to identify networks across different initializations based on their data Jacobians and their number of parameters , and automatically produce the first architecture embeddings independent from the parametrization of the search space . Using our contrastive embeddings , we show that traditional black-box optimization algorithms , without modification , can reach state-of-the-art performance in Neural Architecture Search . As our method provides a unified embedding space , we successfully perform transfer learning between search spaces . Finally , we show the evolution of embeddings during training , motivating future studies into using embeddings at different training stages to gain a deeper understanding of the networks in a search space . 1 INTRODUCTION . Traditionally , the design of state-of-the-art neural network architectures is informed by domain knowledge and requires a large amount of manual work to find the best hyperparameters . However , automated architecture search methods have recently achieved state-of-the-art results on tasks such as image classification , object detection , semantic segmentation and speech recognition ( Ren et al. , 2020 ) . While a significant amount of work has gone into improving the algorithms used for architecture search , there has been only limited work on constructing the space in which these algorithms operate . There exists a vast number of different ways in which the architectures in a search space can be encoded , and the effects of these choices will affect the performance of the search algorithms . Previous embedding methods such as ( Yan et al. , 2020 ) have focused on preserving the edit distance between the computational graph of the architectures in the embedding space . Mellor et al . ( 2020 ) have showed that statistics computed on architectures at initialization , before training , can be used to infer which will perform better after training . Motivated by this , our intent is to automatically learn such statistics at initialization and generate an embedding space based on them , so that the resulting embedding space preserves fundamental properties of the architectures . To achieve our goal , we leverage contrastive learning , that has gathered interest in the computer vision community , and produced various state-of-the-art results ( He et al. , 2020 ; Chen et al. , 2020 ; Caron et al. , 2020 ; Grill et al. , 2020 ; Zbontar et al. , 2021 ) . In contrastive learning , the model learns an informative embedding space through a self-supervised pre-training phase : from the images in a batch , pairs are generated through random transformations , and the model is trained to generate similar ( dissimilar ) embeddings for similar ( dissimilar ) images . In this work , we combine contrastive learning with the theory presented by Wang et al . ( 2016 ) on the Jacobians of networks at initialization , in order to find an embedding space suitable for Neural Architecture Search . The embedding space that we generate is invariant to the search space of origin , allowing us to accomplish transfer learning between different search spaces . 1.1 MOTIVATIONS AND CONTRIBUTIONS . We design a method to produce architecture embeddings using Contrastive Learning and information available from their initial state . Our technique is capable of generating embeddings that are independent from the parametrization of the search space , and evolve during training . We leverage these contrastive embeddings in Neural Architecture Search using traditional black-box optimization algorithms . Moreover , since they provide a unified embedding space across search spaces , we exploit them to perform transfer learning between search spaces . Parametrization-independent embeddings NAS methods promise to outperform random search , however the encoding of the architectures must show some structure for the search algorithm to exploit . These encodings are typically produced by condensing all the parameters used to generate an architecture into a single vector . The method used to generate architectures from a search space thereby implicitly parametrizes it . The parametrization of the search space affects the performance of a NAS algorithm , as noted by e.g . White et al . ( 2020 ) ; Ying et al . ( 2019 ) ; Wang et al . ( 2019 ) . However , when performing architecture search , it is not feasible to test multiple different parametrizations of the search space and evaluate which one performs better : once we have started to evaluate networks in a search space , there is no reason to discard previous evaluations . While the current generation of NAS alleviates the need for experts in the design of architectures , now expert knowledge is needed to build and parametrize a search space compatible with the chosen search algorithm , implying that it is exceedingly difficult to outperform a simple random search . We present in Sec . 4 the first method to create networks embeddings without relying on their parametrization in any search space , through the combination of modern contrastive learning and the theory of data Jacobian matrices for neural architecture search . Evolution of the embeddings during training In Section 4.4 , we show how the embeddings vary during training , noting that the training procedure tends to connect areas with similar final test accuracy together . We hypothesize that this information could enable even more efficient architecture search methods in the future . Leveraging traditional black-box algorithms Existing methods to generate architecture embeddings rely on metrics from their computational graphs to identify similar architectures , either by explicitly trying to preserve the edit distance in the embedding space , or by leveraging more sophisticated methods from graph representation learning . Our method leverages the information contained in the Data Jacobian Matrix of networks at initialization to train a contrastive model . As such , it can produce embeddings that more meaningfully capture the structure of the search space . As a result , traditional black-box algorithms perform well for architecture search , as shown for NAS-Bench-201 ( Dong & Yang , 2020 ) in Section 5.1 . Transfer learning between search spaces Our method provides a unified embedding space , since it does not depend on the parametrization of networks in any search space . We exploit this feature to perform for the first time transfer learning between the two search spaces . In practice , we perform it between the size and the topology spaces in NATS-Bench ( Dong et al. , 2020 ) in Section 5.3 . 2 RELATED WORK . Neural Architecture Search Previous works have attempted to improve network embeddings : Klyuchnikov et al . ( 2020 ) use graph2vec ( Narayanan et al. , 2017 ) to find embeddings such that networks with the same computational graph share the same embeddings , and similarly Yan et al . ( 2020 ) produce embeddings that are invariant to graph isomorphisms . However , the method differs in that this work trains a variational graph isomorphic autoencoder to produce the embeddings . They show that their architecture embeddings arch2vec perform better on downstream architecture search than supervised alternatives , additionally the embeddings that they produce approximately preserves the edit distance of the DAGs in a continuous space . Wei et al . ( 2020 ) uses a contrastive loss to find a low dimensional metric space where the graph edit distances of the original parametrization is approximately preserved . In the absence of dense sampling , all of these works rely on the prior that the edit distance is a good predictor for relative performance . In contrast , our method , learns to find an embedding space based on intrinsic properties of the architectures . It can therefore discover properties about the architectures which are not present in their graph representation . Data Jacobian Methods based on the Jacobians with respect to the input of trained networks have been shown to provide valuable information for knowledge transfer and distillation ( Czarnecki et al. , 2017 ; Srinivas & Fleuret , 2018 ) , as well as analysis and regularization of networks ( Wang et al. , 2016 ) . Neural Tangent Kernel The Jacobian of the network with respect to the parameters is computed for inference with neural tangent kernels ( NTK ) ( Jacot et al. , 2018 ) . Using NTK as a proxy for NAS ( Park et al. , 2020 ) underperforms the neural network Gaussian process ( NNGP ) kernel . The NNGP provides an inexpensive signal for predicting if an architecture exceeds median performance , but it is worse than training for a few epochs in predicting the order of the top performing architectures . Contrastive learning Different techniques have been developed in contrastive learning . He et al . ( 2020 ) train a network with a contrastive loss against a memory bank of negative samples produced by a slowly moving average version of itself . Chen et al . ( 2020 ) remove the memory bank and just consider negative samples from within the same minibatch . Grill et al . ( 2020 ) remove the negative samples completely but stabilize the training by encoding the positive samples using a momentum encoder . Zbontar et al . ( 2021 ) use the correlation matrix between the features of the different augmentations and maximize the similarity of the same feature while minimizing the redundancy between features . 3 BACKGROUND . 3.1 TRADITIONAL ARCHITECTURE EMBEDDINGS . A decision tree is created either implicitly or explicitly to sample networks from a search space . To encode an architecture , one records all choices as the decision tree is traversed into a vector , which is then used as the embedding of the architecture . Without any additional knowledge , a NAS algorithm will assume that all choices in the decision tree have an equal importance on the characteristics of an architecture . One commonly used encoding scheme consists of choosing the operations on nodes along with a binary adjacency matrix connecting them . 3.2 DATA JACOBIAN . Extended Data Jacobian matrices are used by Wang et al . ( 2016 ) to analyze trained neural networks . We ground our work in their theoretical setting , and introduce the relevant concepts below . Multi Layer Perceptrons with ReLU activations are locally equivalent to a linear model : the ReLU after a linear layer can be combined into a single linear layer , where each row in the matrix is replaced by zeros if the output pre-activation is negative . ReLU ( Wx ) = Ŵx , Ŵij = { Wij if ( Wx ) j ≥ 0 0 otherwise Since matrix multiplication is closed , within a neighborhood where the signs of all neurons preactivation is constant , the full network can be replaced by a single matrix . This property can be extended to any model whose layers can be rewritten as linear layers , including convolutional layers and average pooling layers . Max pooling layers also retain this property , and can be treated similarly to ReLU . Therefore , in a local neighbourhood close to x , the full information of a network , f , is contained within its Data Jacobian Matrix ( DJM ) . DJM ( x ) = ∂f ( x ) ∂x and within that neighbourhood f ( x ) = DJM ( x ) x We can evaluate the Data Jacobian Matrix at many different points x to gather information about multiple different neighbourhoods . If we assume the network to have a single output , its DJM is a vector , and we can then stack the DJMs at different points to form the Extended Data Jacobian Matrix ( EDJM ) . If a network has multiple outputs we can sum them to get a single output , which we use to calculate the EDJM . Wang et al . ( 2016 ) use the singular values of the EDJM to compute a score , and empirically show that the score is correlated with the depth of the network , and its model capacity .

Summary: The paper proposes a self-supervised embedding learning method to learn embeddings of various-sized neural network architectures. Each network is first represented as a low rank projection of a Jacobian matrix, where the rows are Jacobians (output-averaged if multivariate) evaluated at various inputs at random initialization time, called EDJM. Since EDJM is random, multiple such representations are treated as positive pairs for contrastive learning. From the contrastively learned embedding, a further dimensionality-reduction stage is optimized to (1) preserve distances and (2) achieve uniform volume for each architecture. The main application of the final embedding is NAS, where the method outperform baselines.