Predict Confidently, Predict Right: Abstention in Dynamic Graph Learning

TL;DR

This paper tackles uncertainty-aware prediction in continuous-time dynamic graphs (CTDGs) by integrating a reject option into temporal graph networks. It introduces a coverage-based abstention framework that jointly optimizes prediction and abstention with a target coverage, via an encoder-decoder CTDG architecture and abstention heads for both link prediction and node classification. A weighted auxiliary loss addresses extreme class imbalance, yielding stronger minority-class performance. Empirical results across four dynamic link-prediction datasets and two dynamic node-classification datasets show that judicious abstention improves AUC and AP, demonstrating higher reliability in risk-sensitive dynamic graph applications.

Abstract

Many real-world systems can be modeled as dynamic graphs, where nodes and edges evolve over time, requiring specialized models to capture their evolving dynamics in risk-sensitive applications effectively. Temporal graph neural networks (GNNs) are one such category of specialized models. For the first time, our approach integrates a reject option strategy within the framework of GNNs for continuous-time dynamic graphs. This allows the model to strategically abstain from making predictions when the uncertainty is high and confidence is low, thus minimizing the risk of critical misclassification and enhancing the results and reliability. We propose a coverage-based abstention prediction model to implement the reject option that maximizes prediction within a specified coverage. It improves the prediction score for link prediction and node classification tasks. Temporal GNNs deal with extremely skewed datasets for the next state prediction or node classification task. In the case of class imbalance, our method can be further tuned to provide a higher weightage to the minority class. Exhaustive experiments are presented on four datasets for dynamic link prediction and two datasets for dynamic node classification tasks. This demonstrates the effectiveness of our approach in improving the reliability and area under the curve (AUC)/ average precision (AP) scores for predictions in dynamic graph scenarios. The results highlight our model's ability to efficiently handle the trade-offs between prediction confidence and coverage, making it a dependable solution for applications requiring high precision in dynamic and uncertain environments.

Figures (2)

• Figure 1: Decision Boundary: (A) is not clear and smooth (B) becomes clearer after abstaining confusing/noisy samples. The examples to be rejected are decided by the objective function, which consists of the losses over unrejected samples and a fixed cost for rejected samples. Minimum value of the objective will be achieved only with proper choice of rejected examples.
• Figure 2: The model overview of classification with rejection for task agnostic temporal graphs. Input to the model is a temporal graph with $t$ timestamps, which is processed through an encoder that samples neighbours from all interactions from previous timestamps $t<t_x$ of these nodes, then aggregates and yields updated temporal node embedding. These embeddings are passed to a Prediction function $f(x)$ to predict downstream tasks such as if a link exists or not for link prediction or if a user is blocked or not for node classification at time $t$. The output embeddings also pass through Abstention function, predict a single score and check if $a(x) > \theta$ threshold i.e. $q(x) = 1$. When it is True, the model abstains, otherwise, it proceeds to make a prediction. The Auxiliary head exposes all samples to the model even with high abstention score for a better training process.