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

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

Keywords: Neural Architecture, Neural Architecture Representation, Task Performance Evaluation

TL;DR: We present ONNX-Bench, a unified benchmark across diverse NAS search spaces, and ONNX-Net, a universal text-based representation and surrogate that enable fast performance prediction of arbitrary neural architectures.

Abstract: Neural architecture search (NAS) automates the design process towards 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 $600$k {$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.

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Submission Number: 11316