Neural Architecture Retrieval

TL;DR

The paper tackles Neural Architecture Retrieval (NAR), aiming to efficiently retrieve architectures with similar designs from large, heterogeneous collections. It introduces a motif-aware representation by decomposing computation graphs into repeated substructures and rebuilding a reduced macro graph, enabling scalable, motif-conscious embeddings. A two-stage pre-training framework first learns robust motif embeddings through motifs-level contrastive learning and then learns macro-graph representations via graph-level contrastive learning and coarse-family classification. Experiments on real-world and NAS-generated datasets show significant improvements over GCN/GAT baselines and demonstrate transferability across domains, supported by a 12k-architecture embedding dataset. This approach provides a practical, scalable solution for navigating vast neural-architecture design spaces and supports rapid design exploration and benchmarking.

Abstract

With the increasing number of new neural architecture designs and substantial existing neural architectures, it becomes difficult for the researchers to situate their contributions compared with existing neural architectures or establish the connections between their designs and other relevant ones. To discover similar neural architectures in an efficient and automatic manner, we define a new problem Neural Architecture Retrieval which retrieves a set of existing neural architectures which have similar designs to the query neural architecture. Existing graph pre-training strategies cannot address the computational graph in neural architectures due to the graph size and motifs. To fulfill this potential, we propose to divide the graph into motifs which are used to rebuild the macro graph to tackle these issues, and introduce multi-level contrastive learning to achieve accurate graph representation learning. Extensive evaluations on both human-designed and synthesized neural architectures demonstrate the superiority of our algorithm. Such a dataset which contains 12k real-world network architectures, as well as their embedding, is built for neural architecture retrieval.

Figures (9)

• Figure 1: The definition of Neural Architecture Retrieval (NAR). This paper explores pre-training an encoder $\mathcal{F}$ to build a neural network embedding database $\mathcal{H}$ based on the architecture designs.
• Figure 2: An illustration of motifs sampling strategy. The graph nodes encode their neighbours in adjacent matrix via an iterative manner to form the encoded node sequence where each node denotes a subgraph and motifs denote the subsequence.
• Figure 3: An illustration of macro graph setup and pre-training in motifs-level and graph-level.
• Figure 4: Coarse-grained and fine-grained classes.
• Figure 5: Visualization of learnt embeddings. The number of dimensions is reduced by t-SNE. (\ref{['fig:exp:tsne:gcn-real']}, \ref{['fig:exp:tsne:gat-real']} and \ref{['fig:exp:tsne:ours-real']} are embeddings of real-world neural architectures, and \ref{['fig:exp:tsne:nas-gcn']}, \ref{['fig:exp:tsne:nas-gat']} and \ref{['fig:exp:tsne:nas-ours']} are embeddings of NAS architectures)