Physics-Aware Distillation of a Grid Foundation Model for Real-Time Edge Inference

Vlad Alexandru Sănduleac; Shahariya Rasheed Rayamarakkarveettil Puthanpurayil; Noman Khan; Remus Teodorescu

cs.LGOct 15, 2026

Physics-Aware Distillation of a Grid Foundation Model for Real-Time Edge Inference

Authors: Vlad Alexandru Sănduleac, Shahariya Rasheed Rayamarakkarveettil Puthanpurayil, Noman Khan, Remus Teodorescu

Organizations: Faculty of Power Engineering, National University of Science and Technology POLITEHNICA Bucharest, Bucharest, Romania · AAU Energy, Aalborg University, Aalborg, Denmark

Abstract

Grid foundation models (GridFMs) are graph transformers pretrained on a masked power-flow objective, and evaluate power flow orders of magnitude faster than an iterative solver. Their size, however, targets data-centre inference, whereas low-inertia operation increasingly asks for inference at the edge, next to the measurement, at batch size one and within a single synchrophasor (PMU) reporting frame. We distil a 20.1M-parameter, 12-layer GridFM teacher into a 0.62M-parameter, 6-layer student, a $32.5\times$ parameter reduction at half the message-passing depth, using an objective that combines the masked power-flow loss with output, physics-residual, latent-state and attention transfer. The physics-residual term, which matches the teacher's power-residual profile, has not previously been applied to typed power-grid graphs. We evaluate on a renewable variant of the IEEE 118-bus benchmark grid in which 26 of the conventional (synchronous) generators are replaced by inverter-based renewables with battery flexibility and standardised reactive-power limits (IEEE 1547), giving $97\,529$ physically valid (AC-OPF-feasible) operating scenarios, and we ablate all eight combinations of the transfer terms, reporting accuracy against ground truth and fidelity to the teacher separately. Distillation lowers student error on all four predicted quantities, by up to $48\%$ on active-power dispatch, and every distilled student is more accurate than its teacher on reactive-power dispatch. The gain is largely carried by the output matching term. The physics-residual term, in the layer-aggregated form used here, is consistently harmful, which we attribute to the spatial averaging it applies to a bus-level residual field. The resulting student is $32.5\times$ smaller in parameter count and has half the message-passing depth, making it a candidate for subsequent edge-device benchmarking.

Figures & tables

	Teacher	Student
Hidden	48	24
Heads	8	4
Layers	12	6
Params	20,069,187	616,707
Reduction	1.0x	32.5x

Table 1: Comparison between full model and reduced parameter student.

Ablations				Quantities ( $\times 10^{3}$ )
				RMSE vs. ground truth				RMSE vs. teacher
$\mathcal{L}_{\mathrm{out}}$	$\mathcal{L}_{\mathrm{physKD}}$	$\mathcal{L}_{\mathrm{lat}}$	$\mathcal{L}_{\mathrm{att}}$	$V_{m}$	$V_{a}$	$P_{g}$	$Q_{g}$	$V_{m}$	$P_{g}$	$Q_{g}$
Teacher				$\mathbf{0.966}$	$\mathbf{7.276}$	$\mathbf{8.531}$	$35.692$	–	–	–
Baseline				$1.596$	$26.334$	$40.686$	$47.236$	$1.547$	$35.921$	$47.627$
$\bullet$				$1.169$	$18.450$	$21.381$	$34.984$	$0.687$	$\mathbf{15.719}$	$10.296$
$\bullet$	$\bullet$			$1.231$	$19.724$	$22.383$	$34.853$	$0.772$	$16.720$	$10.275$

Table 2: Ablation study of the proposed distillation losses. Columns left of the rule mark the active terms;

\mathcal{L}_{\mathrm{GraphKit}}

is present in every row. Teacher is the frozen 20.1M-parameter GridFM and Baseline the 0.62M-parameter student trained without any transfer term. Values are RMSE scaled by

10^{3}

; lower is better and the best entry in each column is in bold. Fidelity columns are undefined for the teacher. All runs use a single seed.

Loss term	$\Delta V_{m}$	$\Delta V_{a}$	$\Delta P_{g}$	$\Delta Q_{g}$
Output KD	$\mathbf{+0.4272}$	$\mathbf{+7.8847}$	$\mathbf{+19.3056}$	$\mathbf{+12.2514}$
Physics KD	$-0.0999\pm 0.0606$	$-1.7491\pm 1.0676$	$-1.7155\pm 1.2020$	$-0.5585\pm 0.4717$
Latent	$-0.0327\pm 0.0583$	$-0.6464\pm 1.1335$	$-1.0418\pm 1.2703$	$+0.6609\pm 0.3840$
Attention	$+0.0078\pm 0.0351$	$-0.7584\pm 0.6116$	$-0.1895\pm 0.5023$	$+0.0487\pm 0.3360$

Table 3: Marginal impact of adding each distillation loss on the test set. Each entry is the mean

\pm

sample standard deviation of the change in RMSE over all configuration pairs differing only in that term, scaled by

10^{3}

; positive values indicate improved accuracy. The spread is over configurations, not over seeds, and all runs use a single seed. Output KD is estimated from the single baseline-to-

\mathcal{L}_{\mathrm{out}}

contrast and so has no spread.

Knowledge Distillation

Knowledge distillation transfers useful behavior or representations from larger or specialized models into more compact or adaptable learners.

See all 821 papers

Sep 24, 2026

GridSFM: A Foundation Model for Solving AC Optimal Power Flow

We introduce GridSFM, a framework that combines a pretrained foundation model across grid topologies with physics-informed fine-tuning for solving AC Optimal Power Flow (AC-OPF) at scale. It is a

15

million parameter physics-inspired graph neural network pretrained across

54

topologies of

500

4{,}000

buses. Our model attains a

2.45\%

zero-shot generation-cost error on a

10{,}000

bus case held-out operating conditions with no degradation as system size grows. Building on this, we pair the pretrained backbone with a physics-informed fine-tuning design based on Newton's method for power flow. With only

100

solved instances, GridSFM adapts to unseen grids up to

10{,}000

buses. We show it out performs single topology, dedicated neural network models that are trained more data, both in terms of cost and solver iterations when deployed as warm starting points. In designing this foundation model, we overcome the fact that the feasible set for AC-OPF can be disconnected. This is an obstruction that prevents any continuous neural network from approximating the solution map. To do so, we lift the problem and relax its constraints with logarithmically penalized slacks. We prove that the resulting elastic feasible set is contractible, that the AC-OPF minimizers remain minimizers of the elastic problem above an explicit penalty threshold, and that projecting an approximate solution back onto the AC-OPF feasible set is well posed. We release all models, data, and code so that the community can build on a shared starting point for AC-OPF.

Luke Bhan, Weiwei Yang, Margaret Capetz +1

University of California, San Diego · Microsoft Research · University of Washington

Neural Surrogate Modeling Optimal Power Flow Physics-Informed GNNs Power Systems

Ask PDF

Mar 27, 2026

Knowledge Distillation for Efficient Transformer-Based Reinforcement Learning in Hardware-Constrained Energy Management Systems

Transformer-based reinforcement learning has emerged as a strong candidate for sequential control in residential energy management. In particular, the Decision Transformer can learn effective battery dispatch policies from historical data, thereby increasing photovoltaic self-consumption and reducing electricity costs. However, transformer models are typically too computationally demanding for deployment on resource-constrained residential controllers, where memory and latency constraints are critical. This paper investigates knowledge distillation to transfer the decision-making behaviour of high-capacity Decision Transformer policies to compact models that are more suitable for embedded deployment. Using the Ausgrid dataset, we train teacher models in an offline sequence-based Decision Transformer framework on heterogeneous multi-building data. We then distil smaller student models by matching the teachers' actions, thereby preserving control quality while reducing model size. Across a broad set of teacher-student configurations, distillation largely preserves control performance and even yields small improvements of up to 1%, while reducing the parameter count by up to 96%, the inference memory by up to 90%, and the inference time by up to 63%. Beyond these compression effects, comparable cost improvements are also observed when distilling into a student model of identical architectural capacity. Overall, our results show that knowledge distillation makes Decision Transformer control more applicable for residential energy management on resource-limited hardware.

Pascal Henrich, Jonas Sievers, Maximilian Beichter +3

Karlsruhe Institute of Technology, Karlsruhe, Germany

Policy Distillation Transformer-Based RL Knowledge Distillation Power Systems

Ask PDF

Jul 15, 2026

MxGPS: Multiplex Graph Transformers for a Power Grid Foundation Model

Single-task fine-tuning of graph neural networks (GNNs) for power grid problems exhibits a systematic failure mode: models that achieve the lowest in-distribution error degrade the most under topology shift. We term this topology overfitting: the tendency of task-specific gradient signals to encode relational structure particular to the training topologies rather than the underlying physics, causing models to fail on unseen grids despite strong in-distribution performance. To expose and address this failure mode, we introduce MxGPS (Multiplex GPS), a multiplex graph transformer that runs K task-specialised GPS branches over a shared node encoder, jointly trained on Static State Estimation (SSE) and AC Power Flow (PF) via a self-supervised pre-training and multi-task fine-tuning protocol, with a cross-branch attention module evaluated in ablation. The joint SSE+PF objective forces the shared encoder to simultaneously satisfy complementary gradient signals, preventing it from overfitting to topology-specific relational structure. Under a 3-fold sliding-window cross-validation spanning four unseen topologies (14-, 24-, 162-, and 300-bus), MxGPS attains 0% boundary violation rate (BVR) on all four zero-shot Power Flow topologies. Critically, models with substantially lower in-distribution PF error degrade by 190% to 1400% under topology shift, whereas MxGPS degrades by only 39%, an inversion that directly implicates topology overfitting as the failure mechanism rather than insufficient model capacity. With only 1.6M parameters (12x fewer than the GridFM reference baseline), MxGPS demonstrates that multi-task joint training is a principled and parameter-efficient mechanism for topology-agnostic generalisation in power grid foundation models.

Charilaos Papaioannou, Ioannis Tsantilas, Dimitris Giannakakos +9

Multi-Task Learning Graph Foundation Models Graph Transformer Power Systems

Ask PDF

Abstract

Knowledge Distillation

Explore similar work

GridSFM: A Foundation Model for Solving AC Optimal Power Flow

Knowledge Distillation for Efficient Transformer-Based Reinforcement Learning in Hardware-Constrained Energy Management Systems

MxGPS: Multiplex Graph Transformers for a Power Grid Foundation Model