Scalability and Sample Efficiency Analysis of Graph Neural Networks for
  Power System State Estimation
Data-driven state estimation (SE) is becoming increasingly important in
modern power systems, as it allows for more efficient analysis of system
behaviour using real-time measurement data. This paper thoroughly evaluates a
phasor measurement unit-only state estimator based on graph neural networks
(GNNs) applied over factor graphs. To assess the sample efficiency of the GNN
model, we perform multiple training experiments on various training set sizes.
Additionally, to evaluate the scalability of the GNN model, we conduct
experiments on power systems of various sizes. Our results show that the
GNN-based state estimator exhibits high accuracy and efficient use of data.
Additionally, it demonstrated scalability in terms of both memory usage and
inference time, making it a promising solution for data-driven SE in modern
power systems.
Ognjen Kundacina, Gorana Gojic, Mirsad Cosovic, Dragisa Miskovic, Dejan Vukobratovic
Scalability and Sample Efficiency Analysis of Graph Neural Networks for Power System State Estimation

###### Abstract

Data-driven state estimation (SE) is becoming increasingly important in modern power systems, as it allows for more efficient analysis of system behaviour using real-time measurement data. This paper thoroughly evaluates a phasor measurement unit-only state estimator based on graph neural networks (GNNs) applied over factor graphs. To assess the sample efficiency of the GNN model, we perform multiple training experiments on various training set sizes. Additionally, to evaluate the scalability of the GNN model, we conduct experiments on power systems of various sizes. Our results show that the GNN-based state estimator exhibits high accuracy and efficient use of data. Additionally, it demonstrated scalability in terms of both memory usage and inference time, making it a promising solution for data-driven SE in modern power systems.

 State Estimation, Graph Neural Networks, Machine Learning, Power Systems, Real-Time Systems

## I Introduction

**Motivation and literature review:** The state estimation (SE) algorithm is a key component of the energy management system that provides an accurate and up-to-date representation of the current state of the power system. Its purpose is to estimate complex bus voltages using available measurements, power system parameters, and topology information [1]. In this sense, the SE can be seen as a problem of solving large, noisy, sparse, and generally nonlinear systems of equations. The measurement data used by the SE algorithm usually come from two sources: the supervisory control and data acquisition (SCADA) system and the wide area monitoring system (WAMS) system. The SCADA system provides low-resolution measurements that cannot capture system dynamics in real-time, while the WAMS system provides high-resolution data from phasor measurement units (PMUs) that enable real-time monitoring of the system. The SE problem that considers measurement data from both WAMS and SCADA systems is formulated in a nonlinear way and solved in a centralized manner using the Gauss-Newton method [1]. On the other hand, the SE problem that considers only PMU data provided by WAMS has a linear formulation, enabling faster, non-iterative solutions.

In this work, we will focus on the SE considering only phasor measurements, described with a system of linear equations [2], which is becoming viable with the increasing deployment of PMUs. This formulation is usually solved using linear weighted least-squares (WLS), which involve matrix factorizations and can be numerically sensitive [3]. To address the numerical instability issues that often arise when using traditional SE solvers, researchers have turned to data-driven deep learning approaches [4, 5]. These approaches, when trained on relevant datasets, are able to provide solutions even when traditional methods fail. For example, in [4], a combination of feed-forward and recurrent neural networks was used to predict network voltages using historical measurement data. In the nonlinear SE formulation, the study [5] demonstrates the use of deep neural networks as fast and quality initializers of the Gauss-Newton method.

Both linear WLS and common deep learning SE methods at its best approach quadratic computational complexity regarding the power system size. To fully utilize high sampling rates of PMUs, there is a motivation to develop SE algorithms with a linear computational complexity. One way of achieving this could be using increasingly popular graph neural networks (GNNs) [6, 7]. GNNs have several advantages when used in power systems, such as permutation invariance, the ability to handle varying power system topologies, and requiring fewer trainable parameters and less storage space compared to conventional deep learning methods. One of the key benefits of GNNs is the ability to perform distributed inference using only local neighbourhood measurements, which can be efficiently implemented using the emerging 5G network communication infrastructure and edge computing [8]. This allows for real-time and low-latency decision-making even in
large-scale networks, as the computations are performed at the edge of the network, closer to the data source, reducing the amount of data that needs to be transmitted over the network. This feature is particularly useful for utilizing the high sampling rates of PMUs, as it can reduce communication delays in PMU measurement delivery that occur in centralized SE implementations.

GNNs are being applied in a variety of prediction tasks in the field of power systems, including fault location [9], stability assessment [10], and load forecasting [11]. GNNs have also been used for power flow problems, both in a supervised [12] and an unsupervised [13] manner. A hybrid nonlinear SE approach [14] combines a model and data-based approach using a GNN that outputs voltages which are used a regularization term in the SE loss function.

**Contributions**: In our previous work [15], we proposed a data-driven linear PMU-only state estimator based on GNNs applied over factor graphs. The model demonstrated good approximation capabilities under normal operating conditions and performed well in unobservable and underdetermined scenarios. This work significantly extends our previous work in the following ways:

* We conduct an empirical analysis to investigate how the same GNN architecture could be used for power systems of various sizes. We assume that the local properties of the graphs in these systems are similar, leading to local neighbourhoods with similar structures which can be represented using the same embedding space size and the same number of GNN layers.
* To evaluate the sample efficiency of the GNN model, we run multiple training experiments on different sizes of training sets. Additionally, we assess the scalability of the model by training it on various power system sizes and evaluating its accuracy, training convergence properties, inference time, and memory requirements.
* As a side contribution, the proposed GNN model is tested in scenarios with high measurement variances, using which we simulate phasor misalignments due to communication delays, and the results are compared with linear WLS solutions of SE.

## II Linear State Estimation with PMUs

The SE algorithm has a goal of estimating the values of the state variables \(\mathbf{x}\), so that they are consistent with measurements, as well as the power system model defined by its topology and parameters. The power system's topology is represented by a graph \(\mathcal{G}=(\mathcal{N},\mathcal{E})\), where \(\mathcal{N}=1,\ldots,n\) is the set of buses and \(\mathcal{E}\subseteq\mathcal{N}\times\mathcal{N}\) is the set of branches.

PMUs measure complex bus voltages and complex branch currents, in the form of magnitude and phase angle [16, Sec. 5.6]. PMUs placed at a bus measure the bus voltage phasor and current phasors along all branches incident to the bus [17]. The state variables are given as \(\mathbf{x}\) in rectangular coordinates, and therefore consist of real and imaginary components of bus voltages. The PMU measurements are transformed from polar to rectangular coordinate system, since then the SE problem can be formulated using a system of linear equations [15]. The solution to this sparse and noisy system can be found by solving the linear WLS problem:

\[\left(\mathbf{H}^{T}\boldsymbol{\Sigma}^{-1}\mathbf{H}\right)\mathbf{x}= \mathbf{H}^{T}\boldsymbol{\Sigma}^{-1}\mathbf{z}, \tag{1}\]

where the Jacobian matrix \(\mathbf{H}\in\mathbb{R}^{m\times 2n}\) is defined according to the partial first-order derivatives of the measurement functions, and \(m\) is the total number of linear equations. The observation error covariance matrix is \(\boldsymbol{\Sigma}\in\mathbb{R}^{m\times m}\), while the vector \(\mathbf{z}\in\mathbb{R}^{m}\) contains measurement values in rectangular coordinate system. The aim of the WLS-based SE is to minimize the sum of residuals between the measurements and the corresponding values that are calculated using the measurement functions [1].

This approach has the disadvantage of requiring a transformation of measurement errors (magnitude and angle errors) from polar to rectangular coordinates, making them correlated, resulting in a non-diagonal covariance matrix \(\boldsymbol{\Sigma}\) and increased computational effort. To simplify the calculation, the non-diagonal elements of \(\boldsymbol{\Sigma}\) are often ignored, which can impact the accuracy of the SE [17]. We can use the classical theory of propagation of uncertainty to compute variances in rectangular coordinates from variances in polar coordinates [18]. The solution to (1) obtained by ignoring the non-diagonal elements of the covariance matrix \(\boldsymbol{\Sigma}\) to avoid its computationally demanding inversion is referred to as the _approximative WLS SE solution_.

In the rest of the paper, we will explore whether using a GNN model trained with measurement values, variances, and covariances labelled with the exact solutions of (1) leads to greater accuracy compared to the approximative WLS SE, which ignores covariances. The GNN model, once trained, scales linearly with respect to the number of power system buses, allowing for lower computation time compared to both the approximate and exact solvers of (1).

## III Methods

In this section, we introduce spatial GNNs on a high-level and describe how can they be applied to the linear SE problem.

### _Spatial Graph Neural Networks_

Spatial GNNs are a type of machine learning models that process graph-structured data by iteratively applying message passing to local subsets of the graph. The goal of GNNs is to transform the inputs from each node and its connections into a higher-dimensional space, creating a \(s\)-dimensional vector \(\mathbf{h}\in\mathbb{R}^{s}\) for each node. GNNs contain \(K\) layers, with each layer representing a single iteration \(k\) of the message passing process. Each GNN layer includes trainable functions, which are implemented as neural networks, such as a message function, an aggregation function, and an update function, as shown in Fig. 1.

The message function calculates the message \(\mathbf{m}_{i,j}\in\mathbb{R}^{u}\) between two node embeddings, the aggregation function combines the incoming messages in a specific way, resulting in an aggregated message \(\mathbf{m_{j}}\in\mathbb{R}^{u}\), and the update function calculates the update to each node's embedding. The message
passing process is repeated a fixed number of times, with the final node embeddings passed through additional neural network layers to generate predictions. GNNs are trained by optimizing their parameters using a variant of gradient descent, with the loss function being a measure of the distance between the ground-truth values and the predictions.

### _State Estimation using Graph Neural Networks_

The proposed GNN model is designed to be applied over a graph with a SE factor graph topology [19], which consists of factor and variable nodes with edges between them. The variable nodes are used to create a \(s\)-dimensional embedding for the real and imaginary parts of the bus voltages, which are used to generate state variable predictions. The factor nodes serve as inputs for measurement values, variances, and covariances. Factor nodes do not generate predictions, but they participate in the GNN message passing process to send input data to their neighbouring variable nodes. To improve the model's representation of a node's neighbourhood structure, we use binary index encoding as input features for variable nodes. This encoding allows the GNN to better capture relationships between nodes and reduces the number of input neurons and trainable parameters, as well as training and inference time, compared to the one-hot encoding used in [15]. The GNN model can be applied to various types and quantities of measurements on both power system buses and branches, and the addition or removal of measurements can be simulated by adding or removing factor nodes. In contrast, applying a GNN to the bus-branch power system model would require assigning a single input vector to each bus, which can cause problems such as having to fill elements with zeros when not all measurements are available and making the output sensitive to the order of measurements in the input vector.

Connecting the variable nodes in the \(2\)-hop neighbourhood of the factor graph topology significantly improves the model's prediction quality in unobservable scenarios [15]. This is because the graph remains connected even when simulating the removal of factor nodes (e.g., measurement loss), which allows messages to be propagated in the entire \(K\)-hop neighbourhood of the variable node. This allows for the physical connection between power system buses to be preserved when a factor node corresponding to a branch current measurement is removed.

The proposed GNN for a heterogeneous graph has two types of layers: one for factor nodes and one for variable nodes. These layers, denoted as \(\mathrm{Layer^{f}}\) and \(\mathrm{Layer^{v}}\), have their own sets of trainable parameters, which allow them to learn their message, aggregation, and update functions separately. Different sets of trainable parameters are used for variable-to-variable and factor-to-variable node messages. Both GNN layers use two-layer feed-forward neural networks as message functions, single layer neural networks as update functions, and the attention mechanism [7] in the aggregation function. Then, a two-layer neural network \(\mathrm{Pred}\) is applied to the final node embeddings \(\mathbf{h}^{K}\) of variable nodes only, to create state variable predictions. The loss function is the mean-squared error (MSE) between the predictions and the ground-truth values, calculated using variable nodes only. All trainable parameters are updated via gradient descent and backpropagation over a mini-batch of graphs. The high-level computational graph of the GNN architecture specialized for heterogeneous augmented factor graphs is depicted in Figure 2.

The proposed model uses an inference process that requires measurements from the \(K\)-hop neighbourhood of each node, allowing for computational and geographical distribution. Additionally, since the node degree in the SE factor graph is limited, the computational complexity for the inference process is constant. As a result, the overall GNN-based SE has a linear computational complexity, making it efficient and scalable for large networks.

## IV Numerical Results

In this section, we conduct numerical experiments to investigate the scalability and sample efficiency of the proposed GNN approach. By varying the power system and training set sizes, we are able to assess the model's memory requirements, prediction speed, and accuracy and compare them to those of traditional SE approaches.

Fig. 1: A GNN layer, which represents a single message passing iteration, includes multiple trainable functions, depicted as yellow rectangles. The number of first-order neighbours of the node \(j\) is denoted as \(n_{j}\).

Fig. 2: Proposed GNN architecture for heterogeneous augmented factor graphs. Variable nodes are represented by circles and factor nodes are represented by squares. The high-level computational graph begins with the loss function for a variable node, and the layers that aggregate into different types of nodes have distinct trainable parameters.
We use the IEEE 30-bus system, the IEEE 118-bus system, the IEEE 300-bus system, and the ACTIVSg 2000-bus system [20], with measurements placed so that measurement redundancy is maximal. For the purpose of sample efficiency analysis, we create training sets containing 10, 100, 1000, and 10000 samples for each of the mentioned power systems. Furthermore, we use validation and test sets comprising 100 samples. These datasets are generated by solving the power flow problem using randomly generated bus power injections and adding Gaussian noise to obtain the measurement values. All the data samples were labelled using the traditional SE solver. An instance of the GNN model is trained on each of these datasets.

In contrast to our previous work, we use higher variance values of \(5\times 10^{-1}\) to examine the performance of the GNN algorithm under conditions where input measurement phasors are unsynchronized due to communication delays [21]. While this is usually simulated by using variance values that increase over time, as an extreme scenario we fix the measurement variances to a high value.

In all the experiments, the node embedding size is set to \(64\), and the learning rate is \(4\times 10^{-4}\). The minibatch size is \(32\), and the number of GNN layers is \(4\). We use the ReLU activation function and a gradient clipping value of \(5\times 10^{-1}\). The optimizer is Adam, and we use mean batch normalization.

### _Properties of Power System Augmented Factor Graphs_

For all four test power systems, we create augmented factor graphs using the methodology described in Section III-B. Fig. 3 illustrates how the properties of the augmented factor graphs, such as average node degree, average path length, average clustering coefficient, along with the system's maximal measurement redundancy, vary across different test power systems.

The average path length is a property that characterizes the global graph structure, and it tends to increase as the size of the system grows. However, as a design property of high-voltage networks, the other graph properties such as the average node degree, average clustering coefficient, as well as maximal measurement redundancy do not exhibit a clear trend of change with respect to the size of the power system. This suggests that the structures of local, \(K\)-hop neighbourhoods within the graph are similar across different power systems, and that they contain similar factor-to-variable node ratio. Consequently, it is reasonable to use the same GNN architecture (most importantly, the number of GNN layers and the node embedding size) for all test power systems, regardless of their size. In this way, the proposed model achieves scalability, as it applies the same set of operations to the local, \(K\)-hop neighbourhoods of augmented factor graphs of varying sizes without having to adapt to each individual case.

### _Training Convergence Analysis_

First, we analyse the training process for the IEEE 30-bus system with four different sizes of the training set. As mentioned in III-B, the training loss is a measure of the error between the predictions and the ground-truth values for data samples used in the training process. The validation loss, on the other hand, is a measure of the error between the predictions and the ground-truth values on a separate validation set. In this analysis, we used a validation set of 100 samples.

The training losses for all the training processes converged smoothly, so we do not plot them for the sake of clarity. Figure 4 shows the validation losses for 150 epochs of training on four different training sets. For smaller training sets, the validation loss decreases initially but then begins to increase, which is a sign of overfitting. In these cases, a common practice in machine learning is to select the model with the lowest validation loss value. As it will be shown in IV-C, the separate test set results for models created using small training sets are still satisfactory. As the number of samples in the training set increases, the training process becomes more stable. This is because the model has more data to learn from and is therefore less prone to overfitting.

Next, in Table I, we present the training results for the other power systems and training sets of various sizes. The numbers in the table represent the number of epochs after which either the validation loss stopped changing or began to increase. Similarly to the experiments on the IEEE 30

Fig. 4: Validation losses for trainings on four different training set sizes.

Fig. 3: Properties of augmented factor graphs along with the system’s measurement redundancy for different test power systems, labelled with their corresponding number of buses.
bus system, the trainings on smaller training sets exhibited overfitting, while others converged smoothly. For the former, the number in the table indicates the epoch at which the validation loss reached its minimum and stopped improving. For the latter, the number in the table represents the epoch when there were five consecutive validation loss changes less than \(10^{-5}\).

Increasing the size of the training set generally results in a lower number of epochs until the validation loss reaches its minimum. However, the epochs until the validation loss reaches its minimum vary significantly between the different power systems. This could be due to differences in the complexity of the systems or the quality of the data used for training.

### _Accuracy Assessment_

Fig. 5 reports the mean squared errors (MSEs) between the predictions and the ground-truth values on 100-sample sized test sets for all trained models and the approximate WLS SE. These results indicate that even the GNN models trained on small datasets outperform the approximate WLS SE, except for the models trained on the IEEE 30-bus system with 10 and 100 samples. These results suggest that the quality of the GNN model's predictions and the generalization capabilities improve as the amount of training data increases, and the models with the best results (highlighted in bold) have significantly smaller MSEs compared to the approximate WLS SE. While we use randomly generated training sets in this analysis, using carefully selected training samples based on historical load consumption data could potentially lead to even better results with small datasets.

### _Inference Time and Memory Requirements_

The plot in Fig. 6 shows the ratio of execution times between WLS SE and GNN SE inference as a function of the number of buses in the system. These times are measured on a test set of 100 samples. As expected, the difference in computational complexity between GNN, with its linear complexity, and WLS, with more than quadratic complexity, becomes apparent as the number of buses increases. From the results, it can be observed that GNN significantly outperforms WLS in terms of inference time on larger power systems.

The number of trainable parameters in the GNN model remains relatively constant, as the number of power system buses increases. The number of input neurons for variable node binary index encoding does grow logarithmically with the number of variable nodes. However, this increase is relatively small compared to the total number of GNN parameters1. This

Fig. 5: Test set results for various power systems and training set sizes.

\begin{table}
\begin{tabular}{|c|c|c|c|} \hline
**Power system** & IEEE 118 & IEEE 300 & ACTIVSg 2000 \\ \hline
**10 samples** & \(61\) & \(400\) & \(166\) \\ \hline
**100 samples** & \(38\) & \(84\) & \(200\) \\ \hline
**1000 samples** & \(24\) & \(82\) & \(49\) \\ \hline
**10000 samples** & \(12\) & \(30\) & \(15\) \\ \hline \end{tabular}
\end{table} TABLE I: Epoch until validation loss minimum for various power systems and training set sizes.
indicates that the GNN approach is scalable and efficient, as the model's complexity does not significantly increase with the size of the power system being analysed.

## V Conclusions

In this study, we focused on thoroughly testing a GNN-based state estimation algorithm in scenarios with large variances, and examining its scalability and sample efficiency. The results showed that the proposed approach provides good results for large power systems, with lower prediction errors compared to the approximative SE. The GNN model used in this approach is also fast and maintains constant memory usage, regardless of the size of the scheme. Additionally, the GNN was found to be an effective approximation method for WLS SE even with a relatively small number of training samples, particularly for larger power systems, indicating its sample efficiency. Given these characteristics, the approach is worthy of further consideration for real-world applications.
