A General Framework for quantifying Aleatoric and Epistemic uncertainty in Graph Neural Networks

In this paper, we address the lack of systematic and explicit incorporation of different sources of uncertainty in GNNs within a Bayesian framework. We formally define the different sources of aleatoric and epistemic uncertainty in GNNs. Specifically, we consider the aleatoric uncertainty arising from (i) imprecise information about the graph structure via probabilistic links and (ii) measurement noise in feature vectors of nodes. We propagate the aleatoric uncertainty through the node embedding layers and classification layers of the GNN model via Assumed Density Filtering (ADF). We quantify the epistemic uncertainty arising from probabilistic parameters of the GNN model with Monte-Carlo sampling. The proposed framework exhibits the following advantages:

• 1.

  Quantification of total uncertainty due to aleatoric and epistemic sources

• 2.

  Propagation of aleaotric uncertainty through all layers of a GNN model

• 3.

  Generic framework that can be applied at the time of fresh training, as well as to pre-trained networks without the need of a computationally expensive retraining process.

We present the different sources of uncertainty in GNNs and discuss the related literature on handling uncertainty in Section 2. We discuss the specific problem definition targeted in this work and the proposed approach to quantify the total uncertainty in GNNs in Section 3. Experiments and results are presented in Section 4 and the article ends with concluding remarks in Section 5.

Aleatoric uncertainty refers to intrinsic randomness of the data due to noisy or inaccurate measurements. In the case of a GNN, the input data is in the form of graphs that model a real-world network. This data consists of the feature vectors $h_{i}$ of the nodes and probabilities/weights $p_{ij}$ of the links. The different sources of aleatoric uncertainty in GNN are as follows:

1. AU1

   Measurement uncertainty associated with feature vectors of nodes $h_{i}$, i.e., the measured feature vectors are considered as being the sum of true feature vectors ($h_{i}^{*}$) and measurement noise ($\epsilon_{i}$) as: $h_{i}=x_{i}+\epsilon_{i}$, with $\epsilon_{i}\sim p(\epsilon)$.

2. AU2

   Structural uncertainty of the graph captured via probabilities of links $p_{ij}$, $i=1,2,3,...,n_{V}$, $j=1,2,3,...,n_{V}$.

Uncertainty AU1 refers to the uncertainty in features of the nodes. The feature vectors can represent physical quantities such as the coverage of a cell tower in a communication network, or the functionalities of a protein in a biological network. For instance, every node (research article) in the citation dataset “Cora” is characterized by a feature vector of size $1433$, where each unit indicates the absence/presence of predefined words from the dictionary. These measured quantities can be uncertain in many scenarios because of the sensing process or imprecise information about the system. For example, in a protein network, the exact functionality of all proteins is rarely accurately known. This type of uncertainty propagates through the layers of the GNN and ultimately affects the model predictions.

Uncertainty AU2, on the other hand, represents the uncertainty in the existence of links. In the protein network scenario, the knowledge of interactions between different proteins and protein complexes is also highly uncertain, which results in probabilistic links between different nodes. Since the interactions are derived through noisy and error-prone lab experiments, each link is associated with an uncertainty value [22]. For instance, a graph with $5$ nodes and binary weights of links has $2^{5}=32$ possible configurations. In a graph with continuous valued weights, these weights influence the extent to which information between any two nodes is exchanged and assimilated. This type of uncertainty therefore results in fundamentally different information exchange and processing through the graph.

A variety of Bayesian methods are used in the literature to handle aleatoric uncertainty in deep neural networks [23, 24, 25, 26]. However, there are very few works for GNN models. The authors in Zhang et al. [20] propose a Bayesian framework using joint estimation of graph structure and GNN parameters. The authors make use of families of parametric random graphs to estimate the structure and parameters. This makes the approach sensitive to the choice of the random graph model and the extent to which the random network can accurately represent the characteristics of the true underlying network. As a result, inferences can be inconsistent for different problems and datasets. Another significant drawback of the technique is that the posterior inference of the graph is carried out solely conditioned on the observed graph. As a result, any information provided by the node features and the training labels is completely disregarded. Therefore, Pal et al. [21] proposed an alternative approach which formulates the posterior inference of the graph in a non-parametric fashion, conditioned on the observed graph, features and training labels. Precisely, they obtain MAP estimate of graph, and conducted all the classification/regression tasks on this estimate. It is argued that MAP estimate handles aleatoric uncertainty of the input graph. However, the approach does not systematically define/quantify the sources of uncertainty and their impact on the predictions. Specifically, the uncertainties AU1 and AU2 are not considered in the framework. We address these shortcomings by explicitly incorporating AU1 and AU2 in our framework. Specifically, ADF is leveraged to propagate the aleatoric uncertainty from the input of the GNN to final node predictions through all the intermediate layers.

Epistemic uncertainty is the scientific uncertainty in the model that exists because of model in-competency to completely explain the underlying process. A GNN model $\Xi(\Theta)$ typically consists of several layers of complex aggregation and combination operations followed by feedforward processing. The different forms of epistemic uncertainty in the context of GNNs are:

1. EU1

   Parametric uncertainty in the GNN model, i.e., the parameters $\Theta$ of the GNN are assumed to be probabilistic with a probability density function $p(\Theta)$

2. EU2

   Uncertainty in activation functions of neurons of the GNN model

Uncertainty EU1 represents the uncertainty of the learnable parameters $\Theta$ of the GNN model, and is represented by by placing a distribution over the neural network weights. However, estimation of the posterior density function of the parameters $p(\Theta|D)$ given the data $D$ is mathematically intractable to compute for deep neural networks and is approximated by different methods. Among these methods, variational inference [27], and sampling-based approaches [28] are the most effective ones. Monte-Carlo sampling methods involve sampling the parameters from a distribution and are generally obtained using an ensemble of neural network predictions. The prediction ensemble could either be generated by differently trained networks [29], or by using dropout at test-time [28].

Similar to aleatoric uncertainty, the literature on handling epistemic uncertainty in GNN is limited. Zhang et al.[20] and Pal et al. [21] are some of the few efforts that generated multiple Monte-Carlo samples by using dropout at test time. To address the problem of huge dependency of Zhang et al.[20] on the assumed random graph model, Hasan et al. [30] introduces a stochastic regularization technique for GNN by adaptive connection sampling. Specifically, it adaptively learns the dropout rate for each layer of GNN. Akin to many of the aforementioned methods, we use the dropout-at-test approach to generate Monte-Carlo samples and estimate the epistemic uncertainty in predictions.

Uncertainty EU2 represents the randomness in the activation function of neurons in the GNN, and is not well considered in the literature. We also restrict our analysis to AU1, AU2 and EU1 in GNN and construct the total uncertainty from the contributions of these components.

We consider a network $G=(V,E)$ as described in Section 2 where nodes are associated with noisy feature vectors (AU1) and links are associated with probabilities (AU2) as follows:

1. 1.

   The noise in the features is considered to be Gaussian with zero mean and known variance. For a node $u$ in the network, we have:

   $$\begin{split}h_{u}=h_{u}^{*}+\epsilon_{u}\\ \epsilon_{u}\sim\mathcal{N}(0,\Sigma_{u})\end{split}$$

   (1)

   where $h$ and $h^{*}$ are the measured and true feature vector respectively, $\epsilon_{u}$ is the noise in feature vector and $\Sigma_{u}$ is a diagonal matrix consisting of the known variances of noise in individual features. The noise in the features of different nodes are assumed to be uncorrelated, i.e., for any two nodes $u$ and $\nu$ in the network, we have:

   $$\mathbb{E}[\epsilon_{u}\epsilon_{\nu}]=0$$

2. 2.

   The probabilities $p_{u\nu}$ of links are assumed to be known a priori.

3. 3.

   The learnable parameters $\Theta$ of the GNN are assumed to be random variables with an unknown probability density function:

   $$\Theta\sim p_{\Theta}(\psi)$$

These sources of uncertainty result in a probabilistic propagation of the feature vectors through the model, and thus result in probabilistic outputs, i.e., $\hat{y}\sim p(y|h,p,\Theta)$. Obtaining the exact distribution of $\hat{y}$ is mathematically intractable. The problem of uncertainty quantification considered here is to systematically obtain the variance in the predictions because of the different sources of uncertainty. In the following subsections, we discuss how these effects are quantified in the proposed framework.

We propose a Bayesian approach to propagate the uncertainty in feature vectors through a GNN while also explicitly incorporating probabilistic links in the system. We achieve this with Assumed Density Filtering (ADF) and moment matching. We consider propagating and matching the mean and variance of the probability density function of outputs of all node embedding and feedforward layers of the model. While this can be achieved with existing approaches for feedforward layers, the following result formalises the result for node embedding layers.

Equations (11) and (12) can be determined analytically for most of the functions used in neural network such as ReLu, sigmoid, convolution, etc. For instance if the function $g$ is ReLu, then the modified mean and variance are [35]:

where, $\sigma_{u}^{(i)}=\sqrt{\biggl{(}v_{u}^{(i)}\biggr{)}}$, $\Phi$ and $\phi$ are the cumulative normal and standard normal distributions, respectively. Basically, Eqs. (13) and (14) are recursive formulae to compute mean $\left(\mu_{u}^{(i)}\right)$ and uncertainty $\left(v_{u}^{(i)}\right)$ of the embeddings, given the parameters of the embedding distribution $q\left(h^{(i-1)}\right)$ in previous layer.

Typically, node embeddings from GNN are fed to feed-forward layers for classification/regression task. Therefore, $\mu_{u}^{(l_{g})}$, $v_{u}^{(l_{g})}$ serve as an input to feed-forward layers, and mean and variance is propagated in a similar way as shown in Gast & Roth [33], Loquercio et al. [14]. In a nutshell, ADF reshapes the forward pass of a GNN to generate not only output predictions $\mu_{u}^{(l)}$, but also their respective aleatoric uncertainties $v_{u}^{(l)}$. This is achieved by considering two values per dimension of both embeddings in GNN layers as well as neural units in feed-forward layers.

Epistemic uncertainty, also known as model uncertainty refers to the model confidence on its prediction. This uncertainty arises because of the single adoption of weights out of many combinations that can attain same loss values on training data. This is usually captured by assuming a probability distribution for neural network weights rather than a scalar value. However, computation of this distribution $p(\omega|X,y)$ is usually intractable. Therefore, MC based approaches have been used to obtain different weight samples by using dropout at test time [36, 28, 37]. Specifically, in our case, epistemic uncertainty is the variance of $M$ MC samples obtained via different dropout masks as shown below:

where, $\{y_{t}\}_{t}^{M}$ is a set of M sampled outputs for different weight instances from the distribution $\omega^{t}\sim q(\omega,\phi)$ and $\hat{y}=\frac{1}{T}\sum_{t}y_{t}$. Authors in Gal et al.[28] have shown that the optimal dropout rate $\phi$ for the computation of $\sigma_{\text{model}}$ is same as training dropout rate.

Total variance of GNN predictions $y$ for a sample node with feature vector $X$ corrupted by noise variance $v^{0}$ can be written as:

The first term $\left(v_{t}^{(L)}\right)$ in denotes aleatoric variance and the second term $\left(\mu_{t}^{(L)}-\hat{\mu}\right)^{2}$ represents the model uncertainty from M MC predictions. $L=l_{g}+l_{f}$ is the total number of layers in GNN. Thus, the first part of total variance captures ensembles of propagated variance and the second part handles the ensembles of mean prediction, thereby, addressing both aleatoric and epistemic uncertainty. The overall algorithm to compute total uncertainty can be summarized in following steps: (i) Transform GNN into a bayesian network by associating mean and variance to each embedding vector and neuron unit; (ii) Obtain $M$ mean and variance predictions by forwarding $\left(X,v^{0}\right)$ to network with weights $\omega^{t}$ sampled from $q(\omega,\phi)$; (iii) Compute output predictions and its variances according to Eq. (16).

We compare the performance of the proposed method with Bayesian Graph Convolutional Network (BGCN) proposed by [21]. This work captures the aleatoric uncertainty through MAP estimation of the network structure and quantifies the epistemic uncertainty with Monte-Carlo sampling. However, uncertainties AU1 and AU2 are not explicitly incorporated. Moreover, the MAP estimation is dependent on accurate knowledge of node features, and hence the method is incapable of handling AU1. Finally, the approach is computationally expensive because of the MAP estimation step. These drawbacks are addressed with our method through explicit incorporation of uncertainty with the added advantage of reduced computational complexity.

All training and evaluation experiments are performed on a computer with Intel i9-4820K processor running at 3.70GHz with 8 cores, one Nvidia RTX 2080 Ti GPU with 12 GB memory, and 64 GB RAM. The implementation is performed with the help of deep graph library in PyTorch. Evaluation is repeated $100$ times and average of metrics are reported for each dataset described in Section 4.4.

GNN is trained with the GraphSAGE algorithm [31]. The detailed architecture of the GNN is as follows: Depth i.e., no. of node embedding modules: 2; no. of neurons in $2$ layers: 64, 32; no. of Multi-layer perceptron (MLP) layers: 3; no. of neurons in MLP layers: 12,8,1; Activation function: Linear (except last layer with softmax); Aggregation function: Mean.

The data is divided into training, validation and testing subsets by randomly masking $70\%$, $10\%$ and $20\%$ of the nodes respectively. Since the core task of the GNN is node classification, the node embeddings are concatenated with feedforward layers to provide class probabilities. The loss function is categorical cross entropy with ADAM optimizer.

The training is carried out in a mini-batch manner. The batch size is set to 50, the learning rate is set to 0.001 with a dropout rate of 0.1. The models are trained 50 epochs in total.

We introduce uncertainty in nodes feature (AU1) by adding Gaussian white noise to the true feature values with zero mean and a known variance as shown in Eq. (1). The proposed method is compared with the baseline with different levels of noise. The variance of noise is also varied to highlight the ability of the proposed method to capture the impact of this noise as it propagates through the model. We introduce uncertainty in links (AU2) with probability of nodes. These probabilities are not available in the datasets for all links in the networks. We therefore perform link prediction in a supervised manner following the approach presented by Zhang & Chen [40] and obtain the probabilities of links. These predicted probabilities are then used for training Bayesian models in the proposed method. We introduce uncertainty in parameters of GNN (EU1) with a Bernoulli distribution of parameters according to Eq. (15).

The results demonstrate the computational efficiency of the proposed approach, the ability to reflect different levels of uncertainty in predictions and generalizability of the approach. The source code to regenerate all the results can be accessed at this link.