Robust Ante-hoc Graph Explainer using Bilevel Optimization

TL;DR

The paper addresses the need for explainable graph predictions in high-stakes domains by proposing RAGE, a robust ante-hoc graph explainer that learns edge-level explanations jointly with a GNN using bilevel optimization. The method optimizes an edge-influence matrix $Z$ so that the adjacency is effectively $A_Z = Z \odot A$, and employs an inner loop to train the GNN parameters $\theta$ and an outer loop to update the explainer parameters $\Phi$ guiding $Z$. This yields compact, discriminative explanations while maintaining or improving predictive accuracy, demonstrated across eight molecular graph datasets, including four real-world benchmarks and three semi-synthetic ground-truth scenarios. RAGE achieves competitive or superior performance in both graph classification and explainability metrics, and shows enhanced reproducibility of predictor behavior when explanations are used to reproduce decisions. The work highlights the potential of bilevel, ante-hoc explainers to provide robust, reproducible insights for graph-based predictions with practical impact in molecular discovery and related domains.

Abstract

Explaining the decisions made by machine learning models for high-stakes applications is critical for increasing transparency and guiding improvements to these decisions. This is particularly true in the case of models for graphs, where decisions often depend on complex patterns combining rich structural and attribute data. While recent work has focused on designing so-called post-hoc explainers, the broader question of what constitutes a good explanation remains open. One intuitive property is that explanations should be sufficiently informative to reproduce the predictions given the data. In other words, a good explainer can be repurposed as a predictor. Post-hoc explainers do not achieve this goal as their explanations are highly dependent on fixed model parameters (e.g., learned GNN weights). To address this challenge, we propose RAGE (Robust Ante-hoc Graph Explainer), a novel and flexible ante-hoc explainer designed to discover explanations for graph neural networks using bilevel optimization, with a focus on the chemical domain. RAGE can effectively identify molecular substructures that contain the full information needed for prediction while enabling users to rank these explanations in terms of relevance. Our experiments on various molecular classification tasks show that RAGE explanations are better than existing post-hoc and ante-hoc approaches.

Figures (8)

• Figure 1: An example of explanations generated by our approach (RAGE) on molecular graphs from the synthetic dataset LACTAM. Subfigure (a) shows an input molecule, subfigure (b) highlights the ground-truth explanation and subfigure (c) illustrates the edge weights assigned by RAGE. Our method correctly identified the subgraph corresponding to the ground-truth lactam functional group.
• Figure 2: (a) Post-hoc models generate explanations for a pre-trained GNN classifier using its predictions. (b) Ante-hoc models, as our approach, learn GNNs and explanations jointly. This enables ante-hoc models to identify GNNs that are both explainable and accurate.
• Figure 3: Illustration of an edge-based ante-hoc explainer that uses bilevel optimization. Explainer generates an explanation graph from the input graph by assigning an influence value to each edge. Edge influences are incorporated to edge weights on the explanation graph, the input of GNN Classifier. The inner problem optimizes GNN Classifier with $T$ iterations, while the outer problem updates Explainer using gradients from inner iterations. The dotted edges in the explanation graph show that they do not influence the classification, while others have different degrees of influence.
• Figure 4: Comparisons between the absolute weights assigned by RAGE and GSAT on ground-truth explanation (foreground) edges and the remaining (background) edges. The plots illustrate the density of edge weights across 3 synthetic molecular datasets. Densities of foreground edge weights are shown in green while those of background edge weights are shown in pink. The density of foreground edges by RAGE on BenLac is barely visible because most foreground weights are close to 1.0.
• Figure 5: Reproducibility of various ante-hoc and post-hoc explainers on Mutagenicity. We extracted explanation subgraphs by keeping the edges with the highest weights and trained a classifier on the resulting dataset comprising these subgraphs. RAGE and GSAT, being ante-hoc, perform better than other post-hoc explainers under such perturbation.