Source-linked AI summary
DTA-GT: Direction- and Topology-Aware Graph Transformer for Neural Network Representation Learning
Yuxiang Zeng, Kun Xie, Yang Wang, Jigang Wen, Xiaocan Li, Guangxing Zhang, Gaogang Xie, Jiannong Cao
TL;DR
Architecture attribute prediction can reduce the cost of neural architecture search, but architecture DAGs encode computational flow in their directed edges. DTA-GT coordinates direction-aware representation learning across the encoder, and the paper reports consistent improvements in accuracy and latency prediction across limited-budget settings.
Problem
Accuracy and latency prediction for neural architecture search is costly, while architecture representations must preserve the computational flow encoded by directed edges.
Source: §1 · pp. 1–2Method
DTA-GT coordinates direction- and topology-aware node initialization, Magnetic spectral encoding, direction-dual structural attention, and direction-sensitive feature updates across the encoder.
Source: §1 · pp. 2–3Results
Across four tasks and training budgets, DTA-GT achieves the highest Kendall’s τ and lowest MAPE, with consistent gains in accuracy and latency prediction.
Source: §1 · p. 3, §5.3 · p. 8, §7 · p. 10Takeaways & Limitations
The results support encoder-wide directional consistency as an effective design principle for architecture-DAG representation.
Source: §1 · p. 2, §7 · p. 10Evaluation focuses on NAS-Bench-101 and NAS-Bench-201; transfer requires aligned operation vocabularies and dependency semantics, and pairwise attention incurs quadratic costs as graph size grows.
Source: §6 · p. 10
Abstract
from arXiv · showhide
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.
1 Introduction
Neural architecture search relies on costly accuracy training and latency profiling, motivating predictors trained from a limited set of evaluated candidates. DTA-GT addresses the challenge of preserving directed computational semantics throughout architecture-DAG representation learning.
Motivation
- Limited-budget predictors estimate unseen candidates because accuracy evaluation typically requires full training and latency evaluation requires hardware measurement or profiling.NAS searches for architectures that meet performance and deployment constraints.Source: §1 · p. 1
Why direction matters
- Reversing one edge can swap operations’ split/merge roles while leaving the undirected graph skeleton unchanged.Directed edges encode computational flow and dependencies, including branch structure and reachability.Source: §1 · p. 2, Figure 1 · p. 2
Research gap
- Existing predictor families model topology and direction only partially, often leaving directional cues local or stage-specific rather than coordinating them across the encoder.Table 1 frames the comparison as design-oriented, not a performance ranking.Source: §1 · p. 2, Table 1 · p. 2
Key idea
- DTA-GT’s encoder-wide directional consistency preserves directed computational semantics across representation learning instead of injecting direction at only one stage.The paper motivates this as a coordinated principle for the main encoder stages.Source: §1 · p. 2
Proposed model
- DTA-GT coordinates direction- and topology-aware node initialization, Magnetic Laplacian encoding, direction-dual structural attention, and a Direction-Sensitive MoE-FFN.These modules address node roles, a global edge-orientation-sensitive spectral prior, pairwise relations, and predecessor/successor distinctions during feature updates.Source: §1 · pp. 2–3
Evaluation overview
- Across limited-budget settings, the paper evaluates accuracy and latency prediction and reports consistent improvements over representative predictors.The evaluation includes ablations, Magnetic spectral controls, and mechanistic diagnostics.Source: §1 · p. 3
2 Related Work
Prior architecture predictors use sequence, graph, or Transformer/hybrid representations, with different ways of encoding topology and direction. The paper identifies a gap in coordinating directed computational semantics beyond local or stage-specific cues.
Cross-paradigm gap
- Across these paradigms, the stated gap is coordinated preservation of direction across topology encoding, local cues, a global direction-aware prior, and direction-aware updates.The comparison concerns representation design, not a performance ranking.Source: §2 · p. 3
Sequence-based Methods
GNN-based Methods
- GNN predictors model graph topology and local edge directions, but higher-order topology and long-range dependencies rely on repeated local message passing.The passage notes that information can be mixed, over-smoothed, or over-squashed, while an explicit global direction-aware prior is usually absent.Source: §2 · p. 3
Transformer-based and Hybrid Methods
- Transformer and hybrid predictors use self-attention for global interactions and inject topology cues, but their directional signals remain mostly local or stage-specific.Examples include topology tokenization, graph-aided attention, relation masks, and bidirectional aggregation.Source: §2 · p. 3
3 Problem Formulation and Model Overview
The formulation treats each candidate architecture as an operation-labeled DAG and learns an attribute predictor from measured candidates. The overview maps that DAG through direction-aware initialization and encoder stages to a graph representation for prediction.
3.1 Problem Formulation
- Each architecture is an operation-labeled DAG G = (V, E, τ), with nodes as operations, directed edges as data-flow dependencies, and τ assigning operation types.The adjacency convention is A_u,v = 1 for an edge u → v.Source: §3.1 · p. 4
- The formulation represents a candidate architecture a in space A as graph G, with its node operations specified by τ.The graph representation carries topology and directed dependencies into the prediction problem.Source: §3.1 · p. 4
- The target y is a scalar attribute, such as accuracy or latency.Source: §3.1 · p. 4
- Given a measured subset of architectures, the predictor estimates attributes for unmeasured candidates in the same space without fully training or profiling each one.Source: §3.1 · p. 4
- The predictor is decomposed into an encoder h_φ that produces z_G and a prediction head p_ψ that maps z_G to the estimated attribute.The formulation writes ˆy = p_ψ(z_G), with ϑ = {ϕ, ψ}.Source: §3.1 · p. 4
- The stated modeling challenge is to preserve operation semantics, topology, and directed computational semantics across the encoder.The formulation warns against allowing directional information to be lost after a single encoding stage.Source: §3.1 · p. 4
3.2 Overall Architecture
- DTA-GT initializes node states by combining NodeEmbed(G) with DirSpecEnc(G), joining operation and local directed-role features with a global edge-orientation-sensitive spectral prior.Source: §3.2 · p. 4
- The encoder applies L layers alternating direction-dual structural attention and direction-sensitive MoE-FFN updates, then uses attention readout to produce z_G for prediction.The relation-aware interaction and predecessor/successor-aware updates operate without altering the original architecture DAG.Source: §3.2 · p. 4
4 Method
DTA-GT preserves directed computational semantics across node initialization, spectral encoding, pairwise attention, and feature updates. Its graph readout and training objective support prediction of architecture attributes and their ranking.
4.1 Direction- and Topology-Aware Node Initialization
- DTA-GT initializes nodes with operation labels, in/out degrees, and forward/reverse depths, capturing computation, local roles, and input-to-output position.The depth pair distinguishes nodes with the same forward depth but different remaining distances to the output.Source: §4.1 · p. 5
4.2 Magnetic Laplacian Spectral Encoding
- Magnetic Laplacian encoding adds a global spectral prior whose complex phase records edge orientation, complementing local directed node features.Real and imaginary spectral channels retain orientation-dependent phase information.Source: §4.2 · p. 5
4.3 Direction-Dual Structural Attention
- Six paired directional relation channels encode local adjacency, directed reachability, and shared successors or predecessors for relation-constrained attention.Within each channel, content-based attention weights admissible keys according to their node representations.Source: §4.3 · pp. 5–6
4.4 Direction-Sensitive Mixture-of-Experts Feed-Forward Network
- A Direction-Sensitive MoE-FFN uses separate self, predecessor, and successor experts with node-wise routing to update features.The routing gate adapts expert weights to the current node state, allowing predecessor/successor balance to vary across architectures and prediction targets.Source: §4.4 · p. 6
4.5 Attention-based Graph Readout
- Attention pooling summarizes final node states without virtual tokens, keeping graph-level prediction separate from the topology used by the encoder.The readout keeps the original architecture node set unchanged.Source: §4.5 · p. 7
4.6 Prediction and Training Objective
- A linear head predicts the target attribute, while training combines MSE calibration, pairwise ranking, and routing regularization.MSE fits attribute values, whereas ranking loss constrains candidate ordering for selection.Source: §4.6 · p. 7
5 Experiments
Across four accuracy and latency prediction tasks and training budgets, DTA-GT leads the reported methods on ranking and attribute estimation. Ablations, spectral controls, and mechanistic diagnostics examine its components and their behavior.
5.1 Experimental Setup
- The experiments compare predictors on NAS-Bench-101 and NAS-Bench-201 for accuracy and latency, reporting Kendall’s τ and MAPE over five-run means.Higher Kendall’s τ and lower MAPE are better.Source: §5.1 · p. 7
5.2 Ablation Analysis
- Under low-budget training, removing direction-aware components at individual encoder stages degrades performance; CLS readout and the MSE-only control also underperform.The ablations include directed versus undirected models, stage-wise removals, readout controls, and objective controls.Source: §5.2 · p. 8, Table 2 · p. 8
5.3 Main Results
- DTA-GT achieves the highest Kendall’s τ and lowest MAPE across all four tasks and training budgets.The tasks span accuracy and latency prediction on NAS-Bench-101 and NAS-Bench-201.Source: §5.3 · p. 8, Table 3 · p. 9
5.4 Magnetic Spectral Encoding Analysis
- Magnetic Laplacian encoding performs best on all reported spectral-control metrics, while symmetric Laplacian encoding gives limited, task-dependent gains.Replacing Magnetic with symmetric Laplacian encoding inside the directed encoder also weakens performance.Source: §5.4 · p. 8, Table 4 · p. 10
5.5 Mechanistic Analysis
- Mechanistic diagnostics show edge-orientation-sensitive spectral responses at q = 0.25, role-aligned MoE routing patterns, and differing enrichment across relation channels.These diagnostics characterize module behavior rather than directly attributing performance gains.Source: §5.5 · p. 9, Figure 3 · p. 10
6 Discussion
The evaluation is limited to NAS-Bench-101 and NAS-Bench-201, and transfer depends on aligned operation vocabularies and dependency semantics. Larger graphs also face quadratic costs, while end-to-end NAS effects remain untested.
6 Discussion
- Evaluation covers NAS-Bench-101 and NAS-Bench-201; transfer requires aligned operation vocabularies and dependency semantics across search spaces.Graph nodes may represent computations at different granularities.Source: §6 · p. 10
- Dense directed-relation masks and pairwise attention incur quadratic storage and attention costs as graph size grows.Source: §6 · p. 10
- The effect of improved predictor ranking on end-to-end NAS remains to be evaluated.Future work proposes integration into NAS pipelines and evaluation under a fixed end-to-end budget.Source: §6 · p. 10
7 Conclusion
DTA-GT applies encoder-wide directional consistency across node initialization, spectral encoding, pairwise interaction, and feature update. It achieves consistent gains on accuracy and latency prediction, with ablations and diagnostics supporting the proposed stages’ contributions and behavior.
7 Conclusion
- DTA-GT preserves directed computational semantics across node initialization, global spectral encoding, pairwise interaction, and feature update.Its components are direction- and topology-aware node initialization, Magnetic Laplacian spectral encoding, direction-dual structural attention, and a Direction-Sensitive MoE-FFN.Source: §7 · p. 10
- Experiments on NAS-Bench-101 and NAS-Bench-201 show consistent gains on accuracy and latency prediction across limited-budget settings.Source: §7 · p. 10
- Ablations, Magnetic spectral controls, and mechanistic diagnostics support the direction-aware stages’ contribution and intended behavior.Source: §7 · p. 10
Version Note
The preprint is provisional and may contain errors; the forthcoming camera-ready NeurIPS proceedings version supersedes it.
Version Note
- The forthcoming camera-ready NeurIPS proceedings version supersedes this provisional preprint.The preprint may contain errors.Source: Version Note · p. 10
A.1 Implementation Platform
The appendix documents the experimental platform, datasets and evaluation protocol, baseline comparisons, spectral settings and aggregation, and mechanistic diagnostics.
A.1 Implementation Platform
- Experiments use a workstation with an Intel Core i9-13900K, 64 GB RAM, and one RTX 4090 GPU; large-DAG scalability experiments use an A800-SXM4-80GB.The NAS-Bench experiments run on Ubuntu 22.04 with PyTorch 2.4.0; large-DAG hardware and timing details are in Appendix F.Source: §A.1 · p. 14
A.2 Datasets, Splits, and Training Budgets
- NAS-Bench-101 contains 423,624 cells (at most 7 nodes and 9 edges), while NAS-Bench-201 contains 15,625 cells (4 nodes and 6 edges).Five shared seeds define varying train/validation partitions; Kendall’s τ and MAPE are computed over the full architecture pool, including training and validation architectures.Source: §A.2 · p. 14
A.3 Fair Comparison and Reproducibility
- All methods share partitions, evaluation targets, budgets, and metrics; DTA-GT uses one fixed default configuration across tasks and budgets.Main comparison and ablation tables report means over five runs; standard deviations appear in Appendix H.Source: §A.3 · p. 14
A.4 Descriptions of Baselines
- Baselines span sequence-, CNN-, GNN-, Transformer-, and hybrid architecture predictors, including methods based on tensor features, graph structure, and global attention.Public implementations are used when available; otherwise predictors are reproduced and tuned within reported or standard ranges under the shared protocol.Source: §A.4 · pp. 14–15
A.5 Code and Data Access
- The public repository provides the DTA-GT implementation, experiment scripts, benchmark data files, and measured NAS-Bench-101 latency labels.Source: §A.5 · p. 15
A.6 Direction-Aware Spectral Encoding Settings
- Magnetic encoding uses q = 0.25 and the first k = 4 eigenvectors, retaining their real and imaginary parts as two-channel spectral inputs.The phase-parameter sensitivity experiment is described in Appendix B.1.Source: §A.6 · p. 15
A.7 Magnetic Spectral Aggregation Details
- Each node’s two-channel spectral components are mapped to hidden representations, attended over spectral modes, flattened, projected, and added to node initialization.Source: §A.7 · p. 16
B Additional Mechanistic Diagnostics
- Mechanistic diagnostics examine orientation sensitivity, task-dependent MoE routing, and attention allocation over predefined relation channels; they complement rather than attribute the main gains.The diagnostics use the second training budget for each task, while the phase-parameter experiment uses the lowest budget.Source: B Additional Mechanistic Diagnostics · p. 16
B.1 Magnetic Phase Sensitivity and Parameter Analysis
Magnetic phase controls make edge-orientation sensitivity explicit, while routing and relation-channel diagnostics show how directional contexts are used. The complexity analysis identifies graph preprocessing and dense pairwise computation as scale-dependent costs, with reuse and support-aware execution as practical responses.
B.1 Magnetic Phase Sensitivity and Parameter Analysis
- q = 0.25 yields mean edge-reversal sensitivity of 0.2398 on NAS-Bench-101 and 0.2065 on NAS-Bench-201; q = 0 gives zero in all four settings.NAS-Bench-101 responses vary across topology-diverse cells, while NAS-Bench-201's fixed encoding topology gives a constant response across operation assignments.Source: §B.1 · p. 16, Figure 4 · p. 17
- With q = 0.25, Kendall’s τ is highest and MAPE lowest across all four tasks among tested phase values.The sensitivity sweep tests q ∈ {0, 0.125, 0.25, 0.5} under each task’s lowest training budget; Table 6 reports means over five runs.Source: §B.1 · p. 17, Table 6 · p. 17
B.2 Detailed MoE Routing by Node Role
- NAS-Bench-201 accuracy shows positive role-aligned routing contrasts for source, branching, converging, and sink nodes: 0.0342, 0.0371, 0.0152, and 0.0439.Across other settings, contrast signs and magnitudes vary with prediction target and node role.Source: §B.2 · pp. 17–18, Table 7 · p. 17, Figure 5 · p. 18
B.3 Detailed Relation-Channel Statistics
- Across second-budget settings, shared-successor and shared-predecessor relations are enriched, while directed-reachability channels remain near proportional to their coverage.On NAS-Bench-101, the shared-successor and shared-predecessor enrichments are approximately 2.20 and 2.70; on NAS-Bench-201, shared-predecessor enrichment reaches 4.67.Source: §B.3 · p. 18, Figure 6 · p. 19
C.1 Time Complexity
- With fixed k, encoder evaluation costs O(nkd^2 + nk^2d + L(n^2d + nd^2)); dense preprocessing gives a conservative first-use bound of O(n^3 + L(n^2d + nd^2)).Dense relation masks require O(n^2) storage, retained eigenvectors O(nk), node states O(nd), and encoder parameters O(Ld^2); a conventional backward pass adds O(L(n^2+nd)) storage.Source: §C.1 · p. 19, §C.2 · p. 19
C Complexity Analysis
- Graph-dependent construction and dense pairwise processing are the identified scale-dependent costs; graph-derived quantities are computed once and reused.Degrees, directed depths, and sparse Magnetic Laplacians take O(n+m) to construct, while dense reachability/common-neighbor construction and full eigendecomposition have O(n^3) worst-case time and O(n^2) storage.Source: §C.3 · p. 19, §C.1 · p. 19
C.3 Design Choices for Practical Scalability
- Sparse retained-mode spectral computation avoids forming a complete spectral basis, while reusable preprocessing and support-indexed attention target repeated work and supported pairs.Support-indexed attention has an O(Sd) term and sparse predecessor/successor aggregation costs O(md); the stated practical motivation is architecture-dependent connectivity.Source: §C.3 · pp. 19–20
D Computational Efficiency
The paper examines a compact DTA-GT variant, parameter choices, and prediction on larger DAGs. The reported results show competitive compact-model prediction, sensitivity to spectral modes and model scale, and strong classification results on MalNet-derived DAGs.
D Computational Efficiency
- DTA-GT-S uses 19.4× fewer parameters than DTA-GT, beats Graphormer in Kendall’s τ in all four settings and MAPE in three, while Graphormer wins MAPE on NAS-Bench-201 accuracy.Inference times are 0.049 ms/sample for DTA-GT-S and 0.050 ms/sample for Graphormer; DTA-GT’s NAS-Bench-201 profile uses 6.99M parameters and averages 0.121 ms/sample and 0.180 h training time.Source: Table 9 · p. 20, D Computational Efficiency · pp. 20–21, Table 8 · p. 20
E Design Rationale and Parameter Sensitivity
- The encoder’s four directional stages preserve computational direction from node initialization through pairwise interaction and feature updates.The stages represent local directed roles, a global orientation-sensitive spectral prior, directed node-pair relations, and adaptive self, predecessor, and successor contexts.Source: E Design Rationale and Parameter… · p. 21
E.1 Retained Spectral Modes
- Four retained spectral modes achieve the highest Kendall’s τ and lowest MAPE in all four tasks; on NAS-Bench-201 accuracy, k = 1 to 4 changes τ from 0.7383 to 0.8015 and MAPE from 0.0481 to 0.0244.The comparison varies k ∈ {1, 2, 4} with other parameters fixed at the shared configuration.Source: §E.1 · p. 22, Table 11 · p. 22
E.2 Model Depth and Width
- Across all four tasks, L = 10 outperforms L = 4, while L = 12 declines on all eight metrics; d = 180 leads tested widths in Kendall’s τ and MAPE.The depth and width sweeps vary one parameter at a time under the lowest training budget of each task.Source: §E.2 · pp. 22–23, Table 12 · p. 23, Table 13 · p. 23
F.1 Predictive Accuracy on MalNet-Derived DAGs
- On MalNet-derived DAGs, DTA-GT achieves 74.1% accuracy and 73.8% macro-F1, leading the compared methods on all four classification metrics.Its accuracy exceeds Graphormer by 4.6 percentage points and BRP-NAS by 17.1 percentage points; all methods use the same derived graphs and partitions.Source: §F.1 · pp. 23–24, Table 14 · p. 24
F.2 Computational Efficiency on Large DAGs
The large-DAG profiles separate graph-input construction from model inference and show that DTA-GT executes on graphs up to 32,768 nodes, while its first-use preprocessing and dense relation storage remain important scaling costs. The discussion identifies larger-graph efficiency and broader adaptation as open directions, and Appendix H supplies repeated-run statistics for the benchmark results.
F.2 Computational Efficiency on Large DAGs
- On the 3,416-node DAG, DTA-GT’s structural-mask construction, spectral construction, and model forward take 11.46 ms, 7.48 ms, and 40.01 ms, totaling 58.95 ms.For repeated predictions of the same architecture, constructed graph inputs can be reused and only model forward runs again.Source: Table 15 · p. 24, §F.2 · p. 24
- 58.95 ms: DTA-GT’s end-to-end latency on the largest profiled real DAG is below PINAT’s 122.56 ms and NN-Former’s 161.20 ms, while BRP-NAS takes 4.24 ms.The comparison includes each method’s required preprocessing and one inference pass.Source: §F.2 · p. 24, Table 16 · p. 25
- 3,125.11 ms: DTA-GT model forward on the 32,768-node generated DAG is below PINAT’s 9,684.47 ms and NAR-Former V2’s 12,036.15 ms, while BRP-NAS records 313.80 ms.NAR-Former and NN-Former run out of memory under their reported settings.Source: Table 17 · p. 25, §F.2 · p. 25
- On the 32,768-node generated DAG, independently measured feature preprocessing, spectral construction, structural-mask construction, and model-forward means sum to 25,601.083 ms for first use.Model forward takes 3,125.117 ms and reaches 42.249 GiB peak GPU memory.Source: Table 18 · p. 25, §F.2 · p. 25
- On a 32,768-node spanning-chain DAG, topology-specific construction of six masks takes 33.340 ms, and the resulting dense Boolean masks occupy 6 GiB.The reported construction time and storage are separate measures.Source: §F.2 · p. 25
G Extended Discussion
- Cross-space adaptation would require compatible operation definitions and dependency semantics; the discussion proposes aligned operation vocabularies and small-subset adaptation for compatible spaces and tasks.It also proposes evaluating predictor-guided architecture selection within a fixed end-to-end NAS evaluation budget.Source: G Extended Discussion · p. 26
- The evaluated direction-dual attention uses six heuristic relation channels for local dependencies, multi-hop reachability, and shared predecessor/successor structure; expanded sets’ predictive value and cost remain unstudied.The discussion suggests hop-specific and directed-motif relations as possible extensions.Source: G Extended Discussion · p. 26
- Directed reachability can grow quadratically even when adjacency is sparse: a chain has n - 1 edges but n(n - 1)/2 reachable ordered pairs, and A800 profiles demonstrate generated-DAG execution up to 32,768 nodes.Source: G Extended Discussion · p. 26
H.2 Ablation Analysis
Appendix H provides repeated-run uncertainty estimates for ablation results and Magnetic-Laplacian controls. The supplied checklist passages also report reproducibility and disclosure details, including incomplete compute accounting and asset-license information.
H.2 Ablation Analysis
- Table 23 reports complete ablation results under the lowest training budget of each task as mean ± standard deviation.The accompanying passage identifies these as uncertainty estimates for the mean-only ablation results in Table 2.Source: §H.2 · p. 28, Table 23 · p. 29
H.3 Magnetic Spectral Controls
- Table 24 reports symmetric- and Magnetic-Laplacian controls under the lowest training budget of each task, including the q = 0 row from the five-run phase-parameter sensitivity experiment.Cells report mean ± standard deviation.Source: Table 24 · p. 29, §H.3 · p. 29
NeurIPS Paper Checklist
- The checklist states that datasets, per-seed shared splits, training settings, and reproduction details are provided in Appendices A–E, and that Appendix H reports five-run means and standard deviations.The public code and data repository is linked in Appendix A.5.Source: NeurIPS Paper Checklist · pp. 31–32
- The checklist notes that hardware and selected resource profiles are reported, but complete per-run and total compute accounting is not provided.The checklist also says existing assets are credited but their licenses and terms of use are not fully enumerated.Source: NeurIPS Paper Checklist · pp. 33–34