🤖 AI Summary
This work addresses the high computational cost of conventional neural architecture search (NAS) performance predictors, which often rely on expensive fine-tuning or intricate architecture representations. The authors propose Code-Oriented Language Model Embeddings (COLE), a method that directly uses raw PyTorch class definition code as input and leverages a frozen off-the-shelf language model to extract architecture embeddings. Coupled with a lightweight regression head, COLE constructs an efficient performance predictor without requiring NAS-specific fine-tuning. Evaluated on the NAS-Bench-201 benchmark, COLE achieves within 1% of the optimal architecture’s accuracy while reducing the evaluation budget by 34% compared to path-based encoding. Furthermore, experiments on CIFAR-100 demonstrate its strong generalization capability and superior search efficiency.
📝 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.