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: The paper considers different aspects of feature-importance-based explanations. In particular, the authors propose a ‘counterfactual training’ approach to deal with observations that are out of the original training distribution (OOD smaples). Moreover, the authors illustrate different approaches to replace input features with a pre-defined baseline, i.e., a common approach to assess the quality of an attribution/explanation method, which however may lead to OOD samples. Finally, the authors propose search algorithms to find binary-valued, feature-based explanations that have been shown to improve relevant test measures.