🤖 machine learning

ONNX-Net: Towards Universal Representations and Instant Performance Prediction for Neural Architectures

This paper introduces ONNX-Net, a unified, text-based neural architecture representation built on the 600k+ sample ONNX-Bench benchmark, which enables a single performance predictor to instantly and accurately evaluate arbitrary neural architectures across diverse search spaces without being restricted to specific cell-based designs.

Original authors: Shiwen Qin, Alexander Auras, Shay B. Cohen, Elliot J. Crowley, Michael Moeller, Linus Ericsson, Jovita Lukasik

Published 2026-08-27

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Shiwen Qin, Alexander Auras, Shay B. Cohen, Elliot J. Crowley, Michael Moeller, Linus Ericsson, Jovita Lukasik

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

In the world of artificial intelligence, the most powerful tools are often neural networks, systems loosely modeled on the human brain that learn to recognize patterns, translate languages, or drive cars. For decades, scientists have been trying to build these networks by hand, tweaking their structure like a mechanic adjusting an engine. But as these systems grew more complex, human intuition became a bottleneck. This gave rise to a field called neural architecture search, where computers are tasked with designing their own networks. The goal is to let machines find the best possible structure for a specific job. However, this process has hit a wall. To know if a design is good, the computer must build it and train it from scratch, a process that can take days or even weeks on massive supercomputers. This makes the search incredibly slow and expensive, limiting how much experimentation can happen.

A team of researchers has now proposed a way to break through this bottleneck, not by building faster computers, but by changing how they talk about the designs. They introduced a new system called ONNX-Net, which acts as a universal translator for neural networks. Instead of forcing every new design into a rigid, pre-defined box, this system reads the network's blueprint as if it were a simple sentence. By converting complex technical diagrams into plain text descriptions, the researchers trained a language model to predict how well a network will perform just by reading that text. The result is a tool that can instantly guess the success of a neural network design without ever having to build or train it, saving vast amounts of time and energy.

The core problem the researchers tackled was that previous attempts to speed up this process were too narrow. Earlier methods worked well only for networks built from small, repeating blocks, like stacking identical Lego bricks. But as scientists began designing more flexible and complex networks, those old methods failed because they could not understand the new, irregular shapes. Furthermore, the data used to train these prediction tools was scattered across different formats, making it impossible to create a single, smart predictor that could handle any type of network. To solve this, the team first built a massive, unified library called ONNX-Bench. They gathered nearly 650,000 different neural network designs from various sources and converted them all into a single, standard file format known as ONNX. This format is like a universal language for neural networks, capable of describing any structure, no matter how complex or unusual.

Once they had this massive collection of standardized designs, the researchers faced the challenge of how to feed this information into a computer program. They realized that the most flexible way to describe a network was not through complex graphs or matrices, but through natural language. They developed a method to translate the technical details of each network—what operations it performs, how its parts are connected, and the specific settings for each component—into a readable text description. Imagine a network as a set of instructions: "Take an input, pass it through a filter, then apply a correction, and finally output a result." The researchers wrote software to generate these instructions for every single one of the 650,000 networks in their library.

They then used this text to train a large language model, a type of artificial intelligence familiar with human language, to act as a performance predictor. Instead of learning to write poetry or answer questions, this model learned to read the technical "sentences" describing neural networks and predict their accuracy. The training was surprisingly efficient. The model learned to make accurate predictions after seeing only a tiny fraction of the data required by previous methods. In fact, it achieved the same level of performance as older systems using less than one percent of the training data. This suggests that by describing the networks in a clear, text-based format, the computer can grasp the essential features of a design much faster than before.

The true test of this new approach was whether it could generalize to designs it had never seen before. The researchers tested the model by training it on networks from one specific family and then asking it to predict the performance of completely different networks from other families. The results were striking. The model successfully predicted the performance of these unseen designs with high accuracy, a capability known as zero-shot transfer. This means the system did not need to be retrained or adjusted for the new types of networks; it simply applied what it learned from the text descriptions to understand the new structures. This is a significant leap forward, as previous tools were often locked into the specific types of networks they were trained on and could not adapt to new, more complex designs.

The researchers also explored exactly what information the model needed to make these predictions. They broke down the text descriptions into different parts, such as the names of the operations, the shapes of the data flowing through them, and the specific parameters set by the designer. They found that including details about how the network parts were connected and the shapes of the data was the most critical factor for success. While the model performed well with just the basic structure, adding these specific details allowed it to make much sharper predictions. This confirms that the text-based approach captures the necessary details of a network's design without getting lost in the noise of irrelevant code variations.

Despite these successes, the researchers are careful to note the boundaries of their work. The system was trained and tested primarily on networks designed for image recognition tasks using a specific dataset called CIFAR-10. While the method is theoretically capable of handling any network that can be converted into their standard format, the current results are limited to this specific domain. The team acknowledges that the diversity of their training data, while vast, still comes from existing benchmarks and does not yet cover every possible type of neural network or task. They also point out that while their text-based method is more flexible than previous graph-based approaches, those older methods can still outperform them in very specific, restricted scenarios.

The release of this new benchmark and the accompanying prediction tool opens a door for more efficient exploration in artificial intelligence. By providing a way to instantly evaluate designs without the heavy cost of training, researchers can now test thousands of ideas in the time it used to take to test a few. The work suggests that the future of designing neural networks may not lie in building faster hardware, but in finding better ways to describe and understand the designs themselves. The researchers hope that by making their tools and data available, they will encourage others to build search methods that are not limited to specific types of networks, but can adapt to the ever-evolving landscape of artificial intelligence.

Technical Summary: ONNX-Net: Towards Universal Representations and Instant Performance Prediction for Neural Architectures

Problem Statement

Neural Architecture Search (NAS) automates the design of high-performing networks but is severely bottlenecked by the computational cost of evaluating candidate architectures. While surrogate models have been developed to approximate performance without full training, existing approaches suffer from two primary limitations:

Search Space Restrictions: Most methods are tied to specific, restrictive search spaces (e.g., cell-based designs) and cannot generalize to more expressive, hierarchical, or heterogeneous search spaces.
Representation Limitations: Current encodings, such as adjacency matrices or raw computational graphs, often fail to scale to sparse, variable-sized graphs or lack sensitivity to fine-grained operator parameters (e.g., kernel sizes). Furthermore, representations based on Python code are often redundant and inconsistent across frameworks, hindering cross-space transfer.

Consequently, there is a lack of a unified, search-space-agnostic representation that can capture both network topology and detailed operator parameters to enable instant performance prediction across diverse architectural families.

Methodology

The authors propose a two-pronged approach consisting of a new benchmark dataset and a universal representation framework.

1. ONNX-Bench: A Unified Benchmark
To address the lack of diverse training data, the authors introduce ONNX-Bench, a large-scale dataset aggregating approximately 650,000 {architecture, accuracy} pairs.

Unification: It consolidates architectures from multiple existing NAS benchmarks (including NAS-Bench-101/201/301, NATS-Bench, TransNAS-Bench-101, hNAS-Bench-201, and einspace) into a single, standardized format based on the Open Neural Network Exchange (ONNX).
Evaluation: All architectures are evaluated on the CIFAR-10 dataset using a consistent training pipeline to ensure that performance differences reflect architectural properties rather than training protocol inconsistencies.
Diversity: The dataset spans both cell-based and hierarchical search spaces, providing the necessary diversity to train predictors that generalize across different architectural patterns.

2. ONNX-Net: Universal Text-Based Representation
The core contribution is ONNX-Net, a surrogate modeling framework that represents neural architectures as natural language text derived from ONNX files.

Encoding Process:
- ONNX files are parsed to reconstruct the computational graph in memory.
- The graph undergoes lossless optimization and is condensed by merging linear chains of operations without branches.
- The condensed graph is converted into a text format where each operation chain is printed on a single line, explicitly capturing:
  - Operation names and input/output indices.
  - Input tensor shapes and names.
  - Operator-specific parameters (e.g., kernel size, padding, strides).
  - Output tensor shapes.
Predictor Architecture: A Large Language Model (specifically a fine-tuned ModernBERT-large) is trained on this text encoding to predict architecture performance (e.g., accuracy) directly. This text-based approach allows the model to naturally capture topology, operators, and fine-grained parameters without redesigning the encoder for each search space.

Key Contributions

ONNX-Bench: The release of a large-scale, open-source benchmark containing ~650k architectures unified under the ONNX format, enabling the training and evaluation of NAS surrogates across diverse architecture families.
ONNX-Net: A universal neural architecture representation and surrogate framework that encodes complete structural and parametric details into a compositional, extensible text format independent of predefined search spaces.
Performance and Efficiency: Demonstration that an LLM-based surrogate using ONNX-Net achieves strong predictive performance and zero-shot transferability. Notably, the approach exhibits exceptional sample efficiency, matching the performance of concurrent methods (like RLM) with less than 1% of the training data (e.g., reaching the same performance level as the fully trained RLM with 150x fewer samples).

Experimental Results

The paper evaluates ONNX-Net across several dimensions:

In-Domain Performance: On NAS-Bench-301, ONNX-Net outperforms RLM and GNN-based predictors, achieving a Kendall's $\tau$ of 0.546 with only 100 training samples, surpassing RLM's 0.528.
Sample Efficiency: When trained on NAS-Bench-101, ONNX-Net demonstrates exceptional sample efficiency. Specifically, in a comparison mirroring the RLM protocol (where RLM is trained on 381k samples to reach a Spearman's $\rho$ of 0.723), ONNX-Net achieves the same performance level as the fully trained RLM after seeing only 2,500 samples (150x fewer data points). While ONNX-Net reaches a Spearman's $\rho$ of 0.919, this higher score is achieved after seeing 150,000 samples.
Zero-Shot Transfer:
- Across Search Spaces: Trained on NAS-Bench-101, the model transfers to NAS-Bench-201 with a Spearman's $\rho$ of 0.747 (50k samples), outperforming FLAN variants and approaching the ensemble performance of GENNAPE, despite using a single predictor and no zero-cost proxies.
- Across Datasets: The model demonstrates cross-task generalization, successfully transferring knowledge from CIFAR-10 to CIFAR-100 and ImageNet16-120, and from CIFAR-10 to unseen NAS datasets within the einspace search space.
Ablation Studies: The text encoding components were analyzed, revealing that adding explicit input information (connectivity and weight shapes) yields the most significant performance gains. While the full encoding (including parameters and output shapes) is generally best, the results suggest a trade-off between sequence length and sample efficiency.

Significance and Claims

The authors claim that ONNX-Net and ONNX-Bench represent a significant step toward search-space-agnostic NAS. By moving away from rigid, cell-based representations and Python code dependencies, the proposed method enables:

Instant Performance Prediction: The ability to evaluate any neural network architecture (convertible to ONNX) instantly without retraining the surrogate for specific spaces.
Generalization: The capacity to generalize across disparate search spaces and tasks with minimal pretraining data, addressing the "bottleneck" of expensive evaluation in expressive search spaces.
Foundation for Future Research: The benchmark serves as a foundation for studying encodings that capture both structural and operator-level details, potentially guiding the development of general search methods and LLM-driven architecture generation.

The paper modestly acknowledges limitations, noting that the benchmark currently relies on pre-existing datasets (mostly low-resolution image classification) and that direct comparison with baselines is challenging due to the novelty of the methodology. However, the results suggest that text-based, ONNX-derived representations offer a highly sample-efficient and flexible alternative to existing graph-based or code-based approaches.