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

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 LLMs (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.

• The paper introduces COLE, which converts neural architectures into PyTorch class definitions, extracts frozen language-model embeddings, and trains a lightweight rank-based surrogate predictor.
• COLE achieves Kendall’s Tau of about 0.71 on einspace with only 220 training architectures and consistently outperforms ONNX and derivation-tree text encodings at larger sample budgets.
• Integrating COLE with BANANAS reduces evaluations needed to approach the best architecture by 13% on CIFAR-10, 34% on CIFAR-100, and 18% on ImageNet16-120, although broader validation is needed.

Overview

This paper introduces COLE (Code-Oriented LM Embeddings), a surrogate-modeling strategy for Neural Architecture Search (NAS) that represents architectures as raw PyTorch class definition text and embeds them using frozen, off-the-shelf LLMs. The central claim is that the fine-tuning overhead required by prior LM-based NAS surrogates is unnecessary: because LMs are pre-trained on large code corpora, native code representations align with their learned distributions and yield competitive performance predictors out of the box. The authors validate this claim on NAS-Bench-201 and einspace, and demonstrate downstream search gains when integrating COLE into the BANANAS algorithm within NASLib (2605.15649).

Methodology

The pipeline consists of three stages. First, native code stringification: architectures are deterministically translated into valid torch.nn.Module subclass definitions. For NAS-Bench-201, cell strings are parsed into Cell classes; for einspace, a recursive transpiler converts derivation trees into PyTorch modules, mapping non-terminal grammar nodes to container classes (SequentialModule, BranchingModule, RoutingModule, ComputationModule) and terminal nodes to atomic primitives. Second, frozen embedding extraction: hidden states from the last layer of a frozen LM are token-wise mean pooled, following evidence that mean pooling outperforms last-token pooling on regression tasks. CodeLlama Python 7B is the default embedder due to its Python-centric pre-training. Third, surrogate head: embeddings are reduced via PCA to 128 dimensions and passed through a three-layer MLP trained with a Pairwise Hinge Loss optimizing Kendall's Tau rank correlation, aligning the surrogate objective with rank-based candidate selection in search.

Notably, the framework is robust to input verbosity: Helper Method, Inline, and Excluded Helper Method modes perform comparably, indicating the LM infers operation semantics without full implementations. The token-efficient Excluded Helper mode is adopted as standard.

Embedding quality analysis

Using stratified subsampled cross-validation over NAS-Bench-201's 15,625 architectures (100 trials per budget), the Base configuration reaches Kendall's Tau of 0.463 at $N=14$ training samples and 0.797 at $N=3516$. Several findings stand out:

• PCA is data-scale dependent: it hurts at very low budgets ($N \le 55$) but becomes critical at scale — at $N=879$, PCA yields $\tau = 0.703$ versus 0.645 without it.
• Rank loss beats MSE for ranking: Pairwise Hinge Loss consistently outperforms MSE on Kendall's Tau (0.703 vs. 0.691 at $N=879$), though MSE achieves lower absolute error.
• Model scaling offers marginal returns: CodeLlama Python 34B improves only marginally over 7B (0.711 vs. 0.703 at $N=879$), and Yi Coder 1.5B is nearly competitive (0.698). Decoder-only models generally outperform encoder-only ones, with CodeBERT lagging substantially (0.659 at $N=879$).

The practical implication is that effective NAS surrogates do not require large models or extensive labeled data; small code-pre-trained models with lightweight heads suffice in low-data regimes.

Comparison against other text encodings

Against ONNX-to-text encodings on a matched 7,400-architecture subset of NAS-Bench-201, PyTorch text encoding is statistically significantly better ($p < 0.05$) for both CodeLlama Python 7B and ModernBERT Large at budgets $N \ge 220$. On einspace, PyTorch code also outperforms derivation tree strings (e.g., $N=3516$0 vs. 0.763 at $N=3516$1 with CodeLlama). A particularly strong claim emerges here: COLE achieves $N=3516$2 on einspace with only 220 training samples, whereas Qin et al.'s fine-tuned ModernBERT approach requires roughly 6.5k architectures to reach $N=3516$3. This supports the paper's argument that matching the LM's pre-training distribution eliminates the need for expensive format-learning fine-tuning.

t-SNE visualizations reinforce this: high-performing architectures separate more cleanly under COLE than under adjacency matrices, path encodings, ONNX-to-text, or derivation tree encodings.

Downstream search efficiency

Integrating COLE into BANANAS (replacing path encodings, keeping the same MLP predictor) across 100 trials per end-task with a 500-evaluation budget produces consistent efficiency gains:

End-task	Top accuracy	Evals to reach within 1% (COLE)	Evals (path)	Reduction
CIFAR-10	91.61%	52	60	13%
CIFAR-100	73.51%	132	200	34%
ImageNet16-120	47.31%	164	200	18%

The CIFAR-100 gap is statistically significant across nearly the entire search trajectory, and ImageNet16-120 shows significant early-stage gains. Since COLE is purely a feature extractor, these gains are not specific to Bayesian Optimization and should transfer to evolutionary surrogate-assisted pipelines as well.

Positioning relative to structural encodings

The authors concede that specialized structural predictors such as FLAN and HyperNAS can achieve higher Kendall's Tau than COLE, and that embedding extraction with a 7B model is orders of magnitude more expensive than computing a path string. Their counterargument rests on generality: structural encodings treat operations as discrete labels, require manual engineering to capture hyperparameters like kernel sizes, and need re-engineering for new search spaces. Code, by contrast, captures topology and hyperparameters uniformly and applies to any architecture expressible as a program. COLE's value proposition is therefore cold-start applicability and search-space agnosticism rather than peak single-space accuracy.

Limitations and open questions

The paper is explicit about several constraints. Downstream search validation was restricted to NAS-Bench-201; whether the reported evaluation-budget reductions hold on larger spaces such as DARTS remains unverified. The comparison against alternative text encodings covers only NAS-Bench-201 and einspace, so the claimed superiority of native code has limited empirical breadth. The ONNX comparison relied on an accuracy-matched subset (7,400 of 15,625 architectures) due to benchmark inconsistencies, which constrains that experiment's scope. Additionally, the einspace corpus comprised 2,837 architectures drawn from nine CIFAR-10 search runs, raising questions about corpus representativeness in an open-ended space. Open questions include whether advanced pooling strategies or head architectures improve results, whether fine-tuning on architecture-specific syntax would compound the frozen-model gains, and whether the unified code representation enables effective joint architecture–hyperparameter optimization or multi-objective (e.g., hardware-aware) NAS.

Conclusion

This work establishes that frozen, pre-trained LMs applied directly to PyTorch source code serve as competitive zero-shot neural architecture embedders, removing both the fine-tuning overhead of prior LM-based surrogates and the rigidity of structural encodings. The strongest quantitative results — a 34% reduction in evaluations to near-optimal CIFAR-100 architectures under BANANAS, and order-of-magnitude sample-efficiency improvements over fine-tuned derivation-tree surrogates on einspace — substantiate the core thesis that input modality alignment with pre-training data matters more than task-specific adaptation. The approach's generality makes it a practical foundation for surrogate-assisted NAS, though broader validation beyond NAS-Bench-201 remains necessary to confirm its scope.