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

TL;DR

The paper tackles the bottleneck of NAS evaluation by introducing ONNX-Bench, a unified ONNX-format benchmark spanning diverse search spaces, and ONNX-Net, a text-based surrogate that uses an ONNX-to-text encoding fed to an LLM to predict architecture performance instantly. The approach emphasizes space-agnostic representations and operator-level detail to enable cross-space generalization and zero-shot transfer, addressing limitations of graph-based encodings that rely on fixed topologies. Empirical results show strong cross-space transfer, competitive zero-shot performance, and improved cross-dataset generalization when training data is diverse, with ablations highlighting the importance of input-related encoding components and encoder-based backbones. The work advances NAS by enabling universal, fast performance prediction across a broad spectrum of architectures, paving the way for more flexible and scalable NAS methods and living benchmarks that span multiple search spaces.

Abstract

Neural architecture search (NAS) automates the design process of high-performing architectures, but remains bottlenecked by expensive performance evaluation. Most existing studies that achieve faster evaluation are mostly tied to cell-based search spaces and graph encodings tailored to those individual search spaces, limiting their flexibility and scalability when applied to more expressive search spaces. In this work, we aim to close the gap of individual search space restrictions and search space dependent network representations. We present ONNX-Bench, a benchmark consisting of a collection of neural networks in a unified format based on ONNX files. ONNX-Bench includes all open-source NAS-bench-based neural networks, resulting in a total size of more than 600k {architecture, accuracy} pairs. This benchmark allows creating a shared neural network representation, ONNX-Net, able to represent any neural architecture using natural language descriptions acting as an input to a performance predictor. This text-based encoding can accommodate arbitrary layer types, operation parameters, and heterogeneous topologies, enabling a single surrogate to generalise across all neural architectures rather than being confined to cell-based search spaces. Experiments show strong zero-shot performance across disparate search spaces using only a small amount of pretraining samples, enabling the unprecedented ability to evaluate any neural network architecture instantly.

Figures (9)

• Figure 1: Overview of our approach: ONNX-Bench contains $\{$architecture, accuracy$\}$ pairs from multiple search spaces in unified ONNX representation; ONNX-Net consists of a robust, universal text encoding and an LLM-based performance predictor.
• Figure 2: Diversity within and between the search spaces in ONNX-Bench, diversity is measured using Jensen-Shannon divergence (in bits).
• Figure 3: Distribution of nodes (left) in log scale and CIFAR-10 accuracy (right) of the search spaces contained in ONNX-Bench.
• Figure 4: The ONNX-Net pipeline for performance prediction of any ONNX-encoded neural network architecture.
• Figure 5: Zero-shot predictor trained on different number of NAS-Bench-101, evaluated on NAS-Bench-201. Avg. Spearman's $\rho$ and standard error over 5 random seeds is reported.