# Scalable Heterogeneous Graph Foundation Models for Data-Driven Optimal Power Flow in Smart Grids

arXiv: https://arxiv.org/abs/2605.23194 · HTML: https://mlsanity.com/paper/2605.23194/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 18 of its 44 references that are in the map.
- First posted on arXiv 22 May 2026.

## Builds on (most cited first, 18 of 18)

- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Overcoming catastrophic forgetting in neural networks](https://mlsanity.com/paper/1612.00796/) (2016)
- [How Attentive are Graph Attention Networks?](https://mlsanity.com/paper/2105.14491/) (2021)
- [Characterizing possible failure modes in physics-informed neural networks](https://mlsanity.com/paper/2109.01050/) (2021)
- [Principal Neighbourhood Aggregation for Graph Nets](https://mlsanity.com/paper/2004.05718/) (2020)
- [Predicting AC Optimal Power Flows: Combining Deep Learning and Lagrangian Dual Methods](https://mlsanity.com/paper/1909.10461/) (2019)
- [The Power Grid Library for Benchmarking AC Optimal Power Flow Algorithms](https://mlsanity.com/paper/1908.02788/) (2019)
- [Optimal Power Flow Using Graph Neural Networks](https://mlsanity.com/paper/1910.09658/) (2019)
- [Physics-Informed Neural Networks for AC Optimal Power Flow](https://mlsanity.com/paper/2110.02672/) (2021)
- [Learning Optimal Power Flow: Worst-Case Guarantees for Neural Networks](https://mlsanity.com/paper/2006.11029/) (2020)
- [Topology-aware Graph Neural Networks for Learning Feasible and Adaptive ac-OPF Solutions](https://mlsanity.com/paper/2205.10129/) (2022)
- [Foundation Models for the Electric Power Grid](https://mlsanity.com/paper/2407.09434/) (2024)
- [PowerFlowNet: Power Flow Approximation Using Message Passing Graph Neural Networks](https://mlsanity.com/paper/2311.03415/) (2023)
- [Heterogeneous Edge-Enhanced Graph Attention Network For Multi-Agent Trajectory Prediction](https://mlsanity.com/paper/2106.07161/) (2021)
- [LUMINA: Foundation Models for Topology Transferable ACOPF](https://mlsanity.com/paper/2603.04300/) (2026)
- [LUMINA: A Grid Foundation Model for Benchmarking AC Optimal Power Flow Surrogate Learning](https://mlsanity.com/paper/2605.02133/) (2026)
- [Towards Systematic Generalization for Power Grid Optimization Problems](https://mlsanity.com/paper/2605.02026/) (2026)

## Abstract (arXiv)

Fast and reliable optimal power flow (OPF) approximation is important for power system operation, yet heterogeneous OPF graph models are often evaluated either with architecture-specific implementations or on limited training corpora. This paper presents a large-scale heterogeneous graph-learning workflow, built on HydraGNN, for data-driven OPF surrogate modeling and graph foundation-model (GFM) development. The workflow preserves buses, generators, loads, shunts, alternating-current (AC) lines, transformers, and device-to-bus relations and provides a common implementation for distributed preprocessing, multi-graphics-processing-unit (GPU) training, hyperparameter optimization (HPO), and downstream adaptation. Using approximately three million heterogeneous graphs from ten Power Grid Library for Benchmarking AC Optimal Power Flow Algorithms (PGLib-OPF) cases spanning 14 to 13,659 buses, we compare six heterogeneous graph neural network (GNN) families under a common HPO campaign on Frontier, resulting in two compact HeteroSAGE and HeteroHEAT models (~1.6-1.7 million parameters) that achieve the lowest observed validation MSE. Strong and weak scaling using up to 1,024 Frontier nodes show that 256 nodes (corresponding to 1,024 AMD MI250X GPUs) provides the best time-to-solution balance for the measured workload. Two downstream tasks evaluate task adaptation on the IEEE 118-bus case: a classification task to distinguish nominal operating conditions from artificially generated extreme-overload conditions and N-1 single-component-outage regression to predict bus voltage magnitude and phase angle. Partial fine-tuning provides the strongest balance between predictive performance and computational requirements.