Task-Adaptive Neural Network Search with Meta-Contrastive Learning

1 Introduction . Neural Architecture Search ( NAS ) aims to automate the design process of network architectures by searching for high-performing architectures with RL [ 76 , 77 ] , evolutionary algorithms [ 43 , 11 ] , parameter sharing [ 6 , 42 ] , or surrogate schemes [ 38 ] , to overcome the excessive cost of trial-and-error approaches with the manual design of neural architectures [ 47 , 23 , 27 ] . Despite their success , existing NAS methods suffer from several limitations , which hinder their applicability to practical scenarios . First of all , the search for the optimal architectures usually requires a large amount of computation , which can take multiple GPU hours or even days to finish . This excessive computation cost makes it difficult to efficiently obtain an optimal architecture for a novel dataset . Secondly , most NAS approaches only search for optimal architectures , without the consideration of their parameter values . Thus , they require extra computations and time for training on the new task , in addition to the architecture search cost , which is already excessively high . For this reason , supernet-based methods [ 8 , 37 ] that search for a sub-network ( subnet ) from a network pretrained on large-scale data , are attracting more popularity as it eliminates the need for additional ⇤Equal contribution . †This work was done while the author was interning at AITRICS . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . training . However , this approach may be suboptimal when we want to find the subnet for a dataset that is largely different from the source dataset the supernet is trained on ( e.g . medical images or defect detection for semiconductors ) . This is a common limitation of existing NAS approaches , although the problem did not receive much attention due to the consideration of only a few datasets in the NAS benchmarks ( See Figure 1 , left ) . However , in real-world scenarios , NAS approaches should search over diverse datasets with heterogeneous distributions , and thus it is important to task-adaptively search for the architecture and parameter values for a given dataset . Recently , MetaD2A [ 31 ] has utilized meta-learning to learn common knowledge for NAS across tasks , to rapidly adapt to unseen tasks . However it does not consider parameters for the searched architecture , and thus still requires additional training on unseen datasets . Given such limitations of NAS and meta-NAS methods , we introduce a novel problem of Neural Network Search ( NNS ) , whose goal is to search for the optimal pretrained networks for a given dataset and conditions ( e.g . number of parameters ) . To tackle the problem , we propose a novel and extremely efficient task-adaptive neural network retrieval framework that searches for the optimal neural network with both the architecture and the parameters for a given task , based on cross-modal retrieval . In other words , instead of searching for an optimal architecture from scratch or taking a sub-network from a single super-net , we retrieve the most optimal network for a given dataset in a task-adaptive manner ( See Figure 1 , right ) , by searching through the model zoo that contains neural networks pretrained on diverse datasets . We first start with the construction of the model zoo , by pretraining state-of-the-art architectures on diverse real-world datasets . Then , we train our retrieval model via amortized meta-learning of a cross-modal latent space with a contrastive learning objective . Specifically , we encode each dataset with a set encoder and obtain functional and topological embeddings of a network , such that a dataset is embedded closer to the network that performs well on it while minimizing the similarity between irrelevant dataset-network pairs . The learning process is further guided by a performance predictor , which predicts the model ’ s performance on a given dataset . The proposed Task-Adaptive Network Search ( TANS ) largely outper- forms conventional NAS/AutoML methods ( See Figure 2 ) , while significantly reducing the search time . This is because the retrieval of a trained network can be done instantly without any additional architecture search cost , and retrieving a task-relevant network will further reduce the fine-tuning cost . To evaluate the proposed TANS , we first demonstrate the sample-efficiency of our model zoo construction method , over construction with random sampling of dataset-network pairs . Then , we show that the TANS can adaptively retrieve the best-fitted models for an unseen dataset . Finally , we show that our method significantly outperforms baseline NAS/AutoML methods on real-world datasets ( Figure 2 ) , with incomparably smaller computational cost to reach the target performance . In sum , our main contributions are as follows : • We consider a novel problem of Neural Network Search , whose goal is to search for the optimal network for a given task , including both the architecture and the parameters . • We propose a novel cross-modal retrieval framework to retrieve a pretrained network from the model zoo for a given task via amortized meta-learning with constrastive objective . • We propose an efficient model-zoo construction method to construct an effective database of dataset-architecture pairs , considering both the model performance and task diversity . • We train and validate TANS on a newly collected large-scale database , on which our method outperforms all NAS/AutoML baselines with almost no architecture search cost and significantly fewer fine-tuning steps . 2 Related Work . Neural Architecture Search Neural Architecture Search ( NAS ) , which aims to automate the design of neural architectures , is an active topic of research nowadays . Earlier NAS methods use non-differentiable search techniques based on RL [ 76 , 77 ] or evolutionary algorithms [ 43 , 11 ] . However , their excessive computational requirements [ 44 ] in the search process limits their practical applicability to resource-limited settings . To tackle this challenge , one-shot methods share the parameters [ 42 , 6 , 35 , 65 ] among architectures , which reduces the search cost by orders of magnitude . The surrogate scheme predicts the performance of architectures without directly training them [ 38 , 75 , 54 ] , which also cuts down the search cost . Latent space-based NAS methods [ 38 , 54 , 67 ] learn latent embeddings of the architectures to reconstruct them for a specific task . Recently , supernetbased approaches , such as OFA [ 8 ] , receive the most attention due to their high-performances . OFA generates the subnet with its parameters by splitting the trained supernet . While this eliminates the need for costly re-training of each searched architecture from scratch , but , it only trains a fixedsupernet on a single dataset ( ImageNet-1K ) , which limits their effectiveness on diverse tasks that are largely different from the training set . Whereas our TANS task-adaptively retrieves a trained neural network from a database of networks with varying architectures trained on diverse datasets . Meta-Learning The goal of meta-learning [ 55 ] is to learn a model to generalize over the distribution of tasks , instead of instances from a single task , such that a meta-learner trained across multiple tasks can rapidly adapt to a novel task . While most meta-learning methods consider few-shot classification with a fixed architecture [ 56 , 20 , 48 , 40 , 33 , 30 ] , there are a few recent studies that couple NAS with meta-learning [ 46 , 34 , 17 ] to search for the well-fitted architecture for the given task . However , these NAS approaches are limited to small-scale tasks due to the cost of roll-out gradient steps . To tackle this issue , MetaD2A [ 31 ] proposes to generate task-dependent architectures with amortized meta-learning , but does not consider parameters for the searched architecture , and thus requires additional cost of training it on unseen datasets . To overcome these limitations , our method retrieves the best-fitted architecture with its parameters for the target task , by learning a cross-modal latent space for dataset-network pairs with amortized meta-learning . Neural Retrieval Neural retrieval aims to search for and return the best-fitted item for the given query , by learning to embed items in a latent space with a neural network . Such approaches can be broadly classified into models for image retrieval [ 21 , 14 , 66 ] and text retrieval [ 73 , 9 , 63 ] . Crossmodal retrieval approaches [ 32 , 74 , 58 ] handle retrieval across different modalities of data ( e.g . image and text ) , by learning a common representation space to measure the similarity across the instances from different modalities . To our knowledge , none of the existing works is directly related to our approach that performs cross-modal retrieval of neural networks given datasets . 3 Methodology . We first define the Neural Network Search ( NNS ) problem and propose a meta-contrastive learning framework to learn a cross-modal retrieval space . We then describe the structural details of the query and model encoders , and an efficient model-zoo construction strategy . 3.1 Problem Definition . Our goal in NNS is to search for an optimal network ( with both architectures and parameters ) for a dataset and constraints , by learning a cross-modal latent space over the dataset-network pairs . We first describe the task-adaptive network retrieval with amortized meta-contrastive learning . 3.1.1 Meta-Training . We assume that we have a database of neural networks ( model zoo ) pretrained over a distribution of tasks p ( ⌧ ) with each task ⌧ = { D⌧ , M⌧ , s⌧ } , where D⌧ denotes a dataset , M⌧ denotes a neural network ( or model ) trained on the dataset , and s⌧ denotes a set of auxiliary constraints for the network to be found ( e.g . number of parameters or the inference latency ) . Then , our goal is to learn a cross-modal latent space for the dataset-network pairs ( D⌧ , M⌧ ) while considering the constraints s⌧ over the task distribution p ( ⌧ ) , as illustrated in Figure 1 . In this way , a meta-trained model from diverse ( D⌧ , M⌧ ) pairs , will rapidly generalize on an unseen dataset D̃ ; D⌧ \ D̃ = ; by retrieving a well-fitted neural network on the dataset that satisfies the conditions s⌧ . Task-Adaptive Neural Network Retrieval To learn a cross-modal latent space for dataset-network pairs over a task distribution , we first introduce a novel task-adaptive neural network retrieval problem . The goal of task-adaptive retrieval is to find an appropriate network M⌧ given the query dataset D⌧ for task ⌧ . To this end , we need to calculate the similarity between the dataset-network pair ( D⌧ , M⌧ ) 2 Q⇥M , with a scoring function f that outputs the similarity between them as follows : max ✓ , X ⌧2p ( ⌧ ) f ( q , m ) , q = EQ ( D ⌧ ; ✓ ) and m = EM ( M⌧ ; ) , ( 1 ) where EQ : Q ! Rd is a query ( dataset ) encoder , EM : M ! Rd is a model encoder , which are parameterized with the parameter ✓ and respectively , and fsim : Rd ⇥ Rd ! R is a scoring function for the query-model pair . In this way , we can construct the cross-modal latent space for dataset-network pairs over the task distribution with equation 1 , and use this space to rapidly retrieve the well-fitted neural network in response to the unseen query dataset . We can learn such a cross-modal latent space of dataset-network pairs for rapid retrieval by directly solving for the above objective , with the assumption that we have the query and the model encoder : Q and M . However , we further propose a contrastive loss to maximize the similarity between a dataset and a network that obtains high performance on it in the learned latent space , and minimize the similarity between irrelevant dataset-network pairs , inspired by Faghri et al . [ 19 ] , Engilberge et al . [ 18 ] . While existing works such as Faghri et al . [ 19 ] , Engilberge et al . [ 18 ] target image-to-text retrieval , we tackle the problem of cross-modal retrieval across datasets and networks , which is a nontrivial problem as it requires task-level meta-learning . Retrieval with Meta-Contrastive Learning Our meta-contrastive learning objective for each task ⌧ 2 p ( ⌧ ) consisting of a dataset-model pair ( D⌧ , M⌧ ) 2 Q ⇥ M , aims to maximize the similarity between positive pairs : fsim ( q , m+ ) , while minimizing the similarity between negative pairs : fsim ( q , m ) , where m+ is obtained from the sampled target task ⌧ 2 p ( ⌧ ) and m is obtained from other tasks 2 p ( ⌧ ) ; 6= ⌧ , which is illustrated in Figure 3 . This meta-contrastive learning loss can be formally defined as follows : Lm ( ⌧ ; ✓ , ) = L ( fsim ( q , m+ ) , fsim ( q , m ) ; ✓ , ) ( 2 ) q = EQ ( D ⌧ ; ✓ ) , m+ = EM ( M ⌧ ; ) , m = EM ( M ; ) . We then introduce L for the meta-contrastive learning : max 0 @ 0 , ↵ log exp ( fsim ( q , m + ) ) exp ⇣P 2p ( ⌧ ) , 6=⌧ fsim ( q , m ) ⌘ 1 A , ( 3 ) where ↵ 2 R is a margin hyper-parameter and the score function fsim is the cosine similarity . The contrastive loss promotes the positive ( q , m+ ) embedding pair to be close together , with at most margin ↵ closer than the negative ( q , m ) embedding pairs in the learned cross-modal metric space . Note that , similar to this , we also contrast the query with its corresponding model : Lq ( fsim ( q+ , m ) , fsim ( q , m ) ) , which we describe in the supplementary material in detail . With the above ingredients , we minimize the meta-contrastive learning loss over a task distribution p ( ⌧ ) , defined with the model Lm and query Lq contrastive losses , as follows : min , ✓ X ⌧2p ( ⌧ ) Lm ( ⌧ ; ✓ , ) + Lq ( ⌧ ; ✓ , ) . ( 4 ) Meta-Performance Surrogate Model We propose the meta-performance surrogate model to predict the performance on an unseen dataset without directly training on it , which is highly practical in real-world scenarios since it is expensive to iteratively train models for every dataset to measure the performance . Thus , we meta-train a performance surrogate model a = S ( ⌧ ; ) over a distribution of tasks p ( ⌧ ) on the model-zoo database . This model not only accurately predicts the performance a of a network M⌧ on an unseen dataset D⌧ , but also guides the learning of the cross-modal retrieval space , thus embedding a neural network closer to datasets that it performs well on . Specifically , the proposed surrogate model S takes a query embedding q⌧ and a model embedding m⌧ as an input for the given task ⌧ , and then forwards them to predict the accuracy of the model for the query . We train this performance predictor S ( ⌧ ; ) to minimize the mean-square error loss Ls ( ⌧ ; ) = ( s⌧acc S ( ⌧ ; ) ) 2 between the predicted accuracy S ( ⌧ ; ) and the true accuracy s⌧acc for the model on each task ⌧ , which is sampled from the task distribution p ( ⌧ ) . Then , we combine this objective with retrieval objective in equation 4 to train the entire framework as follows : min , ✓ , X ⌧2p ( ⌧ ) Lm ( ⌧ ; ✓ , ) + Lq ( ⌧ ; ✓ , ) + · Ls ( ⌧ ; ) , ( 5 ) where is a hyper-parameter for weighting losses .

Summary: This paper proposes and tackles the following task: given an unseen dataset, selecting the optimal pre-trained neural network from a model zoo of networks which will obtain the best fine-tuned performance on the new dataset. To solve this problem they propose a novel contrastive learning method to learn a shared embedding space for the models and datasets which minimizes the distance of the dataset embeddings and model embeddings of models that can be well adapted to the datasets. They propose encoding the properties of a pretrained model using a combination of a topological and functional encoding and demonstrate the effectiveness of this encoding scheme. They also propose a novel method for efficient model-zoo construction and build a dataset and train their algorithm on a collection of Kaggle datasets. They demonstrate that they can achieve very high performance on unseen datasets with few training steps.