PaperlayerView exact revision on arXiv

Source-linked AI summary

Node-wise Feature Encoding for Neural Performance Prediction

Matthew Grenier, William Hammer, Andrew Heuer, Nikhil Krishna, Yi Wang, Ramtin Zand

arXiv:2608.27794v1cs.LG

TL;DR

Hardware-aware neural architecture search needs accurate latency and energy prediction, yet existing graph predictors largely omit node-level computational cost. The paper introduces FeatureFormer, which combines explicit computational encodings with gated graph attention, and NNEQ, an energy-consumption dataset. FeatureFormer achieves state-of-the-art results across latency and energy prediction, while the encoding also improves existing predictors with negligible trade-offs.

  • Problem

    Existing graph-based performance predictors largely lack explicit node-level computational-cost information, although accurate latency and energy prediction matters for resource-constrained edge deployment.

  • Method

    FeatureFormer combines gated graph self-attention with node-level encodings of FLOPs, parameter counts, and a memory-operation proxy, alongside global computational features.

  • Results

    FeatureFormer achieves state-of-the-art latency and energy prediction, improving existing predictors with negligible trade-offs across the reported evaluation settings.

  • Takeaways & Limitations

    Node-level computational encodings improve prediction accuracy and reliability for FeatureFormer and several recent predictors, supporting their broader applicability.

Abstract

from arXiv · show

As neural networks are increasingly deployed on resource constrained edge devices, accurate prediction of latency and energy is critical for efficient neural architecture search. Existing GNN and transformer based predictors achieve strong results but largely ignore node-level computational cost, limiting their ability to model performance critical operations. To address this, we introduce FeatureFormer, a neural performance predictor that incorporates explicit node-wise encodings of FLOPs, parameter counts, and memory proxies within a gated graph attention architecture. We also present NNEQ, a new large-scale energy consumption dataset that enables unified evaluation of latency and energy prediction. Extensive experiments demonstrate that FeatureFormer achieves state-of-the-art performance across both metrics, including challenging out-of-domain settings. Finally, we show that the proposed encoding is broadly applicable and consistently improves existing predictors with negligible overhead.

1 Introduction

Hardware-aware NAS must optimize latency, energy, and memory under edge-device constraints, but repeated measurements are expensive. The paper introduces FeatureFormer and NNEQ to improve prediction and evaluation of these performance metrics.

  • Edge deployment requires neural architectures to satisfy strict latency, energy-consumption, and memory constraints.
  • Predictive models replace repeated on-device latency and energy measurements, accelerating hardware-aware NAS.
  • Prior graph-based predictors omit explicit node-level FLOP counts and parameter sizes, treating costly and lightweight operations similarly.
  • FeatureFormer adds operation-derived computational attributes and gated attention, concentrating attention on operations that dominate inference cost.
  • FeatureFormer combines architectural and encoding innovations to achieve state-of-the-art prediction of inference latency and energy consumption.
  • NNEQ extends NNLQ with energy-consumption measurements, while the proposed encoding is reported to improve existing methods with negligible trade-offs.

2 Related Work

Neural performance predictors have progressed from lightweight and lookup-based models to GNNs and graph transformers. FeatureFormer addresses their remaining encoding gap by adding computational cost directly to node features while retaining topological information.

  • Earlier predictors used lightweight models, lookup tables, or simple performance proxies for latency and energy estimation.
  • GNN-based predictors became a dominant paradigm by aggregating information from neighboring nodes in graph-structured architectures.
  • Graph transformers capture long-range dependencies through global attention, but distant nodes may have little relevance to execution time.
  • NAR-Former V2 and NN-Former refine topology modeling with masked attention over adjacent or parallel nodes.
  • Despite improved topological modeling, prior methods retain encoding strategies that give individual nodes little explicit computational-cost information.

3 Methods

FeatureFormer encodes each node's operation, attributes, and computational proxies before modeling architectural relationships with gated graph self-attention. The resulting representations are aggregated with global computational features for latency or energy regression.

  • Feature Encoding Stage: FeatureFormer represents each node using operation type, runtime-relevant attributes, and computational information.Attributes include output shape and, for structured layers, kernel size, stride, and padding.
  • Feature Encoding Stage: Sinusoidal encodings map node attributes and computational values from scalar inputs into higher-dimensional representations.The encoding is applied to arbitrary floating-point inputs, including layer attributes and computational proxies.
  • Feature Encoding Stage: Computational proxies include FLOPs, parameter counts, and output-shape products plus parameter counts as a memory-operation proxy.These values are concatenated with operation and attribute encodings to form the node representation.
  • Relationship Encoding Stage: The relationship encoder repeats four blocks, each combining gated graph self-attention with a two-layer ReLU feed-forward network.This depth matches the number of message-passing layers used in NN-Former.
  • Relationship Encoding Stage: Four attention heads use forward, backward, and two sibling-based graph views to model interactions among operations.The views are represented by A, A^T, A^T A, and AA^T, respectively.
  • Relationship Encoding Stage: A multiplicative sigmoid gate nonlinearly combines attention-head outputs by selectively emphasizing or suppressing graph views.The resulting node representations are summed, concatenated with global computational features, and passed to a regression head.

4 Experiments

The experiments evaluate FeatureFormer for latency and energy prediction under out-of-domain and in-domain protocols, using NNLQ and the newly introduced NNEQ dataset. FeatureFormer achieves state-of-the-art results across both metrics, while adding node-level computational encoding also improves prior predictors.

  • Experimental Setup: FeatureFormer is evaluated against prior predictors for latency and energy in out-of-domain and in-domain settings.The out-of-domain protocol trains on nine model families and tests on the held-out family; in-domain experiments use separate training and testing models from each family.
  • Datasets: NNEQ extends the NNLQ setup with energy measurements for CNN model families, enabling comparable evaluation of latency and energy prediction.Both datasets contain readings from 20,000 models across 10 model families, while NNEQ is designed for energy prediction.
  • Out-of-Domain Results: 1.15 percentage points lower MAPE and 2.88 percentage points higher Acc(10%) than NN-Former establish FeatureFormer as the best out-of-domain latency predictor.FeatureFormer also substantially outperforms prior methods on several individual families, although it trails NNLP and NAR-Former V2 on VGG.
  • Out-of-Domain Results: 0.85 percentage points lower MAPE and 3.93 percentage points higher Acc(10%) than the next best method establish state-of-the-art out-of-domain energy prediction.The competing-method ranking differs between NNEQ’s MAPE and Acc(10%) results.
  • Encoding Experiments: Adding Fcomp substantially improves MAPE and Acc(10%) for almost every model and task, with NNLP gaining 8.04 percentage points in average Acc(10%) on NNLQ.The improvement is consistent across models and tasks except for a slight decline for NN-Former on NNEQ, while FeatureFormer remains best overall.

5 Ablations

The ablations isolate the contributions of gating, masking, feature segments, computational attributes, and model size. Results support multiplicative gating, specialized masking, combined encodings, and architectural advantages beyond parameter count.

  • Effect of Gating: Removing the multiplicative gate increases MAPE by 0.41 percentage points and reduces Acc(10%) by 1.77 percentage points on NNEQ.The ablation indicates that gating improves both average accuracy and prediction reliability.
  • Masking Variations: Global attention consistently underperforms, while specialized masking outperforms shared masking in the four-head FeatureFormer configuration.The result supports assigning different attention masks to different heads for greater head-wise specialization.
  • Encoding Components: Each encoding segment alone is unsatisfactory, whereas the Fattr-plus-Fcomp configuration performs best on average and the full encoding wins on nearly every family except NasBench201.The full configuration’s overall comparison is mainly affected by a large NasBench201 improvement for the alternative configuration.
  • Computational Features: Each individual computational feature improves over the no-feature baseline, while combinations of two features generally outperform single-feature variants.The parameters-plus-memory combination is the stated exception among the tested combinations.
  • Model Size Comparison: FeatureFormer consistently outperforms parameter-matched NN-Former, including when FeatureFormer is downscaled and NN-Former is scaled up.Both models use node-level static features, and NN-Former is augmented with FeatureFormer’s encoding for the comparison.

6 Conclusion and Future Work

FeatureFormer combines gated graph self-attention with explicit node-level computational encoding for latency and energy prediction. Results support the encoding across methods, datasets, evaluation settings, and out-of-domain scenarios.

  • FeatureFormer combines gated graph self-attention with explicit node-level encodings of FLOPs, parameter counts, and memory-operation proxies.
  • The encoding improves prediction accuracy and reliability for both latency and energy tasks.
  • Applying the encoding to several recent predictors also improves their results, suggesting it is not specific to one architecture.
  • FeatureFormer performs competitively across multiple datasets and evaluation settings, including out-of-domain scenarios.
  • Node-level computational encodings are suggested as a direction for improving neural performance prediction and supporting future hardware-aware NAS.

A.1 Training Hyperparameters

FeatureFormer uses a four-layer GGSA architecture with fixed model dimensions and a 222-dimensional input representation. Training follows a warm-up, cosine-decay, and AdamW optimization schedule, while prior methods use their released configurations except in Fcomp studies.

  • Architecture: FeatureFormer uses four layers, each combining a GGSA block with a standard ReLU feed-forward network.The model uses dmodel = 512 and dffn = 2048.
  • Input representation: The default input vector has 222 dimensions, including 32 Fop dimensions, 160 Fattr dimensions, and 30 additional dimensions.
  • Optimization: Training uses batch size 16 for 75 epochs, with learning-rate warm-up followed by cosine decay and AdamW optimization.The learning rate increases from 1 × 10−4 to 1 × 10−3 during the first 10% of training steps.
  • Baseline evaluation: Prior methods are evaluated with their released implementations and hyperparameters, except Fcomp experiments, which use a small grid search over Fop and Fcomp sizes.The grid search is intended to give each model a fair opportunity to benefit from Fcomp.

A.2 Training Cost

FeatureFormer experiments require substantial computational resources, with each single-seed model run taking approximately two hours on a 40 GB A100 GPU. The reported results are estimated to consume about 2,000 A100 GPU-hours overall.

  • Per-run cost: Approximately 2 hours are required to train one FeatureFormer model with a single random seed on a 40 GB A100 GPU.
  • Experimental scale: A table containing 10 models evaluated with 10 seeds each would require 100 individual training runs.
  • Total cost: Approximately 2,000 A100 GPU-hours were required to collect all results reported in the paper.

B Implementation Details

The implementation clarifies how GGSA gating combines differently masked attention heads and connects this mechanism to improved accuracy. Additional analyses compare latency measurements from NNLQ with energy measurements from NNEQ across families and datasets.

  • GGSA implementation: GGSA gating nonlinearly combines outputs from differently masked attention heads.The paper provides pseudocode for the layer in figure 6.
  • GGSA implementation: The added gating flexibility improves accuracy, as demonstrated in table 14.
  • Dataset analysis: Figure 7 examines the correlation between NNLQ latency measurements and NNEQ energy measurements, while table 11 reports per-family and per-dataset statistics.

D Expanded Results

The expanded results provide full ablation and dataset-specific evaluations of node-level computational information, gating, masking, and feature choices. They also distinguish latency and energy as related but distinct prediction tasks.

  • Fcomp Ablations: Tables 12 and 13 report out-of-domain performance when node-level computational information is added to various models on NNLQ and NNEQ.Each model is evaluated over 10 independent trials.
  • GGSA Design: Tables 14 and 15 evaluate FeatureFormer’s gating mechanism and alternative attention-head masking schemes on out-of-domain NNEQ.Gating results use 10 trials, while masking results use 3 trials with 4 attention heads.
  • Dataset Comparison: NNEQ ground-truth energy and NNLQ ground-truth latency are neither fully linear nor fully proportional across individual models.The comparison emphasizes that latency and energy constitute distinct predictive tasks.
  • Dataset Statistics: Dataset statistics are summarized by model family, with an “All” aggregate covering all families within each dataset.The statistics are computed over ground-truth values for all samples in each family.
  • Feature Analysis: Tables 16 and 17 examine the importance and composition of Fop, Fattr, and Fcomp for FeatureFormer on out-of-domain NNLQ.Mem denotes a memory-operation proxy formed from multiplied output shape plus parameter count; results use 3 trials.
Loading 2608.27794v1…