RISE: Radius of Influence based Subgraph Extraction for 3D Molecular Graph Explanation

TL;DR

This work tackles explainability for 3D GNNs in molecular graphs, where edges are defined by distance cut-offs and learning captures distance-dependent interactions. It introduces RISE, a radius-of-influence framework that reformulates explanations as directed proximity graphs and optimizes a per-atom radius mask without relaxing discrete masks, enabling exact budgeting and interpretable, atom-centered subgraphs. Across QM9 and GEOM datasets and with both invariant and equivariant backbones, RISE outperforms state-of-the-art baselines and yields subgraphs that align with chemical bonds, demonstrating both higher fidelity explanations and chemical interpretability. By grounding explanations in the physics of spatial interactions and enforcing exact budgets, RISE provides a faithful, physically meaningful framework for explaining 3D molecular GNN predictions.

Abstract

3D Geometric Graph Neural Networks (GNNs) have emerged as transformative tools for modeling molecular data. Despite their predictive power, these models often suffer from limited interpretability, raising concerns for scientific applications that require reliable and transparent insights. While existing methods have primarily focused on explaining molecular substructures in 2D GNNs, the transition to 3D GNNs introduces unique challenges, such as handling the implicit dense edge structures created by a cut-off radius. To tackle this, we introduce a novel explanation method specifically designed for 3D GNNs, which localizes the explanation to the immediate neighborhood of each node within the 3D space. Each node is assigned an radius of influence, defining the localized region within which message passing captures spatial and structural interactions crucial for the model's predictions. This method leverages the spatial and geometric characteristics inherent in 3D graphs. By constraining the subgraph to a localized radius of influence, the approach not only enhances interpretability but also aligns with the physical and structural dependencies typical of 3D graph applications, such as molecular learning.

Figures (9)

• Figure 1: Comparison with existing approaches. Existing approaches require a relaxation from binary masks to soft continuous masks, leading to inconsistencies between the optimized masks and the explanatory binary masks. These inconsistencies not only compromise explanation fidelity quantitatively but also produce chemically uninterpretable results, making the model explanation itself a black-box. On the other hand, RISE introduces a novel framework for extracting explanatory substructures based on atomic radii of influence that aim to find localized regions that capture the most important interactions decisive to the prediction under some budgets. Given an appropriate small budget, RISE can precisely extract chemical bonds and chemical bonds only. For instance, for Ethane ($\text{CH}_3\text{CH}_3$) in the figure, the radii of influence from our experiments assign the C of interest with a radius of $1.532$ and the H of interest with a radius of $1.171$. Under similar radii of influence for C atoms and H atoms, respectively, RISE extracts the precise chemical bonds and chemical bonds only: C-H ($1.171 >1.095$); C-C ($1.532 > 1.530$); all other edges have a distance greater than $1.532$ and will be masked out by RISE. Note the unit of distance is $\mathring{A}$, which is used throughout the paper unless otherwise specified.
• Figure 2: (a): 2D representation of $\text{C}_8\text{H}_{18}$---Nodes are atoms, and edges are chemical bonds. No geometric information; typically a small number of edges. (b): 3D representation of $\text{C}_8\text{H}_{18}$---Nodes are atoms with spatial locations. Edges are constructed with a specified cut-off distance, resulting in a dense graph. (c): 3D representation of $\text{C}_8\text{H}_{18}$ with all non-bonding edges removed.
• Figure 3: (a): Original 3D graphs constructed based on a common cut-off distance; this is the approach taken in most 3D GNNs schnetdimenetSphereNetSEGNN. The edges are bidirectional and dense. (b): Explanatory substructure identified by finding the radii of influence. The radii of influence are optimized (the circles shrink dynamically; see an illustration in Fig. \ref{['fig: shrinking_circles']} in Appendix \ref{['append: shrinking_circles']}) such that the critical messages that are most relevant to the prediction task are preserved. The edges are directed and more sparse.
• Figure 4: Explanatory substructure produced from experiments by different explanation methods on the Ethane molecule ($\text{CH}_3\text{CH}_3$) in the QM9 dataset QM9_dataset. The same budget (number of edges) is used for different explainers. It is obvious that only RISE yields chemically interpretable results that conform to interpretable chemical structures. It should be noted that, under larger budgets, baseline methods yield a more "chaotic" set of edges that are uninterpretable at all, whereas RISE identifies atomic regions of influence, enabling chemical interpretation. This underscores the need for interpretable 3D methods like RISE.
• Figure 5: The visualization of the quantitative results of $\alpha$ of the QM9 dataset on SEGNN when randomly masking 10% edges in different annuli. It shows that the influence of edge masking has a significant correlation with the distances, i.e., masking short edges will cause larger perturbation than long edges. Moreover, the small variance of each annulus demonstrates the comparable importance among edges with similar distances.

Theorems & Definitions (2)

• Definition 3.1
• Remark 3.2