TEFormer: A Topology-Enhanced Transformer for Architecture Performance Prediction

Han Ji, Yanan Sun

Keywords: neural architecture search, performance predictor

Abstract: Evaluating architecture performance is a crucial step in neural architecture search (NAS) but remains computationally expensive. Performance predictors offer an efficient alternative by learning from a limited set of architecture-performance pairs. However, previous predictors tend to oversimplify the topological structure of neural architectures using adjacency matrices, node depths, or computation flow, which fail to fully capture topological features of architectures, leading to poor generalization. To address this limitation, we propose TEFormer, a Topology-Enhanced Transformer that integrates both local and global topological information beneficial to performance prediction. Specifically, we employ a topology-aware flow encoding module that incorporates local topological characteristics via a learnable structural encoding and a flow-based encoder. At the global level, we design a hierarchical attention mechanism to jointly model intra-flow and inter-flow interactions within the architecture. To further improve generalization, we propose an architecture augmentation strategy that synthesizes additional samples by interpolating similar architectures in the latent space. Extensive experiments on computer vision, graph learning, and automatic speech recognition tasks demonstrate that TEFormer consistently outperforms state-of-the-art predictors and exhibits superb performance across diverse search spaces.

Primary Area: other topics in machine learning (i.e., none of the above)

Submission Number: 15363