Uncertainty Modeling in Graph Neural Networks via Stochastic Differential Equations

TL;DR

This paper introduces LGNSDE, a stochastic differential equation framework for graph neural networks that effectively models uncertainty, providing theoretical guarantees and demonstrating improved robustness and out-of-distribution detection in experiments.

Contribution

It develops a novel SDE-based approach for uncertainty quantification in graph neural networks, combining Bayesian and stochastic methods with theoretical guarantees.

Findings

LGNSDE effectively quantifies uncertainty in graph data.

The framework improves robustness to noise and out-of-distribution detection.

Theoretical analysis confirms stability and variance bounds of the model.

Abstract

We propose a novel Stochastic Differential Equation (SDE) framework to address the problem of learning uncertainty-aware representations for graph-structured data. While Graph Neural Ordinary Differential Equations (GNODEs) have shown promise in learning node representations, they lack the ability to quantify uncertainty. To address this, we introduce Latent Graph Neural Stochastic Differential Equations (LGNSDE), which enhance GNODE by embedding randomness through a Bayesian prior-posterior mechanism for epistemic uncertainty and Brownian motion for aleatoric uncertainty. By leveraging the existence and uniqueness of solutions to graph-based SDEs, we prove that the variance of the latent space bounds the variance of model outputs, thereby providing theoretically sensible guarantees for the uncertainty estimates. Furthermore, we show mathematically that LGNSDEs are robust to small perturbations in the input, maintaining stability over time. Empirical results across several benchmarks demonstrate that our framework is competitive in out-of-distribution detection, robustness to noise, and active learning, underscoring the ability of LGNSDEs to quantify uncertainty reliably. Code is available at \href{https://github.com/Richard-Bergna/GraphNeuralSDE}{\texttt{github.com/Richard-Bergna/GraphNeuralSDE}}.