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 a new Task-Adaptive Neural network search framework (TANS). Such a task needs to retrieve a good network in well pre-trained dataset-network sets. In order to obtain a task-adaptive neural network, the amortized meta-learning method is employed. Meanwhile, this meta-learning framework is trained with contrastive learning to maximize the similarity between the positive dataset-network pairs and minimize the similarity between the negative pairs. The experiments show that TANS can rapidly retrieve a well-fitted network for unseen datasets.