PaperlayerView exact revision on arXiv

Source-linked AI summary

Interpreting Graph Neural Networks for NLP With Differentiable Edge Masking

Michael Sejr Schlichtkrull, Nicola De Cao, Ivan Titov

arXiv:2010.00577v3cs.CLcs.LGstat.ML

TL;DR

GNNs support NLP models with structural information, but their predictions and information flow are difficult to interpret. GRAPHMASK learns amortized, differentiable hard gates that identify unnecessary edges at each layer. Applied to question answering and semantic role labeling, it enables analysis of retained edges and paths while removing many edges without deteriorating model performance.

  • Problem

    GNNs have seen little interpretability work focused on which graph components contribute to NLP predictions, despite the need to understand their encoded linguistic information.

  • Method

    GRAPHMASK learns an amortized erasure function with sparse stochastic gates to predict which edges at each layer can be discarded without changing model predictions.

  • Results

    The method analyzes GNNs for question answering and semantic role labeling, identifying relied-upon edge types and paths while retaining predictions after substantial edge removal.

  • Takeaways & Limitations

    GRAPHMASK provides scalable, model-agnostic, and faithful analyses of how GNN edges and paths influence NLP predictions.

  • Takeaways & Limitations

    Erasure search remains infeasible in practical scenarios, and the method’s stated interpretation criteria include identifying relevant paths, tractability, and faithfulness.

Abstract

from arXiv · show

Graph neural networks (GNNs) have become a popular approach to integrating structural inductive biases into NLP models. However, there has been little work on interpreting them, and specifically on understanding which parts of the graphs (e.g. syntactic trees or co-reference structures) contribute to a prediction. In this work, we introduce a post-hoc method for interpreting the predictions of GNNs which identifies unnecessary edges. Given a trained GNN model, we learn a simple classifier that, for every edge in every layer, predicts if that edge can be dropped. We demonstrate that such a classifier can be trained in a fully differentiable fashion, employing stochastic gates and encouraging sparsity through the expected $L_0$ norm. We use our technique as an attribution method to analyze GNN models for two tasks -- question answering and semantic role labeling -- providing insights into the information flow in these models. We show that we can drop a large proportion of edges without deteriorating the performance of the model, while we can analyse the remaining edges for interpreting model predictions.

1 INTRODUCTION

The paper introduces GRAPHMASK, a scalable post-hoc method for interpreting GNN predictions by identifying edges that can be discarded without changing outputs. It addresses erasure search’s tractability and hindsight-bias problems through amortized, differentiable edge masking.

  • GNN interpretability is needed to understand encoded linguistic information, support trust and bias discovery, and simplify error analysis.
  • Erasure search provides discrete, faithful attributions but is computationally infeasible for practical GNNs.Approximations that remove one feature at a time remain prohibitively expensive and can underestimate contributions because of saturation.
  • GRAPHMASK learns an erasure function that predicts whether each edge at each layer should be retained, producing interpretable subgraphs whose omitted edges do not influence predictions.The method uses sparse stochastic gates to enable gradient-based optimization.
  • Amortizing parameter learning over training data makes masking decisions from graph representations without access to final predictions or gold labels, addressing hindsight bias.Per-example optimization can aggressively prune non-superfluous edges when alternative smaller subgraphs yield similar predictions.
  • The authors apply GRAPHMASK to semantic role labeling and multi-hop question answering to analyze edge and path information flow.They present the method as applicable potentially to any end-to-end neural model containing a GNN.

2 RELATED WORK

Related interpretability methods for GNNs include per-example erasure, gradients, attention or gate scores, perturbation methods, and feature-decomposition approaches. The paper positions GRAPHMASK as a scalable alternative that uses a trained masking model and hard edge decisions.

  • GNNExplainer learns a soft edge-erasure function separately for each example, lacks a guarantee that gated edges do not affect predictions, and is vulnerable to hindsight bias.
  • Attention or gate scores are not universally available and may scale messages rather than filter them, making them unreliable importance measures for GNN edges.
  • Shapley-style decomposition is computationally prohibitive for realistic GNN graphs, while LIME’s local models do not trivially identify useful paths or long-distance structure.
  • Perturbation-based interpretability methods equate feature importance with prediction sensitivity to perturbations and may use a secondary trained model.
  • Related differentiable masking work applies sparse stochastic gates and L0 regularization to transformer tokens, while concurrent GNN work uses mutual information and local binary-concrete classifiers.

3 METHOD

GRAPHMASK learns a shared edge-erasure function that identifies sparse, prediction-preserving subgraphs across GNN layers. It combines amortized gating with a differentiable expected-L0 objective to make interpretation tractable and reduce hindsight bias.

  • GraphMask mechanism: Each edge is assigned a binary retain-or-replace decision, with masked messages replaced by learned baseline vectors rather than simply removed.This replacement accounts for the sensitivity of GNNs to changes in graph structure.
  • Amortized edge masking: Instead of optimizing edge gates separately for each example, GRAPHMASK learns one erasure function across datapoints and applies it to unseen examples.The function uses information available to the original model when computing each message, preventing look-ahead.
  • Toy-task comparison: The toy example contrasts amortized GRAPHMASK with erasure search, GNNExplainer, non-amortized GRAPHMASK, integrated gradients, and an information bottleneck approach.The figure evaluates whether methods identify black and blue edges in a graph-counting task.
  • Inference and interpretation: After training, GRAPHMASK runs the original model to obtain representations, computes gates for every edge and layer, and executes the sparsified model with original parameters frozen.If the sparsified model preserves the original prediction, masked messages are interpreted as superfluous.
  • GraphMask objective: GRAPHMASK seeks informative subgraphs at every layer that minimize retained edges while preserving the original model’s prediction within a tolerance.The objective measures output divergence and penalizes the number of unmasked edges.
  • Differentiable optimization: The method optimizes an otherwise nondifferentiable objective using stochastic Hard Concrete gates, which permit exact zeros and low-variance reparameterized gradients.Attribution scores are based on the expectation that sampled masks remain non-zero.

4 SYNTHETIC EXPERIMENT

The synthetic star-graph experiment tests whether GRAPHMASK identifies the edges that determine a counting prediction. Only amortized GRAPHMASK approximately recovers the gold-standard attribution, while several alternatives exploit smaller but unfaithful subgraphs.

  • Experimental setup: The task predicts whether the graph contains more edges of colour x or y, with all x and y edges useful and every other edge superfluous.Examples contain 6 to 12 leaves, and a one-layer R-GCN perfectly classifies them.
  • Baselines: GRAPHMASK is compared with erasure search, integrated gradients, an information bottleneck approach, and GNNExplainer.Thresholds convert soft gates or normalized scores into hard edge decisions for evaluation.
  • Results: Only amortized GRAPHMASK approximately replicates the gold standard, whereas erasure search, GNNExplainer, and non-amortized GRAPHMASK recall only some non-superfluous edges.These methods can reach perfect model performance using an unfaithful low-penalty subgraph.
  • Interpretation: Amortization prevents hindsight bias by avoiding example-specific overfitting to alternative subgraphs that preserve the prediction.The method learns a parametrized erasure function across a training dataset.

5 QUESTION ANSWERING

GRAPHMASK analyzes a multi-hop question-answering GNN by retaining a sparse, stable set of edges while preserving performance. The retained structure reveals layer-specific use of edge types, query-centered paths, and symmetric message exchange.

  • Model and task: The WikiHop model answers query sentences by selecting an entity from context documents, using mention nodes connected by four edge types.The model combines a query-reading BiLSTM with three shared-parameter R-GCN layers.
  • Overall sparsity: GRAPHMASK retains 27% of edges with a −0.4% accuracy change, and retained-edge decisions show high agreement across five random seeds (κ = 0.65).Most retained edges occur in the bottom layer.
  • Faithfulness: Randomly dropping 25% of the retained edges greatly harms performance, indicating that the retained subset is not interchangeable noise.The result contrasts with the small accuracy change when GRAPHMASK constructs the sparsified model.
  • Layer allocation: Without amortization, only 0.4% of retained edges are in the bottom layer and 91.0% are in the top layer, unlike the amortized attribution.Removing the bottom layer causes a −26% accuracy drop, compared with −7% for the top layer.
  • Edge types: COMPLEMENT edges dominate useful edges in the bottom layer but are always superfluous later, suggesting an initial context-pooling propagation step.COREF edges are retained mainly in harder cases not handled by surface MATCH overlap.
  • Query connections: At least one query-entity edge is retained in 92.7% of cases, including 84.1% in the bottom layer, despite query representations being concatenated to every node.This suggests the GNN still uses graph connections involving query mentions.
  • Information flow: Query-originating paths account for 11.8% of retained edges at layer 0, 42.7% at layer 1, and 73.8% at the top layer.The predicted answer is targeted by a retained edge in 99.7% of examples, while query connectivity is similar for predicted and average candidates.
  • Symmetry: Inverse edge pairs are jointly retained or discarded, with inverse retention among retained edges of 49%, 98%, and 79% in layers 0, 1, and 2.This pattern reflects the symmetric relations used by the model.

6 SEMANTIC ROLE LABELING

GRAPHMASK reveals which dependency edges and paths semantic role labeling models retain, while preserving performance with substantial edge removal. Its analyses differ across model architectures, predicate types, and path distances.

  • Paths: Table 3 reports percentages of paths with 0, 1, or 2 retained edges by path length and predicate type for both models.The LSTM+GNN model permits at most one retained edge per path because it uses a single GNN layer.
  • Model performance: The masked LSTM+GNN model changes performance by only −0.62% F1 while retaining substantially fewer edges.The supplied passage does not state the exact retained-edge proportion.
  • Model performance: 42% of edges can be removed under scalar gate thresholding, yet performance decreases by 16.1% F1.This contrasts learned hard-edge attribution with directly inspecting the model’s scalar gates.
  • Dependency types: NMOD dominates for nominal predicates, whereas SBJ and OBJ play the largest roles for verbal predicates.Several frequent dependency relations are entirely superfluous, including P, NAME, CO-ORD, CV, CONJ, HYPH, SUFFIX, and POSTHON.
  • Dependency types: For the LSTM-GNN model, 88% of retained edges point to predicted roles, while the remaining 12% mostly point to other predicates’ arguments.These proportions characterize where retained graph information is directed.
  • Paths: Reliance on retained paths decreases with predicate distance for nominal predicates in LSTM+GNN but increases with distance in the GNN-only model.Direct predicate-role connections are practically always retained, while longer paths are often useful at lower rates.

7 CONCLUSION

The paper concludes that GRAPHMASK provides faithful, scalable, model-agnostic interpretation for GNNs and identifies how edges and paths influence predictions. Its NLP applications illustrate analyses of edge types and information flow, while broader domains remain future work.

  • Conclusion: GRAPHMASK learns differentiable hard gates for every message and amortises them over training data.This supports faithfulness to the studied model and scalability to modern GNNs.
  • Conclusion: GRAPHMASK identifies both how individual edges and paths influence predictions in GNN models.The method was applied to an SRL model and a QA model from prior work.
  • Conclusion: Applications to NLP reveal which edge types the analysed models rely on and how they use paths when predicting.The paper presents these analyses as examples of what GRAPHMASK enables.
  • Conclusion: Applying GRAPHMASK to domains beyond NLP is left for future work.The paper’s demonstrated applications focus on graph-based NLP models with linguistic and structural priors.

A ERASURE FUNCTION ARCHITECTURE

The erasure-function architecture constructs edge representations and predicts differentiable hard gates, with a matrix-form alternative for faster GNN implementations. Hard Concrete sampling enables differentiable sparsity optimization, while training uses staged gate introduction and task-specific optimization choices.

  • Edge representations: Each edge representation q^(k)_u,v is formed by concatenating the source embedding, target embedding, and message representation.The supplied passages identify Layer Normalization in the architecture and define q^(k)_u,v through concatenation.
  • Efficient computation: The matrix-form alternative computes gates from source and target vertex embeddings using a bilinear product rather than enumerating all messages.It is designed for faster but less expressive adjacency-matrix GNN formulations.
  • Efficient computation: Computing edge representations for every possible vertex pair would raise complexity from O(V + E) to O(V^2).This is problematic for large graphs, motivating a faster matrix-form gate computation.
  • Hard Concrete gates: Hard Concrete assigns mass to exactly 0 and 1 while permitting differentiable reparameterized sampling.This turns the L0 objective into an expectation whose gradient can be estimated by Monte Carlo without REINFORCE or bias.
  • Hard Concrete gates: Samples are generated by applying a sigmoid-based stochastic transformation and then stretching and rectifying it into [0, 1].The location and temperature parameters control distribution skew and concentration before rectification.
  • Training: Training progressively adds gates from the top layer downward and optimizes sparsity under a development-set performance constraint.The implementation uses separate optimizers for GRAPHMASK parameters and the Lagrangian multiplier.
  • Training: Each training iteration requires at most twice the computation time of an equivalent iteration using the investigated model.GRAPHMASK performs one forward pass to compute gate values and a backward pass through the sparsified model.

D DATASETS

The experiments use CoNLL-2009 for semantic role labeling and WikiHop for question answering. The WikiHop dataset statistics are reported from prior work, while the supplied passages provide predicate counts for CoNLL-2009.

  • Semantic role labeling: CoNLL-2009 contains 179,014 training predicates, 6,390 validation predicates, and 10,498 test predicates.The dataset is used for the semantic role labeling experiments.
  • Question answering: WikiHop is used for question answering with preprocessing from De Cao et al. (2019).Table 4 reports its numbers of candidates and documents per sample and document length.

E SYNTHETIC TASK MODEL

The synthetic-task model uses a one-layer R-GCN with MLP components, while integrated gradients assigns edge attributions through scalar message gates. The approach has important multi-layer limitations, especially for long-distance information and graph-degree-changing baselines.

  • Model architecture: The synthetic task uses a one-layer R-GCN whose vertex embeddings begin with concatenated one-hot encodings of x and y.An initial MLP constructs zeroth-layer vertex embeddings before graph convolution.
  • Model architecture: Messages are aggregated by sum-pooling, and predictions are produced by a one-hidden-layer MLP from the centroid embedding.The R-GCN states have dimensionality 50, while the MLP hidden states have dimensionality 100.
  • Integrated gradients: Integrated gradients assigns edge attributions by multiplying each layer-specific message by a scalar gate and interpolating from gate value 0.This treats edges as transitioning between fully absent and fully present through partially present states.
  • Limitations: In multi-layer GNNs, integrated gradients can systematically underestimate information that travels through long-distance paths.Partially present upper-layer edges affect gradient flow and attribution to lower-layer edges during interpolation.
  • Limitations: The zero-vector baseline may be inappropriate for general GNNs because removing edges changes graph degree statistics.The synthetic task is constructed so that changes in the number of leaves alter the centroid degree, avoiding this issue for its perfect-scoring model.
  • Comparison method: The information-bottleneck adaptation uses one gate per message, predicts gate logits conditionally, and promotes sparsity through a readout-bottleneck strategy.This adapts a hidden-state attribution method by replacing dimension-specific gates with a single message gate.

H IMPLEMENTATION INVARIANCE FOR GNNS

The paper compares graph attributions across functionally equivalent networks and examines which edge types are retained or kept. Its real-world evaluation reports that masking superfluous edges preserves performance, whereas removing retained edges harms it.

  • Functionally equivalent networks: Functionally equivalent networks can receive the same graph input while producing different edge attributions after equivalent downstream MLP implementations.The setup uses a sum-aggregation GNN without an activation function followed by two functionally equivalent networks.
  • Edge-type analysis: Retained-edge distributions and edge-keeping probabilities are separated by predicate type and dependency direction.The figure distinguishes nominal and verbal predicates, with flow either from the head or toward the head, while excluding rare or almost-always-superfluous edge types.
  • Masking evaluation: Dropping edges marked superfluous by GRAPHMASK does not impact performance, but randomly dropping retained edges significantly hurts the three real-world models.Even randomly removing 25% of the retained edges causes a significant performance decrease.

K BASELINE PERFORMANCE ON QUESTION ANSWERING

On question answering, baseline attribution methods concentrate importance in upper layers even though ablations show that the bottom layer is more consequential. GRAPHMASK is evaluated alongside these baselines through layer-wise attribution analysis.

  • Layer-wise attribution: GNNExplainer and Integrated Gradients assign low attribution to the first two layers and rely primarily on the top layer.This pattern appears in the QA model’s layer-separated attribution analysis.
  • Layer ablations: -26% performance follows from dropping the bottom QA layer, compared with -7% after dropping the top layer.These ablations conflict with the baselines’ attribution pattern, which suggests low importance for the bottom layers.
  • GNNExplainer pathology: GNNExplainer can select specific top-layer edges adjacent to the predicted answer that reproduce the original prediction.This behavior mirrors answer-span-only explanations reported as a pathology of erasure search on textual data.
  • Figure 9: Figure 9 compares, by QA-model layer, the mean percentage of messages assigned attribution scores above specified levels by three techniques.The techniques are GNNExplainer, Integrated Gradients, and GRAPHMASK.
  • Integrated Gradients pathology: Integrated Gradients systematically underestimates long-distance information that must traverse half-open pseudo-gates in other layers.The resulting scores prevent meaningful comparisons of attribution between layers.
  • Faithfulness comparison: The total-attribution allocation across layers again gives GNNExplainer and Integrated Gradients low scores for the bottom layer.The discrepancy with the empirical loss from excluding that layer indicates that these baselines are unlikely to be faithful.
Loading 2010.00577v3…