Energy-based Epistemic Uncertainty for Graph Neural Networks

TL;DR

This paper introduces GEBM, an energy-based model that enhances uncertainty estimation in Graph Neural Networks by aggregating energy across structural levels, improving robustness and out-of-distribution detection.

Contribution

The paper proposes GEBM, a novel energy-based approach that combines structural uncertainty levels in GNNs and offers a simple post hoc method for better uncertainty quantification.

Findings

GEBM achieves superior separation of in-distribution and out-of-distribution data.

It outperforms existing methods on 6 out of 7 anomaly types.

The approach is effective across various datasets and shifts.

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.