The Out-of-Distribution Problem in Explainability and Search Methods for Feature Importance Explanations

1 Introduction . Estimating feature importance ( FI ) is a common approach to explaining how learned models make predictions for individual data points [ 51 , 46 , 34 , 57 , 36 , 11 ] . FI methods assign a scalar to each feature of an input representing its “ importance ” to the model ’ s output , where a feature may be an individual component of an input ( such as a pixel or a word ) or some combination of components . Alongside these methods , many approaches have been proposed for evaluating FI estimates ( also known as attributions ) [ 41 , 1 , 13 , 22 , 20 , 67 ] . Many of these approaches use test-time input ablations , where features marked as important are removed from the input , with the expectation that the model ’ s confidence in its original prediction will decline if the selected features were truly important . 1All supporting code for experiments in this paper is publicly available at https : //github.com/ peterbhase/ExplanationSearch . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . For instance , according to the Sufficiency metric [ 13 ] , the best FI explanation is the set of features which , if kept , would result in the highest model confidence in its original prediction . Typically the top-k features would be selected according to their FI estimates , for some specified sparsity level k. Hence , the final explanation e is a k-sparse binary vector in { 0 , 1 } d , where d is the dimensionality of the chosen feature space . For an explanation e and a model f that outputs a distribution over classes p ( y|x ) = f ( x ) , Sufficiency can be given as : Suff ( f , x , e ) = f ( x ) ŷ − f ( Replace ( x , e ) ) ŷ where ŷ = argmaxy f ( x ) y is the predicted class for x and the Replace function replaces features in x with some uninformative feature at locations corresponding to 0s in the explanation e. The Replace function plays a key role in the definition of such metrics because it defines the counterfactual input that we are comparing the original input with . Though FI explanations are often presented without mention of counterfactuals , all explanations make use of counterfactual situations [ 38 ] , and FI explanations are no exception . The only way we can understand what makes some features “ important ” to a particular model prediction is by reference to a counterfactual input which has its important features replaced with a user-specified ( uninformative ) feature . In this paper , we study several under-explored dimensions of the problem of finding good explanations according to test-time ablation metrics including Sufficiency and a related metric , Comprehensiveness , with a focus on natural language processing tasks . We describe three primary contributions below . First , we argue that standard FI explanations are heavily influenced by the out-of-distribution ( OOD ) nature of counterfactual model inputs , which results in socially misaligned explanations . We use this term , first introduced by Jacovi and Goldberg [ 24 ] , to describe a situation where an explanation communicates a different kind of information than the kind that people expect it to communicate . Here , we do not expect the model prior or random weight initialization to influence FI estimates . This is a problem insofar as FI explanations are not telling us what we think they are telling us . We propose a training algorithm to resolve the social misalignment , which is to expose models to counterfactual inputs during training , so that counterfactuals are not out-of-distribution at test time . Second , we systematically compare Replace functions , since this function plays an important role in evaluating explanations . To do so , we remove tokens from inputs using several Replace functions , then measure how OOD these ablated inputs are to the model . We compare methods that remove tokens entirely from sequences of text [ 41 , 13 ] , replace token embeddings with the zero embedding or a special token [ 34 , 57 , 1 , 65 , 57 ] , marginalize predictions over possible counterfactuals [ 68 , 30 , 63 ] , and edit the input attention mask rather than the input text . Following our argument regarding the OOD problem ( Sec . 4 ) , we recommend the use of some Replace functions over others . Third , we provide several novel search-based methods for identifying FI explanations . While finding the optimal solution to argmaxe Suff ( f , x , e ) is a natural example of binary optimization , a problem for which search algorithms are a common solution [ 43 , 49 , 3 ] , we are aware of only a few prior works that search for good explanations [ 19 , 47 , 15 ] . We introduce our novel search algorithms for finding good explanations by making use of general search principles [ 43 ] . Based on experiments with two Transformer models and six text classification datasets ( including FEVER , SNLI , and others ) , we summarize our core findings as follows : 1 . We propose to train models on explanation counterfactuals and find that this leads to greater model robustness against counterfactuals and yields drastic differences in explanation metrics . 2 . We find that some Replace functions are better than others at reducing counterfactual OOD-ness , although ultimately our solution to the OOD problem is much more effective . 3 . We introduce four novel search-based methods for identifying explanations . Out of all the methods we consider ( including popular existing methods ) , the only one that consistently outperforms random search is the Parallel Local Search ( PLS ) that we introduce , often by large margins of up to 20.8 points . Importantly , we control for the compute budget used by each method . 2 Related Work . Feature Importance Methods . A great number of methods have been introduced for FI estimation , drawing upon local approximation models [ 46 , 47 , 31 ] , attention weights [ 25 , 61 , 66 ] , model gradients [ 51 , 50 , 57 , 52 ] , and model-based feature selection [ 4 , 2 , 62 , 42 , 8 , 11 ] . While search approaches are regularly used to solve combinatorial optimization problems in machine learning [ 49 , 5 , 3 , 16 , 40 ] , we know of only a few FI methods based on search [ 19 , 47 , 15 ] . Fong and Vedaldi [ 19 ] perform gradient descent in explanation space , while Ribeiro et al . [ 47 ] search for probably-sufficient subsets of the input ( under a perturbation distribution ) . In concurrent work , Du and Xu [ 15 ] propose a genetic search algorithm for identifying FI explanations . We introduce several novel search algorithms for finding good explanations , including ( 1 ) a gradient search similar to Fong and Vedaldi [ 19 ] , a ( 2 ) local heuristic search inspired by an adversarial attack method [ 16 ] , ( 3 ) a global heuristic search , and ( 4 ) a parallel local search ( PLS ) making use of general search principles [ 43 ] . Choice of Replace Function . Past evaluations of explanation methods typically remove tokens or words from the text entirely [ 41 , 13 ] or replace them with fixed feature values [ 22 , 64 ] . Methods for creating explanations also use several distinct approaches , including ( 1 ) replacing token embeddings with the zero vector [ 34 , 57 , 1 ] , ( 2 ) using a special token [ 65 , 57 ] , ( 3 ) marginalizing predictions over a random variable representing an unobserved token [ 68 , 30 , 63 , 27 ] , and ( 4 ) adversarially selecting counterfactual features [ 23 ] . Sturmfels et al . [ 55 ] carry out a case study involving a few Replace functions for a vision model , which they compare via test-time ablations with image blurring techniques , though the case study is not intended to be a full comparison of methods . Haug et al . [ 21 ] assess a number of Replace functions for explanation methods used with tabular data , but they compare between functions to use when generating explanations , rather than when evaluating explanations , for which they offer no recommendation . In addition to evaluating Replace functions from the above works , we also consider setting attention mask values for individual tokens to 0 . The Out-of-distribution Problem of Explanations . Many papers have expressed concerns over how removing features from an input may result in counterfactuals that are out-of-distribution to a trained model [ 65 , 57 , 19 , 7 , 22 , 30 , 23 , 63 , 56 , 27 , 45 , 21 , 48 , 28 , 59 ] . In response to the problem , some have proposed to marginalize model predictions over possible counterfactual inputs [ 68 , 30 , 63 , 27 ] , use counterfactuals close to real inputs [ 7 , 48 ] , weight their importance by their closeness to real inputs [ 45 ] , or to adversarially select counterfactual features rather than use any user-specified features [ 23 ] . Others reject the whole notion of test-time ablations , preferring metrics based on train-time ablations [ 22 ] . Jethani et al . [ 28 ] propose a specialized model for evaluating explanations that is trained on counterfactual inputs in order to make them in-distribution , but since the evaluation model is distinct from the model used to solve the task , explanation metrics calculated using this evaluation model may not reflect the faithfulness of explanations to the task model . In concurrent work , Vafa et al . [ 59 ] independently adopt a solution equivalent to our Counterfactual Training with an Attention Mask Replace function , an approach which we empirically justify in Sec . 5 . In general , prior works make arguments for their approach based on intuition or basic machine learning principles , such as avoiding distribution shift . In Sec . 4 , we give a more detailed argument for preferring in-distribution counterfactuals on the basis of social alignment , a concept introduced by Jacovi and Goldberg [ 24 ] , and we propose a solution to the OOD problem . Our solution allows for test-time evaluation of explanations of a particular model ’ s decisions for individual data points , unlike similar proposals which either evaluate large sets of explanations all at once [ 22 ] or use a separate model trained specifically for evaluation rather than the blackbox model [ 28 ] . 3 Problem Statement . Feature Importance Metrics . The problem we are investigating is to find good feature importance explanations for single data points , where explanation are evaluated under metrics using test-time ablations of the input . In this context , an explanation for an input in a d dimensional feature space is a binary vector e ∈ { 0 , 1 } d , which may be derived from discretizing an FI estimate v ∈ Rd . We consider two primary metrics , Sufficiency and Comprehensiveness [ 13 ] . Sufficiency measures whether explanations identify a subset of features which , when kept , lead the model to remain confident in its original prediction for a data point . Comprehensiveness , meanwhile , measures whether an explanation identifies all of the features that contribute to a model ’ s confidence in its prediction , such that removing these features from the input lowers the model ’ s confidence . The Sufficiency metric for an explanation e ∈ { 0 , 1 } d and model p ( y|x ) = fθ ( x ) is given as Suff ( fθ , x , e , s ) = fθ ( x ) ŷ − fθ ( Replaces ( x , e ) ) ŷ ( 1 ) where ŷ = argmaxy f ( x ) y is the predicted class for x , and the Replaces function retains a proportion s of the input features ( indicated by e ) while replacing the remaining features with some user-specified feature . In order to control for the explanation sparsity , i.e . the proportion s of tokens in the input that may be retained , we average Sufficiency scores across sparsity levels in { .05 , .10 , .20 , .50 } , meaning between 5 % and 50 % of tokens in the input are retained [ 13 ] . A lower Sufficiency value is better , as it indicates that more of the model ’ s confidence is explained by just the retained features ( increasing fθ ( Replace ( x , e ) ) ŷ ) . Similarly , Comprehensiveness is given as Comp ( fθ , x , e , s ) = fθ ( x ) ŷ − fθ ( Replaces ( x , e ) ) ŷ but with sparsity values in { .95 , .90 , .80 , .50 } , as we are looking to remove features that are important to the prediction ( while keeping most features ) . A higher Comprehensiveness value is better , as it indicates that the explanation selects more of the evidence that contributes to the model ’ s confidence in its prediction ( resulting in a lower fθ ( Replace ( x , e ) ) ŷ ) . Overall Objective . Finally , our overall Sufficiency and Comprehensiveness objectives are given by averaging Suff ( or Comp ) scores across several sparsity levels . With a model p ( y|x ) = f ( x ) , a single data point x with d features , and a set of sparsity levels S , the Sufficiency objective is optimized by obtaining a set E = { ei } |S|i=1 with one explanation per sparsity level as argmax E 1 |S| |S|∑ i=1 Suff ( f , x , ei , si ) s.t . ei ∈ { 0 , 1 } d and ∑ d e ( d ) i ≤ ceiling ( si · d ) where the ceiling function rounds up the number si · d of tokens to keep . When optimizing for Comprehensiveness , we use the Comp and argmin functions , and the inequality is flipped . In general , we will optimize this objective using a limited compute budget , further described in Sec . 6.2 .

Summary: This paper investigates feature important methods in NLP. The authors argue that counterfactual inputs should be in-distribution, and propose training on counterfactual inputs to reduce this distribution shift. They also investigate the best approach to producing counterfactuals and the best search method for finding FI explanations, finding that one of their proposed methods does best.