Task-Adaptive Neural Network Search with Meta-Contrastive Learning

TL;DR

The paper introduces Task-Adaptive Neural Network Search (TANS), a meta-learning framework that efficiently finds optimal pretrained networks for new datasets from a model zoo, outperforming existing NAS methods in speed and accuracy.

Contribution

It proposes a novel meta-learning approach using contrastive loss to select pretrained networks tailored to specific datasets, addressing limitations of traditional NAS methods.

Findings

TANS outperforms baseline NAS methods on ten real-world datasets.

It retrieves high-performing networks with fewer training steps.

The method reduces total cost of obtaining task-specific networks.

Abstract

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.