🤖 machine learning

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

This paper proposes Code-Oriented LM Embeddings (COLE), a low-cost strategy that leverages off-the-shelf language models to represent neural architectures as PyTorch code, enabling competitive surrogate-assisted Neural Architecture Search without specialized fine-tuning and significantly reducing the evaluation budget required to find optimal models.

Original authors: Pranav Somu, Advay Balakrishnan, Stepan Kravtsov, Aaron McDaniel, Jason Zutty

Published 2026-05-18

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Pranav Somu, Advay Balakrishnan, Stepan Kravtsov, Aaron McDaniel, Jason Zutty

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Finding the Best Recipe is Expensive

Imagine you are a chef trying to find the absolute best recipe for a cake. You have a massive cookbook with millions of possible combinations of ingredients (flour, sugar, eggs, spices) and methods (baking time, temperature).

In the world of Artificial Intelligence, this is called Neural Architecture Search (NAS). The "recipes" are neural network designs, and the "taste test" is training the AI to see how well it works. The problem is that "tasting" every single recipe takes forever and costs a fortune in computer power. You can't just bake every cake in the book.

To solve this, scientists use a Surrogate (or a "taste predictor"). This is like a food critic who looks at a recipe on paper and guesses how good the cake will taste, so you don't have to actually bake it. The goal is to find a critic who is fast, cheap, and very accurate.

The Old Way: Translating Recipes into a Foreign Language

Previously, to get a critic to understand a recipe, scientists had to translate the complex code of the neural network into a strange, simplified text format (like turning a recipe into a list of chemical formulas or a weird diagram).

Think of it like this: You have a master chef (a Large Language Model, or LM) who knows how to cook because they have read millions of real cookbooks and recipes. But instead of showing them the actual recipe, you hand them a translation written in a language they've never seen before (like "ONNX-to-text" or "derivation trees").

To make the chef understand this new language, you have to spend weeks fine-tuning them—essentially re-educating them from scratch. This is slow, expensive, and defeats the purpose of having a fast "taste predictor."

The New Idea: Just Show Them the Real Recipe

The authors of this paper, Pranav Somu and his team, had a simple insight: Neural networks are already written in code. They look like Python programming scripts.

Since the "master chefs" (Large Language Models like CodeLlama) were already trained on trillions of lines of real computer code, they already know how to read these recipes perfectly. They don't need to be re-educated.

The Solution (COLE):
Instead of translating the recipe into a weird format, they just feed the raw Python code directly into the pre-trained AI.

No Fine-Tuning: They use the AI "as is" (frozen). It's like hiring a chef who already knows how to read your specific cookbook without needing a training seminar.
The Embedding: The AI reads the code and turns it into a mathematical "fingerprint" (an embedding) that captures the essence of the recipe.
The Predictor: A small, simple calculator (a regression head) looks at that fingerprint and predicts how good the architecture will perform.

What They Found

The team tested this idea using two different "cookbooks" (search spaces) and found some surprising results:

Raw Code Wins: When they fed the AI the actual Python code, it predicted performance much better than when they fed it the weird, translated text formats. It's like the chef saying, "I understand the real recipe; I don't need you to translate it into chemical formulas for me."
No Training Needed: They proved you don't need to spend weeks fine-tuning the AI. The "off-the-shelf" models worked great immediately.
Faster Search: When they used this method to search for the best AI architecture (using an algorithm called BANANAS), they found the best designs much faster.
- For one task (CIFAR-100), they reached the top 1% of the best possible designs using 34% fewer computer evaluations than the old method. It's like finding the best cake recipe by tasting only 66 cakes instead of 100.

Why This Matters

This approach is like giving a universal translator to a search engine. Because every neural network can be written as code, and code is a universal language that these AI models already speak fluently, this method works for almost any type of architecture without needing custom engineering.

In short: Instead of forcing a smart AI to learn a new, weird language to understand neural networks, the authors just spoke to it in the language it already knows: Code. This saves time, money, and gets better results.

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

Problem Statement

Neural Architecture Search (NAS) automates the discovery of high-performing neural architectures but faces a significant computational bottleneck: evaluating the fitness of candidate architectures is expensive. To mitigate this, researchers employ surrogate models (performance predictors) to approximate fitness without full training. However, existing surrogate approaches suffer from two main limitations:

Representation Rigidity: Structural encodings (e.g., adjacency matrices, path encodings) are often tailored to specific search spaces (like cell-based DAGs) and struggle to generalize to novel architectures or capture complex hyperparameters (e.g., kernel sizes) without custom engineering.
Fine-Tuning Overhead: Recent attempts to use Language Models (LMs) as universal regressors require converting architectures into non-standard text formats (e.g., derivation tree strings or ONNX-to-text) and subsequently fine-tuning the LM on these foreign representations. This "warm-up" phase negates the efficiency gains of using a surrogate, as it requires expensive pre-training before the search can even begin.

The authors posit that these inefficiencies arise because current methods ignore the fact that neural architectures naturally exist as executable code. Since LMs are pre-trained on massive corpora of code, they possess an inherent "inductive bias" for understanding architectural logic that is misaligned when fed artificial text encodings.

Methodology: COLE

The authors propose Code-Oriented LM Embeddings (COLE), a strategy that leverages frozen, pre-trained Language Models as zero-shot feature extractors for NAS.

1. Native Code Stringification

Instead of converting architectures into abstract strings, COLE maps the search space directly into PyTorch class definitions (torch.nn.Module).

NAS-Bench-201: The authors parse standard graph topology strings and convert them into valid Python classes. They investigated three verbosity levels:
- Helper Method: Defining complex operations as separate classes.
- Inline: Defining logic within the class using nn.Sequential.
- Excluded Helper: Referencing helper classes without defining their implementation (most token-efficient).
einspace: A deterministic transpiler converts context-free grammar derivation trees directly into executable PyTorch code, mapping non-terminal nodes (e.g., Branching, Routing) to custom container classes and terminal nodes to atomic primitives.

2. Frozen Embedding Extraction

The system utilizes frozen pre-trained LMs (specifically CodeLlama Python 7B) to process the raw code text.

No Fine-Tuning: The LM parameters are frozen entirely, eliminating the computational cost of domain-specific pre-training.
Pooling: The hidden state activations of the last layer are aggregated using Mean Pooling across the token sequence to generate a fixed-dimensional embedding vector.

3. Surrogate Head

The high-dimensional embeddings (e.g., 4096 dimensions) are processed by a lightweight regression pipeline:

Dimensionality Reduction: Principal Component Analysis (PCA) reduces the embedding to 128 dimensions to prevent overfitting in low-data regimes.
Regression: A 3-layer Multi-Layer Perceptron (MLP) predicts architecture performance.
Loss Function: The model is optimized for Kendall's Tau rank correlation using a Pairwise Hinge Loss. This prioritizes the correct ordering of candidates (crucial for evolutionary search algorithms) over predicting exact absolute accuracy values.

Key Contributions

Superiority of Native Code: The authors demonstrate that representing architectures as raw PyTorch code yields higher predictive performance (Kendall's Tau) than other text-based encodings (ONNX-to-text, derivation trees) when using frozen LMs.
Elimination of Fine-Tuning: COLE proves that computationally expensive fine-tuning is not strictly necessary for effective NAS embeddings; off-the-shelf LMs act as competitive feature extractors "out-of-the-box."
Integration with SA-NAS: Integrating COLE into the BANANAS surrogate-assisted search algorithm significantly improves search efficiency compared to using standard structural path encodings.
Practical Tuning Guidelines: The study documents effective configurations, such as the use of "Excluded Helper" mode for token efficiency and PCA for dimensionality reduction, showing that the LM can infer operation semantics even without full implementation details in the input string.

Experimental Results

The framework was evaluated on NAS-Bench-201 and the einspace search space.

Embedding Quality:
- On NAS-Bench-201, COLE (using CodeLlama 7B) achieved a Kendall's Tau of 0.703 with 879 training samples, outperforming ONNX-to-text encodings (0.678) and derivation tree strings (0.763 in einspace vs. 0.778 for PyTorch code).
- Model Selection: Code-centric models (CodeLlama, Yi Coder) consistently outperformed general encoder-only models (ModernBERT, CodeBERT). Interestingly, smaller models (1.5B) performed comparably to larger 7B/34B models, suggesting massive scale is not strictly required.
- Loss Function: Pairwise Hinge Loss yielded better rank correlation than standard Mean Squared Error (MSE), despite MSE having lower absolute error.
Search Efficiency (BANANAS):
- Replacing path encodings with COLE in the BANANAS algorithm reduced the evaluation budget required to reach near-optimal performance.
- CIFAR-100: COLE reached within 1% of the fittest architecture in 132 evaluations, a 34% reduction compared to the 200 evaluations required by path encodings.
- CIFAR-10: Achieved a 13% reduction (52 vs. 60 evaluations).
- ImageNet16-120: Achieved an 18% reduction (164 vs. 200 evaluations).
- These improvements were statistically significant ( $p < 0.05$ ) across most of the search trajectory.

Significance and Claims

The paper claims that COLE establishes a versatile and efficient foundation for advancing NAS by leveraging the existing knowledge of code in pre-trained LMs.

Search-Agnostic: Since any neural architecture can be represented as code, COLE does not require re-engineering encoding schemes for new search spaces, unlike structural encodings.
Cold-Start Capability: By removing the need for fine-tuning, COLE enables "cold-start" searches where the surrogate is effective immediately, addressing the efficiency paradox of previous LM-based NAS methods.
Generalizability: The approach unifies topology and hyperparameter representation in a single format (code), allowing the LM to interpret both structural flow and parameter values (e.g., kernel_size=3) naturally.

The authors conclude that while their current evaluation is limited to specific benchmarks (NAS-Bench-201, einspace), the methodology offers a practical, low-cost alternative to complex GNNs or fine-tuned transformers, particularly for resource-constrained NAS scenarios.