Uncertainty Modeling in Graph Neural Networks via Stochastic Differential Equations

Dear Reviewer 9vdg, thank you for your thoughtful and detailed review. We appreciate your recognition of the clarity and strength of our model description, the originality of our theoretical results, and the potential impact of our work on uncertainty quantification for graph-structured learning. We hope to have carefully addressed your concerns and made significant updates to the manuscript based on your feedback. We believe that our responses and the revised submission should address your concerns and demonstrate the strengths and contributions of our work. If not, please let us know so that we can address any remaining concerns.

Find below more details on our responses.

Some of the cited works (particularly Calvo-Ordoñez et al., 2024, and Xu et al., 2022) leverage SDEs in a very similar way...

We appreciate the reviewer's comment and the chance to clarify our unique contributions. While Calvo-Ordoñez et al. (2024) and Xu et al. (2022) use SDEs to model weight trajectories in Bayesian Neural Networks, our work focuses on modeling uncertainty in the latent space of Graph Neural Networks. Our approach is fundamentally different in that it addresses the specific problem of learning on graph-structured data.

Quantifying the uncertainty in the latent space of GNNs allows us to provide theoretical guarantees and motivation on well-posedness, robustness, and the connection between latent space variance and output uncertainty. Which, to us, is not translatable to a framework where one uses SDEs in weight space. Please, also note that these contributions are further supported by strong empirical results in tasks like out-of-distribution detection, noise perturbation and active learning.

The choice of a constant drift and diffusion function in the OU prior is not sufficiently explored...

We thank the reviewer for pointing this out. The choice of a constant drift and diffusion function in the Ornstein-Uhlenbeck (OU) prior is analogous to using a standard Gaussian prior in Bayesian models (e.g. BNNs, GPs, etc) to reflect the lack of informative prior knowledge – an uninformative prior.

In response to the reviewer’s suggestion, we have empirically tested the impact of different constant values for these terms and included the results in the revised manuscript in Table 10 and 11 in the new Appendix B.3. For instance, the table at this link (Table 10) evaluates different prior means, and this link (Table 11) shows results for different prior diffusion constants. Our findings indicate that the prior mean has minimal effect, as all experiments converge to similar accuracy, while the prior diffusion constant plays a more significant role, with values below 5 yielding the best results. These findings, now included in the appendix, demonstrate that the model is robust to such variations and that our chosen constants do not compromise performance.

The choice of a constant drift function of zero in the prior OU process reflects the standard uninformative prior commonly adopted in Bayesian frameworks. However, our framework is flexible and supports both constant and non-constant priors, similar to Gaussian Process models. For instance, one could define a prior with a drift function such as f(x, t) = 2x + t, or even utilize another neural network as the drift function, introducing time-dependent dynamics into the process. This flexibility enables the incorporation of more informative priors tailored to specific tasks, which could potentially enhance model performance. But this is beyond the point of our paper.

Similarly, a constant diffusion term is standard practice (Calvo-Ordoñez et al. (2024); Xu et al. (2022); Li et al. (2020)). Parametrizing the diffusion with a learnable function (e.g., a neural network) often causes the diffusion term to collapse to zero, a well-documented issue in the literature (see, e.g., Li et al. (2020) or Xu et al. (2022), Section 2.1), since this would lead to the highest training accuracy. Furthermore, as noted in Calvo-Ordoñez et al. (2024), using the same diffusion function for both the prior and posterior is necessary to keep the KL divergence finite, further supporting the choice of a constant diffusion function.

Thank you again for the review and all the comments. We would like to kindly ask if we have addressed all of your concerns. If so, could you increase your score accordingly? If there are any remaining concerns, please let us know so that we can address them promptly.

Dear Reviewer kdL1, thank you for your detailed feedback and constructive comments. We are glad that you found the topic and approach interesting and appreciate your recognition of the model's potential and its evaluation of uncertainty estimation. We believe the updates and clarifications provided in our response and the revised manuscript directly address your concerns and demonstrate the unique strengths and contributions of our work. We hope these additions encourage you to reconsider your rating. Otherwise, please let us know so that we can address your concerns promptly.

Now, find below our responses:

The paper lacks comparisons with SOTA works...

We thank the reviewer for this helpful feedback. We agree that these comparisons are important and have added new experiments in the revised manuscript. Specifically, we now include results for 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). GNNSafe, as an energy-based method, directly addresses the reviewer’s suggestion to include such approaches. The results are summarized in anonymized links: Cora OOD and Amazon OOD which are table 16 and 17 in the Appendix.

LGNSDE performed well in these comparisons, consistently ranking second-best across most metrics and achieving the best performance in some cases. For example, in the Amazon OOD experiment, it achieved the highest AUPR (Out) score of 97.78. These additions provide a more comprehensive evaluation. We hope that we have addressed the reviewer’s concerns.

The method shows significant similarities to Lin et al. (2024), who also propose an SDE-based GNN...

Thank you for the suggestion. We have now clarified the differences between LGNSDE and GNSD (Lin et al., 2024) in more detail in the paper. Specifically:

GNSD Lin et al. (2024) models message passing (information propagation in the graph) using a stochastic diffusion equation, similar to GRAND (Chamberlain et al., 2021), which uses an ODE for the same purpose. This approach integrates stochasticity directly into the graph diffusion process via a Q-Wiener process, enabling uncertainty propagation during node feature updates. Our model LGNSDE, in contrast, models the latent space of the graph as evolving stochastically under an SDE, while employing a standard GCN for message passing. Resulting in two fundamentally different approaches.

As mentioned above, we have included the GNSD (Lin et al. (2024) and other baselines, see Amazon OOD, and Cora OOD.

Can you elaborate on how the 'framework effectively quantifies uncertainty' based on Proposition 1?...

We thank the reviewer for pointing this out. We revised the statement in the updated manuscript to better reflect our intended meaning and it now reads as follows: “... This ensures that the output variance is bounded by the prior variance of the latent space, providing a controlled measure of uncertainty”. This means that the variance of the output will not explode, since the variance has to be smaller or equal to the prior variance.

Formally, the key point is that, in a Bayesian framework, the posterior variance $\text{Var}(\mathbf{H}(t))$ cannot exceed the prior variance. This ensures that the posterior is bounded by the prior, and as a result, the output variance $\text{Var}(\hat{\mathbf{y}}(t))$ is also bounded by the prior variance of the latent space through the relationship $\text{Var}(\hat{\mathbf{y}}(t)) \leq L_h^2 \text{Var}(\mathbf{H}(t))$ This property ensures that the uncertainty remains controlled and cannot grow arbitrarily large. We hope this clarification resolves the confusion and thank you again for the question.

We would like to start by thanking the reviewer for highlighting the strengths of our work, including the innovation of the idea, the elegance of the mathematical framework, and the potential demonstrated by our results. We hope that our detailed responses and the updates in the revised manuscript reaffirm your satisfaction with the paper. If any questions or concerns arise, please do not hesitate to let us know, and we will address them promptly.

Please find below our responses to all your questions and comments ⤵️

I would imagine the result is heavily dependent upon the integration methods...

We appreciate the observation made by the reviewer and agree with the claim. As noted in GRAND (Chamberlain et al., 2021), numerical integration methods play a significant role in the performance of diffusion-based models. Following their approach, we experimented with several numerical integration schemes, including SRK and Euler-Maruyama, and arrived at similar conclusions. Specifically, we found SRK to consistently perform the best, which aligns with GRAND’s findings (the stochastic variant of RK4) and is also well-supported in the literature as an effective method for solving stochastic differential equations (e.g., Kloeden & Platen, 1992; Milstein, 1995). Their results were a key motivator for our choice of SRK as the default method in our experiments.

The accuracy of this model is not as performant as many cheaper variants.

We thank the reviewer for this comment. While the accuracy of our model is comparable to standard approaches, it does achieve slightly better results in most cases. However, the primary objective of our work was not to maximize accuracy but to develop a model capable of effectively quantifying uncertainty. In this regard, our model demonstrates strong performance in tasks such as out-of-distribution detection, noise perturbation handling, and active learning. The observed improvement in accuracy is an additional benefit but not the primary focus of our contributions.

I would love to see its speed and memory benchmark, as I imagine it to be quite expensive.

For benchmarks of time and memory complexity, please refer to complexity and components, or see Table 15 and 16 in the Appendix C of the paper. This bottleneck is due to the SDE solver, which need multiple samples to estimate the true solution just like BNNs. However, these samples could be parallelised in future work, helping with the time complexity. The memory complexity on the other hand, is one of the main advantages as it becomes constant $O(1)$ when employing the adjoint method (Chen et al., 2018).

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.

Thank you for the question. The choice of backbone is indeed important, and we have explored this aspect in our experiments. Specifically, we tested a Graph Attention Network (GAT) as a backbone instead of the GCN used in the original model. The results can be viewed here or Table 12 in the appendix C. We can see that on the Cora dataset, LGNSDE-GAT outperformed LGNSDE-GCN, achieving an AUROC of 0.8249 compared to 0.7614, making it the best-performing variant so far. However, on the Citeseer dataset, LGNSDE-GAT underperformed relative to LGNSDE-GCN. For a comparison of the entropy distributions on the Cora dataset, see this figure (Figure 4 in the Appendix). The reviewer is right in noting that our proposed framework is flexible and allows the use of any graph model as a backbone, enabling adjustments to optimize performance depending on the dataset. Thank you again for the question.