TLDR: This research introduces ONNX-Bench, a unified benchmark of over 600,000 neural network architectures in the ONNX format, evaluated on CIFAR-10. Building on this, ONNX-Net proposes a novel method to represent any neural architecture as a natural language description. This text-based encoding, combined with an LLM-based performance predictor, enables instant and universal performance prediction across diverse and previously unseen search spaces, significantly accelerating Neural Architecture Search (NAS) by overcoming limitations of traditional, search-space-specific methods.
Neural Architecture Search (NAS) is a powerful field aiming to automate the design of neural networks, potentially leading to architectures that surpass human-designed ones. However, this process has long been hampered by a significant bottleneck: the immense computational cost of evaluating the performance of each candidate architecture. Traditional NAS methods often rely on restrictive search spaces or specialized graph encodings that limit their flexibility and ability to generalize across different types of neural networks.
A new research paper, titled “ONNX-NET: TOWARDS UNIVERSAL REPRESENTATIONS AND INSTANT PERFORMANCE PREDICTION FOR NEURAL ARCHITECTURES” by Shiwen Qin, Alexander Auras, Shay B. Cohen, Elliot J. Crowley, Michael Moeller, Linus Ericsson, and Jovita Lukasik, introduces a novel approach to overcome these limitations. The researchers present ONNX-Bench, a comprehensive benchmark dataset, and ONNX-Net, a method for universally representing neural architectures and instantly predicting their performance. You can read the full paper here: ONNX-NET Research Paper.
ONNX-Bench: A Unified Dataset for Diverse Architectures
The first key contribution is ONNX-Bench, an extensive benchmark dataset designed to unify various neural network architectures. It comprises over 600,000 pairs of architecture designs and their corresponding performance accuracies. What makes ONNX-Bench unique is its use of the Open Neural Network Exchange (ONNX) format, a widely adopted standard for representing neural networks. This unified format allows for consistent evaluation and comparison of networks from diverse search spaces, including popular cell-based and more expressive hierarchical designs.
All architectures within ONNX-Bench are evaluated on the CIFAR-10 dataset using a standardized training pipeline. This consistency ensures that any observed performance differences are due to the architectural design itself, rather than variations in training protocols. The dataset’s sheer size and diversity, encompassing networks with varying complexities and performance levels, make it an ideal foundation for developing performance predictors that can generalize across a broad spectrum of neural network designs.
ONNX-Net: Text-Based Encoding and LLM-Powered Prediction
Building upon ONNX-Bench, the researchers developed ONNX-Net, a system that transforms neural network architectures into natural language descriptions, which are then used to predict performance. The core idea is to convert the ONNX file, which represents a network as a computational graph of operations and tensors, into a human-readable text format.
To achieve this, the ONNX graph undergoes several optimizations, such as removing less important nodes and merging common subgraphs, to create a condensed representation. This optimized graph is then converted into a textual description that captures the network’s topology, specific operations, and fine-grained parameters like kernel sizes or strides. This text-based encoding is highly flexible, capable of describing any neural architecture regardless of its complexity or specific layer types.
Once an architecture is encoded as text, a large language model (LLM) is fine-tuned to act as a performance predictor. This LLM learns to map the textual description of a neural network to its expected accuracy. This surrogate model can then instantly infer the performance of a new, unseen architecture without the need for time-consuming training, significantly accelerating the NAS process.
Also Read:
- ZEROSHOTOPT: A Pretrained Model for Efficient Black-Box Optimization Without Tuning
- DiffuSpec: Accelerating LLM Inference with Diffusion Language Models
Promising Results and Future Directions
The experiments conducted by the team demonstrate the strong capabilities of ONNX-Net. The surrogate model showed competitive performance when generalizing across different search spaces, even when trained on a subset of spaces and tested on an entirely excluded one. Notably, ONNX-Net achieved strong zero-shot transfer performance, meaning it could accurately predict the performance of architectures from a new search space without any prior training on that specific space, outperforming existing methods, especially with limited training data.
The research also explored the importance of different components in the text encoding, finding that input information (connectivity and weight shapes) is particularly crucial. Furthermore, encoder-based LLMs like ModernBERT generally outperformed decoder-based models for zero-shot transfer tasks.
ONNX-Bench and ONNX-Net represent a significant step towards more flexible and efficient neural architecture search. By providing a unified benchmark and a universal, text-based representation, this work paves the way for developing general search methods and performance predictors that are not confined to specific architectural styles or search spaces. Future work aims to expand ONNX-Bench with even more diverse architectures and to enable surrogates to consider the dataset as context for prediction, further enhancing their versatility.
