Uncertainty-Aware Robust Learning on Noisy Graphs

The field of graph learning has witnessed significant advancements in recent years, fueled by the remarkable performance of graph neural networks (GNNs) (gcn, 1, 2). GNNs leverage message-passing techniques to enable efficient information exchange among nodes in a graph, leading to improved embeddings. Among the various graph-based learning tasks, node classification, especially in a semi-supervised setting, stands out as a prominent and widely applicable problem that has greatly benefited from GNNs. The objective of node classification is to learn high-quality node embeddings and make predictions about their labels using a graph that has only a subset of nodes labeled.

However, similar to other deep neural networks, GNNs are susceptible to the influence of noisy input, including both noisy features and graph structure (li2022recent, 3, 4, 5). This issue becomes even more critical in the low-data setting, where the performance can be greatly impaired by noisy observations when only a limited number of labeled nodes are available (li2022recent, 3, 4, 6). For example, in a social network, new users may inconsistently express their preferences or engage with content that aligns with their actual interests due to limited options available at that time. Such noises would introduce uncertainty and incorrect information about the user’s preferences into the training data. With only a few labeled samples available, the GNN heavily relies on these noisy observations to learn the user’s preferences and generalize them to make less accurate recommendations. Hence, the GNN’s capacity to capture the underlying patterns of different nodes can be significantly compromised when the input graph is noisy.

In this paper, we propose a novel uncertainty-aware graph learning framework based on distributionally robust optimization (DRO) (duchi2021learning, 7, 8, 9, 10). The proposed Distributionally Robust Graph Learning (DRGL) framework significantly improves node embedding quality and enhances model performance, particularly in the presence of noisy data that hinders the direct observation of key graph features. To tackle this challenge, we first allow the underlying node feature patterns to vary within a pre-defined family of distributions and seek the least favorable distribution (LFD) in the worst-case scenario. It is important to note that the uncertainty arising from data noise is represented by an uncertainty set (NEURIPS2018_a08e32d2, 9, 11, 12), characterized by a Wasserstein ball (NEURIPS2018_a08e32d2, 9, 11, 13, 14). Then, by minimizing the risk associated with this worst-case distribution, we can discover the most robust node embeddings that effectively mitigate the impact of uncertainty. To seamlessly integrate this minimax formulation into graph learning, we also leverage differentiable optimization techniques (cvx, 15, 16) and develop a tractable end-to-end learning algorithm. Furthermore, the resulting minimax solution also provides a means to estimate the potential uncertainty associated with predictions, allowing decision-makers to assess the likelihood of different outcomes. This information is particularly valuable in high-stakes scenarios where incorrect decisions can have severe consequences. Lastly, our numerical results demonstrate that the proposed learning strategy significantly enhances the quality of node embeddings and achieves promising predictive performance on semi-supervised node classification when compared to state-of-the-art methods.

Our approach handles noisy graphs based on the DRO (duchi2021learning, 7, 8), which aims to enhance the robustness of GNNs by minimizing the worst-case risk in the presence of data noise. Only a handful of recent research attempts, such as (zhang2021robust, 11, 27), have embraced similar concepts and investigated the potential applications of DRO in the context of graph-structured data. For instance, Sadeghi et al. (sadeghi2021distributionally, 27) use GNNs to encode the input graph and employ DRO for model training. However, their work is mainly focused on the noisy node features, while we consider a more general scenario involving both noisy features and edges. Furthermore, their method solves for worst-case distribution and updates model parameters separately at each iteration (sadeghi2021distributionally, 27), while our approach offers an end-to-end approach updating parameters with gradient-based learning techniques. Additionally, our approach provides uncertainty quantification under Least Favorable Distributions (LFDs). This capability has meaningful applications in uncertainty-sensitive scenarios (ABDAR2021243, 28, 29), such as molecular classification tasks in graphs (ryu2019bayesian, 30, 31).

It is important to note that these node embeddings obtained through minimizing standard loss functions (e.g., sum of squared errors) may not accurately reflect the true feature patterns of the graph due to the presence of noise in the observed data.

Our approach assumes that node embeddings share the same label $m$ adhere to an underlying distribution ($\xi_{i}\sim P_{m}\in\mathcal{P}_{m},\forall i:y_{i}=m$), within an uncertainty set $\mathcal{P}_{m}$ encompassing all potential distributions $P_{m}$. However, obtaining a direct observation of this distribution is challenging due to inaccuracies in the nodal or topological information, and any changes in node embeddings will influence their corresponding distributions. The key idea of our proposed framework is to find the most robust node embeddings represented by $\theta$ that minimize the worst-case risk over the uncertainty set $\mathcal{P}_{m}$ in the probability simplex $\Delta_{M}=\{\pi\in\mathbb{R}_{+}^{M}:\sum_{m=1}^{M}\pi_{m}=1\}$. This results in the definition of our distributionally robust minimax problem:

where $\Psi$ is the risk function, and we define the risk of a classifier $\pi$ as the summation of error probabilities under each class, i.e., $\Psi\left(\pi;P_{1},\ldots,P_{M}\right):=\sum_{m=1}^{M}\mathbb{E}_{\xi\sim P_{m}}\left[1-\pi_{m}(\xi)\right]$. We note that the optimal solution $P_{1}^{*},\dots,P_{M}^{*}$ to the inner maximization problem is known as the least favorable distributions (LFDs) in statistics literature. The risk associated with these distributions is considered the worst-case risk.

As shown in Figure 2, we choose the uncertainty set $\mathcal{P}_{m}$ to be a Wasserstein ball of radius $\vartheta_{m}$ centered at the empirical distribution $\widehat{P}_{m}$:

where $\mathscr{P}(\Xi)$ denotes the set of all probability distributions on $\Xi$. The $\mathcal{W}_{1}$ is defined as $\mathcal{W}_{1}\left(P,P^{\prime}\right):=\min_{\gamma}\mathbb{E}_{\left(\xi,\xi^{\prime}\right)\sim\gamma}\left[c\left(\xi,\xi^{\prime}\right)\right]$, where $c(u,v)$ is some cost function transferring from $u$ to $v$, $c(u,v)\geq 0$. The empirical distribution $\widehat{P}_{m}$ is represented by the Dirac point mass, denoted as:

Here, $\delta$ refers to the Dirac delta function, $|\cdot|$ represents the cardinality of a set, and $\mathbbm{1}$ denotes the indicator function.

As the original problem (1) entails an intractable infinite dimensional functional optimization, we follow (NEURIPS2018_a08e32d2, 9, 10) and introduce the proposition below. This proposition reformulates the problem into a computationally tractable convex optimization problem due to our careful selection of the risk function and uncertainty sets.

We apply the mini-batch stochastic gradient descent as summarized in Algorithm 1. It is necessary that each mini-batch, which is provided to the convex solver, includes at least one training sample from every class. This requirement is put in place to maintain the integrity of the optimization and its ability to generate valid solutions at all times.

To assess the robustness of our method, we introduce random noise into these data sets in the following two ways: (1) We add Gaussian noise $\epsilon\sim\mathcal{N}(0,\sigma)$ to the node feature matrix $X$, where $\sigma$ is set proportionally to the standard deviation of the bag-of-words representation of all nodes in each graph; (2) We randomly remove a certain percentage $r$ of edges in the graph. We repeat each test three times and calculate the average accuracy.

Since our framework is designed to be compatible with various models, we choose to experiment with Graph Convolutional Networks (GCNs) (gcn, 1) and Graph Attention Networks (GATs) (gat, 2) as our graph encoders due to their popularity as competitive baselines. We also evaluate our framework using two commonly used classifiers (output layers), including a shallow (2-layer) neural network with Softmax output and a weighted $k$-NN. To fit the models, we use the Adam optimizer with a learning rate of $10^{-4}$ for all experiments. A more detailed experimental setup can be found in Appendix A.

Next we assess the performance of DRGL in the missing edge setting with $K=5$. Table 2 shows that both GCN and GAT equipped with DRGL exhibit substantial improvement, particularly in scenarios where a higher percentage of edges are removed. This showcases the efficacy of DRGL in enhancing the resilience of GNNs when facing challenges associated with missing edges.

We also conduct experiments in the standard few-shot setting using GCN (xu2021dp, 37, 38), where no noise is added to the data. In this setting, each node class is provided with a limited number of labeled samples, along with the original features and edges. The results presented in Table 3 demonstrate that using DRGL in conjunction with GCN yields the highest average test accuracy across all data sets when compared to baseline methods. This indicates that while the DRGL is primarily designed for noisy graphs, it also enhances model performance in standard few-shot scenarios. Additional experimental results can be found in Appendix B.

Figure 3 presents a comparison between the learned node embeddings and those generated by the vanilla GCN in various scenarios. We can observe that the between-class distance demonstrates a slight increase in contrast to the vanilla GCN. Conversely, the intra-class distance shows a minor reduction, as denoted by the denser distribution of dots with matching colors. It is worth mentioning that such differences in distribution can become more prominent in higher dimensional embedding spaces. Although seemingly subtle, this shift in the distribution of embeddings plays a pivotal role in significantly improving the accuracy of the classification outcomes, indicated by the numbers displayed in the lower right corner.

Figure 4 showcases the visual representation of the LFDs generated by DRGL. The shades of grey in the visualization are obtained using kernel density estimation, where darker shades indicate a higher level of uncertainty between the different categories. These LFDs serve as a valuable tool for uncertainty quantification (UQ), enabling us to pinpoint the data points that are susceptible to the greatest impact in worst-case scenarios. Let $\tilde{p}^{*}=[\tilde{p}_{1}^{*},\dots,\tilde{p}_{M}^{*}]$ denote the predicted probabilities of a classifier built on learned LFDs of a graph. The predictive uncertainty can be expressed using entropy $H(\tilde{p}^{*})=-\sum_{m=1}^{M}\tilde{p}_{m}^{*}\log\tilde{p}_{m}^{*}$ (6769090, 39, 40, 41). The result highlights the improved capability of our approach in capturing and quantifying uncertainty compared with the vanilla GCN method.

To address the challenges posed by noisy graphs, we proposed a novel GNN-agnostic framework that enhances the robustness of node embeddings and predictive performance. Particularly, our proposed framework improves the model robustness by accounting for the uncertainties raised by the data noise within the graph, which shows substantial improvements over state-of-the-art baselines on different benchmark datasets.