Towards Generalization of Graph Neural Networks for AC Optimal Power Flow

TL;DR

This paper tackles the computational bottleneck of ACOPF by introducing HH-MPNN, a hybrid model that unifies heterogeneous graph neural networks with a scalable transformer to handle topology changes and grid scaling. By explicitly modeling power-system components as distinct node/edge types and using effective resistance-based positional encodings, the approach achieves $<1\%$ optimality gaps on default topologies and $<3\%$ in zero-shot N-1 cases across grids from 14 to 2,000 buses, while delivering $10^3$ to $10^4\times$ speedups over interior-point solvers. Pre-training on smaller grids further improves performance on larger grids, and fine-tuning on a target grid yields substantial gains in accuracy and generalization, including size generalization and reduced data-generation costs. The results suggest a practical, generalizable ML framework for real-time power system operations, with implications for topology adaptability and scalable data efficiency.

Abstract

AC Optimal Power Flow (ACOPF) is computationally expensive for large-scale power systems, with conventional solvers requiring prohibitive solution times. Machine learning approaches offer computational speedups but struggle with scalability and topology adaptability without expensive retraining. To enable scalability across grid sizes and adaptability to topology changes, we propose a Hybrid Heterogeneous Message Passing Neural Network (HH-MPNN). HH-MPNN models buses, generators, loads, shunts, transmission lines and transformers as distinct node or edge types, combined with a scalable transformer model for handling long-range dependencies. On grids from 14 to 2,000 buses, HH-MPNN achieves less than 1% optimality gap on default topologies. Applied zero-shot to thousands of unseen topologies, HH-MPNN achieves less than 3% optimality gap despite training only on default topologies. Pre-training on smaller grids also improves results on a larger grid. Computational speedups reach 1,000x to 10,000x compared to interior point solvers. These results advance practical, generalizable machine learning for real-time power system operations.

Figures (6)

• Figure 1: The proposed HH-MPNN architecture with detailed message passing mechanism showing how information flows between different node and edge types. Skip connections and bidirectional updates enhance model expressivity. K=5 message passing layers are used.
• Figure 2: Heterogeneous versus Homogeneous graph representation of the power grid
• Figure 3: Comparison of average optimality gaps in the base and zero-shot experiments. Optimality gaps for the proposed HH-MPNN remain below $3\%$ when models trained on only the full topology are applied to numerous N-1 topologies.
• Figure 4: Optimality gaps for high-impact N-1 contingencies. Naive: full-topology model assuming default topology. Zero-shot: full-topology model with correct N-1 topology. Base: model trained on N-1 data. Zero-shot sometimes performs worse than the naive baseline due to distribution shifts, while the N-1 trained model achieves <4% gaps across all systems.
• Figure 5: Each bar shows the number of unique topologies in the N-1 test subset for each grid. There is exponential growth in the number of grid topologies with increasing grid size.