NeuralMAE: Data-Efficient Neural Architecture Predictor with Masked Autoencoder

Predictor-based Neural Architecture Search (NAS) offers a promising solution for enhancing the efficiency of traditional NAS methods. However, it is non-trivial to train the predictor with limited architecture evaluations for efficient NAS. While current approaches typically focus on better utilizing the labeled architectures, the valuable knowledge contained in unlabeled data remains unexplored. In this paper, we propose a self-supervised transformer-based model that effectively leverages unlabeled data to learn meaningful representations of neural architectures, reducing the reliance on labeled data to train a high-performance predictor. Specifically, the predictor is pre-trained with a masking strategy to reconstruct input features in both latent and raw data spaces. To further enhance its representative capability, we introduce a multi-head attention-masking mechanism that guides the model to attend to different representation subspaces from both explicit and implicit perspectives. Extensive experimental results on NAS-Bench-101, NAS-Bench-201 and NAS-Bench-301 demonstrate that our predictor requires less labeled data and achieves superior performance compared to existing predictors. Furthermore, when combined with search strategies, our predictor exhibits promising capability in discovering high-quality architectures.