GNN's Uncertainty Quantification using Self-Distillation

TL;DR

This paper introduces a self-distillation-based method for efficiently quantifying the predictive uncertainty of Graph Neural Networks, improving trustworthiness in clinical applications without high computational costs.

Contribution

It proposes a novel self-distillation approach that captures GNN uncertainty more precisely by weighting classifiers differently, avoiding the need for multiple trained models.

Findings

Effectively captures GNN uncertainty in clinical datasets

Achieves similar performance to ensemble and MC Dropout methods

Improves uncertainty quantification for out-of-distribution detection

Abstract

Graph Neural Networks (GNNs) have shown remarkable performance in the healthcare domain. However, what remained challenging is quantifying the predictive uncertainty of GNNs, which is an important aspect of trustworthiness in clinical settings. While Bayesian and ensemble methods can be used to quantify uncertainty, they are computationally expensive. Additionally, the disagreement metric used by ensemble methods to compute uncertainty cannot capture the diversity of models in an ensemble network. In this paper, we propose a novel method, based on knowledge distillation, to quantify GNNs' uncertainty more efficiently and with higher precision. We apply self-distillation, where the same network serves as both the teacher and student models, thereby avoiding the need to train several networks independently. To ensure the impact of self-distillation, we develop an uncertainty metric that captures the diverse nature of the network by assigning different weights to each GNN classifier. We experimentally evaluate the precision, performance, and ability of our approach in distinguishing out-of-distribution data on two graph datasets: MIMIC-IV and Enzymes. The evaluation results demonstrate that the proposed method can effectively capture the predictive uncertainty of the model while having performance similar to that of the MC Dropout and ensemble methods. The code is publicly available at https://github.com/tailabTMU/UQ_GNN.