Task-Adaptive Neural Network Search with Meta-Contrastive Learning
This work addresses the inefficiency of conventional NAS methods for tasks dissimilar from pre-training datasets, offering a more adaptive and cost-effective solution for machine learning practitioners.
The paper tackles the problem of finding the optimal pretrained neural network for a new dataset and constraints from a model zoo, proposing a task-adaptive search framework that uses meta-contrastive learning to match datasets with high-performing networks. The results show that the method retrieves networks that outperform existing NAS/AutoML baselines with significantly fewer training steps, reducing the cost of obtaining task-optimal networks.
Most conventional Neural Architecture Search (NAS) approaches are limited in that they only generate architectures without searching for the optimal parameters. While some NAS methods handle this issue by utilizing a supernet trained on a large-scale dataset such as ImageNet, they may be suboptimal if the target tasks are highly dissimilar from the dataset the supernet is trained on. To address such limitations, we introduce a novel problem of \emph{Neural Network Search} (NNS), whose goal is to search for the optimal pretrained network for a novel dataset and constraints (e.g. number of parameters), from a model zoo. Then, we propose a novel framework to tackle the problem, namely \emph{Task-Adaptive Neural Network Search} (TANS). Given a model-zoo that consists of network pretrained on diverse datasets, we use a novel amortized meta-learning framework to learn a cross-modal latent space with contrastive loss, to maximize the similarity between a dataset and a high-performing network on it, and minimize the similarity between irrelevant dataset-network pairs. We validate the effectiveness and efficiency of our method on ten real-world datasets, against existing NAS/AutoML baselines. The results show that our method instantly retrieves networks that outperform models obtained with the baselines with significantly fewer training steps to reach the target performance, thus minimizing the total cost of obtaining a task-optimal network. Our code and the model-zoo are available at https://github.com/wyjeong/TANS.