Reproduce “Node-wise Feature Encoding for Neural Performance Prediction” on OpenResearch

About this paper

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

FeatureFormer improves neural-network latency and energy prediction by embedding FLOPs, parameter counts, and a memory-operation proxy directly into individual graph-node representations. Across leave-one-family-out tests, it achieves 7.24% latency MAPE and 6.63% energy MAPE, outperforming the strongest reported graph-transformer baseline, while a new NNEQ dataset broadens evaluation beyond latency.

Key insights

• Explicit node-level cost information helps graph predictors focus on operations that are likely to dominate hardware execution, rather than relying only on graph position or operation identity.
• The gains transfer to existing predictors, indicating that computational feature encoding is a modular representation improvement rather than an advantage restricted to the full FeatureFormer architecture.
• Latency and energy are related but not interchangeable: their measurements are not consistently proportional, so they require separate prediction targets and evaluation.

Summary by alphaXiv.