Energy-based Epistemic Uncertainty for Graph Neural Networks

TL;DR

The paper tackles epistemic uncertainty in graph neural networks by introducing GEBM, an energy-based model that aggregates energy signals at independent, local, and group graph scales via diffusion. A Gaussian regularizer on logits ensures the induced data density is integrable, and the framework admits an evidential interpretation that yields robust predictions under distribution shifts. Empirically, GEBM achieves state-of-the-art out-of-distribution detection and the best average performance across multiple distribution shifts while preserving backbone accuracy. The work also clarifies theoretical connections between EBMs and evidential learning, and discusses limitations and broader impact in the graph learning setting.

Abstract

In domains with interdependent data, such as graphs, quantifying the epistemic uncertainty of a Graph Neural Network (GNN) is challenging as uncertainty can arise at different structural scales. Existing techniques neglect this issue or only distinguish between structure-aware and structure-agnostic uncertainty without combining them into a single measure. We propose GEBM, an energy-based model (EBM) that provides high-quality uncertainty estimates by aggregating energy at different structural levels that naturally arise from graph diffusion. In contrast to logit-based EBMs, we provably induce an integrable density in the data space by regularizing the energy function. We introduce an evidential interpretation of our EBM that significantly improves the predictive robustness of the GNN. Our framework is a simple and effective post hoc method applicable to any pre-trained GNN that is sensitive to various distribution shifts. It consistently achieves the best separation of in-distribution and out-of-distribution data on 6 out of 7 anomaly types while having the best average rank over shifts on \emph{all} datasets.

Figures (8)

• Figure 1: Overview of the Graph Energy-based Model (GEBM). Graph-agnostic energy (uncertainty) of a trained GNN $f_\theta(\mathcal{G})$ is first regularized to mitigate overconfidence and then aggregated at a local, cluster, and structure-independent scale by interleaving energy marginalization and graph diffusion. While group energy marginalizes before diffusion, local energy is fine-grained and can pick up conflicting evidence. GEBM assigns high uncertainty to several anomaly types simultaneously.
• Figure 2: Accuracy of evidential inference at different regularization $\gamma$. GEBM performs on-par with evidential methods at increasingly severe perturbations.
• Figure 3: Logit-based, Gaussian, and regularized energy at different magnitudes. Logit-EBMs become overconfident while the corrected energy is eventually dominated by its regularizer.
• Figure 4: Energy of different types for anomalies of increasing severity on synthetic data. We vary the size of an o.o.d. cluster on real data (left), insert per-node anomalies to the energies of an SBM (middle), and increase the heterophily in an SBM (right).
• Figure 5: Confidence (Maximum Softmax Probability and negative energy) for different GNNs at increasing distribution shift severity.

Theorems & Definitions (11)

• Proposition 4.1
• Theorem 4.2
• Corollary 4.3
• Theorem 4.4
• proof
• proof
• proof
• proof
• proof
• proof