How do Decisions Emerge across Layers in Neural Models? Interpretation with Differentiable Masking (2020)

Abstract: Attribution methods assess the contribution of inputs to the model prediction. One way to do so is erasure: a subset of inputs is considered irrelevant if it can be removed without affecting the prediction. Though conceptually simple, erasure’s objective is intractable and approximate search remains expensive with modern deep NLP models. Erasure is also susceptible to the hindsight bias: the fact that an input can be dropped does not mean that the model ‘knows’ it can be dropped. The resulting pruning is over-aggressive and does not reflect how the model arrives at the prediction. To deal with these challenges, we introduce Differentiable Masking. DiffMask learns to mask-out subsets of the input while maintaining differentiability. The decision to include or disregard an input token is made with a simple model based on intermediate hidden layers of the analyzed model. First, this makes the approach efficient because we predict rather than search. Second, as with probing classifiers, this reveals what the network ‘knows’ at the corresponding layers. This lets us not only plot attribution heatmaps but also analyze how decisions are formed across network layers. We use DiffMask to study BERT models on sentiment classification and question answering.

## introduction deep neural networks have become standard tools in nlp demonstrating impressive improvements over traditional approaches on many tasks @xcite . their power typically comes at the expense of interpretability, which may prevent users from trusting predictions @xcite , makes it hard to detect model or data deficiencies (gururangan et al., 2018; kaushik and 1 source code available at https://github.com/ nicola-decao/diffmask gated input model model with gated input figure 1 : diffmask: hidden states up to layer ℓ from a model (top) are fed to a classifier g that predicts a mask z. we use this to mask the input and re-compute the forward pass @xcite . the classifier g is trained to mask the input as much as possible without changing the output (minimizing a divergence d ). @xcite or verify that a model is fair and does not exhibit harmful biases @xcite . these challenges have motivated work on interpretability, both in nlp and generally in machine learning; see @xcite and jacovi and @xcite for reviews. in this work, we study post hoc interpretability where the goal is to explain the prediction of a trained model and to reveal how the model arrives at the decision. this goal is usually approached with attribution methods @xcite @xcite , which explain the behavior of a model by assigning relevance to inputs. one way to perform attribution is to use erasure where a subset of features (e.g., input tokens) is considered irrelevant if it can be removed without affecting the model prediction @xcite . the advantage of erasure is that it is conceptually simple and optimizes a well-defined objective. this contrasts with most other attribution methods which rely on heuristic rules to define feature salience; for example, attention-based attri-(f) our diffmask non-amortized. figure 2 : question answering token attribution: (b) and (c), are misleading (i.e., not faithful) as they attribute the prediction mostly to the answer span itself (underlined). our method (d) reveals that the model pays attention to other named entities and the predicate 'practice' in both sentences. predictions of the path-based methods (a) are more spread-out. exact search (e) as well as approximate search (f) leads to pathological attributions. despite its conceptual simplicity, subset erasure is not commonly used in practice. first, it is generally intractable, and beam search @xcite or leave-one-out estimates @xcite are typically used instead. these approximations may be inaccurate. for example, leaveone-out can underestimate the contribution of features due to saturation @xcite . more importantly, even these approximations remain very expensive with modern deep (e.g., bertbased; @xcite models, as they require multiple computation passes through the model. second, the method is susceptible to the hindsight bias: the fact that a feature can be dropped does not mean that the model 'knows' that it can be dropped and that the feature is not used by the model when processing the example. this results in over-aggressive pruning that does not reflect what information the model uses to arrive at the decision. the issue is pronounced in nlp tasks (see @xcite , though it is easier to see on an artificial example (figure 3a ). a model is asked to predict if there are more 8s than 1s in the sequence. the erasure attributes the prediction to a single 8 digit, as this reduced example yields the same decision as the original one. however, this does not reveal what the model was relying on: it has counted digits 8 and 1 as otherwise, it would not have achieved the perfect score on the test set. we propose a new method, differentiable masking (diffmask), which overcomes the aforementioned limitations and results in attributions that are more informative and help us understand how the model arrives at the prediction. diffmask relies on learning sparse stochastic gates (a.k.a., masks), guaranteeing that the information from the maskedout inputs does not get propagated while maintaining end-to-end differentiability without having to resort to reinforce @xcite . the decision to include or disregard an input token is made with a simple model based on intermediate hidden layers of the analyzed model (see figure 1 ). first, this amortization circumvents the need for combinatorial search making the approach efficient at test time. second, as with probing classifiers @xcite , this reveals whether the network 'knows' at the corresponding layer what input tokens can be disregarded. during training inputs are truly masked whenever we sample zeros. after training, attribution scores correspond to the expectation of sampling non-zeros. contributions we introduce diffmask, a technique addressing limitations of attribution-based methods (especially erasure and its approximations), and demonstrate that it is stable and faithful to the analyzed models. we then use this technique to analyze bert-based models fined-tuned on sentiment classification and question answering. ## method we aim to understand how a trained model processes an input (i.e., a sequence of embedded to-kens) to produce an output (e.g., a vector of class probabilities). first, for an input @xmath0, . . . , x n , we obtain the output @xmath1, . . . , h (l) , where h (0) = x. we then probe the model using a shallow interpreter network which takes hidden states up to a certain layer ℓ and outputs a binary mask @xmath2, . . . , z n indicating which input tokens are necessary and which can be disregarded. to assess whether the masked input @xmath3, . . . , xn is sufficient, we re-feed the model with it and compute the output ŷ = f (x). as long as ŷ approximates the original output y well, we deem the inputs masked by z unnecessary. masking, however, as in multiplication by zero, makes a strong assumption about the geometry of the feature space, in particular, it assumes that the zero vector bears no information. instead, we replace some of the inputs by a learned baseline vector b, i.e., @xmath4 see figure 1 for an overview. the interpreter model consists of l + 1 classifiers, the ℓth of which conditions on the stack of hidden states up to h (ℓ) to predict binary 'votes' that is, in order to deem x i unnecessary, it is sufficient to do so based on any subset of hidden states up until h (ℓ) . clearly, there is no direct supervision to estimate the parameters φ of the probe and the baseline b, thus we borrow erasure's objective: namely, we train the probe to mask-out as many input tokens as possible constrained to keeping f (x) ≈ f (x). since often, the output of f parameterizes a likelihood (e.g., a categorical distribution), we formulate the constraint in terms of a divergence d between the two functions' outputs. we cast this, rather naturally, in the language of constrained optimization. objective a practical way to minimize the number of non-zeros predicted by g is minimizing the l 0 'norm'. 2 thus, our l 0 loss is defined as the total number of positions that are not masked: ## experiments the goal of this work is to uncover a faithful interpretation of an existing model, i.e. revealing, as accurately as possible, the process by which the model arrives at the prediction. human-provided labels, such as human rationales @xcite , will not help us in demonstrating this, as humans cannot judge if an interpretation is faithful (jacovi and goldberg, 2020). more precisely, human-provided labels do not show how the model behaves -e.g., annotations of what parts of the input are relevant for solving a particular task do not constitute a guarantee that a model relies on those parts more than others when making a prediction. when we evaluate an attribution method by comparing its outputs with human annotations, we are not measuring whether it provides faithful attributions but only if they are plausible according to humans. this goes against our goals as we aim to use the interpretation method to detect model deficiencies, which are usually cases where the model does not behave like humans. the ground-truth explanations of how a model makes certain predictions depend not only on the data but also on the model, and, unfortunately, are generally not known for real tasks and with complex models. this makes the evaluation and comparison of attribution methods non-trivial. our strategy is to i) show the effectiveness of diffmask in a controlled setting (i.e., a toy task) where ground-truth is available; ii) test the effectiveness of our relaxation for learning discrete masks (on a real model for sentiment classification); and iii) demonstrate that the method is stable and models behave the same when masking is ap- ## related work while we motivated our approach through its relation to erasure, an alternative way of looking at our approach is considering it as a perturbationbased method. this recently introduced class of attribution methods @xcite @xcite , instead of erasing input, injects noise. besides back-propagation and attention-based methods discussed in the introduction, another class of interpretation methods @xcite @xcite builds on prior work in cooperative game theory (e.g., shapley value of @xcite . these methods are not trivial to apply to a new model, as they are architecturespecific. their hierarchical versions (e.g., @xcite also make a strong assumption about the structure of interaction (e.g., forming a tree) which may affect their faithfulness. also @xcite share some similarities to our work as they also do amortization but use the gumbel softmax trick @xcite to approximate minimal subset selection. they assume that the subset contains exactly k elements where k is a hyperparameter. moreover, their explainer is a separate model predicting input subsets, rather than a 'probe' on top of the model's hidden layers, and hence cannot be used to reveal how decisions are formed across layers. transformed-based models. for example, tenney et al. (2019) and van aken et al. (2019) probed bert layers for a range of linguistic tasks, while hao et al. (2019) analyzed the optimization surface. rogers et al. (2020) provides a comprehensive overview of recent bert analysis papers. there is a stream of work on learning interpretable models by means of extracting latent rationales @xcite . some of the techniques underlying diffmask are related to that line of work. they employ stochastic masks to learn an interpretable model, which they train by minimizing a downstream loss subject to constraints on l 0 , whereas we employ stochastic masks to interpret an existing model, and for that, we minimize l 0 subject to constraints on that model's output distribution. in our very recent work @xcite , we also employ stochastic masks and l 0 regularization for analyzing graph neural networks. we learn which edges are relevant in multi-hop question answering and graph-based semantic role labeling @xcite . ## conclusion we have introduced a new post hoc interpretation method which learns to completely remove subsets of inputs or hidden states through masking. we circumvent an intractable search by learning an end-to-end differentiable prediction model. to overcome the hindsight bias problem, we probe the model's hidden states at different depths and amortize predictions over the training set. faithfulness is validated in a controlled experiment pointing more clearly to some flaws of other attribution methods. we used our method to study bert-based models on sentiment classification and question answering. diffmask sheds light on what different layers 'know' about the input and where information about the prediction is stored in different layers.