Graph transformer improves neural architecture predictions using direction and topology
DTA-GT: Direction- and Topology-Aware Graph Transformer for Neural Network Representation Learning
Machine LearningArtificial IntelligenceNeural and Evolutionary Computing
Summary
Predicting how well a neural network design will perform is difficult because these designs are graphs with edges that show the computation order. The authors created a new graph-based method called DTA-GT that carefully considers both the direction of these edges and the overall graph structure throughout its process. This approach helps the model better understand the flow of computation, leading to more accurate estimates of accuracy and speed for different network designs. Their experiments show that this method outperforms previous models when only limited data is available.
What this means in practice
- •For automated machine learning teams: Estimate neural network model accuracy and latency more accurately using fewer evaluated designs in architecture search workflows.
- •For network hardware engineers: Improve latency predictions for hardware deployment of neural networks with better consideration of model computational flow.
Authors
Yuxiang Zeng, Kun Xie, Yang Wang, Jigang Wen, Xiaocan Li, Guangxing Zhang, Gaogang Xie, Jiannong Cao
Abstract
Architecture attribute predictors reduce the cost of neural architecture search by estimating accuracy or latency from a small set of evaluated candidates. The key difficulty is that neural architectures are operation-labeled DAGs: edge direction defines computational flow, but existing predictors often treat direction as a local structural cue. We propose DTA-GT, a Direction- and Topology-Aware Graph Transformer based on encoder-wide directional consistency, which preserves directed computational semantics across node initialization, global spectral encoding, pairwise interaction, and feature update. DTA-GT realizes this principle with direction- and topology-aware node initialization, Magnetic Laplacian spectral encoding, direction-dual structural attention, and a Direction-Sensitive MoE-FFN. Across accuracy and latency prediction on NAS-Bench-101 and NAS-Bench-201, DTA-GT consistently outperforms representative sequence-, GNN-, Transformer-, and hybrid predictors under limited-budget settings. Ablations and Magnetic spectral controls show that the improvements come from coordinated direction-aware stages and edge-orientation-sensitive spectral information, while mechanistic diagnostics confirm that the modules exhibit the intended directional behavior. These results support encoder-wide directional consistency as an effective design principle for architecture DAG representation.