Accurate and Robust Feature Importance Estimation under Distribution Shifts

We consider the setup where we build a predictive model $\mathcal{F}(\Theta)$ which takes as input a sample $\mathbf{x}\in\mathbb{R}^{d}$ with $d$ features and produces the output $\hat{\mathbf{y}}\in\mathbb{R}^{k}$ of dimensionality $k$. Note that this setup is deliberately unspecific in terms of both model architecture and data modality as PRoFILE is agnostic to either. Given a training set $\{(\mathbf{x}_{i},\mathbf{y}_{i})\}_{i=1}^{N}$, we optimize for the parameters $\Theta$ using the loss function $\mathcal{L}:\mathbf{y}\times\hat{\mathbf{y}}\rightarrow s$, where $s\in\mathbb{R}$. In other words, $\mathcal{L}$ measures the discrepancy between the true and predicted outputs using a pre-specified error metric. Examples include categorical cross-entropy for classification or mean-squared error for regression.

While our approach does not need access to the training data or specifics of the training procedure while generating explanations, similar to any post-hoc interpretability method, our approach requires the training of an auxiliary network $\mathcal{G}(\Phi;\Theta)$ that takes the same input $\mathbf{x}$ and produces the output $\hat{s}\approx\mathcal{L}(\mathbf{y},\mathcal{F}(\mathbf{x}))$. The objective of this network is to directly estimate the fidelity for the prediction that $\mathcal{F}$ makes for $\mathbf{x}$, which we will use in order to construct our post-hoc expalantions without any additional re-training. Note that, the loss estimates implicitly provide information about the inherent uncertainties; for example, in Ash et al. (2020), the gradients of loss estimates have been used to capture the model uncertainties. We define the auxiliary objective $\mathcal{L}_{aux}:s\times\hat{s}\rightarrow\mathbb{R}$, in order to train the parameters $\Phi$ of model $\mathcal{G}$. As showed in Figure 1(top), the loss estimator $\mathcal{G}$ uses the latent representations from different stages of $\mathcal{F}$ (e.g., every layer in the case of an FCN or every convolutional block in a CNN) to estimate $\hat{s}$. We use a linear layer along with non-linear activation (ReLU in our experiments) to transform each of the latent representations from $\mathcal{F}$ and they are finally concatenated to predict the loss. During training, the gradients from both the losses are used to update the parameters $\Theta$ of model $\mathcal{F}$.

Since the proposed feature estimation strategy relies directly on the quality of the loss estimator, the choice of the loss function $\mathcal{L}_{aux}$ is crucial. In particular, our approach (see Figure 1(bottom)) is based on ranking input features using the loss values obtained by masking those features. Consequently, we expect the loss estimator to preserve the ordering of samples (based on their losses), even if the original scale is discarded. Here, we explore the use of two different objectives to minimize the discrepancy between true and estimated losses. As we will show in our empirical studies, both these objectives are highly effective at revealing the input feature importance.

Given the loss estimator $\mathcal{G}$, we estimate the feature importance using a Granger causality-based objective, similar to Schwab and Karlen (2019). The Humean definition of causality adopted by Granger Granger (1969) postulates that a causal relationship exists between random variables $x_{j}$ and $y$, i.e., $x_{j}\rightarrow y$, if we can better predict using all available information than the case where the variable $x_{j}$ was excluded. This definition is directly applicable to our setting since it satisfies the key assumptions of Granger causality analysis, our data sample $\mathbf{x}$ contains all relevant variables required to predict the target and $\mathbf{x}$ temporally precedes $\mathbf{y}$. Mathematically,

where $\epsilon$ denotes the model error. For a sample $\mathbf{x}$, we can compute this objective for each feature $j$ to construct the explanation. As showed in Figure 1(bottom), we use the loss estimator to measure the predictive model’s error in the presence and absence of a variable $x_{j}$ to check if $x_{j}$ causes the predicted output. There are a variety of strategies that can be adopted to construct $\mathbf{x}\setminus\{j\}$. In the simplest form, we can mask the chosen feature by replacing it with zero or a pre-specified constant. However, in practice, one can also adopt more sophisticated masking strategies that take into account the underlying data distribution Janzing et al. (2013); Štrumbelj, Kononenko, and Šikonja (2009). Interestingly, our approach is agnostic to the masking strategy and the loss estimator can be used to compute the causal objective in Eqn.(4) for any type of masking. In contrast, existing approaches such as CXPlain requires re-training of the explanation model for the new masking strategy.

Since the loss estimator is jointly trained with the main predictor, our approach does not require any additional adversarial training as done in Lakkaraju, Arsov, and Bastani (2020) to ensure that the explanations are consistent with the black-box. Furthermore, existing works in the active learning literature have found that the loss function Yoo and Kweon (2019) or its gradients Ash et al. (2020) effectively capture the inherent uncertainties in a model and hence can be used for selecting informative samples. Using a similar argument, we show that even though our causal objective is similar to CXPlain, our approach more effectively generalizes to even complex distribution shifts where CXPlain fails.

We consider a suite of synthetic and real-world datasets to evaluate the fidelity and robustness of our approach. For the fidelity comparison study under standard testing conditions, we use the: (a) UCI Handwritten Digits dataset, (b) OpenML benchmarks Vanschoren et al. (2013), Kropt, Letter Image Recognition, Pokerhand and RBF datasets and (c) Cifar-10 image classification dataset Krizhevsky and Hinton (2009). The dimensionality of the input data ranges from $10$ to $64$ and the total number of examples varies between $1797$ and $13750$ for different benchmarks. For each of the UCI and OpenML datasets, we utilized $\sim 90\%$ of the data while for Cifar-10, we used the prescribed dataset of $50$K RGB images of size $32$$\times$$32$ for training our proposed model.

For the robustness study, we used the following datasets: (a) Synthetic dataset: In order to study the impact of distribution shifts on explanation fidelity, we constructed synthetic data based on correlation and variance shifts to the data generation process defined using a multi-variate normal distribution. More specifically, we generated multiple synthetic datasets of $5$K samples, where the number of covariates was randomly varied between $10$ and $50$. In each case, the samples were drawn from $\mathcal{N}(\boldsymbol{\mu},\mathbf{\Sigma})$, where $\mu_{ii}=\alpha,\Sigma_{ii}=1$ and $\Sigma_{ij}=\beta$ and the values $\alpha$, $\beta$ (the correlation between any two variables) were randomly chosen from the uniform intervals $[-2,2]$ and $[-1,1]$ respectively. The label for each sample was generated using their corresponding quantiles (i.e., defining classes separated by nested concentric multi-dimensional spheres). To generate correlation shifts, we created new datasets following the same procedure, but using a different correlation $\bar{\beta}=\beta+\delta_{\beta}$. Here, $\delta_{\beta}$ was randomly drawn from the uniform interval $[-0.2,0.2]$. Next, we created a third dataset to emulate variance shifts, wherein we changed the variance $\Sigma_{ii}=\Sigma_{ii}+\kappa$ and $\kappa$ was drawn from the uniform interval $[0.25,0.75]$. While the predictive model was trained only using the original dataset, the explanations were evaluated using both correlation- and variance-shifted datasets. We generated $10$ different realizations with this process and report the explanation fidelity metrics averaged across the $10$ trials; (b) Cifar10 to Cifar10-C Hendrycks and Dietterich (2019): This is a popular benchmark for distribution-shift studies, wherein we train the predictive model and loss estimator using the standard Cifar10 dataset and generate explanations for images from the Cifar10-C dataset containing wide-variety of natural image corruptions; and (c) MNIST-USPS: In this case, we train the predictive model using only the MNIST handwritten digits dataset LeCun, Cortes, and Burges (2010) and evaluate the explanations on the USPS dataset Hull (1994) at test time.

We compared PRoFILE  against the following baseline methods that are commonly adopted to produce sample-level explanations. All baseline methods considered belong to the class of post-hoc explanation strategies which aim to construct interpretable models that can approximate the functionality of any black-box predictor.

To evaluate the explanation fidelity, we utilize the commonly used difference in log-odds metric, which is a measure of change in prediction when $k\%$ of the most relevant features in the input data are masked.

Here $\text{log-odds}(\text{p})=\text{log}(\frac{\text{p}}{1-\text{p}})$ and $\text{p}_{\text{ref}}$ is the reference prediction probability of the original data and $\text{p}_{\text{masked}}$ refers to the prediction probability when a subset of features are masked. A higher value for $\Delta\text{log-odds}$ implies higher fidelity of the feature importance estimation. More specifically, for: (a) Non-Image Datasets. We sort the feature attribution scores obtained from the explainability method (PRoFILE and baselines) and zero mask the top $k\%$ important features in the input sample to evaluate the metric, and (b) Image Datasets. We use the SLIC Achanta et al. (2012) segmentation algorithm to generate superpixels, which are then used to compute the feature importance scores. For CXPlain and DeepSHAP, we aggregate the pixel-level feature importance scores to estimate attributions for each superpixel.

For all non-imaging datasets, the black-box model was a 5 layer MLP with ReLU activations, each fully-connected (FC) layer in the loss estimator contained $16$ units. In the case of Cifar-10, we used the standard ResNet-18 architecture, and the loss estimator used outputs from each residual blocks (with fully connected layers containing $128$ hidden units). Finally, for the MNIST-USPS experiment, we used a $3-$layer CNN with $2$ FC layers. The loss estimator was designed to access outputs from the first $4$ layers of the network and utilized FC layers with $16$ units each. All networks were trained using the ADAM optimizer with a learning rate $0.001$ and batch size $128$.