🤖 machine learning

MxGPS: Multiplex Graph Transformers for a Power Grid Foundation Model

MxGPS introduces a parameter-efficient multiplex graph transformer that leverages multi-task joint training on Static State Estimation and AC Power Flow to overcome "topology overfitting," achieving robust zero-shot generalization across unseen power grid topologies where single-task models fail.

Original authors: Charilaos Papaioannou, Ioannis Tsantilas, Dimitris Giannakakos, Vasilis Michalakopoulos, Sotiris Pelekis, Vangelis Marinakis, Arsam Aryandoust, Antonello Monti, Ricardo J. Bessa, Perdo P. Vergara, Joc

Published 2026-07-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Charilaos Papaioannou, Ioannis Tsantilas, Dimitris Giannakakos, Vasilis Michalakopoulos, Sotiris Pelekis, Vangelis Marinakis, Arsam Aryandoust, Antonello Monti, Ricardo J. Bessa, Perdo P. Vergara, Jochen Cremer, Elissaios Sarmas

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine the electrical grid as a giant, living city of energy. In this city, power plants are the factories, homes and businesses are the consumers, and the wires connecting them are the busy streets. To keep the lights on and the city safe, engineers have to constantly solve a massive, moving puzzle: figuring out exactly how much electricity is flowing through every street and what the voltage (the "pressure" of the electricity) is at every house. Traditionally, they use heavy, slow math calculators to solve this. But as the city grows and adds unpredictable elements like solar panels and wind turbines, those old calculators get too slow. Enter Artificial Intelligence, specifically a type called Graph Neural Networks (GNNs). Think of GNNs as a super-smart student who learns the layout of the city's streets by studying maps. They can predict the flow of electricity almost instantly. However, there's a catch: if this student only studies one specific city map, they might memorize the exact shape of the streets rather than learning the actual rules of traffic. When they try to drive in a brand new city with a different layout, they get lost. This paper tackles that exact problem: how to teach an AI to understand the rules of electricity so it can handle any city, not just the one it studied.

The researchers behind this paper, led by Charilaos Papaioannou and colleagues, discovered a sneaky failure mode they call "topology overfitting." Imagine a student who memorizes the answer key for a specific math test. If the test questions change slightly (like a different city layout), the student fails miserably, even though they got a perfect score on the original practice test. The authors found that standard AI models for power grids do exactly this: they get great scores on the grids they were trained on, but when the grid's shape changes (a "topology shift"), their performance crashes. They memorized the specific connections of the training grid rather than the underlying physics of how electricity actually behaves.

To fix this, the team built a new model called MxGPS (Multiplex Graph Transformers). Instead of training one AI to do just one job, they created a "multiplex" system. Think of it like a school with a shared library (the shared encoder) but two different specialized classrooms. One classroom is dedicated to Static State Estimation (SSE), which is like a detective trying to figure out the true state of the city even when some sensors are broken or noisy. The other classroom is for Power Flow (PF), which is like a traffic controller predicting exactly how electricity moves through the wires.

Here is the magic trick: these two classrooms share the same library. The AI is forced to learn a single, shared understanding of the city that satisfies both the detective and the traffic controller at the same time. Because the detective needs to be robust against noise and the traffic controller needs to follow strict physical laws, the shared AI cannot just memorize the specific street layout of one city. It has to learn the universal rules that work for both jobs. The researchers also added a "physics check" (based on the Bus Injection Model) that acts like a strict teacher, ensuring the AI's answers always obey the laws of physics, no matter which task it is doing.

The results were striking. When they tested these models on four completely new, unseen city layouts (ranging from small 14-bus grids to massive 300-bus grids), the old single-task models fell apart. The best-performing single-task models saw their errors skyrocket by anywhere from 190% to 1400% when faced with a new grid shape. In contrast, MxGPS only saw its error increase by 39%. Even more impressively, MxGPS achieved a 0% boundary violation rate on all four new grids. This means it never predicted a voltage level that was unsafe or outside the legal limits, whereas the other models frequently made dangerous predictions.

The paper also compared MxGPS to a massive reference model called GridFM. While GridFM is powerful, it is huge, with 20.1 million parameters. MxGPS, using this clever multi-task approach, achieved similar or better results with only 1.6 million parameters—that's 12 times smaller. This proves that you don't need a giant, expensive brain to solve these problems; you just need a brain that learns the right lessons by doing multiple jobs at once. The authors suggest that this "multi-task joint training" is a smarter, more efficient way to build foundation models for power grids, ensuring they can handle the unpredictable, changing world of modern energy without needing to be retrained for every new city layout.

Technical Summary: MxGPS – Multiplex Graph Transformers for a Power Grid Foundation Model

1. Problem Statement: Topology Overfitting in Single-Task GNNs

The paper identifies a systematic failure mode in applying Graph Neural Networks (GNNs) to power grid problems, termed topology overfitting. While single-task fine-tuned GNNs often achieve the lowest in-distribution (ID) error on tasks like Power Flow (PF) and Static State Estimation (SSE), they degrade catastrophically when applied to unseen grid topologies (zero-shot transfer).

The authors argue that task-specific gradient signals in single-task models encode relational structures particular to the training topologies rather than the underlying physical laws. Consequently, models with superior ID performance (e.g., GNS, GPS) exhibit massive performance decay (190%–1400% increase in error) under topology shifts. Furthermore, single-task fine-tuning for SSE often leads to "magnitude collapse," where models predict mean voltage values to minimize error, failing to provide genuine state estimation despite low Mean Absolute Error (MAE).

2. Methodology: MxGPS Architecture

To address topology overfitting, the authors propose MxGPS (Multiplex Graph Transformers), a foundation model architecture designed for multi-task learning.

Architecture: MxGPS consists of a shared node encoder and $K$ $K$ task-specialised Graph GPS (GatedGCN + Transformer) branches. In this work, $K=2$ $K = 2$ , handling SSE and PF simultaneously.
- Shared Encoder: Maps raw bus features and structural positional encodings (Random Walk Structural Encoding and Laplacian Positional Encoding) into a common physical embedding.
- Task Branches: Each branch processes the shared embedding through its own GPS layers, conditioned by a learnable task token ( $t_k$ ) to steer specialization.
- Cross-Branch Attention (Optional): An ablation study evaluates a module where branches exchange per-node representations after each GPS layer to mutual regularize each other. The authors find that for $K=2$ , the independent branch configuration (MxGPS-ind) is superior, suggesting cross-branch attention may be more valuable as the task portfolio expands.
Training Protocol:
1. Phase A (Pre-training): A single GPS model is pre-trained on the Masked Grid Task (MGT) using Scaled Cosine Error (SCE) and a layered physics loss to learn general grid representations.
2. Phase B (Fine-tuning): The pre-trained weights are transferred to all $K$ branches. A linear probe phase freezes the shared encoder while training task tokens and heads, followed by full fine-tuning.
Loss Function: The model is jointly trained on SSE and PF using a composite loss:
- Supervised Loss: Dimension-weighted MSE for task-specific predictions.
- Physics-Informed Loss: A shared Bus Injection Model (BIM) residual loss applied simultaneously to all branches, enforcing AC power balance constraints (Eqs. 2–3) as a soft constraint.
- Boundary Loss: A penalty for voltage magnitudes violating operational limits.

3. Key Contributions

Characterization of Topology Overfitting: The paper demonstrates that lower in-distribution error correlates with higher degradation under topology shift. Models like GNS and GPS degrade by 790% and >1300% respectively, whereas MxGPS degrades by only 39%.
MxGPS Architecture: A multiplex graph transformer that uses multi-task joint training as a structural inductive bias. By forcing the shared encoder to satisfy complementary gradient signals from SSE and PF simultaneously, the model cannot overfit to topology-specific relational structures.
Zero-Shot Generalization: MxGPS achieves 0% Boundary Violation Rate (BVR) on all four unseen zero-shot topologies (14-, 24-, 162-, and 300-bus systems), a metric where single-task baselines fail significantly (up to 37% BVR).
Parameter Efficiency: MxGPS achieves these results with only 1.6M parameters, which is 12 $\times$ fewer than the GridFM reference baseline (GNS-kit), while maintaining comparable or superior zero-shot accuracy.

4. Experimental Results

Experiments were conducted using a 3-fold sliding-window cross-validation protocol across IEEE test systems (14 to 300 buses).

Power Flow (PF):
- In-Distribution: Single-task models (GNS, GPS) achieved lower MAE but high BVR (7–9%). MxGPS had higher ID MAE (~0.021) but near-zero BVR (≤0.33%).
- Zero-Shot: MxGPS maintained the lowest or near-lowest error on unseen topologies (e.g., 0.0230 MAE on 162-bus) with 0% BVR. In contrast, single-task baselines saw error increases of 190% to 1400%.
Static State Estimation (SSE):
- Single-task models (GCN, GAT, GPS) exhibited "magnitude collapse" in SSE, predicting mean voltages (slope < 0.1) despite low MAE.
- MxGPS avoided this collapse (slope $\approx$ 0.92, $R^2 \approx$ 0.98) due to the joint gradient signal, providing genuine state estimation.
Efficiency: MxGPS-ind offers a Pareto-optimal balance, achieving the best zero-shot PF accuracy among 0% BVR models at 12 $\times$ lower parameter cost and 5 $\times$ lower inference latency than GNS-kit.

5. Significance and Claims

The paper claims that multi-task joint training is a principled and parameter-efficient mechanism for achieving topology-agnostic generalization in power grid foundation models. The central finding is that the "failure" of single-task models is not a lack of capacity, but a structural tendency to overfit to training topologies. By introducing a shared physics-informed objective across tasks, MxGPS forces the encoder to learn representations grounded in physical laws rather than specific graph structures.

The authors position MxGPS as a step toward robust foundation models that can be deployed on grids absent from the training portfolio. They note limitations, including the current scope of IEEE systems (up to 300 buses) and the need for scalable attention mechanisms for production-scale grids, but assert that the multi-task approach successfully mitigates the identified failure mode of topology overfitting.