AbstainGNN: Teaching Graph Neural Networks to Abstain for Graph Classification

2026-05-28

AI Summary

This work addresses the unreliability of graph neural networks under high uncertainty, which jeopardizes decision-making in safety-critical applications. The authors propose a rejection-aware graph classification framework that, for the first time, leverages PAC-Bayesian theory to formally characterize the trade-off between classification error and rejection cost. Building upon this theoretical foundation, they devise a joint optimization objective and a two-stage training strategy—comprising prediction warm-up and rejection calibration—to effectively integrate graph structural information with a theoretically grounded rejection mechanism. Evaluated on five benchmark datasets, the proposed method consistently outperforms existing rejection approaches, achieving substantially higher classification accuracy at equivalent rejection rates.

Abstract

Graph classification is a core task in graph data mining with widespread real-world applications. Recent advances in graph neural networks (GNNs) have led to substantial performance improvements for graph classification. However, existing GNNs are typically forced to make predictions even under high uncertainty or unknown conditions, resulting in unreliable decisions that can severely impact downstream tasks, particularly in safety-critical scenarios. To address this critical limitation, we propose AbstainGNN, a novel and theory-driven framework for graph classification with abstention, which enables GNNs to reject uncertain predictions instead of producing incorrect decisions. Specifically, AbstainGNN explicitly models both the predictive function and the abstention function, allowing for effective utilization of graph structural information. Moreover, unlike existing heuristic abstention methods, we theoretically characterize the trade-off between classification errors and rejection costs from a PAC-Bayesian generalization perspective, and derive a unified learning objective for model optimization. Guided by this theoretical insight, we further develop an efficient two-stage training strategy consisting of predictive function warm-start and abstention function calibration. Extensive experiments on five benchmark datasets show that AbstainGNN outperforms existing abstention methods, achieving superior classification performance under the same rejection rates.

Problem

• graph classification
• abstention
• uncertainty
• graph neural networks
• reliable decision-making

Innovation

• abstention learning
• graph neural networks
• PAC-Bayesian theory
• graph classification
• uncertainty-aware prediction