Interpreting Graph Neural Networks for NLP With Differentiable Edge Masking

1 INTRODUCTION . Graph Neural Networks ( GNNs ) have in recent years been shown to provide a scalable and highly performant means of incorporating linguistic information and other structural biases into NLP models . They have been applied to various kinds of representations ( e.g. , syntactic and semantic graphs , co-reference structures , knowledge bases linked to text , database schemas ) and shown effective on a range of tasks , including relation extraction ( Zhang et al. , 2018 ; Zhu et al. , 2019 ; Sun et al. , 2019a ; Guo et al. , 2019 ) , question answering ( Sorokin & Gurevych , 2018 ; Sun et al. , 2018 ; De Cao et al. , 2019 ) , syntactic and semantic parsing tasks ( Marcheggiani & Titov , 2017 ; Bogin et al. , 2019 ; Ji et al. , 2019 ) , summarisation ( Fernandes et al. , 2019 ) , machine translation ( Bastings et al. , 2017 ) and abusive language detection in social networks ( Mishra et al. , 2019 ) . While GNNs often yield strong performance , such models are complex , and it can be difficult to understand the ‘ reasoning ’ behind their predictions . For NLP practitioners , it is highly desirable to know which linguistic information a given model encodes and how that encoding happens ( Jumelet & Hupkes , 2018 ; Giulianelli et al. , 2018 ; Goldberg , 2019 ) . The difficulty in interpreting GNNs represents a barrier to such analysis . Furthermore , this opaqueness decreases user trust , impedes the discovery of harmful biases , and complicates error analysis ( Kim , 2015 ; Ribeiro et al. , 2016b ; Sun et al. , 2019b ; Holstein et al. , 2019 ) . The latter is a particular issue for GNNs , where seemingly small implementation differences can make or break models ( Zaheer et al. , 2017 ; Xu et al. , 2019 ) . In this work , we focus on post-hoc analysis of GNNs . We are interested especially in developing a method for understanding how the GNN uses the input graph . As such , we seek to identify which edges in the graph the GNN relies on , and at which layer they are used . We formulate some desiderata for an interpretation method , seeking a technique that is : 1. able to identify relevant paths in the input graph , as paths are one of the most natural ways of presenting GNN reasoning patterns to users ; 2. sufficiently tractable to be applicable to modern GNN-based NLP models ; 3. as faithful ( Jacovi & Goldberg , 2020 ) as possible , providing insights into how the model truly arrives at the prediction . A simple way to perform interpretation is to use erasure search ( Li et al. , 2016 ; Feng et al. , 2018 ) , an approach wherein attribution happens by searching for a maximal subset of features that can be entirely removed without affecting model predictions . The removal guarantees that all information about the discarded features is ignored by the model . This contrasts with approaches which use heuristics to define feature importance , for example attention-based methods ( Serrano & Smith , 2019 ; Jain & Wallace , 2019 ) or back-propagation techniques ( Bach et al. , 2015 ; Sundararajan et al. , 2017 ) . They do not guarantee that the model ignores low-scoring features , attracting criticism in recent years ( Nie et al. , 2018 ; Sixt et al. , 2019 ; Jain & Wallace , 2019 ) . The trust in erasure search is reflected in the literature through other methods motivated as approximations of erasure ( Baehrens et al. , 2010 ; Simonyan et al. , 2014 ) , or through new attribution techniques evaluated using erasure search as ground truth ( Serrano & Smith , 2019 ; Jain & Wallace , 2019 ) . Applied to GNNs , erasure search would involve searching for the largest subgraph which can be completely discarded . Besides faithfulness considerations and conceptual simplicity , discrete attributions would also simplify the comparison of relevance between paths ; this contrasts with continuous attribution to edges , where it is not straightforward to extract and visualise important paths . Furthermore , in contrast to techniques based on artificial gradients ( Pope et al. , 2019 ; Xie & Lu , 2019 ; Schwarzenberg et al. , 2019 ) , erasure search would provide implementation invariance ( Sundararajan et al. , 2017 ) . This is important in NLP , as models commonly use highly parametrised decoders on top of GNNs ( e.g. , Koncel-Kedziorski et al . ( 2019 ) ) . While arguably satisfying criteria ( 1 ) and ( 3 ) in our desiderata , erasure search unfortunately fails on tractability . In practical scenarios , it is infeasible , and even approximations , which remove one feature at a time ( Zintgraf et al. , 2017 ) and underestimate their contribution due to saturation ( Shrikumar et al. , 2017 ) , remain prohibitively expensive . Our GRAPHMASK aims at meeting the above desiderata by achieving the same benefits as erasure search in a scalable manner . That is , our method makes easily interpretable hard choices on whether to retain or discard edges such that discarded edges have no relevance to model predictions , while remaining tractable and model-agnostic ( Ribeiro et al. , 2016a ) . GRAPHMASK can be understood as a differentiable form of subset erasure , where , instead of finding an optimal subset to erase for every given example , we learn an erasure function which predicts for every edge 〈u , v〉 at every layer k whether that connection should be retained . Given an example graph G , our method returns for each layer k a subgraph G ( k ) S such that we can faithfully claim that no edges outside G ( k ) S influence the predictions of the model . To enable gradient-based optimization for our erasure function , we rely on sparse stochastic gates ( Louizos et al. , 2018 ; Bastings et al. , 2019 ) . In erasure search , optimisation happens individually for each example . This can result in a form of overfitting where even non-superfluous edges are aggressively pruned because a similar prediction could be made using an alternative smaller subgraph ; we refer to this problem as hindsight bias . Because our interpretation method relies on a parametrised erasure function rather than an individual per-edge choice , we can address this issue by amortising parameter learning over a training dataset through a process similar to the readout bottleneck introduced in Schulz et al . ( 2020 ) . In other words , the decision to drop or keep an edge is made based on the information available in the network ( i.e. , representation of the graph nodes ) without having access to the final prediction ( or to the gold standard ) . As we demonstrate in Section 4 , this strategy avoids hindsight bias . Contributions Our contributions are as follows : • We present a novel interpretation method for GNNs , applicable potentially to any end-toend neural model which has a GNN as a component.1 • We demonstrate using artificial data the shortcomings of the closest existing method , and show how our method addresses those shortcomings and improves faithfulness . • We use GRAPHMASK to analyse GNN models for two NLP tasks : semantic role label- ing ( Marcheggiani & Titov , 2017 ) and multi-hop question answering ( De Cao et al. , 2019 ) . 2 RELATED WORK . Several recent papers have focused on developing interpretability techniques for GNNs . The closest to ours is GNNExplainer ( Ying et al. , 2019 ) , wherein a soft erasure function for edges is learned individually for each example . Unlike our method ( and erasure search ) , GNNExplainer can not as such guarantee that gated edges do not affect predictions . Furthermore , as we show in our experiments ( Section 4 ) , separate optimisation for each example results in hindsight bias and compromises faithfulness . Pope et al . ( 2019 ) ; Xie & Lu ( 2019 ) explore gradient-based methods , including gradient heatmaps , Grad-CAM , and Excitation Backpropagation . Similarly , Schwarzenberg et al . ( 2019 ) ; Baldassarre & Azizpour ( 2019 ) ; Schnake et al . ( 2020 ) apply Layerwise Relevance Propagation ( Bach et al. , 2015 ) to the GNN setting . These methods represent an alternative to GRAPHMASK , but as we have noted their faithfulness is questionable ( Nie et al. , 2018 ; Sixt et al. , 2019 ; Jain & Wallace , 2019 ) , and the lack of implementation invariance ( Sundararajan et al. , 2017 ) is problematic ( see Appendix H ) . Furthermore , significant engineering is still required to develop these techniques for certain GNNs , e.g . networks with attention as the aggregation function ( Veličković et al. , 2018 ) . Another popular approach is to treat attention or gate scores as a measure of importance ( Serrano & Smith , 2019 ) . However , even leaving questionable faithfulness ( Jain & Wallace , 2019 ) aside , many GNNs use neither gates nor attention . For those that do ( Marcheggiani & Titov , 2017 ; Veličković et al. , 2018 ; Neil et al. , 2018 ; Xie & Grossman , 2018 ) , such scores are , as we demonstrate in Section 6 , not necessarily informative , as gates can function to scale rather than filter messages . Outside of graph-specific methods , one line of research involves decomposing the output into a part attributed to a specific subset of features and a part attributed to the remaining features ( Shapley , 1953 ; Murdoch et al. , 2019 ; Singh et al. , 2019 ; Jin et al. , 2020 ) . For GNNs , the computational cost for realistic use cases ( e.g . the thousands of edges per example in De Cao et al . ( 2019 ) ) is prohibitive . LIME ( Ribeiro et al. , 2016b ) like us relies on a trained erasure model , but interprets local models in place of global models . Local models can not trivially identify useful paths or long-distance dependent pairs of edges , and as also pointed out in Ying et al . ( 2019 ) LIME can not be easily applied for large general graphs . Similarly , it is unclear how to apply integrated gradients ( Sundararajan et al. , 2017 ) to retrieve relevant paths , especially for deep GNNs operating in large graphs . Masking messages in GRAPHMASK can be equivalently thought of as adding a certain type of noise to these messages . Therefore , GRAPHMASK can be categorised as belonging to the recently introduced class of perturbation-based methods ( Guan et al. , 2019 ; Taghanaki et al. , 2019 ; Schulz et al. , 2020 ) which equate feature importance with sensitivity of the prediction to the perturbations of that feature . The closest to our model is Schulz et al . ( 2020 ) , wherein the authors like us apply a secondary , trained model to predict the relevancy of a feature in a given layer . Unlike us , this trained model has ‘ look-ahead ’ , i.e . access to layers above the studied layer , making their model vulnerable to hindsight bias . Their approach uses soft gates on individual hidden state dimension to interpolate between hidden states , Gaussian noise in order to detect important features for CNNs on an image processing task , and makes independent Gaussian assumptions on the features to derive their objective . We adapted their method to GNNs and used it as a baseline in our experiments . In our very recent work ( De Cao et al. , 2020 ) we have introduced a similar differentiable masking approach to post-hoc analysis for transformers . We used sparse stochastic gates and L0 regularisation to determine which input tokens can be dropped , conditioning on various hidden layers . Concurrently to this paper , Luo et al . ( 2020 ) have also developed an interpretability technique for GNNs relying on differentiable edge masking . Their approach uses a mutual information objective like GNNExplainer , along with local binary concrete classifiers as in GRAPHMASK . 1Source code available at https : //github.com/MichSchli/GraphMask .

Summary: This paper proposes an interpretation method for graph neural networks (GNNs) by learning parameterized, differentiable edge masks. The proposed method uses a single-layer NN classifier to predict  whether an edge can be dropped. It further uses L0 norm to encourage sparsity and amortizes parameter learning over a training dataset to avoid the hindsight bias. The method is experimented on a toy dataset (to demonstrate its faithfulness) and two real NLP tasks (multi-hop question answering and semantic role labeling). There are some interesting discussions based on the edge pruning results of these 2 tasks.