Uncertainty Modeling in Graph Neural Networks via Stochastic Differential Equations

Authors: richard  bergna
sergio  calvo-ordo ñez
felix l opolka
pietro  liò
jose  miguel hernandez-lobato
Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-08-28 22:09:17

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.

Main paper: 
 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., 2016 , Quach et al., 2007 , Mandelzweig and Tabakin, 2001 , Cardelli, 2008 , Buckdahn et al., 2011 , Cvijovic et al., 2014] .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., 2020] , and GCN ensembles [Lin et al., 2022] . 
methodology: 
 Inspired by Graph Neural ODEs [Poli et al., 2019] and Latent SDEs [Li et al., 2020] , we now introduce our model: Latent Graph Neural SDEs -LGNSDEs (Figure 1 ), which use SDEs to define prior and approximate posterior stochastic trajectories for H(t) [Xu et al., 2022] . Furthermore,LGNSDEs can be viewed as the continuous representations of existing discrete architectures (A.4).Figure 1 : The diagram shows the evolution of one of the nodes of the input graph in latent space, 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.LGNSDEs are designed to capture the stochastic latent evolution of H(t) on graph-structured data.We use an Ornstein-Uhlenbeck (OU) prior process, represented bydH(t) = FG(H(t), t) dt + GG(H(t), t) dW(t),where we set the drift and diffusion functions, F G and G G , to constants and consider them hyperparameters. Moreover, dW(t) is a Wiener process. The approximate posterior is defined asEQUATIONwhere F G is parameterized by a GCN with ϕ 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 G G (H(t)) dW(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. 2024 , Archambeau et al. 2007 ].Let Y be a collection of target variables, e.g., class labels, for some of the graph nodes. Given Y we train our model with variational inference, with the ELBO computed asLELBO(ϕ) = E log p(Y|H(t1)) - t 1 t 0 1 2 ∥v(H(u), ϕ, θ, G)∥ 2 2 du ,where the expectation is approximated over trajectories of H(t) sampled from the approximate posterior SDE, andv = G G (H(t)) -1 [F G,ϕ (H(u), u) -F G,θ (H(u), u)].To sample H(t) from the approximate posterior, we integrate the SDE in Eq. 1:H(t1) = H(t0) + t 1 t 0 F G,ϕ (H(u), u) du + t 1 t 0 GG(H(u), u) dW(u),where H(t 0 ) are the node-wise features X in in the graph G. In practice, this is not feasible since the posterior drift F G,ϕ is parametrised by a neural network. We numerically solve this integral with a standard Stochastic Runge-Kutta method [Rößler, 2010] . We then use a Monte Carlo approximation to get the expectation of H(t 0 ) and approximate the posterior predictive distribution asp(Y * |G, Xin, Y) ≈ 1 N N n=1 p (Y * |Hn(t1), G) ,where  
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:Proposition 1. Under assumptions 1-3, there exists a unique mild 1 solution to an LGNSDE of the formdH(t) = F G (H(t), t, θ) dt + G G (H(t), t) dW(t),whose variance bounds the variance of the model output ŷ(t) as:Var(ŷ(t)) ≤ L 2 h Var(H(t)),where L 2 h is the Lipschitz constant of the readout layer. This ensures that the output variance is bounded by the prior variance of the latent space, providing a controlled measure of 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.Proposition 2. Under assumptions 1-3, consider two initial conditions H 0 and H0 = H 0 + δH(0), where δH(0) ∈ R n×d is a small perturbation in the initial node features with ∥δH(0)∥ F = ϵ. Assume that H 0 is taken from a compact set H ⊆ R n×d . Then, the deviation between the solutions H(t) and H(t) of the LGNSDE with these initial conditions remains bounded across time t 2 , specificallyE[∥H(t) -H(t)∥ F ] ≤ ϵe (L f + 1 2 L 2 g )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.  
experiments: 
 We evaluate LGNSDEs on 5 datasets (see A.2 for details on these datasets and hyperparameters), we compare it to GNODE [Poli et al., 2019] , GCN Kipf and Welling [2016] , Bayesian GCN (BGCN) [Hasanzadeh et al., 2020] , and an ensemble of GCNs [Lin et al., 2022] .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. 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 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.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|X i ) = -C c=1 p(y = c|X i ) log p(y = c|X i ), where p(y = c|X i ) represents the probability of input X i belonging to class c. Entropy quantifies the uncertainty in the model's predicted probability distribution over C classes for a given input 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. AU-ROC 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. 
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. 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 F G and G G . Finally, Var(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 y(t) = h(H(t)). 
 
 Appendix: Assumption 1. The drift and diffusion functions F G and G G satisfy the following Lipschitz conditions:∥F G (H 1 (t), t, θ) -F G (H 2 (t), t, θ)∥ F ≤ L f ∥H 1 (t) -H 2 (t)∥ F (2) ∥G G (H 1 (t), t) -G G (H 2 (t), t)∥ F ≤ L g ∥H 1 (t) -H 2 (t)∥ F (3) for all H 1 , H 2 ∈ R n×d , t ∈ [0, T ],and some constants L f and L g .Assumption 2. The drift and diffusion functions F G and G G satisfy a linear growth condition:∥F G (H(t), t, θ)∥ 2 F + ∥G G (H(t), t)∥ 2 F ≤ K(1 + ∥H(t)∥ 2 F ),for all H ∈ R n×d , t ∈ [0, T ], and some constant K.Assumption 3. The variance of the initial conditions, H(0) = H 0 , of the dynamical system is bounded:E[∥H 0 ∥ 2 F ] < ∞.A.3 Proofs Proposition 1. Under assumptions 1-3, there exists a unique mild 3 solution to an LGNSDE of the form dH(t) = F G (H(t), t, θ) dt + G G (H(t), t) dW(t), whose variance bounds the variance of the model output ŷ(t) as:Var(ŷ(t)) ≤ L 2h Var(H(t)), where L 2 h is the Lipschitz constant of the readout layer. This ensures that the output variance is bounded by the prior variance of the latent space, providing a controlled measure of uncertainty.Proof. Using Theorem 1 in Lin et al. [2024] , it follows that the Lipschitz conditions of F G and G G ensure the existence and uniqueness of a mild solution 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:E t 0 G G (H(u), u)dW(u) 2 F = E t 0 ∥G G (H(u), u)∥ 2 F du .Under the Lipschitz condition on G G , we can bound the variance of H(t) as follows:Var(H(t)) = t 0 ∥G G (H(u), u)∥ 2 F du.If G G is bounded, i.e., ∥G G (H(u), u)∥ F ≤ M for some constant M , then Var(H(t)) ≤ M 2 t. This shows that the variance of the latent state H(t) is bounded and grows linearly with time, capturing the aleatoric uncertainty introduced by the stochastic process.Finally, assuming that the model output y(t) is a function of the latent state H(t), y(t) = h(H(t)), where h : R n×d → R n×p is a smooth function, we can apply Itô's Lemma as follows:dy(t) = h ′ (H(t)) [F G (H(t), t, θ) dt + G G (H(t), t) dW(t)] + 1 2 h ′′ (H(t))G G (H(t), t) 2 dt.For the variance of y(t), we focus on the term involving G G (H(t), t) dW(t):Var(y(t)) = t 0 tr h ′ (H(u)) ⊤ G G (H(u), u)G G (H(u), u) ⊤ h ′ (H(u)) du.Using the Cauchy-Schwarz inequality for matrix norms, we can bound this as follows:tr h ′ (H(u)) ⊤ G G (H(u), u)G G (H(u), u) ⊤ h ′ (H(u)) ≤ ∥h ′ (H(u))∥ 2 F ∥G G (H(u), u)∥ 2 F .Therefore, if h is Lipschitz continuous with constant L h , then:Var(y(t)) ≤ L 2 h t 0 ∥G G (H(u), u)∥ 2 F du = L 2 h Var(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.Proposition 2. Under assumptions 1-3, consider two initial conditions H 0 and H0 = H 0 + δH(0), where δH(0) ∈ R n×d is a small perturbation in the initial node features with ∥δH(0)∥ F = ϵ. Assume that H 0 is taken from a compact set H ⊆ R n×d . Then, the deviation between the solutions H(t) and H(t) of the LGNSDE with these initial conditions remains bounded across time t4 , specificallyE[∥H(t) -H(t)∥ F ] ≤ ϵe (L f + 1 2 L 2 g )t .3 A mild solution to an SDE is expressed via an integral equation involving the semigroup generated by the linear operator and represents a weaker notion of the solution.Proof. Consider two solutions H 1 (t) and H 2 (t) of the GNSDE with different initial conditions. Define the initial perturbation as δH(0) where H 1 (0) = H 0 + δH(0) and H 2 (0) = H 0 , with ∥δH(0)∥ F = ϵ.The difference between the two solutions at any time t is given by δH(t) = H 1 (t) -H 2 (t). The dynamics of δH(t) are:d(δH(t)) = [F G (H 1 (t), t, θ) -F G (H 2 (t), t, θ)] dt + [G G (H 1 (t), t) -G G (H 2 (t), t)] dW(t).Applying Itô's Lemma to tr(δH(t) ⊤ δH(t)), we obtain:d(tr(δH(t) ⊤ δH(t))) = 2tr δH(t) ⊤ [F G (H 1 (t), t, θ) -F G (H 2 (t), t, θ)] dt + 2tr δH(t) ⊤ [G G (H 1 (t), t) -G G (H 2 (t), t)] dW(t) + tr [G G (H 1 (t), t) -G G (H 2 (t), t)] ⊤ [G G (H 1 (t), t) -G G (H 2 (t), t)] dt.Taking the expected value, the stochastic integral term involving dW(t) has an expectation of zero due to the properties of the Brownian motion. Thus, we have:E[d(tr(δH(t) ⊤ δH(t)))] = E 2tr(δH(t) ⊤ [F G (H 1 (t), t, θ) -F G (H 2 (t), t, θ)]) dt + E[∥G G (H 1 (t), t) -G G (H 2 (t), t)∥ 2 F ] dt.Here, the second term arises from the variance of the diffusion term, as captured by Itô's Lemma. Solving this using Gronwall's inequality gives:E[tr(δH(t) ⊤ δH(t))] ≤ tr(δH(0) ⊤ δH(0))e (2L f +L 2 g )t .Since ∥δH(0)∥ F = ϵ, we conclude that:E[∥δH(t)∥ F ] ≤ ϵe (L f + 1 2 L 2 g )t . 5Hence, the deviation in the output remains bounded under small perturbations to the initial conditions, providing robustness guarantees.A Here, ∆t represents a fixed time step and ∆W j is a Brownian increment, normally distributed with mean zero and variance ∆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.A mild solution to an SDE is expressed via an integral equation involving the semigroup generated by the linear operator and represents a weaker notion of the solution.Note that while the bound is exponential in t, in practice, the time horizon is usually constrained to a limited range, such as t ∈ [0, 1]. Within this interval, the exponential factor does not grow excessively. Furthermore, in practice, L 2 g = 0 and L f is easily controllable.Note that while the bound is exponential in t, in practice, the time horizon is usually constrained to a limited range, such as t ∈ [0, 1]. Within this interval, the exponential factor does not grow excessively. Furthermore, in practice, L 2 g = 0 and L f is easily controllable.Note that the second term (stochastic part) can be omitted as the first term dominates. 



Reviews: review: title: title: The paper introduces Latent Graph Neural Stochastic Differential Equations (LGNSDEs), a new approach for modeling uncertainty in graph-structured data.
main_review: main_review: The work is well explained, This novel method combines the concepts from graph neural networks and stochastic differential equations.

Strengths:- 
1) Novelty, supported with enough evidence,
2) The paper demonstrates the model's capability in identifying out-of-distribution samples, 

Weakness:- 
1) It would be better if you can provide the ablations with different hyperparameters, and how it affects the model performance.
2) No mention of computational efficiency
3) I Would like to see the image fonts a bit bigger.[Figure 2]
title: title: The paper presents a novel approach (Latent Graph Neural Stochastic Differential Equations) to uncertainty quantification in graph neural networks. However, a more justification of why these specific approach is well-suited for GNNs is expected to provide and a few mathematical annotations could be improved to enhance clarity with explanations for readers. Overall the theoretical foundation and empirical results look promising and make this a valuable contribution.
main_review: main_review: # Quality

Pros:
- Theoretical foundation: The authors provide a solid mathematical basis for their Latent Graph Neural Stochastic Differential Equations (LGNSDEs) model, including theoretical guarantees (Lemmas 1 and 2).
- Empirical evaluation: The experiment compares LGNSDEs against several baseline models (GCN, GNODE, BGCN, Ensemble) across multiple datasets.
- Performance: LGNSDEs show superior or comparable performance in terms of AUROC, AURC and accuracy.

Cons:
- Mathematical clarity: Some equations and notations lack clear explanations, which results difficulty of understanding for the readers.
- Justification for GNNs: A more detailed explanation of why the proposed SDE formulations are particularly well-suited for Graph Neural Networks should be provided.

# Clarity

Pros:
- The paper is generally well-structured and follows a logical flow.

Cons:
- Most of equations are not numbered. Only 1 question is indexed.
- The drift and diffusion functions are mentioned as constants, consider hyperparameters first, but these functions are then parametrized later depending on a few inputs.
- Some mathematical symbols are not well explained. For example, in the ELBO equation:
$L_{ELBO}(\phi) = \mathbb{E}\left[\log p(Y|H(t_1)) - \int_{t_0}^{t_1} \frac{1}{2} |v(H(u), \phi, \theta, G)|_2^2 du\right]$
The meaning of $v(H(u), \phi, \theta, G)$ is not clearly defined.
- The memory and time complexity is not explained how both are calculated in the appendix. 

# Originality

Pros:
- Integration of SDEs with GNNs for uncertainty quantification is innovative and addresses a gap in the field. The ability to quantify both epistemic and aleatoric uncertainties in graph neural networks is promising.
- Theoretical contributions: The provided theoretical guarantees (Lemmas 1 and 2) are novel to the field of uncertainty in GNNs.

# Significance

Pros:
- Novel integration of SDEs with GNNs for uncertainty quantification
- Ability to capture both epistemic and aleatoric uncertainties
- Strong theoretical foundation with proven guarantees