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

2026-05-21

AI Summary

This work addresses the limited generalization and scalability of existing learning-based optimal power flow (OPF) surrogate models, which often neglect the heterogeneous graph structure of power grids. We propose a scalable heterogeneous graph neural network framework based on HydraGNN, implemented with distributed preprocessing, training, and fine-tuning on the ORNL Frontier supercomputer, coupled with DeepHyper for efficient hyperparameter optimization. For the first time, we develop a foundation model tailored to large-scale heterogeneous power grid graphs, enabling cross-topology transfer with minimal adaptation cost—requiring only partial or head-only fine-tuning. Evaluated on ten PGLib-OPF datasets spanning 14 to 13,659 buses, our compact models (∼1.6–1.7M parameters) significantly improve few-shot accuracy, training stability, and convergence speed in both feasibility classification and N-1 contingency regression tasks.

Abstract

Fast and reliable optimal power flow (OPF) approximation is essential for reliable smart-grid operation, yet many learning-based surrogates either flatten the native heterogeneous structure of power networks, target a limited set of grid topologies, or lack scalable infrastructure for graph foundation model (GFM) training. This paper presents a scalable heterogeneous graph neural network (GNN) workflow, built on HydraGNN, for data-driven OPF surrogate modeling and OPF-GFM development. The workflow preserves the distinct node and edge types of power grids -- buses, generators, loads, shunts, AC lines, transformers, and device-to-bus couplings -- and supports distributed preprocessing, training, hyperparameter optimization (HPO), and downstream fine-tuning on leadership-class supercomputers. Using three million heterogeneous graph instances spanning ten PGLib-OPF cases, from 14 to 13,659 buses, we conduct DeepHyper-driven HPO on the ORNL Frontier supercomputer. The campaign identifies compact models ($\sim$1.6--1.7M parameters) with the lowest validation losses. Downstream experiments on feasibility classification and N-1 contingency regression show that fine-tuning pretrained OPF GFM improves low-data accuracy, stabilizes training, accelerates convergence, and reduces adaptation cost when partial or head-only fine-tuning is used.

Problem

• optimal power flow
• heterogeneous graph
• graph foundation model
• smart grids
• scalability

Innovation

• heterogeneous graph neural network
• graph foundation model
• optimal power flow
• scalable GNN workflow
• distributed hyperparameter optimization