🤖 machine learning

AbstainGNN: Teaching Graph Neural Networks to Abstain for Graph Classification

This paper introduces AbstainGNN, a novel, theory-driven framework that equips Graph Neural Networks with the ability to abstain from uncertain predictions by modeling both classification and rejection functions, theoretically optimizing their trade-off via a PAC-Bayesian perspective, and demonstrating superior performance on benchmark datasets through an efficient two-stage training strategy.

Original authors: Xixun Lin, Zhiheng Zhou, Zhengyin Zhang, Yancheng Chen, Shuai Zhang, Ge Zhang, Shichao Zhu, Lixin Zou, Chuan Zhou, Peng Zhang, Shirui Pan, Yanan Cao

Published 2026-06-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xixun Lin, Zhiheng Zhou, Zhengyin Zhang, Yancheng Chen, Shuai Zhang, Ge Zhang, Shichao Zhu, Lixin Zou, Chuan Zhou, Peng Zhang, Shirui Pan, Yanan Cao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are hiring a team of expert detectives (Graph Neural Networks, or GNNs) to solve cases by looking at complex maps of clues (graphs). These detectives are incredibly smart and usually get the right answer. However, they have a dangerous habit: they never admit when they are confused. Even when the clues are messy, contradictory, or missing, they will still force a guess. In high-stakes situations—like checking if a new medicine is safe or if a chemical is toxic—this "guessing anyway" can lead to disastrous mistakes.

The paper "AbstainGNN" proposes a solution: teaching these detectives the wisdom to know when not to speak.

Here is a simple breakdown of how they did it, using everyday analogies:

1. The Problem: The "Know-It-All" Detective

Current AI models are like detectives who are afraid to say "I don't know." If a case is unclear, they still pick a suspect.

The Risk: If the AI is wrong about a toxic chemical, it might approve it for testing, wasting money or hurting people.
The Goal: The authors wanted to build a system where the AI can say, "This graph is too confusing; I will abstain (refuse to guess) and let a human expert handle it."

2. The Solution: A Two-Person Team

The authors created AbstainGNN, which acts like a two-person detective squad working together:

Detective A (The Predictor): This is the standard AI that looks at the graph and tries to guess the answer (e.g., "Toxic" or "Safe").
Detective B (The Referee): This is a new, special AI trained specifically to watch Detective A. Its only job is to decide: "Is Detective A confident enough to make a call, or is this a bad guess?"

If the Referee thinks the confidence is too low, the whole team says "Abstain" and passes the case to a human.

3. How They Trained Them: The "Warm-Up" and "Calibration"

Training this team wasn't as simple as just telling them to be careful. The authors used a clever two-stage training strategy:

Stage 1: The Warm-Up (Getting Detective A Ready)
Before teaching the team to say "I don't know," they first trained Detective A to be very good at solving clear cases. They let the AI practice until it was in a "small-gradient region."
- Analogy: Imagine a student studying for a test. You don't ask them to grade their own confidence until they have actually learned the material. First, they just study hard to get the basics down.
Stage 2: Calibration (Teaching the Referee)
Once Detective A was stable, they trained Detective B (the Referee). They showed the Referee examples where Detective A was right and where it was wrong, teaching the Referee to spot the difference.
- The "Global Cluster" Trick: A common problem in training is that if you only look at a small group of students (a "batch"), you might get a biased view of the class. To fix this, AbstainGNN keeps a global scorecard of what "correct answers" look like for the whole class, not just the current group. This ensures the Referee is fair and accurate.

4. The Theory: Why It Works

The paper isn't just a "try it and see" experiment; they did the math to prove why it works.

The "Clustering" Concept: They proved that for an AI to know when to abstain, all the "Toxic" graphs must look very similar to each other (a tight cluster), and all the "Safe" graphs must look very similar to each other.
The Magic: By forcing the AI to group similar graphs tightly together (reducing "intra-class variance"), the AI naturally becomes more confident when it's right and less confident when it's looking at something weird. This mathematical proof ensures the AI doesn't just guess randomly; it knows exactly when to stop.

5. The Results: Smarter, Safer Decisions

The authors tested AbstainGNN on five different real-world datasets (like chemical structures and social networks).

The Scorecard: They measured two things: Coverage (how many cases the AI tried to solve) and Risk (how many mistakes it made).
The Win: AbstainGNN consistently made fewer mistakes than other methods when they were allowed to skip the hard cases.
- Example: On the MUTAG dataset (chemicals), AbstainGNN reduced the risk of error by 16.8% compared to the next best method, all while skipping the same number of difficult cases.

Summary

AbstainGNN is like teaching a super-smart AI to have the humility to say, "I'm not sure."

It uses a two-stage training process: first learn the facts, then learn to judge confidence.
It uses a mathematical rule to ensure similar things look similar, making confidence scores trustworthy.
It outperforms existing methods by making fewer dangerous mistakes in safety-critical fields like chemistry and biology.

The paper concludes that by giving AI the power to abstain, we can make these systems much more reliable for real-world, high-stakes decisions.

Technical Summary: AbstainGNN

Problem Statement

Graph classification is a fundamental task in graph data mining with applications in molecular property prediction, material analysis, and social network analysis. While Graph Neural Networks (GNNs) have achieved significant performance improvements, they suffer from a critical limitation: they are typically forced to make predictions even under conditions of high uncertainty, ambiguous structural evidence, or weak class separation. In safety-critical domains (e.g., drug discovery or financial risk assessment), such forced decisions can lead to severe downstream errors.

Existing "Learning with Abstention" (or classification with rejection) methods, which allow models to defer uncertain predictions to human experts, have primarily been developed for computer vision. These methods often fail to capture the rich topological and relational information inherent in graph-structured data. Furthermore, most existing approaches rely on heuristic designs lacking a solid theoretical foundation specifically for graph domains, hindering the development of principled abstention mechanisms for graph classification.

Methodology

The authors propose AbstainGNN, a novel, theory-driven framework that explicitly equips GNNs with the ability to abstain from making predictions on uncertain graph samples. The framework consists of two primary components: a predictive function $f$ for classification and an abstention function $g$ for determining whether to predict or reject.

Theoretical Foundation

The core of AbstainGNN is derived from a PAC-Bayesian generalization perspective. The authors theoretically characterize the trade-off between classification errors and rejection costs.

Generalization Bound: Theorem 4.1 establishes that the generalization performance of AbstainGNN is upper-bounded by the empirical margin-based loss and the intra-class variance of graph-level representations.
Key Insight: Minimizing the intra-class variance of graph-level representations is crucial for improving performance in graph classification with abstention.

Learning Objective

Based on the theoretical analysis, the authors derive a unified learning objective:
$\mathcal{L}_{total} = \mathcal{L}_{ce} + \lambda_r \mathcal{L}_{icv}$
Where $\mathcal{L}_{ce}$ is the standard cross-entropy loss, and $\mathcal{L}_{icv}$ is a novel regularization loss designed to reduce the intra-class variance of graph-level representations. Theoretical analysis (Theorem 4.2) demonstrates that minimizing this objective leads to a contraction of the intra-class variance, tightening the generalization bound.

Implementation Strategy

To optimize this objective efficiently, AbstainGNN employs a two-stage training strategy:

Predictive Function Warm-start: The predictive function $f$ is pre-trained using only $\mathcal{L}_{ce}$ to drive the model into a small-gradient region. This stabilizes the optimization process before introducing the variance regularization.
Abstention Function Calibration:
- Global Cluster Adjustment: To mitigate estimation bias caused by batch-wise updates, the model maintains global class clusters ( $\mu_y$ ) throughout training using a moving average strategy. This ensures accurate estimation of the intra-class variance.
- Calibration: The abstention function $g$ (implemented as a GNN) is trained to predict whether the predictive function $f$ 's output is correct. This casts confidence calibration as a binary classification task, aligning confidence scores with prediction accuracy.

Key Contributions

Problem Definition: This work is the first to investigate the problem of graph classification with abstention, enabling GNNs to reject uncertain or unknown graph samples to improve reliability in safety-critical domains.
Theoretical Analysis: The authors derive PAC-Bayesian generalization bounds for AbstainGNN, revealing that minimizing intra-class variance is theoretically necessary for effective abstention. This leads to a unified learning objective.
Model Implementation: Guided by theory, the authors develop an efficient implementation featuring a two-stage training strategy and a global class-cluster adjustment mechanism to alleviate estimation bias.
Experimental Verification: Comprehensive experiments on five real-world benchmark datasets demonstrate the method's superiority over existing abstention baselines.

Experimental Results

The authors evaluated AbstainGNN on five benchmark datasets: PROTEINS, MUTAG, NCI1 (chemical graphs), and IMDB-BINARY, COLLAB (social graphs). They compared against state-of-the-art baselines including SR, MC-Dropout, Deep Gamblers, SAT, CCL-SC, NCwR, and GraphPPD.

Performance: AbstainGNN consistently outperformed existing methods across all datasets and coverage levels.
Risk Reduction: Under the same rejection rates, the model achieved significant relative risk reductions:
- 16.8% on MUTAG
- 15.8% on IMDB-BINARY
- 9.8% on PROTEINS
- 5.9% on COLLAB
- 4.8% on NCI1
Ablation Studies: Removing the warm-start strategy or the intra-class variance regularization ( $\mathcal{L}_{icv}$ ) consistently degraded performance, confirming the necessity of both components.
Efficiency: Runtime analysis showed that AbstainGNN incurs only moderate computational overhead compared to baselines, maintaining an acceptable balance between effectiveness and efficiency.

Significance and Claims

The paper claims that AbstainGNN addresses a significant gap in graph data mining by providing the first principled framework for graph classification with abstention. By moving beyond heuristic designs to a theory-driven approach grounded in generalization bounds, the method offers a more reliable decision-making mechanism for GNNs. The authors emphasize that their approach is particularly valuable for safety-critical applications where the cost of an incorrect decision outweighs the cost of abstention. The work validates that explicitly modeling graph structural information and reducing intra-class variance are key to achieving superior coverage-risk trade-offs in graph classification tasks.