Language Embedding Meets Dynamic Graph: A New Exploration for Neural Architecture Representation Learning

TL;DR

This paper introduces LeDG-Former, a novel framework combining language embeddings and dynamic graph learning to improve neural architecture representation, enabling zero-shot hardware prediction and surpassing state-of-the-art benchmarks.

Contribution

LeDG-Former uniquely integrates language-based semantic embeddings with dynamic graph transformers for neural architecture modeling, addressing hardware info and structural representation limitations.

Findings

Achieves state-of-the-art results on NNLQP benchmark.

First to enable cross-hardware latency prediction.

Outperforms previous methods on NAS-Bench datasets.

Abstract

Neural Architecture Representation Learning aims to transform network models into feature representations for predicting network attributes, playing a crucial role in deploying and designing networks for real-world applications. Recently, inspired by the success of transformers, transformer-based models integrated with Graph Neural Networks (GNNs) have achieved significant progress in representation learning. However, current methods still have some limitations. First, existing methods overlook hardware attribute information, which conflicts with the current trend of diversified deep learning hardware and limits the practical applicability of models. Second, current encoding approaches rely on static adjacency matrices to represent topological structures, failing to capture the structural differences between computational nodes, which ultimately compromises encoding effectiveness. In this paper, we introduce LeDG-Former, an innovative framework that addresses these limitations through the synergistic integration of language-based semantic embedding and dynamic graph representation learning. Specifically, inspired by large language models (LLMs), we propose a language embedding framework where both neural architectures and hardware platform specifications are projected into a unified semantic space through tokenization and LLM processing, enabling zero-shot prediction across different hardware platforms for the first time. Then, we propose a dynamic graph-based transformer for modeling neural architectures, resulting in improved neural architecture modeling performance. On the NNLQP benchmark, LeDG-Former surpasses previous methods, establishing a new SOTA while demonstrating the first successful cross-hardware latency prediction capability. Furthermore, our framework achieves superior performance on the cell-structured NAS-Bench-101 and NAS-Bench-201 datasets.