Towards Code-Oriented LM Embeddings for Surrogate-Assisted Neural Architecture Search

TL;DR

This paper introduces COLE, a low-cost, code-based embedding method leveraging language models to improve neural architecture search efficiency without extensive fine-tuning.

Contribution

It proposes a novel code-oriented embedding strategy using off-the-shelf language models for surrogate modeling in NAS, reducing overhead and enhancing search performance.

Findings

Raw code inputs outperform other text encodings in predictive accuracy.

COLE improves surrogate-assisted search efficiency in NAS-Bench-201.

Replacing structural encodings with COLE reduces evaluation budget by 34%.

Abstract

Developing effective surrogates (performance predictors) for Neural Architecture Search (NAS) typically requires expensive fine-tuning or the engineering of complex representations. We propose a low-cost embedding strategy that leverages the inductive bias of Language Models (LMs) to eliminate these overheads. By representing architectures as PyTorch class definition text, we demonstrate that off-the-shelf LMs act as competitive feature extractors without NAS-specialized fine-tuning. The final predictor is constructed by passing the extracted Code-Oriented LM Embeddings (COLE) through a lightweight regression head. We also investigate strategies to improve embedding quality and utilization. Our experiments on the NAS-Bench-201 and einspace search spaces reveal that raw code inputs yield higher predictive performance than other text-based encodings (e.g., ONNX-to-text encodings) when using frozen LMs. We also observe COLE drives superior surrogate-assisted search using the BANANAS algorithm in NAS-Bench-201. When optimizing for CIFAR-100 performance, replacing structural path encodings with COLE for architecture representation allows for a 34% decrease in the evaluation budget required to reach within 1% of the fittest architecture in the search space (by test accuracy). As any neural architecture can be represented as code, these findings establish COLE as a versatile and efficient foundation for advancing NAS.