Source-linked AI summary
The Out-of-Distribution Problem in Explainability and Search Methods for Feature Importance Explanations
Peter Hase, Harry Xie, Mohit Bansal
TL;DR
Feature-importance explanations and their evaluation often use test-time feature ablations, but OOD counterfactuals can make them socially misaligned. The paper proposes counterfactual training, compares replacement functions, and introduces search methods; only Parallel Local Search consistently beats random search in its experiments.
Problem
OOD counterfactuals can make feature-importance explanations and their metrics depend on model priors and random seeds rather than only on learned evidence.
Source: §4 · p. 4Method
The paper trains models on randomly generated counterfactual inputs, compares five feature-replacement approaches, and introduces four search-based explanation methods.
Source: §1 · p. 2, §6.1 · p. 8Results
Only Parallel Local Search consistently outperforms random search, by up to 20.8 points, with compute budgets controlled.
Source: §1 · p. 2Takeaways & Limitations
Counterfactual Training is proposed to align train and test distributions, while Parallel Local Search is the only evaluated method that consistently beats random search.
Source: §7 · p. 10Feature-importance explanations can lose information through sparsity, and the authors cannot guarantee that Counterfactual Training eliminates all risks.
Source: §8 · p. 10
Abstract
from arXiv · showhide
Feature importance (FI) estimates are a popular form of explanation, and they are commonly created and evaluated by computing the change in model confidence caused by removing certain input features at test time. For example, in the standard Sufficiency metric, only the top-k most important tokens are kept. In this paper, we study several under-explored dimensions of FI explanations, providing conceptual and empirical improvements for this form of explanation. First, we advance a new argument for why it can be problematic to remove features from an input when creating or evaluating explanations: the fact that these counterfactual inputs are out-of-distribution (OOD) to models implies that the resulting explanations are socially misaligned. The crux of the problem is that the model prior and random weight initialization influence the explanations (and explanation metrics) in unintended ways. To resolve this issue, we propose a simple alteration to the model training process, which results in more socially aligned explanations and metrics. Second, we compare among five approaches for removing features from model inputs. We find that some methods produce more OOD counterfactuals than others, and we make recommendations for selecting a feature-replacement function. Finally, we introduce four search-based methods for identifying FI explanations and compare them to strong baselines, including LIME, Anchors, and Integrated Gradients. Through experiments with six diverse text classification datasets, we find that the only method that consistently outperforms random search is a Parallel Local Search (PLS) that we introduce. Improvements over the second-best method are as large as 5.4 points for Sufficiency and 17 points for Comprehensiveness. All supporting code for experiments in this paper is publicly available at https://github.com/peterbhase/ExplanationSearch.
1 Introduction
Feature-importance explanations and their evaluation commonly use test-time feature ablations, but out-of-distribution counterfactuals can make them socially misaligned. The paper proposes counterfactual training, compares replacement functions, and introduces search methods for finding explanations.
1 Introduction
- The paper studies test-time ablation metrics, including Sufficiency and Comprehensiveness, which assess feature importance through changes in model confidence.FI methods assign scalar importance estimates to input features, and evaluation often removes features marked important.Source: §1 · pp. 1–2
- OOD counterfactuals can make explanations socially misaligned because model priors and random initialization influence FI estimates in unintended ways.The proposed training approach exposes models to counterfactual inputs to address this problem.Source: §1 · p. 2
- Counterfactual Training exposes models to counterfactual inputs during training, increasing robustness and substantially changing explanation metrics.The authors report that this approach leads to greater robustness against counterfactuals.Source: §1 · p. 2
- Some Replace functions reduce counterfactual OOD-ness better than others, though the proposed OOD solution is more effective.Source: §1 · p. 2
- Only Parallel Local Search consistently outperforms random search, by up to 20.8 points, under controlled compute budgets.The comparison includes the paper’s search methods and popular existing methods.Source: §1 · p. 2
2 Related Work
Prior work spans diverse feature-importance methods and replacement strategies, but search-based FI methods and systematic comparisons of replacement functions remain limited. The paper introduces additional search algorithms and evaluates replacement functions for explanation evaluation.
2 Related Work
- Earlier FI approaches use local approximation, attention weights, gradients, and model-based feature selection, while only a few use search.The paper introduces gradient, local heuristic, global heuristic, and Parallel Local Search methods.Source: §2 · pp. 2–3
- Prior replacement approaches remove words, use fixed values, zero embeddings, special tokens, marginalization, or adversarially selected counterfactual features.The paper also considers setting attention-mask values to zero.Source: §2 · p. 3
- Earlier replacement-function assessments address vision or explanation generation, whereas the paper compares functions for evaluating explanations and motivates in-distribution counterfactuals through social alignment.The paper’s approach allows test-time evaluation of individual explanations for a particular task model, unlike ROAR and EVAL-X.Source: §2 · p. 3
3 Problem Statement
The paper formulates explanation finding as optimizing Sufficiency or Comprehensiveness for a single input using binary feature-selection vectors. Both objectives average scores across sparsity levels, with opposite preferences for the metric values.
3 Problem Statement
- An explanation is a binary vector over features; Sufficiency measures confidence when selected features are retained, while Comprehensiveness measures confidence loss when contributing features are removed.Source: §3 · p. 3
- Sufficiency subtracts confidence after replacing unselected features from confidence in the original predicted class; lower scores are better.The selected features are retained at a specified sparsity level, and remaining features are replaced.Source: §3 · pp. 3–4
- Comprehensiveness uses sparsity levels that remove most features, and higher scores indicate that the explanation selects evidence contributing to model confidence.Source: §3 · p. 4
- The overall objective averages Sufficiency or Comprehensiveness across sparsity levels and seeks one explanation per level under a limited compute budget.For Comprehensiveness, the optimization uses Comp and arg min, with the inequality flipped.Source: §3 · p. 4
4 The Out-of-Distribution Problem in Explanations
The paper argues that OOD counterfactuals make feature-importance explanations and metrics reflect unintended model factors rather than only learned evidence. Counterfactual Training aligns train and test distributions by including counterfactual inputs during training.
4 The Out-of-Distribution Problem in Explanations
- The section develops the case against OOD counterfactuals and proposes a training-based solution assessed in later experiments.Source: §4 · p. 4
- OOD counterfactuals are inputs not drawn from a model’s training or finetuning distribution, and feature replacement can create such inputs.The paper notes that a MASK token may not have appeared in language-model training data.Source: §4 · p. 4
- People expect FI explanations to reflect how a model learned to interpret features as evidence, but OOD inputs allow model priors and random seeds to influence explanations.The paper illustrates this with predictions for a MASK token absent from training data.Source: §4 · p. 5
- The argument also applies to explanation metrics: OOD counterfactuals make scores sensitive to model priors and random seeds rather than feature removal alone.Source: §4 · p. 5
- Counterfactual Training exposes models to counterfactual inputs, making their predictions on observed and counterfactual data in-distribution.The authors report an average test-accuracy drop of 0.7 points across six datasets and improved robustness to counterfactual inputs.Source: §4 · p. 5
5 Analysis of Counterfactual Input OOD-ness
The study compares how five replacement functions affect counterfactual OOD-ness and tests whether Counterfactual Training improves model robustness to those inputs. Attention Mask and MASK Token perform best without Counterfactual Training, while training on counterfactuals largely erases differences between replacement functions.
Experiment Design
- The comparison fixes the removed features across replacement functions and uses model accuracy as a proxy for how in- or out-of-distribution ablated inputs are.Differences in accuracy at a given feature sparsity are attributed to the replacement function; the same comparison is used for Counterfactual-Trained and Standard models.Source: §5 · p. 6
- The study tests BERT-Base and RoBERTa-Base on SNLI and SST-2, removing random tokens at proportions of 0.2, 0.5, and 0.8.Ten random seeds are used for each Standard model; additional Counterfactual-Trained BERT models on SNLI use ten seeds and exclude Marginalize because it is too expensive for training.Source: §5 · p. 6
Results for Replace functions
- Without Counterfactual Training, Attention Mask and MASK Token are the best replacement functions, with the best exceeding the third-best by up to 5.48 points.Marginalize is regularly more than 10 points worse than the best method; the authors recommend Attention Mask or MASK Token for standard models.Source: §5 · pp. 6–7
- Counterfactual Training improves robustness to counterfactual inputs by up to 22.9 points and nearly erases differences among replacement functions.A 2.23-point difference remains between Attention Mask and Zero Embedding when 80% of tokens are removed; the authors recommend Counterfactual Training with Attention Mask, MASK Token, or Slice Out for FI explanations.Source: §5 · p. 7
6 Explanation Methods and Experiments
The paper compares salience and search-based feature-importance explanations under a fixed compute budget, across six text-classification datasets. Parallel Local Search is the only method that consistently beats Random Search, while Counterfactual Training can lower metric scores even as the authors recommend it for socially aligned explanations.
6.1 Explanation Methods
- Ordered Search enumerates explanations by a linear score initialized from LIME salience, allocating 25% of compute to scoring and 75% to enumeration.Other approaches include discrete Gradient Search, Taylor-based beam search, and Exhaustive Search, which can identify optimal explanations for sequences typically no longer than 10 tokens.Source: §6.1 · p. 8
- PLS runs parallel local searches from random initial explanations, accepting a proposed neighboring explanation only when it improves the objective.The method uses ten parallel runs; search methods use Attention Mask replacement and generally search explanations at the maximum allowable sparsity, or minimum for Comprehensiveness.Source: §6.1 · pp. 7–8
6.2 Experimental Setup
- The benchmark uses six datasets and 500 test points per dataset; query tokens are not replaced in datasets that pair queries with documents.Ten BERT-Base seeds are trained per dataset for Standard and Counterfactual-Trained models, with Counterfactual Training using Attention Mask replacement.Source: §6.2 · p. 8
6.3 Main Results
- PLS is the only method to consistently outperform Random Search; its Sufficiency gains reach 7.6 points over Random Search, and its Comprehensiveness gains reach 17 points over Random Search.For Comprehensiveness, PLS is best in 9 of 10 comparisons, 7 statistically significant; Table 1 presents explanation metrics across methods and datasets.Source: Table 1 · p. 9, §6.3 · pp. 9–10
6.2 Experimental Setup
- Each explanation method is limited to a fixed number of forward and backward passes, typically 1,000 per explanation.The limit controls compute across methods that otherwise require very different numbers of model passes.Source: §6.2 · pp. 8–9
6.3 Main Results
- LIME is the strongest salience method, but Random Search beats it on Sufficiency in 9 of 10 comparisons, by up to 21.5 points.LIME performs better than Random Search on Comprehensiveness for three of five datasets with Standard models and four of five with Counterfactual-Trained models.Source: §6.3 · p. 10
- The authors recommend PLS with Counterfactual-Trained models, although Counterfactual Training often yields worse Sufficiency and Comprehensiveness scores than Standard training.With Random Search, Comprehensiveness scores are worse for all datasets by up to 24.16 points; the authors attribute the metric differences to counterfactual OOD-ness and favor Counterfactual Training for socially aligned explanations.Source: §6.3 · p. 10
7 Conclusion
The paper argues for addressing OOD counterfactuals, recommends some replacement functions over others, and introduces Parallel Local Search (PLS), the only evaluated method that consistently outperforms random search.
7 Conclusion
- PLS is the only evaluated method that consistently outperforms random search; the paper also presents Counterfactual Training and recommendations for choosing Replace functions.These contributions address OOD counterfactual inputs and feature-importance explanation search.Source: §7 · p. 10
8 Broader Impacts and Limitations
Feature-importance explanations can help developers inspect model behavior and users assess decisions, while risks include misuse, over-reliance, and imperfect faithfulness. Counterfactual Training helps align explanations with expected information, but does not guarantee eliminating seed and model-prior influence.
8 Broader Impacts and Limitations
- Feature-importance explanations can help developers check which features models use and users confirm or dispute whether decisions relied on acceptable features.The authors caution that these explanations are not data summaries or proof that a decision was right; sparse explanations can lose information.Source: §8 · p. 10
- Counterfactual Training helps align explanations with expected information, but cannot guarantee eliminating random-seed and model-prior influence.The authors say explanations for the considered models therefore may remain imperfectly socially aligned.Source: §8 · p. 10
A Method Implementation and Hyperparameter Tuning Details
The appendix details replacement functions and the implementation and tuning of explanation methods, including local surrogate approaches, gradients, and search procedures.
A.1 Replace Functions
- Marginalization estimates predictions over imputed counterfactual tokens using Monte Carlo samples; 10 samples performed best for robustness in the reported tuning, though the margin was small.The method remains over ten times slower than other Replace functions because it requires many masked-language-model forward passes.Source: §A.1 · pp. 16–17
A.2 Explanation Methods
- LIME fits a locally weighted, regularized surrogate using binary-vector perturbations, with forward selection for at most six features and ridge regression otherwise.The perturbation kernel weights by negative cosine similarity; final salience scores use regression weights.Source: §A.2 · p. 17
- Vanilla Gradient uses gradients of predicted probability with respect to token embeddings, while Integrated Gradients scores tokens relative to a baseline and satisfies completeness.For Integrated Gradients, token scores sum to the probability difference between the input and baseline.Source: §A.2 · p. 17
- Gradient Search optimizes a continuous state with differentiable attention-mask replacement and Gumbel-Softmax sampling, then treats state values as token salience scores.An ℓ2 penalty encourages expected sparsity near ceiling(0.05 · L).Source: §A.2 · p. 18
- Taylor Search uses a first-order loss-change heuristic to choose swaps between retained and hidden tokens, then applies beam search of width 3.Its search space contains k-sparse explanations, and the Comprehensiveness objective uses the opposite optimization direction.Source: §A.2 · p. 18
- Parallel Local Search proposes unseen explanations by random walks and uses 10 parallel runs; the appendix also describes Ordered Search and tuning the shared search space.Its deterministic update accepts a new state only when its objective improves; exactly k-sparse explanations were preferred for Random Search and used by other search methods.Source: §A.2 · pp. 17, 19
A.3 Model Training Details and Experiment Runtimes
The experiments train BERT-Base or RoBERTa-Base models across six datasets and report training and explanation-evaluation runtimes. The appendix also documents dataset handling and preprocessing.
A.3 Model Training Details and Experiment Runtimes
- Models use BERT-Base or RoBERTa-Base with AdamW at a 1e-5 learning rate, and ten random seeds are trained for every experiment.Training uses 20 epochs for Standard models in the counterfactual OOD analysis, 10 for Counterfactual-Trained models, and 10 for explanation evaluation.Source: §A.3 · p. 19
- For 500 Evidence Inference points, explanation runtimes range from 2.5 minutes for Vanilla Gradient to 17 hours for Anchors under a 1000-pass budget.LIME takes 2.5 hours and Parallel Local Search takes 5 hours in the same reported setting.Source: §A.3 · p. 19
B.1 Data Preprocessing
- For preprocessing, input text is split on spaces into words, each word is tokenized independently and concatenated, and inputs longer than 512 tokens are truncated by shortening the document.Source: §B.1 · p. 20
B.2 Analysis of Counterfactual Input OOD-ness Details
The appendix details how counterfactual OOD-ness was evaluated and how explanation methods’ compute budgets were allocated. The evaluation used random token masks and a subset of training data on validation splits.
B.2 Analysis of Counterfactual Input OOD-ness Details
- The OOD evaluation sampled 10 random masks per data point and used the majority prediction across the resulting ablated inputs, except for Marginalize, which used one explanation per input.It used 10k training points per task and validation splits to motivate method-design choices.Source: §B.2 · p. 20
- Table 3 reports model accuracies for the counterfactual input OOD-ness analysis.The supplied passage identifies the table’s topic but does not provide its values.Source: Table 3 · p. 21
B.3 Compute Budget Details
- The standard compute budget is 1,000 forward and backward passes per final metric value, counting both explanation creation and evaluation.Some methods compute Sufficiency and Comprehensiveness while finding an explanation; others evaluate metrics afterward.Source: §B.3 · p. 21
- LIME uses 996 forward passes for an explanation and 4 for metrics, while Integrated Gradients uses 498 forward and backward passes and 4 forward passes for metrics.Vanilla Gradient uses one forward and backward pass; Taylor Search uses at most 1,000 forward passes, and the other search methods use 1,000 total.Source: §B.3 · p. 21
C Additional Results
Additional results show that PLS remains competitive across search budgets and model settings, while its performance over time differs from Random and Ordered Search. The appendix also reports accuracy comparisons and checks whether metric trends change under weight-of-evidence scoring.
SST-2 Results
- On SST-2, PLS exactly identifies the optimal Sufficiency explanation and comes within .01 of optimal Comprehensiveness, while LIME trails search methods by 2.86–3.86 points.Most considered search methods perform as well as Exhaustive Search on sequences short enough to exhaustively search.Source: C Additional Results · p. 21
Search Method Performance Over Time
- On FEVER and SNLI, PLS outperforms Random Search early and remains preferable at step 1,000, with statistically significant differences at p < .05.On FEVER, PLS continues improving while Random and Ordered Search plateau by 1,000 steps.Source: C Additional Results · pp. 21–22, Figure 3 · p. 22
Results with Reduced Search Budget
- With a reduced 250-step search budget, search methods outperform salience methods in 15 of 24 comparisons (8 favoring LIME), suggesting they can outperform salience methods with far less compute.Source: C Additional Results · p. 22
Counterfactual-Trained Model Accuracy
- Accuracy differences between Standard and Counterfactual-Trained models are typically below 1 point and average about 0.7 points across datasets.The authors describe this as a small but consistent gap and identify its causes as a topic for future work.Source: C Additional Results · p. 22
Results with Reduced Search Budget
- Search methods are preferable to LIME in 15 of 24 reduced-budget comparisons, while LIME is preferable in 8; PLS is preferable to Random Search 7 times versus once.LIME may be preferable when many sparsity levels, tight compute, and Comprehensiveness are priorities; otherwise the authors favor PLS.Source: C Additional Results · p. 22, Table 5 · p. 23
Results with RoBERTa Models
- With RoBERTa-Base models on a dataset subset, PLS is best in each condition and consistently outperforms Random Search, unlike LIME.The comparison uses the main-paper compute budget and 10 Counterfactual-Trained RoBERTa-Base models per dataset.Source: C Additional Results · p. 23, Table 6 · p. 23
Weight of Evidence Metrics
- Weight-of-evidence and difference-in-probability metrics show no notable trend differences; one significance-test contrast has p = .052 versus p = .030.The passage reports that the difference between those p-values is negligible.Source: C Additional Results · p. 24, Table 7 · p. 24
D Discussion
The discussion argues that counterfactual-trained models may be preferable when explanations should reflect learned data rather than model priors, and distinguishes search algorithms from the objectives they optimize. It also describes practical constraints on matching training and explanation distributions and on comparing computational cost.
D Discussion
- Counterfactual-trained models are preferable when explanations should reflect what the model learned from data rather than its prior or random seed, even if their Sufficiency and Comprehensiveness scores are worse.The authors caution against treating standard models as easier to explain when the difference arises from behavior on out-of-distribution inputs.Source: D Discussion · p. 24
- The discussion of search responds to De Cao et al. [12], who argue that pruning can discard features a model uses, citing evidence that a model counts all 8s and 1s.The authors counter that the search objective and the criterion of identifying features “used by the model” are distinct.Source: D Discussion · pp. 24–25
- Search is not inherently the problem when an explanation is unsatisfactory; the objective may omit preferences such as sparsity or features considered “used by the model.”The authors argue that the latter criterion must be formalized before searching under its corresponding objective.Source: D Discussion · pp. 24–25
- Matching training inputs to explanations produced by explanation methods remains prohibitively expensive, potentially requiring up to 1000x more training forward passes.The resulting training distribution is non-stationary, and limited experiments sometimes produced Sufficiency and Comprehensiveness scores similar to those of standard models.Source: D Discussion · p. 25
- Training on replaced inputs need not imply changing their labels: the model should predict from the visible evidence while accounting for removed features.The authors expect such predictions to be appropriately calibrated to the overall data distribution.Source: D Discussion · p. 25
- Forward and backward passes provide a rate-limiting compute measure, but practitioners should also compare wall-clock time because batching can change runtimes.In the experiments, batching makes PLS efficient relative to Gradient Search, which does not batch its forward or backward passes.Source: D Discussion · p. 25
Related papers
This paper cites
- BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding2018
- A Unified Approach to Interpreting Model Predictions2017
- HuggingFace's Transformers: State-of-the-art Natural Language Processing2019
- Axiomatic Attribution for Deep Networks2017
- Virtual Adversarial Training: A Regularization Method for Supervised and Semi-Supervised Learning2017
- A Benchmark for Interpretability Methods in Deep Neural Networks2018