PaperlayerView exact revision on arXiv

Source-linked AI summary

Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism

Siqi Miao, Miaoyuan Liu, Pan Li

arXiv:2201.12987v3cs.LG

TL;DR

Graph learning needs interpretations that identify predictive graph patterns while avoiding unstable or spuriously correlated features. GSAT uses stochastic attention grounded in the information bottleneck principle to constrain information and select task-relevant subgraphs. Across eight datasets, it improves interpretation AUC by up to 20%↑ and prediction accuracy by up to 5%↑ over state-of-the-art methods.

  • Problem

    Graph-learning interpretations can be unstable or spuriously correlated, while inherently interpretable models raise concerns about prediction-accuracy trade-offs.

  • Method

    GSAT injects stochasticity into attention to block label-irrelevant information and uses reduced stochasticity to select label-relevant subgraphs.

  • Results

    20%↑ in interpretation AUC and 5%↑ in prediction accuracy are the maximum improvements GSAT achieves over state-of-the-art methods across 8 datasets.

  • Takeaways & Limitations

    GSAT provides inherently interpretable and well-generalizable graph models while avoiding potentially biased sparsity, continuity, or connectivity assumptions.

  • Takeaways & Limitations

    Post-hoc interpretation methods may overfit label-irrelevant or spuriously correlated features and can be sensitive to the pretrained model, motivating joint training.

Abstract

from arXiv · show

Interpretable graph learning is in need as many scientific applications depend on learning models to collect insights from graph-structured data. Previous works mostly focused on using post-hoc approaches to interpret pre-trained models (graph neural networks in particular). They argue against inherently interpretable models because the good interpretability of these models is often at the cost of their prediction accuracy. However, those post-hoc methods often fail to provide stable interpretation and may extract features that are spuriously correlated with the task. In this work, we address these issues by proposing Graph Stochastic Attention (GSAT). Derived from the information bottleneck principle, GSAT injects stochasticity to the attention weights to block the information from task-irrelevant graph components while learning stochasticity-reduced attention to select task-relevant subgraphs for interpretation. The selected subgraphs provably do not contain patterns that are spuriously correlated with the task under some assumptions. Extensive experiments on eight datasets show that GSAT outperforms the state-of-the-art methods by up to 20%$\uparrow$ in interpretation AUC and 5%$\uparrow$ in prediction accuracy. Our code is available at https://github.com/Graph-COM/GSAT.

1. Introduction

Graph learning needs interpretations that reveal predictive patterns without sacrificing accuracy or stability. GSAT addresses these concerns with an information-bottleneck attention mechanism and reports stronger interpretation and prediction performance across eight datasets.

  • Motivation: Scientific graph-learning applications need models that reveal patterns inducing predictions, not only accurate predictions.Examples include discovering predictive patterns in scientific data and identifying molecular functional groups to guide experiments.
  • Problem: GNNs’ nonlinear entanglement of irregular graph features makes their predictive patterns difficult to identify.Post-hoc methods therefore search pretrained models for influential input subgraphs.
  • Problem: Inherently interpretable models face concerns about prediction-interpretability trade-offs and potentially unfaithful attention weights.These concerns have limited investigation of inherent interpretability for graph learning.
  • Approach: GSAT injects stochasticity into attention to constrain information flow, retain label-relevant information, and select relevant subgraphs through reduced stochasticity.It is grounded in the information bottleneck principle and avoids potentially biased size or connectivity assumptions.
  • Implications: GSAT can remove spurious correlations under certain assumptions and may improve both interpretation and prediction accuracy when initialized from a pretrained model.Visualizations show performance on varied-size digit subgraphs and disconnected motif subgraphs.
  • Results: 20%↑ in interpretation AUC and 5%↑ in prediction accuracy are the largest reported improvements over state-of-the-art methods across 8 datasets.GSAT also achieves state-of-the-art molhiv performance among models without manually designed expert features.

2. Preliminaries

The paper introduces graph-learning notation, label-relevant subgraphs, GNNs, attention mechanisms, and mutual-information feature selection. It motivates information constraints because graph features and interpretable counterparts occupy irregular spaces.

  • Graph and GNN Basics: An attributed graph is represented as G = (A, X), with adjacency matrix A, node attributes X, node set V, and edge set E.The paper focuses on graph-level prediction from IID graph-label samples.
  • Core Concepts: A label-relevant subgraph is a subgraph of the input graph that mostly supports the graph label.The supplied passage introduces the term but does not provide its completed formal definition.
  • Attention: Attention mechanisms in GNNs may produce low-fidelity interpretations because multiple edge weights must be combined over irregular graph structures.GSAT belongs to the attention category that learns unnormalized weights between 0 and 1.
  • Graph and GNN Basics: GNNs encode graph-structured data into node or graph representations by initializing node states from attributes and aggregating neighbor representations.Graph representations are commonly formed by aggregating node representations.
  • Related Concepts: L2X selects a fixed number of regular-space features by maximizing mutual information with labels, whereas GSAT uses information constraints for irregular graph features.Directly applying L2X to graph features may yield subpar performance.

3. Graph Learning Interpretation via GIB

Graph information bottleneck (GIB) selects label-relevant subgraphs without imposing potentially biased size or connectivity constraints. The section also identifies optimization, overfitting, and initialization problems in post-hoc interpretation methods.

  • GIB-based objective: Graph inputs contain critical subgraphs with varying sizes, irregular structures, and potentially scattered functional groups.Molecular examples range from small groups such as -Cl to large groups such as -C10H9.
  • GIB-based objective: GIB selects subgraphs by maximizing I(GS; Y) subject to I(GS; G) ≤γ, rather than imposing size or connectivity constraints.The selected subgraph inherits only the most indicative information from G for predicting Y.
  • Post-hoc limitations: Post-hoc methods optimize a subgraph selector while keeping the pre-trained predictor fixed, so they may not jointly optimize the information-constrained model.Their procedure is described as a single-step projection that minimizes the induced mutual-information decrease.
  • Post-hoc limitations: A fixed predictor may misestimate the information in selected subgraphs because their joint distribution with labels differs from that of the original graphs.This can mislead selector optimization and prevent selection of subgraphs that genuinely indicate Y.
  • Post-hoc limitations: Aggressive post-hoc optimization can select features that distinguish training labels but are irrelevant or spuriously correlated at the population level.On Mutag, further loss reduction after initially good explanations substantially degraded interpretation performance, complicating stopping decisions without ground-truth explanations.
  • Post-hoc limitations: Post-hoc interpretability can vary substantially with the pre-trained model initialization, motivating joint training of fθ ◦gφ under GIB.Figure 5 reports interpretation performance and training losses across random seeds and pre-trained models with different seeds.

4. Stochastic Attention Mechanism for GIB

GSAT formulates graph subgraph selection as an information bottleneck, using stochastic attention to suppress task-irrelevant information while retaining label-relevant structure. Its stochasticity-reduced attention supplies interpretations, and theory connects the objective to spurious-correlation removal and guaranteed interpretability under stated assumptions.

  • 4.1. A Tractable Objective for GIB: GSAT learns an extractor that defines a distribution over subgraphs, injecting stochasticity to block label-irrelevant information while retaining label-relevant information for prediction.The extractor gφ is represented as Pφ(GS|G), and the predictor models Y from the extracted graph.
  • 4.1. A Tractable Objective for GIB: The GSAT objective combines information about labels with a penalty on information retained from the input graph through a variational GIB bound.The method introduces variational distributions for P(Y|GS) and the marginal P(GS), then optimizes the resulting bound.
  • 4.2. GSAT and Stochastic Attention Mechanism: Stochastic attention samples each edge weight from a Bernoulli distribution and forms the selected adjacency matrix by entry-wise masking of the input adjacency matrix.The Gumbel-softmax reparameterization makes gradients with respect to edge probabilities computable.
  • 4.3. The Interpretation Mechanism of GSAT: Interpretation uses edges with reduced stochasticity, ranking them by p_uv and selecting the highest-ranked edges under an optional budget.Classification pressure reduces stochasticity on task-relevant subgraphs, while the regularization encourages stochasticity broadly.
  • 4.3. The Interpretation Mechanism of GSAT: GSAT avoids sparsity, size, and connectivity constraints, and moving the marginal-regularization parameter r away from zero can provide more robust interpretations.The paper reports that this remains advantageous even when label-relevant subgraphs satisfy those assumptions, with potentially greater gains for disconnected or varying-size patterns.
  • 4.5. Guaranteed Spurious Correlation Removal: Under theorem assumptions, the optimal GIB solution is G*S, which determines Y and excludes spurious backdoor correlations from the selected interpretation.The assumptions include a deterministic invertible relation with independent randomness and a subgraph G*S contained in G.
  • 4.5. Guaranteed Spurious Correlation Removal: With a pre-trained model, GSAT can improve both interpretation and prediction accuracy, while experiments show it almost never harms the original prediction performance.It also often achieves better interpretation than training the GNN from scratch.

5. Other Related Works

Related work largely interprets GNNs post hoc, while fewer methods pursue intrinsic interpretation. The cited intrinsic alternatives can rely on graph transformations or assumptions that make them more complicated or restrictive than GSAT.

  • Post-hoc interpretation methods: Most GNN interpretation methods are post hoc and search for prediction-influencing subgraphs using combinatorial or other specialized procedures.Examples include reinforcement learning, Monte Carlo tree search, graphical models, Granger causality, and HSIC-based approaches.
  • Intrinsic interpretation methods: Intrinsic interpretation has received less attention than post-hoc interpretation in graph learning.The paper identifies DIR as a recent intrinsic method designed to avoid overfitting spurious correlations and capture invariant rationales.
  • Intrinsic interpretation methods: DIR iteratively breaks graphs into subgraphs and assembles them back into graphs during training, making it more complicated than GSAT.This comparison concerns training procedure complexity rather than interpretation quality.

6. Experiments

Experiments evaluate GSAT across datasets, baselines, interpretability, prediction, invariant learning, and ablations. GSAT generally improves interpretation while maintaining or improving prediction performance.

  • Baselines and Setup: GSAT is evaluated against post-hoc interpreters, inherently interpretable models, prediction backbones, and invariant-learning methods across eight datasets.Interpretation uses explanation ROC AUC, while prediction uses classification ROC AUC or accuracy; results are averaged over 10 random-seed tests.
  • Interpretability Results: 9%↑ average and up to 20%↑ interpretation AUC improvements are reported over baselines, with smaller variance indicating more stable interpretations.Fine-tuning pre-trained models as GSAT* can further improve interpretation performance.
  • Prediction Results: GSAT significantly outperforms GIN on Spurious-Motif and remains comparable on other datasets, supporting interpretation without prediction decay.With PNA, GSAT improves performance by about 1−5% on the first three datasets; GSAT* reaches SOTA on molhiv without expert knowledge.
  • Invariant Learning Results: GSAT significantly outperforms invariant-learning methods for removing spurious correlations without using causality analysis.The comparison uses the backbone model from the prior invariant-learning work.
  • Ablation Study: Removing stochasticity or the regularization term significantly reduces performance, with deterministic attention causing the largest drop.Stochastic attention without regularization performs better than baselines but worse than GSAT and has large variance.
  • Ablation Study: Information regularization consistently outperforms sparsity-driven ℓ1-norm regularization, while results often worsen when r approaches 0.The best performance often occurs when r ∈[0.5, 0.9].

7. Conclusion

The conclusion presents GSAT as an information-bottleneck attention mechanism for interpretable graph learning. It uses stochasticity to suppress irrelevant information and reduced stochasticity to identify relevant subgraphs.

  • Conclusion: GSAT injects stochasticity into attention to block label-irrelevant information and selects label-relevant subgraphs through reduced attention stochasticity.The approach is grounded in the information bottleneck principle.
  • Conclusion: GSAT avoids sparsity, continuity, and other potentially biased assumptions about the selected graph patterns without performance decay.Its rationale constrains information rather than imposing a specific subgraph structure.
  • Conclusion: GSAT can remove spurious correlations and thereby support better model generalization.The paper also identifies a potential optimization issue affecting post-hoc interpretation methods.
  • Variational Bounds: The GIB objective is optimized using variational bounds for its mutual-information terms.The method introduces variational approximations for the conditional label distribution and graph-summary marginal.

C.1. Details of the Datasets

The dataset suite spans molecular, synthetic motif, superpixel, sentiment, and OGBG molecular graph tasks. Several datasets provide explanation labels, while others support prediction-only evaluation.

  • Datasets: Mutag uses molecular graphs with atom and chemical-bond structure, and labels -NO2 and -NH2 as ground-truth explanations.Each graph has a binary mutagenic-effect label.
  • Datasets: BA-2Motifs attaches house-like or five-node cycle motifs to Barabási-Albert graphs, with the motif determining the class label and explanation.The two motif types correspond to the two classes.
  • Datasets: Spurious-Motif separates label-determining motifs from bases that are spuriously correlated with labels, with correlation controlled by b = 0.5, 0.7, or 0.9.Testing randomly attaches motifs and bases to assess overfitting to the spurious correlation.
  • Datasets: MNIST-75sp converts images into superpixel graphs, where nonzero-pixel nodes are explanations and explanation subgraphs have different sizes.Edges are based on spatial distance between superpixel centers.
  • Datasets: Graph-SST2 introduces degree shifts by training and validating on graphs with higher average node degree, while explanation labels are unavailable.The study therefore evaluates prediction and provides interpretation visualizations.
  • Datasets: OGBG molecular datasets lack ground-truth explanation labels, so GSAT is evaluated on prediction performance only.The suite includes molhiv, molbace, molbbbp, molclintox, moltox21, and molsider.

C.2.3. BASELINE INTERPRETABLE METHODS/MODELS

The baseline setup covers post-hoc and intrinsic interpretable methods, with dataset-specific attention choices and controlled training randomness. IB-subgraph is substantially more computationally expensive than GSAT.

  • Interpretability Baselines: Interpretability baselines include GNNExplainer, PGExplainer, GraphMask, DIR, and IB-subgraph.The set compares post-hoc methods with inherently interpretable models.
  • Baseline Configuration: Post-hoc methods use recommended or tuned hyperparameters, while real-world post-hoc tuning lacks a clear metric.GNNExplainer is tuned using validation interpretation ROC AUC.
  • Baseline Configuration: DIR uses recommended settings for most datasets but tunes its causal-ratio parameter for Ba-2Motif and Mutag and its α parameter in the stated setting.The α grid ranges from 10 to 0.000001.
  • Efficiency: IB-subgraph requires approximately 40 hours for 100 epochs on Spurious-Motif and 150 hours on OGBG-Molhiv, compared with approximately 15 minutes for GSAT on OGBG-Molhiv.IB-subgraph uses five random seeds because of its inefficiency.
  • Attention Configuration: Node-level attention is used for Mutag, MNIST-75sp, and OGBG-Mol datasets, while edge attention is used for the other datasets.The choice reflects observed usefulness for molecular datasets and larger graphs.

C.4. Further Supplementary Experiments

Supplementary experiments test GSAT on disconnected critical subgraphs and compare it with explanation baselines and design variants. Across these evaluations, GSAT continues to outperform competing methods, while ablations examine β, stochasticity, and its information constraint.

  • Disconnected critical subgraphs: GSAT and GraphMask are evaluated on graphs with disconnected critical subgraphs, generated by attaching two or three house motifs to BA-model base graphs.The number of house motifs determines the graph class, and both methods use the Ba-2Motifs training settings.
  • Interpretation visualizations: Figures 9 and 10 visualize GSAT-discovered label-relevant subgraphs for Ba-2Motifs and Mutag, respectively.The Mutag visualizations show the -NO2 and -NH2 ground-truth explanations.
  • Baseline comparisons: GSAT still significantly outperforms PGExplainer and GNNExplainer in interpretation ROC AUC.Table 6 directly compares GSAT with the interpretation ROC AUC reported by Luo et al. (2020).
  • Baseline comparisons: GSAT greatly outperforms DIR when paired with the backbone model used in DIR.Tables 4 and 7 provide the direct comparisons.
  • Ablation studies: Ablations study β, stochasticity, and the information constraint in Eq. (9), observing the same trends discussed in Sec. 6.3.Table 8 uses PNA as the backbone, while Figure 8 evaluates the information constraint on Spurious-Motif b = 0.7 and b = 0.9.

D. Interpretation Visualization

GSAT visualizations show label-relevant subgraphs across multiple datasets, with edge transparency encoding normalized attention weights. The figures also identify dataset-specific graph classes, sentiment groups, or raw-versus-attention views where applicable.

  • Attention-weight visualization: Edge transparency represents GSAT’s normalized attention weights, obtained by rescaling learned edge weights to [0, 1].For each graph, the minimum and maximum learned edge weights define the rescaling range.
  • Dataset visualizations: Figures 11–13 visualize GSAT-discovered subgraphs for Spurious-Motif at b = 0.5, 0.7, and 0.9, with pink nodes marking ground-truth explanations.Each row represents a graph class.
  • Dataset visualizations: Figure 14 visualizes GSAT-discovered label-relevant subgraphs for OGBG-Molhiv, with each row representing a graph class.
  • Dataset visualizations: Figure 15 visualizes Graph-SST2 subgraphs, with negative-sentiment sentences in the top two rows and positive-sentiment sentences in the bottom two.
  • Dataset visualizations: Figure 16 contrasts MNIST-75sp raw images in the first row with GSAT’s normalized attention weights in the second row.
Loading 2201.12987v3…