CR-LSO: Convex Neural Architecture Optimization in the Latent Space of Graph Variational Autoencoder with Input Convex Neural Networks

TL;DR

CR-LSO addresses the challenge of gradient-based NAS by embedding neural architectures into a continuous latent space via a graph variational autoencoder (G-VAE) and enforcing a convex architecture-performance mapping with an input convex neural network (ICNN) predictor. The joint training with a semi-supervised GNN predictor enables learning from unlabeled architectures, while the ICNN regularization makes gradient optimization more reliable. Empirical results on NAS-Bench-101, NAS-Bench-201, and NAS-Bench-301 demonstrate competitive performance and improved sample efficiency across budgets, with clearer clustering of high-performing architectures in the convex latent space. This approach blends graph-based generative modeling with convex optimization principles to deliver scalable, gradient-based NAS with robust performance gains.

Abstract

In neural architecture search (NAS) methods based on latent space optimization (LSO), a deep generative model is trained to embed discrete neural architectures into a continuous latent space. In this case, different optimization algorithms that operate in the continuous space can be implemented to search neural architectures. However, the optimization of latent variables is challenging for gradient-based LSO since the mapping from the latent space to the architecture performance is generally non-convex. To tackle this problem, this paper develops a convexity regularized latent space optimization (CR-LSO) method, which aims to regularize the learning process of latent space in order to obtain a convex architecture performance mapping. Specifically, CR-LSO trains a graph variational autoencoder (G-VAE) to learn the continuous representations of discrete architectures. Simultaneously, the learning process of latent space is regularized by the guaranteed convexity of input convex neural networks (ICNNs). In this way, the G-VAE is forced to learn a convex mapping from the architecture representation to the architecture performance. Hereafter, the CR-LSO approximates the performance mapping using the ICNN and leverages the estimated gradient to optimize neural architecture representations. Experimental results on three popular NAS benchmarks show that CR-LSO achieves competitive evaluation results in terms of both computational complexity and architecture performance.

Figures (9)

• Figure 1: The conceptual visualization of the proposed CR-LSO. By using a graph variational autoencoder, CR-LSO transforms the discrete search space of NAS into a continuous latent space. Simultaneously, CR-LSO utilizes the guaranteed convexity of the ICNN to regularize the learning process of latent space, so as to obtain a convex architecture performance mapping, which makes the gradient-based LSO more effective.
• Figure 2: The PCA visualizations of architecture representations in the unconstrained latent space (LS) and the convexity regularized latent space (CR-LS) on NAS-Bench-101. The architectures with blue colors own higher rankings than those with red colors. Best
• Figure 3: The cosine similarity of architecture representations in the unconstrained latent space (LS) and the convexity regularized latent space (CR-LS) in NAS-Bench-201. Left: The architecture similarity among the optimal 100 architectures. Right: The architecture similarity between the optimal 100 and the worst 100 architectures.
• Figure 4: Two-dimensional t-SNE projection of latent representations of architectures searched by different methods. Architectures with lower rankings (darker colors) own higher performances on NAS-Bench-301.
• Figure 5: Visual comparison of architecture performance prediction with different VAE encoders.