Uncertainty Modeling in Graph Neural Networks via Stochastic Differential Equations

## Abstract

Abstract We address the problem of learning uncertainty-aware representations for graph-structured data. While Graph Neural Ordinary Differential Equations (GNODE) are effective in learning node representations, they fail to quantify uncertainty. To address this, we introduce Latent Graph Neural Stochastic Differential Equations (LGNSDE), which enhance GNODE by embedding randomness through Brownian motion to quantify uncertainty. We provide theoretical guarantees for LGNSDE and empirically show better performance in uncertainty quantification.

## 1 Introduction

Before the widespread of neural networks and the boom in modern machine learning, complex systems in various scientific fields were predominantly modelled using differential equations. Stochastic Differential Equations (SDEs) were the standard approach to incorporating randomness. These methods were foundational across disciplines such as physics, finance, and computational biology (Hoops et al., ; Quach et al., ; Mandelzweig and Tabakin, ; Cardelli, ; Buckdahn et al., ; Cvijovic et al., ).

In recent years, Graph Neural Networks (GNNs) have become the standard for graph-structured data due to their ability to capture relationships between nodes. They are widely used in social network analysis, molecular biology, and recommendation systems. However, traditional GNNs cannot reliably quantify uncertainty. Both aleatoric (inherent randomness in the data) and epistemic (model uncertainty due to limited knowledge) are essential for decision-making, risk assessment, and resource allocation, making GNNs less applicable in critical applications.

To address this gap, we propose Latent Graph Neural Stochastic Differential Equations (LGNSDE), a method that perturbs features during both the training and testing phases using Brownian motion noise, allowing for handling noise and aleatoric uncertainty. We also assume a prior SDE latent space and learn a posterior SDE using a GNN. This Bayesian approach to the latent space allows us to quantify epistemic uncertainty. As a result, our model can capture and quantify both epistemic and aleatoric uncertainties. More specifically, our contributions are as follows:

- •
We introduce a novel model class combining SDE robustness with GNN flexibility for handling complex graph-structured data, which quantifies both epistemic and aleatoric uncertainties.
- •
We provide theoretical guarantees demonstrating our model’s ability to provide meaningful uncertainty estimates and its robustness to perturbations in the inputs.
- •
We empirically show that Latent GNSDEs demonstrate exceptional performance in uncertainty quantification, outperforming Bayesian GCNs (Hasanzadeh et al., ), and GCN ensembles (Lin et al., ).

## 2 Methodology

Inspired by Graph Neural ODEs (Poli et al., ) and Latent SDEs (Li et al., ), we now introduce our model: Latent Graph Neural SDEs $-$ LGNSDEs (Figure 1), which use SDEs to define prior and approximate posterior stochastic trajectories for $\mathbf{H}(t)$ (Xu et al., ). Furthermore, LGNSDEs can be viewed as the continuous representations of existing discrete architectures (A.4).

Figure: Figure 1: The diagram shows the evolution of input data in latent space $\mathbf{H}(t)$ through an SDE, with sample paths (purple) and confidence bands representing variance. At three timesteps, we visualize graph embeddings, where nodes (white and orange) become more separable over time due to the influence of the vector field. The inset axes represent latent dimensions, while the purple and yellow background highlights the magnitude and direction of the vector field guiding the latent dynamics.
Refer to caption: x1.png

### 2.1 Model Definition

LGNSDEs are designed to capture the stochastic latent evolution of $\mathbf{H}(t)$ on graph-structured data. We use an Ornstein-Uhlenbeck (OU) prior process, represented by

$$ $\text{d}\mathbf{H}(t)=\mathbf{F}_{\mathcal{G}}(\mathbf{H}(t),t)\,\text{d}t+ \mathbf{G}_{\mathcal{G}}(\mathbf{H}(t),t)\,\text{d}\mathbf{W}(t),$ $$

where we set the drift and diffusion functions, $\mathbf{F}_{\mathcal{G}}$ and $\mathbf{G}_{\mathcal{G}}$, to constants and consider them hyperparameters. Moreover, $d\mathbf{W}(t)$ is a Wiener process. The approximate posterior is defined as

$$ $\text{d}\mathbf{H}(t)=\mathbf{F}_{\mathcal{G}}(\mathbf{H}(t),t,\phi)\,\text{d} t+\mathbf{G}_{\mathcal{G}}(\mathbf{H}(t),t)\,\text{d}\mathbf{W}(t),$ (1) $$

where $\mathbf{F}_{\mathcal{G}}$ is parameterized by a GCN with $\phi$ representing the learned weights of the neural network. The drift function mainly determines the dynamics of the evolution of the latent state, while the diffusion term $\mathbf{G}_{\mathcal{G}}(\mathbf{H}(t))\,\text{d}\mathbf{W}(t)$ introduces stochastic elements. With the need to keep the Kullback-Leibler (KL) divergence bounded, we set the diffusion functions of the prior and posterior to be the same [Calvo-Ordonez et al. ; Archambeau et al. ].

Let $\mathbf{Y}$ be a collection of target variables, e.g., class labels, for some of the graph nodes. Given $\mathbf{Y}$ we train our model with variational inference, with the ELBO computed as

$$ $\mathcal{L}_{\text{ELBO}}(\phi)=\mathbb{E}\left[\log{p(\mathbf{Y}|\mathbf{H}(t _{1}))}-\int_{t_{0}}^{t_{1}}\frac{1}{2}\left\|v(\mathbf{H}(u),\phi,\theta, \mathcal{G})\right\|_{2}^{2}\,\text{d}u\right],$ $$

where the expectation is approximated over trajectories of $\mathbf{H}(t)$ sampled from the approximate posterior SDE, and $v=\mathbf{G}_{\mathcal{G}}(\mathbf{H}(t))^{-1}[\mathbf{F}_{\mathcal{G},\phi}(
\mathbf{H}(u),u)-\mathbf{F}_{\mathcal{G},\theta}(\mathbf{H}(u),u)].$

To sample $\mathbf{H}(t)$ from the approximate posterior, we integrate the SDE in Eq. 1:

$$ $\displaystyle\mathbf{H}(t_{1})$ $\displaystyle=\mathbf{H}(t_{0})+\int_{t_{0}}^{t_{1}}\mathbf{F}_{\mathcal{G}, \phi}(\mathbf{H}(u),u)\,\text{d}u+\int_{t_{0}}^{t_{1}}\mathbf{G}_{\mathcal{G}} (\mathbf{H}(u),u)\,\text{d}\mathbf{W}(u),$ $$

where $\mathbf{H}(t_{0})$ are the node-wise features $\mathbf{X}_{\text{in}}$ in the graph $\mathcal{G}$. In practice, this is not feasible since the posterior drift $\mathbf{F}_{\mathcal{G},\phi}$ is parametrised by a neural network. We numerically solve this integral with a standard Stochastic Runge-Kutta method (Rößler, ). We then use a Monte Carlo approximation to get the expectation of $\mathbf{H}(t_{0})$ and approximate the posterior predictive distribution as

$$ $\displaystyle p(\mathbf{Y}^{*}|\mathcal{G},\mathbf{X}_{\text{in}},\mathbf{Y})$ $\displaystyle\approx\frac{1}{N}\sum_{n=1}^{N}p\left(\mathbf{Y}^{*}|\mathbf{H}_ {n}(t_{1}),\mathcal{G}\right),$ $$

where $\mathbf{H}_{1}(t_{1}),\ldots,\mathbf{H}_{N}(t_{1})$ are samples drawn from the approximate posterior $p(\mathbf{H}(t_{1})|\mathbf{Y},\mathbf{X}_{\text{in}},\mathcal{G})$.

Following Poli et al. (), we use a similar encoder-decoder setup. Our encoding focuses solely on the features of individual nodes, while the graph structure remains unchanged. Finally, we remark that the memory and time complexity are $\mathcal{O}(|\mathcal{E}|d+L)$ and $\mathcal{O}(L)$ respectively, where $L$ is the number of SDE solver steps, $\mathcal{E}$ is the number of edges in the graph and $d$ is the dimension of the features.

## 3 Theoretical Guarantees

In this section, we present key results on the stability and robustness of our framework under mild assumptions (Appendix A.2). Firstly, we address the fundamental question of whether our proposed models provide meaningful uncertainties. By showing that the variance of the latent representation bounds the model output variance, we highlight the ability of LGNSDEs to capture and quantify inherent uncertainty in the system. The latent representation is the underlying structure from which the model’s output is generated, i.e. the uncertainty in the latent space directly influences the uncertainty in predictions. We formalize this in the following lemma:

###### Lemma 1 .

Under assumptions 1-3, there exists a unique strong solution to an LGNSDE of the form

$$ $d\mathbf{H}(t)=\mathbf{F}_{\mathcal{G}}(\mathbf{H}(t),t,\boldsymbol{\theta})\, dt+\mathbf{G}_{\mathcal{G}}(\mathbf{H}(t),t)\,d\mathbf{W}(t),$ $$

whose variance bounds the variance of the model output $\mathbf{y}(t)$ as:

$$ $\text{Var}(\mathbf{y}(t))\leq L^{2}_{h}\text{Var}(\mathbf{H}(t)),$ $$

hence providing a meaningful measure of the total uncertainty.

We now demonstrate the robustness of our framework under small perturbations in the initial conditions. By deriving explicit bounds on the deviation between the perturbed and unperturbed solutions over time, we show that the model’s output remains stable.

###### Lemma 2 .

Under assumptions 1-3, consider two initial conditions $\mathbf{H}_{0}$ and $\tilde{\mathbf{H}}_{0}=\mathbf{H}_{0}+\delta\mathbf{H}(0)$, where $\delta\mathbf{H}(0)\in\mathbb{R}^{n\times d}$ is a small perturbation in the initial node features with $\|\delta\mathbf{H}(0)\|_{F}=\epsilon$. Assume that $\mathbf{H}_{0}$ is taken from a compact set $\mathcal{H}\subseteq\mathbb{R}^{n\times d}$. Then, the deviation between the solutions $\mathbf{H}(t)$ and $\tilde{\mathbf{H}}(t)$ of the LGNSDE with these initial conditions remains bounded across time $t$, specifically

$$ $\mathbb{E}[\|\mathbf{H}(t)-\tilde{\mathbf{H}}(t)\|_{F}]\leq\epsilon e^{(L_{f}+ \frac{1}{2}L_{g}^{2})t}.$ $$

In summary, we show analytically that our framework effectively quantifies uncertainty and maintains robustness under small perturbations of the input. First, we confirm that the model’s output variance is controlled and directly linked to the variance of the latent state. Second, we provide a bound on the deviation between solutions with perturbed initial conditions, ensuring stability over time. The proofs can be found in Appendix A.

###### Lemma 1 .

###### Lemma 2 .

## 4 Experiments

Figure: Figure 2: Top: Entropy distributions comparing correct and incorrect model predictions on the CORA dataset. Higher entropy is expected for incorrect predictions. Bottom: Entropy distributions comparing OOD samples with in-distribution samples in the CORA dataset.
Refer to caption: x2.png

We evaluate LGNSDEs on 5 datasets (see A.2 for details on these datasets and hyperparameters), we compare it to GNODE (Poli et al., ), GCN Kipf and Welling (), Bayesian GCN (BGCN) (Hasanzadeh et al., ), and an ensemble of GCNs (Lin et al., ).

The results in Table 3 demonstrate that LGNSDE consistently ranks as either the best or second-best model across most datasets in terms of Micro-AUROC (Area Under the Receiver Operating Characteristic), AURC (Area Under the Risk Coverage), and accuracy. This indicates that LGNSDE effectively handles model uncertainty, successfully distinguishing between classes (AUROC), maintaining low risk while ensuring confident predictions (AURC), and delivering high accuracy.

The top figure 2 shows the entropy distributions of the models for correct and incorrect predictions. Note that most models display similar mean entropy for both correct and incorrect predictions. Notably, our model stands out with the largest difference in entropy, with incorrect predictions having 35% more entropy compared to correct predictions, a larger gap than observed in other models.

### 4.1 Out of Distribution Detection

**Table 1: AUROC (Mean ± Std) and AURC (Mean ± Std) for OOD Detection across datasets. Red denotes the best-performing model, and blue denotes the second-best-performing model.**
| Metric | Model | Cora | Citeseer | Computers | Photo | Pubmed |
| --- | --- | --- | --- | --- | --- | --- |
| AUROC (↑) | GCN | 0.7063 ± 0.0569 | 0.7937 ± 0.0366 | 0.7796 ± 0.0271 | 0.8578 ± 0.0136 | 0.6127 ± 0.0351 |
| GNODE | 0.7398 ± 0.0677 | 0.7828 ± 0.0465 | 0.7753 ± 0.0795 | 0.8473 ± 0.0158 | 0.5813 ± 0.0242 |  |
| BGCN | 0.7193 ± 0.0947 | 0.8287 ± 0.0377 | 0.7914 ± 0.1234 | 0.7910 ± 0.0464 | 0.5310 ± 0.0472 |  |
| ENSEMBLE | 0.7031 ± 0.0696 | 0.8190 ± 0.0375 | 0.8292 ± 0.0338 | 0.8352 ± 0.0059 | 0.6130 ± 0.0311 |  |
| LGNSDE (Our) | 0.7614 ± 0.0804 | 0.8258 ± 0.0418 | 0.7994 ± 0.0238 | 0.8707 ± 0.0099 | 0.6204 ± 0.0162 |  |
| AURC (↓) | GCN | 0.0220 ± 0.0049 | 0.0527 ± 0.0075 | 0.0072 ± 0.0013 | 0.0076 ± 0.0006 | 0.3227 ± 0.0266 |
| GNODE | 0.0184 ± 0.0053 | 0.0545 ± 0.0110 | 0.0070 ± 0.0029 | 0.0097 ± 0.0015 | 0.3357 ± 0.0309 |  |
| BGCN | 0.0208 ± 0.0091 | 0.0458 ± 0.0071 | 0.0064 ± 0.0047 | 0.0108 ± 0.0034 | 0.3714 ± 0.0317 |  |
| ENSEMBLE | 0.0215 ± 0.0061 | 0.0487 ± 0.0072 | 0.0041 ± 0.0011 | 0.0081 ± 0.0003 | 0.3277 ± 0.0265 |  |
| LGNSDE (Our) | 0.0168 ± 0.0070 | 0.0479 ± 0.0109 | 0.0061 ± 0.0011 | 0.0068 ± 0.0008 | 0.3205 ± 0.0135 |  |

We evaluate the models’ ability to detect out-of-distribution (OOD) data by training them with one class left out of the dataset. This introduces an additional class in the validation and test sets that the models have not encountered during training. The goal is to determine if the models can identify this class as OOD. We analyze the entropy, $H(y|\mathbf{X}_{i})=-\sum_{c=1}^{C}p(y=c|\mathbf{X}_{i})\log p(y=c|\mathbf{X}_
{i}),$ where $p(y=c|\mathbf{X}_{i})$ represents the probability of input $\mathbf{X}_{i}$ belonging to class $c$. Entropy quantifies the uncertainty in the model’s predicted probability distribution over $C$ classes for a given input $\mathbf{X}_{i}$.

Figure 2 shows the test entropy distribution for in-distribution (blue) and out-of-distribution (red) data. For each test sample, predictions were made over $C-1$ classes, excluding the left-out class. The OOD class exhibits higher entropy, indicating greater uncertainty. While most models show similar entropy distributions for both data types, our LGNSDE model achieves a clear separation, with a 50% higher mean entropy for OOD data compared to in-distribution data. Other models show less than a 10% difference between the two distributions.

Table 1 presents the AUROC and AURC scores for OOD detection across multiple datasets. AUROC evaluates the model’s ability to differentiate between in-distribution and out-of-distribution (OOD) samples, with higher scores indicating better discrimination. AURC measures the risk of misclassification as coverage increases, where lower values are preferred. The LGNSDE model (ours) consistently achieves the best AUROC and AURC scores across most datasets, indicating its superior performance in accurately identifying OOD samples and minimizing the risk of misclassification.

## 5 Conclusions and Future Work

In conclusion, LGNSDEs outperform the tested models, opening a new avenue for uncertainty quantification in graph data. In future work, stronger benchmarks should be included for a more comprehensive evaluation. Additionally, Neural SDEs face challenges with time and memory complexity. Further work should explore more scalable sampling methods to address these limitations.

## Appendix A Theoretical Remarks

### A.1 Notation

Let $\mathcal{G}=(\mathcal{V},\mathcal{E})$ denote a graph with node set $\mathcal{V}$ and edge set $\mathcal{E}$. The node feature matrix at time $t$ is $\mathbf{H}(t)\in\mathbb{R}^{n\times d}$, where $n$ is the number of nodes and $d$ is the feature dimension. The evolution of $\mathbf{H}(t)$ is described by a Graph Neural Stochastic Differential Equation, with drift function $\mathbf{F}_{\mathcal{G}}(\mathbf{H}(t),t,\boldsymbol{\theta})$ and diffusion function $\mathbf{G}_{\mathcal{G}}(\mathbf{H}(t),t)$. Here, $\mathbf{F}_{\mathcal{G}}$ depends on the graph $\mathcal{G}$, the node features $\mathbf{H}(t)$, time $t$, and parameters $\boldsymbol{\theta}$. The diffusion function $\mathbf{G}_{\mathcal{G}}$ depends on $\mathcal{G}$ and $\mathbf{H}(t)$ but not on $\boldsymbol{\theta}$, as in practice, this is usually a constant function. The randomness is introduced through the Brownian motion $\mathbf{W}(t)$.

The constants $L_{f}$ and $L_{g}$ are Lipschitz constants for the drift and diffusion functions, respectively, ensuring the existence and uniqueness of the solution to the GNSDE. The linear growth condition is controlled by a constant $K$, preventing unbounded growth in $\mathbf{F}_{\mathcal{G}}$ and $\mathbf{G}_{\mathcal{G}}$. Finally, $\text{Var}(\mathbf{H}(t))$ represents the variance of the node features, capturing the aleatoric uncertainty in the system, which is also reflected in the variance of the model output $\mathbf{y}(t)=\mathbf{h}(\mathbf{H}(t))$.

### A.2 Technical Assumptions

###### Assumption 1 .

The drift and diffusion functions $\mathbf{F}_{\mathcal{G}}$ and $\mathbf{G}_{\mathcal{G}}$ satisfy the following Lipschitz conditions:

$$ $\displaystyle\|\mathbf{F}_{\mathcal{G}}(\mathbf{H}_{1}(t),t,\boldsymbol{\theta })-\mathbf{F}_{\mathcal{G}}(\mathbf{H}_{2}(t),t,\boldsymbol{\theta})\|_{F}$ $\displaystyle\leq L_{f}\|\mathbf{H}_{1}(t)-\mathbf{H}_{2}(t)\|_{F}$ (2) $\displaystyle\|\mathbf{G}_{\mathcal{G}}(\mathbf{H}_{1}(t),t)-\mathbf{G}_{ \mathcal{G}}(\mathbf{H}_{2}(t),t)\|_{F}$ $\displaystyle\leq L_{g}\|\mathbf{H}_{1}(t)-\mathbf{H}_{2}(t)\|_{F}$ (3) $$

for all $\mathbf{H}_{1},\mathbf{H}_{2}\in\mathbb{R}^{n\times d}$, $t\in[0,T]$, and some constants $L_{f}$ and $L_{g}$.

###### Assumption 2 .

The drift and diffusion functions $\mathbf{F}_{\mathcal{G}}$ and $\mathbf{G}_{\mathcal{G}}$ satisfy a linear growth condition:

$$ $\|\mathbf{F}_{\mathcal{G}}(\mathbf{H}(t),t,\boldsymbol{\theta})\|_{F}^{2}+\| \mathbf{G}_{\mathcal{G}}(\mathbf{H}(t),t)\|_{F}^{2}\leq K(1+\|\mathbf{H}(t)\|_ {F}^{2}),$ $$

for all $\mathbf{H}\in\mathbb{R}^{n\times d}$, $t\in[0,T]$, and some constant $K$.

###### Assumption 3 .

The variance of the initial conditions, $\mathbf{H}(0)=\mathbf{H}_{0}$, of the dynamical system is bounded: $\mathbb{E}[\|\mathbf{H}_{0}\|_{F}^{2}]<\infty$.

###### Assumption 1 .

###### Assumption 2 .

###### Assumption 3 .

### A.3 Proofs

###### Lemma 3 .

Under assumptions 1-3, there exists a unique strong solution to a Graph Neural SDE of the form:

$$ $d\mathbf{H}(t)=\mathbf{F}_{\mathcal{G}}(\mathbf{H}(t),t,\boldsymbol{\theta})\, dt+\mathbf{G}_{\mathcal{G}}(\mathbf{H}(t),t)\,d\mathbf{W}(t),$ $$

whose variance bounds the variance of the model output $\mathbf{y}(t)$ as:

$$ $\text{Var}(\mathbf{y}(t))\leq L^{2}_{h}\text{Var}(\mathbf{H}(t)),$ $$

hence providing a meaningful measure of the total uncertainty.

###### Proof.

Using standard results from the theory of SDEs Pham [] applied to graph SDEs, it follows that the Lipschitz conditions of $\mathbf{F}_{\mathcal{G}}$ and $\mathbf{G}_{\mathcal{G}}$ ensure the existence and uniqueness of a strong solution $\mathbf{H}(t)$ to the GNSDE.

Now, consider the stochastic part of the variance of the solution. By applying the Itô isometry, we can compute the expectation of the Frobenius norm of the stochastic integral:

$$ $\mathbb{E}\left[\left\|\int_{0}^{t}\mathbf{G}_{\mathcal{G}}(\mathbf{H}(u),u)d \mathbf{W}(u)\right\|_{F}^{2}\right]=\mathbb{E}\left[\int_{0}^{t}\|\mathbf{G}_ {\mathcal{G}}(\mathbf{H}(u),u)\|_{F}^{2}du\right].$ $$

Under the Lipschitz condition on $\mathbf{G}_{\mathcal{G}}$, we can bound the variance of $\mathbf{H}(t)$ as follows:

$$ $\text{Var}(\mathbf{H}(t))=\int_{0}^{t}\|\mathbf{G}_{\mathcal{G}}(\mathbf{H}(u) ,u)\|_{F}^{2}du.$ $$

If $\mathbf{G}_{\mathcal{G}}$ is bounded, i.e., $\|\mathbf{G}_{\mathcal{G}}(\mathbf{H}(u),u)\|_{F}\leq M$ for some constant $M$, then $\text{Var}(\mathbf{H}(t))\leq M^{2}t$. This shows that the variance of the latent state $\mathbf{H}(t)$ is bounded and grows linearly with time, capturing the aleatoric uncertainty introduced by the stochastic process.

Finally, assuming that the model output $\mathbf{y}(t)$ is a function of the latent state $\mathbf{H}(t)$, $\mathbf{y}(t)=\mathbf{h}(\mathbf{H}(t))$, where $\mathbf{h}:\mathbb{R}^{n\times d}\rightarrow\mathbb{R}^{n\times p}$ is a smooth function, we can apply Itô’s Lemma as follows:

$$ $dy(t)=h^{\prime}(\mathbf{H}(t))\left[\mathbf{F}_{\mathcal{G}}(\mathbf{H}(t),t, \boldsymbol{\theta})\,dt+\mathbf{G}_{\mathcal{G}}(\mathbf{H}(t),t)\,d\mathbf{W }(t)\right]+\frac{1}{2}h^{\prime\prime}(\mathbf{H}(t))\mathbf{G}_{\mathcal{G}} (\mathbf{H}(t),t)^{2}\,dt.$ $$

For the variance of $\mathbf{y}(t)$, we focus on the term involving $\mathbf{G}_{\mathcal{G}}(\mathbf{H}(t),t)\,d\mathbf{W}(t)$:

$$ $\text{Var}(\mathbf{y}(t))=\int_{0}^{t}\text{tr}\left(\mathbf{h}^{\prime}( \mathbf{H}(u))^{\top}\mathbf{G}_{\mathcal{G}}(\mathbf{H}(u),u)\mathbf{G}_{ \mathcal{G}}(\mathbf{H}(u),u)^{\top}\mathbf{h}^{\prime}(\mathbf{H}(u))\right)du.$ $$

Using the Cauchy-Schwarz inequality for matrix norms, we can bound this as follows:

$$ $\text{tr}\left(\mathbf{h}^{\prime}(\mathbf{H}(u))^{\top}\mathbf{G}_{\mathcal{G }}(\mathbf{H}(u),u)\mathbf{G}_{\mathcal{G}}(\mathbf{H}(u),u)^{\top}\mathbf{h}^ {\prime}(\mathbf{H}(u))\right)\leq\|\mathbf{h}^{\prime}(\mathbf{H}(u))\|_{F}^{ 2}\|\mathbf{G}_{\mathcal{G}}(\mathbf{H}(u),u)\|_{F}^{2}.$ $$

Therefore, if $\mathbf{h}$ is Lipschitz continuous with constant $L_{h}$, then:

$$ $\text{Var}(\mathbf{y}(t))\leq L_{h}^{2}\int_{0}^{t}\|\mathbf{G}_{\mathcal{G}}( \mathbf{H}(u),u)\|_{F}^{2}du=L^{2}_{h}\text{Var}(\mathbf{H}(t)).$ $$

Hence, under the Lipschitz continuity and boundedness assumptions for the drift and diffusion functions, the solution to the GNSDE exists and is unique, and its output variance serves as a meaningful measure of aleatoric uncertainty.
∎

###### Lemma 4 .

Under assumptions 1-3, consider two initial conditions $\mathbf{H}_{0}$ and $\tilde{\mathbf{H}}_{0}=\mathbf{H}_{0}+\delta\mathbf{H}(0)$, where $\delta\mathbf{H}(0)\in\mathbb{R}^{n\times d}$ represents a small perturbation in the initial node features with $\|\delta\mathbf{H}(0)\|_{F}=\epsilon$. Assume that $\mathbf{H}_{0}$ is taken from a compact set $\mathcal{H}\subseteq\mathbb{R}^{n\times d}$. Then, the deviation between the solutions $\mathbf{H}(t)$ and $\tilde{\mathbf{H}}(t)$ of the Graph Neural SDE with these initial conditions remains bounded across time $t$, specifically:

$$ $\mathbb{E}[\|\mathbf{H}(t)-\tilde{\mathbf{H}}(t)\|_{F}]\leq\epsilon e^{(L_{f}+ \frac{1}{2}L_{g}^{2})t}.$ $$

###### Proof.

Consider two solutions $\mathbf{H}_{1}(t)$ and $\mathbf{H}_{2}(t)$ of the GNSDE with different initial conditions. Define the initial perturbation as $\delta\mathbf{H}(0)$ where $\mathbf{H}_{1}(0)=\mathbf{H}_{0}+\delta\mathbf{H}(0)$ and $\mathbf{H}_{2}(0)=\mathbf{H}_{0}$, with $\|\delta\mathbf{H}(0)\|_{F}=\epsilon$.

The difference between the two solutions at any time $t$ is given by $\delta\mathbf{H}(t)=\mathbf{H}_{1}(t)-\mathbf{H}_{2}(t)$. The dynamics of $\delta\mathbf{H}(t)$ are:

$$ $d(\delta\mathbf{H}(t))=\left[\mathbf{F}_{\mathcal{G}}(\mathbf{H}_{1}(t),t, \boldsymbol{\theta})-\mathbf{F}_{\mathcal{G}}(\mathbf{H}_{2}(t),t,\boldsymbol{ \theta})\right]dt+\left[\mathbf{G}_{\mathcal{G}}(\mathbf{H}_{1}(t),t)-\mathbf{ G}_{\mathcal{G}}(\mathbf{H}_{2}(t),t)\right]d\mathbf{W}(t).$ $$

Applying Itô’s Lemma to $\text{tr}(\delta\mathbf{H}(t)^{\top}\delta\mathbf{H}(t))$, we obtain:

$$ $\displaystyle d(\text{tr}(\delta\mathbf{H}(t)^{\top}\delta\mathbf{H}(t)))$ $\displaystyle=2\text{tr}\left(\delta\mathbf{H}(t)^{\top}\left[\mathbf{F}_{ \mathcal{G}}(\mathbf{H}_{1}(t),t,\boldsymbol{\theta})-\mathbf{F}_{\mathcal{G}} (\mathbf{H}_{2}(t),t,\boldsymbol{\theta})\right]\right)dt$ $\displaystyle+2\text{tr}\left(\delta\mathbf{H}(t)^{\top}\left[\mathbf{G}_{ \mathcal{G}}(\mathbf{H}_{1}(t),t)-\mathbf{G}_{\mathcal{G}}(\mathbf{H}_{2}(t),t )\right]d\mathbf{W}(t)\right)$ $\displaystyle+\text{tr}\left(\left[\mathbf{G}_{\mathcal{G}}(\mathbf{H}_{1}(t), t)-\mathbf{G}_{\mathcal{G}}(\mathbf{H}_{2}(t),t)\right]^{\top}\left[\mathbf{G} _{\mathcal{G}}(\mathbf{H}_{1}(t),t)-\mathbf{G}_{\mathcal{G}}(\mathbf{H}_{2}(t) ,t)\right]\right)dt.$ $$

Taking the expected value, the stochastic integral term involving $d\mathbf{W}(t)$ has an expectation of zero due to the properties of the Brownian motion. Thus, we have:

$$ $\displaystyle\mathbb{E}[d(\text{tr}(\delta\mathbf{H}(t)^{\top}\delta\mathbf{H} (t)))]$ $\displaystyle=\mathbb{E}\left[2\text{tr}(\delta\mathbf{H}(t)^{\top}[\mathbf{F} _{\mathcal{G}}(\mathbf{H}_{1}(t),t,\boldsymbol{\theta})-\mathbf{F}_{\mathcal{G }}(\mathbf{H}_{2}(t),t,\boldsymbol{\theta})])\right]\,dt$ $\displaystyle+\mathbb{E}[\|\mathbf{G}_{\mathcal{G}}(\mathbf{H}_{1}(t),t)- \mathbf{G}_{\mathcal{G}}(\mathbf{H}_{2}(t),t)\|_{F}^{2}]\,dt.$ $$

Here, the second term arises from the variance of the diffusion term, as captured by Itô’s Lemma. Using the Lipschitz bounds for $\mathbf{F}_{\mathcal{G}}$ and $\mathbf{G}_{\mathcal{G}}$, we obtain:

$$ $\mathbb{E}[d(\text{tr}(\delta\mathbf{H}(t)^{\top}\delta\mathbf{H}(t)))]\leq(2L _{f}\mathbb{E}[\text{tr}(\delta\mathbf{H}(t)^{\top}\delta\mathbf{H}(t))]+L_{g} ^{2}\mathbb{E}[\text{tr}(\delta\mathbf{H}(t)^{\top}\delta\mathbf{H}(t))])\,dt.$ $$

Rewriting this as a differential inequality:

$$ $\frac{d}{dt}\mathbb{E}[\text{tr}(\delta\mathbf{H}(t)^{\top}\delta\mathbf{H}(t) )]\leq(2L_{f}+L_{g}^{2})\mathbb{E}[\text{tr}(\delta\mathbf{H}(t)^{\top}\delta \mathbf{H}(t))].$ $$

Solving this using Gronwall’s inequality gives:

$$ $\mathbb{E}[\text{tr}(\delta\mathbf{H}(t)^{\top}\delta\mathbf{H}(t))]\leq\text{ tr}(\delta\mathbf{H}(0)^{\top}\delta\mathbf{H}(0))e^{(2L_{f}+L_{g}^{2})t}.$ $$

Since $\|\delta\mathbf{H}(0)\|_{F}=\epsilon$, we conclude that:

$$ $\mathbb{E}[\|\delta\mathbf{H}(t)\|_{F}]\leq\epsilon e^{(L_{f}+\frac{1}{2}L_{g} ^{2})t}.$ $$

Hence, the deviation in the output remains bounded under small perturbations to the initial conditions, providing robustness guarantees.
∎

###### Lemma 3 .

###### Proof.

###### Lemma 4 .

###### Proof.

### A.4 GNSDE as a Continuous Representation of Graph ResNet with Stochastic Noise Insertion

Consider a Graph Neural Stochastic Differential Equation (GNSDE) represented as:

$$ $d\mathbf{H}(t)=\mathbf{F}_{\mathcal{G}}(\mathbf{H}(t),t)dt+\mathbf{G}_{ \mathcal{G}}(\mathbf{H}(t),t)d\mathbf{W}(t),$ $$

where $\mathbf{H}(t)\in\mathbb{R}^{n\times d}$, $\mathbf{F}_{\mathcal{G}}(\mathbf{H}(t),t)$, and $\mathbf{G}_{\mathcal{G}}(\mathbf{H}(t),t)$ are matrix-valued functions, and $\mathbf{W}(t)$ is a Brownian motion. The numerical Euler-Maruyama discretization of this GNSDE can be expressed as

$$ $\frac{\mathbf{H}(t_{j+1})-\mathbf{H}(t_{j})}{\Delta t}\approx\mathbf{F}_{ \mathcal{G}}(\mathbf{H}(t_{j}),t_{j})+\frac{\mathbf{G}_{\mathcal{G}}(\mathbf{H }(t_{j}),t_{j})\Delta\mathbf{W}_{j}}{\Delta t},$ $$

which simplifies to

$$ $\mathbf{H}_{j+1}=\mathbf{H}_{j}+\mathbf{F}_{\mathcal{G}}(\mathbf{H}_{j},t_{j}) \Delta t+\mathbf{G}_{\mathcal{G}}(\mathbf{H}_{j},t_{j})\Delta\mathbf{W}_{j}.$ $$

Here, $\Delta t$ represents a fixed time step and $\Delta\mathbf{W}_{j}$ is a Brownian increment, normally distributed with mean zero and variance $\Delta t$. This numerical discretization is analogous to a Graph Residual Network (Graph ResNet) with a specific structure, where Brownian noise is injected at each residual layer. Therefore, the Graph Neural SDE can be interpreted as a deep Graph ResNet where the depth corresponds to the number of discretization steps of the SDE solver.

## Appendix B Details of The Experimental Setup

**Table 2: Uniform Hyperparameters of the LGNSDE, GNODE, Bayesian GCN, GCN, Ensemble of GCN models for all datasets.**
| Parameter | GNSDE | GNODE | Other |
| --- | --- | --- | --- |
| Datasets | All | All | All |
| $t_{1}$ | 1 | 1 | n/a |
| Hidden Dimensions | 64 | 64 | 64 |
| Learning Rate | 0.01 | 0.01 | 0.01 |
| Optimizer | adam | adam | adam |
| Method | SRK | SK4 | n/a |
| Dropout | 0.2 | 0.2 | 0.2 |
| Diffusion g | 1.0 | n/a | n/a |

**Table 3: AUROC (Mean ± Std), AURC (Mean ± Std), and Accuracy (Mean ± Std) for all datasets. Red denotes the best-performing model, and blue denotes the second-best-performing model.**
| Metric | Model | Cora | Citeseer | Computers | Photo | Pubmed |
| --- | --- | --- | --- | --- | --- | --- |
| MICRO-AUROC (↑) | GCN | 0.9654 ± 0.0050 | 0.9173 ± 0.0068 | 0.9680 ± 0.0016 | 0.9905 ± 0.0003 | 0.9006 ± 0.0139 |
| GNODE | nan ± nan | 0.9146 ± 0.0063 | 0.9569 ± 0.0067 | 0.9885 ± 0.0007 | 0.8857 ± 0.0203 |  |
| BGCN | 0.9571 ± 0.0092 | 0.9099 ± 0.0090 | 0.9421 ± 0.0097 | 0.9489 ± 0.0189 | 0.7030 ± 0.1331 |  |
| ENSEMBLE | 0.9635 ± 0.0031 | 0.9181 ± 0.0062 | 0.9669 ± 0.0025 | 0.9886 ± 0.0004 | 0.8785 ± 0.0163 |  |
| LGNSDE (Our) | 0.9667 ± 0.0036 | 0.9111 ± 0.0072 | 0.9691 ± 0.0032 | 0.9909 ± 0.0004 | 0.9007 ± 0.0091 |  |
| AURC (↓) | GCN | 0.9966 ± 0.0007 | 0.9966 ± 0.0011 | 0.9994 ± 0.0005 | 0.9987 ± 0.0015 | 0.9994 ± 0.0004 |
| GNODE | nan ± nan | 0.9967 ± 0.0011 | 0.9994 ± 0.0004 | 0.9998 ± 0.0001 | 0.9915 ± 0.0163 |  |
| BGCN | 0.9972 ± 0.0004 | 0.9963 ± 0.0010 | 0.9994 ± 0.0002 | 0.9989 ± 0.0005 | 0.9996 ± 0.0004 |  |
| ENSEMBLE | 0.9970 ± 0.0012 | 0.9967 ± 0.0012 | 0.9994 ± 0.0002 | 0.9989 ± 0.0006 | 0.9996 ± 0.0005 |  |
| LGNSDE (Our) | 0.9970 ± 0.0003 | 0.9971 ± 0.0005 | 0.9995 ± 0.0003 | 0.9997 ± 0.0002 | 0.9995 ± 0.0005 |  |
| Accuracy (↑) | GCN | 0.8105 ± 0.0173 | 0.7258 ± 0.0137 | 0.8098 ± 0.0048 | 0.9116 ± 0.0021 | 0.7570 ± 0.0229 |
| GNODE | nan ± nan | 0.7235 ± 0.0159 | 0.7911 ± 0.0098 | 0.9053 ± 0.0032 | 0.7577 ± 0.0231 |  |
| BGCN | 0.7897 ± 0.0261 | 0.7013 ± 0.0196 | 0.7114 ± 0.0333 | 0.7124 ± 0.0968 | 0.4581 ± 0.1846 |  |
| ENSEMBLE | 0.8038 ± 0.0105 | 0.7108 ± 0.0166 | 0.8070 ± 0.0055 | 0.9070 ± 0.0019 | 0.7299 ± 0.0218 |  |
| LGNSDE (Our) | 0.8113 ± 0.0128 | 0.7120 ± 0.0119 | 0.8247 ± 0.0103 | 0.9169 ± 0.0021 | 0.7595 ± 0.0168 |  |

## Peer Reviews

confidence: 3
contribution: 4
initial_rating: 10
presentation: 4
questions: - Have you tried varying the backbones of the GNN and is there any performance change? I think rewiring the graph structure might have a huge impact.
soundness: 4
strengths: - The idea is sound and innovative.
- The mathematical framework is simply elegant.
- The robustness of this model is clearly demonstrated by theoretical and experimental results.
- This framework can be of high utility to the community.
summary: This paper generalizes GNODE to its stochastic counterpart, using SDEs to derive uncertainty-aware representations for graph-structured data. Theoretical and experimental characterization of the framework, named LGNSDE, showed its robustness and uncertainty quantification capability.
weaknesses: - I would imagine the result is heavily dependent upon the integration methods. See the extensive study conducted in GRAND: Graph Neural Diffusion, Chamberlain et al. 2021. 
- The accuracy of this model is not as performant as many cheaper variants.
- I would love to see its speed and memory benchmark, as I imagine it to be quite expensive.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Can you elaborate on how the 'framework effectively quantifies uncertainty' based on Proposition 1? It is not immediately clear to me how this bounded output variance translates to effectively quantified uncertainty.


While the topic and approach are interesting and the method appears potentially promising, I have concerns about the evaluation and the presentation, particularly the lack of differentiation from existing methods. If additional data and clarifications are provided, I would be willing to reconsider my rating.


--- 
Post rebuttal: experimental evaluation seems good now.
soundness: 2
strengths: 1. The paper demonstrates good empirical performance.
2. The authors evaluate uncertainty estimation through OOD detection, noise perturbation, and active learning.
summary: The paper proposes a new GNN architecture version based on stochastic differential equations to improve uncertainty estimation.
weaknesses: 1. The paper lacks comparisons with SOTA works. Some strong methods, such as GPN (Stadler et al., 2021) and GNSD (Lin et al., 2024), are mentioned but not evaluated, despite their better empirical performance compared to the baselines used. Additionally, other high-performing methods that utilize energy variants are not tested.
2. The method shows significant similarities to Lin et al. (2024), who also propose an SDE-based GNN. While the authors acknowledge the difference with one sentence in the related work section, I believe the similarities and differences should be further discussed in detail, such as by comparing the frameworks mathematically and conducting a specific study (even on a synthetic dataset) to highlight the unique merits of their approach. But the method is not compared against at all.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: - Is it possible to allow a non-constant drift function in the prior OU process? Would a similar method or extension allow this?
- What is the effect upon training and convergence for different choices of hyperparameters, and in particular, for different choices of constant drift and diffusion functions? Has an ablation study been performed for this proposed model?
- How does the walltime of learning with an LGNSDE compare to walltime for the other tested methods?
- Are there any interesting limitations on the practicality of this method, and in your experience, what can allow for learning the best results?
soundness: 2
strengths: - The LGNSDE model description is clear, concise, and intuitive.
- The two theoretical results are important and original, providing a good deal of strength to the proposed methodology.
- The experiments clearly show this model has significant potential in delivering upon the promise of uncertainty quantification for graph-structured learning problems.
summary: This paper introduces a new model (LGNSDE) for learning on graphs with uncertainty quantification. It leverages an Ornstein-Uhlenbeck prior and a posterior with drift parameterized by a GCN model that is trained through variational inference. This SDE approach is latent as it operates on node feature embeddings. Two novel theoretical results are shown, providing a bound on model output variance and a bound on solutions under node feature perturbations. Six graph models are then compared in various experiments on five standard graph datasets, including studies of OOD detection, test noise, and active learning. The proposed model performs favorably across these experiments.
weaknesses: - Some of the cited works (particularly Calvo-Ordoñez et al., 2024, and Xu et al., 2022) leverage SDEs in a very similar way, but do not consider the problem of learning on graphs. While the application to graphs is creative, the model definition is somewhat limited in novelty.
- The choice of a constant drift and diffusion function in the OU prior is not sufficiently explored. It would be great if there was mention of why this is a reasonable restriction if this is indeed the case.
- The experimental section would benefit from increased clarity and specificity, especially in regard to the structure and setup of experiments. Choices such as number of epochs and early stopping criteria, if any, are not explicitly stated.
- It is stated in Section 5 that hyperparameters that achieved the highest validation accuracy were chosen. It is unclear to me which validation accuracy is used here, and knowing the exact grid search setup would be ideal for improved reproducibility.
- Table 7 shows that uniform hyperparameters were chosen across all models; if the hyperparameter choice was performed via search for each model independently, the comparison results might be more fair.
- The most impactful perceived weakness of this paper was the lack of experimental comparison with the models referenced as related works, particularly those referenced in Section 6. The other methods used in comparison do not follow a stochastic approach, limiting the ability for a fair comparison. I would be extremely interested to see a comparative study against the referenced GNSD (Lin et al., 2024) method, for example.
- For OOD detection, comparison against methods built specifically for this task on graphs, like GNNSafe (Wu et al., 2023) would be enlightening.
- One of the theoretical claims was a resistance to small perturbations in graph structure. This might lead to a compelling experiment, but this claim is not tested.

There are two small writing nitpicks:
- In the first line of the introduction, "Before the widespread of neural networks" is missing a word. Perhaps use "widespread success" or similar?
- In Section 6, "Hence, constructing a different method." does not read as a complete sentence, and should probably be folded into the line before it.

Meta-review: This paper investigates uncertainty-aware representations in graph neural networks using stochastic differential equations (SDEs). The authors provide theoretical results for bounding the model output variance and bounding solutions under node feature perturbations. Experiments are conducted to validate the effectiveness of the proposed structure.

During the initial reviews, all reviewers acknowledged the novelty of the proposed method and the soundness of its theoretical results but raised concerns about the lack of empirical comparisons with state-of-the-art methods and insufficient discussion on the methodological differences with other SDE-based approaches. The authors successfully addressed these issues by including more comprehensive experiments and improving the writing.

While the performance of the proposed method is not particularly impressive compared to standard approaches, the reviewers agree that the method offers meaningful characterizations of uncertainty in graph data learning. Congratulations to the authors on this great work.

Decision: Accept (Spotlight)