Uncertainty Modeling in Graph Neural Networks via Stochastic Differential Equations

Authors: Richard Bergna, Sergio Calvo Ordoñez, Felix Opolka, Pietro Lio, José Miguel Hernández-Lobato
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: graph neural networks, stochastic differential equations, uncertainty quantification, bayesian machine learning

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 perturbations, and active learning, underscoring the ability of LGNSDEs to quantify uncertainty reliably.

Decision: accept (spotlight)
Reviewer scores: 10, 6, 6

Reviews and discussion:
<thread>
creation date:
1729650237644

thread_type:
official_review

<thread_object>
creator: reviewer_9vdg

<content>
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.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730579247299

thread_type:
official_review

<thread_object>
creator: reviewer_kdl1

<content>
summary:
The paper proposes a new GNN architecture version based on stochastic differential equations to improve uncertainty estimation.

strengths:
1. The paper demonstrates good empirical performance.
2. The authors evaluate uncertainty estimation through OOD detection, noise perturbation, and active learning.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730694532427

thread_type:
official_review

<thread_object>
creator: reviewer_sz1c

<content>
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.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732126704504

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Comments and Main Elements of our Rebuttal

comment:
Dear reviewers, thank you for your constructive and thoughtful comments.

We are grateful for your feedback, which has helped us improve the clarity and scope of our work. We appreciate that **Reviewer kdL1** highlighted our strong empirical performance and evaluation of uncertainty estimation. **Reviewer 9vdg** recognized the clear model description, the originality of our theoretical results, and the potential for uncertainty quantification. **Reviewer SZ1C** commended the soundness and innovation of the idea, as well as the robustness demonstrated by the theoretical and experimental results.

**The main elements of our rebuttal are as follows:**

➡️ **Expanded Baseline Comparisons.** One major concern was the lack of sufficient baselines. We now include OOD detection comparisons with state-of-the-art methods such as GNSD (Lin et al., 2024), GPN (Stadler et al., 2021), GKDE (Zhao et al., 2020), MSP (Hendrycks & Gimpel, 2016), GNNSafe (Wu et al., 2023), GRAND (Chamberlain et al., 2021), GREAD (Choi et al., 2023), Mahalanobis (Lee et al., 2018), and ODIN (Liang et al., 2018). Our model consistently ranks among the top 1–3 across datasets. Updated results are provided [here (Amazon OOD)](https://anonymous.4open.science/r/ICLR2024Rebutal-3528/amazon_new_OOD.png) and [here (Cora OOD)](https://anonymous.4open.science/r/ICLR2024Rebutal-3528/cora_new_OOD.png), as well as in **Appendix B** of the updated manuscript. We hope these additional baselines address the reviewers’ concerns and highlight the effectiveness of our approach to quantifying uncertainty in graph data.

➡️ **Empirical Study of Prior Choices.** In response to comments on prior drift and diffusion functions, we conducted a study showing the impact of different choices. Results for [drift](https://anonymous.4open.science/r/ICLR2024Rebutal-3528/prior_mean_vs_acc_table.png) and [diffusion](https://anonymous.4open.science/r/ICLR2024Rebutal-3528/sigma_vs_acc.png), or in **Appendix B**, demonstrate robustness to drift constant variations and identify optimal ranges for diffusion constants.

➡️ **Improved Clarity and Reproducibility.** We have made clarifying edits throughout the manuscript to improve comparisons, theoretical explanations, and experimental setups. Additionally, all experiments will be fully reproducible in the final version with the provided code.

We hope these updates address your concerns and strengthen our submission. If you still have any concerns, please let us know so that we can address them. Please find detailed responses below and an updated version of our manuscript attached to the submission.
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.