Year: 2019

Abstract:
Many times, working with data starts with writing code to build a model from the features x to the response y. Answering questions about the data can require understanding what parts of the features x influence the response y. We show that using the KL-divergence within a randomization test discovers important features, while allowing for the reuse of model code. We call this method AMI-CRT. AMI-CRT requires running models for each randomization. We develop a faster variant called FAST-AMI-CRT, and show it is robust to errors in the randomization. Questions of feature importance can also be asked at the level of an individual sample. We provide an example to show that optimal predictive models are insufficient for instance-wise feature selection. We show that the estimators from FAST-AMI-CRT can also be reused to find important features in a particular instance. We evaluate our method on several simulation experiments, on a genomic dataset, a clinical dataset for hospital readmission, and on a subset of classes in ImageNet. Our methods outperform several baselines in various simulated datasets, identifies biologically significant genes, selects the most important predictors of a hospital readmission event, and identifies distinguishing image regions in an image-classification task.

heading: 1 INTRODUCTION
text: Model interpretation techniques aim to select features important for a response by reducing models (sometimes locally) to be human interpretable. However, the phrase model interpretation can be a bit of a misnomer. Any interpretation of a model must be imbued to the model by the population distribution that provides the data to train the model. In this sense, interpreting a model should be viewed as understanding the population distribution of data through the lens of a model. Existing methods for understanding the population distributions only work with particular models fit to the population, particular choices of test statistic, or particular auxiliary models for interpretation (Ribeiro et al., 2016; Lundberg and Lee, 2017). Such structural restrictions limit the applicability of these methods to a smaller class of population distributions. To be able to work in a black-box manner, feature selection methods can use models but must not require a particular structure in models used in selection processes.
Understanding the population distribution can be phrased as assessing whether a response is independent of a feature given the rest of the features; this test is called a conditional randomization test (Candes et al., 2018). Conditional randomization tests require test statistics. Test statistics like linear model coefficients (Barber et al., 2015) or correlation may miss dependence between the response and outcome. To avoid missing relationships between variables, we develop the notion of a proper test statistic. Proper test statistics are those whose power increases to one as the amount of data increases. Conditional independence implies the conditional-joint factorizes into conditionalmarginals. Measuring the divergence between these distributions yields a proper test statistic. Of the class of integral probability metrics (Müller, 1997) and f -divergences (Csiszár, 1964), the KLdivergence simplifies estimation and allows for reuse of the model structures and code from the standard task of predicting the response from the features. Using the KL-divergence in this context has a natural interpretation; it is a measure of the additional information each feature provides about the outcome over the rest. This measure of information is known as the additional mutual information (AMI) (Ranganath and Perotte, 2018).
Our proposed procedure is called the additional mutual information conditional randomization test (AMI-CRT). AMI-CRT uses regressions to simulate data from the null for each feature and compares the additional mutual information (AMI) of the original data to the AMI of the simulations from
the null to assess significance. AMI-CRT works with any regression with probabilistic outputs like cross-entropy-trained neural networks. Training many regressions for each sample from the null can require substantial computation. While this is an embarrassingly parallel computation, we develop FAST-AMI-CRT that only requires a single model trained from null data. FAST-AMI-CRT uses an average of the model from the null with the model from the original data. We show this mixture guards against both variance in model training and poor estimation of the null. Though simple, AMI-CRT outperforms popular procedures for feature importance on a wide variety of simulated data, hospital records, and biological data.
Working with data sometimes requires interpreting individual datapoints. For example, a doctor may benefit from knowing which features for a particular patient relate to their health. The process of identifying features at a datapoint-level is called instance-wise feature selection (Ribeiro et al., 2016; Lundberg and Lee, 2017; Gimenez and Zou, 2019). We identify an issue in instance-wise feature selection, where even features selected using the true population distribution do not yield the features that were used to generate the response of an instance. The crux of this disparity is that the response generation process, conditional on the features, may use randomness to select features. We provide an example to demonstrate where instance-wise feature selection can go awry. We develop sufficient conditions for instance-wise feature selection to avoid this issue. The same regression estimates from AMI-CRT can be used to estimate feature-importances with minimal computational overhead, resulting in a method we term additional mutual information instance-wise feature selection (AMIIW). We demonstrate AMI-IW on multiple simulations and image data. Across all of these tasks AMI-IW outperforms popular baselines.
Related Work. Permutation tests (Fisher, 1937) provide a test for marginal independence between each feature and the outcome. However, they fail to test conditional independence, which is required when covariates are dependent on each other. To address this, solutions like Sure Independence Screening (Barut et al., 2016; Fan and Lv, 2008) and Conditional Randomization Tests (Barber et al., 2015; Candes et al., 2018) have been proposed. These outline frameworks for conditional independence testing. However, they often make linearity or additive noise assumptions about the data generating distribution. Furthermore, they require the choice of a test statistic to capture some notion of conditional independence. The user of such frameworks is often burdened with the task of choosing this test statistic, which may require strong assumptions about the data generating distribution. Extending this approach to neural networks, Lu et al. (2018) propose a fully connected network whose weights are used as a test statistic. Though moving beyond linear models, their method is specific to fully connected networks. Tansey et al. (2018) propose holdout randomization tests (HRTs) that use empirical loss as a test-statistic. The loss they use for continuous-valued distributions of response is the mean-squared-error (MSE), which may ignore higher order dependencies between the response and features. Using AMI inside an HRT would capture these higher order dependencies. In our results, we adapt our test-statistic AMI to HRTs and show that this produces better calibration than other choices of loss. HRTs provide computational speed-ups over CRTs. However this speedup comes at the cost of robustness to poor estimations of the null feature distribution. We demonstrate empirically that FAST-AMI-CRT is robust to such poor estimations.
Beyond understanding the population distribution, some tasks require interpreting a population distribution on the level of an individual datapoint. Methods that test for conditional independence work under distributional notions of feature selection, but are not designed to identify the relevant features for a particular sample. To address this issue of “instance-wise feature selection,” several methods have been proposed, including local perturbations (Simonyan et al., 2013; Sundararajan et al., 2017; Ribeiro et al., 2016) and fitting simpler auxiliary models to explain the predictions of a large model (Chen et al., 2018; Lundberg and Lee, 2017; Yoon et al., 2019; Turner, 2016; Štrumbelj and Kononenko, 2014; Shrikumar et al., 2017). Our instance-wise work is most similar to that of Burns et al. (2019), who repurpose the HRT framework to perform instance-wise feature selection, or Gimenez and Zou (2019), who define a conditional randomization test (CRT) procedure for subsets of the feature space. In general, however, the conditions under which instance-wise feature selection with predictive models may be possible are not well developed. We address this issue by first identifying a set of sufficient conditions under which instance-wise feature selection is always possible. We then show how estimators used in AMI-CRT can be repurposed for use in an instance-wise setting, yielding a procedure called the AMI-IW.

heading: 2 PROPER TESTS FOR FEATURE SELECTION
text: Practitioners of machine learning use feature selection to identify important features for their predictive task. One way to filter out important features is to find those that improve predictions given the rest of the features. This can be formalized through conditional independence. Let xj be the jth feature of x and let x−j be all features but the jth one. The goal is to discover a set S such that ∀xj 6∈ S,xj ⊥ y | x−j , where independence is with respect to the true population distribution q. The only knowledge about q comes from a finite set of samples DN := {(x(i),y(i))}Ni=1 sampled from the population. This means that it is impossible to assess exact conditional independence. Therefore, in the finite sample setting, we must formulate a statistical hypothesis test.
A conditional randomization test (CRT) (Candes et al., 2018) defines a hypothesis test for conditional independence. For the jth feature, CRTs first compute a test statistic t using the N samples of data DN . CRTs place this statistic in a null distribution where samples of the jth feature xj are replaced by samples of x̃j |x−j which by construction satisfy x̃j ⊥ y |x−j . Letting D̃j,N be a dataset where {x(i)j }Ni=1 has been replaced by {x̃ (i) j }Ni=1, the p-value for a CRT is
pj(DN ) = E ∀i=1...N
x̃ (i) j ∼q(xj |x−j=x (i) −j)
[ 1 ( t(DN ) ≤ t(D̃j,N ) )] = E
∀i=1...N x̃
(i) j ∼q(xj |x−j=x (i) −j)
[ 1 ( t(DN )− t(D̃j,N ) ≤ 0 )] ,
(1)
Under smoothness constraints, the p-value is uniform under the null because it computes the cumulative distribution function of the test statistic under the null. While CRTs provide a general method for conditional independence testing, they leave several components including the choice of test statistic unspecified.

heading: 2.1 CHOOSING THE RIGHT TEST STATISTIC
text: Imagine a test statistic t(·) = t({x(i)j ,y(i)}Ni=1) that uses only a feature xj and the outcome y. Any p-values computed using this test statistic would be meaningless when testing for conditional independence, as t never considers the remaining features x−j . Therefore, particular choices for test statistics limit what can be tested. To address this, we introduce the concept of a proper test statistic.
Definition 1. Proper Test Statistic: A test statistic t(DN ) is proper if p-values produced by the statistic converge to 0 when the null must be rejected, and are uniformly distributed otherwise. Using t in Equation (1), this is:
pj (DN ) d−−−−→
N→∞ { Uniform(0, 1) if xj ⊥ y | x−j 0 with probability 1 if xj 6⊥ y | x−j , (2)
where d−→ indicates a convergence in distribution. Under the alternate hypothesis, which in the case of feature selection is xj 6⊥ y | x−j , the power to reject the null hypothesis must be 1, implying pj → 0. A proper test statistic requires that Equation (2) must hold for all distributions of y,x.
Proper tests statistics in a CRT select the features in S as the data grows. Definition 1 mirrors the concept of a scoring rule (Gneiting and Raftery, 2007), which measures the calibration of a probabilistic prediction by a model. A proper scoring rule is one such that the highest expected score is obtained by a model that uses the true probability distribution to make predictions.
Divergences are proper test statistics. Conditional independence means the conditional distribution r factorizes:
r(xj ,y | x−j) = r(xj | x−j)r(y | x−j). (3)
Divergences measure the closeness between two distributions. A divergence is zero when the two distributions are the same and positive otherwise. Computing any divergence K, like an integral probability metric (Müller, 1997) or an f -divergence (Csiszár, 1964), between the left hand side and right hand side of Equation (3) would be a proper test statistic. LetK(a, b) ≥ 0 with equality holding only when a is equal in distribution to b, then a proper test statistic Kj(r) := Er(x−j)[K(r(xj ,y | x−j), r(xj | x−j)r(y | x−j))]. A consistent estimator of this quantity is a proper test statistic (see
Appendix B.1). Casting conditional independence testing as divergence estimation reduces this test to fitting univariate regressions that can reuse pre-developed model code from the features to the response.
Define the resampling distribution q̃j = q̃j(x̃j |x−j)q(y,x−j). Using a divergence in a CRT requires estimates of the following conditional distributions: q(xj | x−j), q(y | x), q̃j(y | x̃j ,x−j), and q(y | x−j). The first distribution q(xj | x−j) is required for any CRT. The next distribution q(y | x) corresponds to the standard task of building a good regression model. The third distribution q̃j(y | x̃j ,x−j) requires a regression model with corrupted inputs. This regression can reuse the model structure and code from the standard regression task q(y | x). However, the last distribution q(y | x−j) could require development of new model structures. For example, if x is an image, a good model for q(y | x) could be a convolutional neural network. If the conditioning set x−j is a subregion of that image, the convolutional neural network used for q(y | x) would need to be modified for different padding and filter sizes. This means new models could be needed for each x−j . In the next section, we show that the KL-divergence removes the need for estimating this distribution, and therefore only requires the piece needed for all CRTs, q(xj | x−j), and model code to fit the response from the features.

heading: 2.2 AMI-CRT
text: Using the KL-divergence as a test statistic in Equation (1) requires the difference δj :
δj =Eq(x−j)[DKL(q(xj ,y | x−j), q(xj | x−j)q(y | x−j))] − Eq̃j(x−j)[DKL(q̃j(x̃j ,y | x−j), q̃j(x̃j | x−j)q̃j(y | x−j))]
=Eq(x−j)[Eq(xj ,y|x−j)[log q(xj ,y | x−j)− log(q(xj | x−j)q(y | x−j))] − Eq̃j(x̃j ,y|x−j)[log q̃j(x̃j ,y | x−j)− log(q̃j(x̃j | x−j)q̃j(y | x−j))]]
=Eq(x−j)[Eq(xj ,y|x−j)[log q(y | x)− log q(y | x−j)] − Eq̃j(x̃j ,y|x−j)[log q̃j(y | x−j)− log q̃j(y | x−j)]]
=Eq(x−j)[Eq(xj ,y|x−j)[log q(y | x)]− Eq̃j(x̃j ,y|x−j)[log q̃j(y | x̃j ,x−j)]] =Eq(x,y)[log q(y | x)]− Eq(x−j ,y)Eq̃j(x̃j |x−j)[log q̃j(y | x̃j ,x−j)]].
The second equality above follows from q(x−j) = q̃j(x−j), and the fourth from q(y | x−j) = q̃j(y | x−j). This simplification of δj means that for the computation of the average KL-divergence as a test statistic, the distribution q(y | x−j) is unecessary, thereby reducing computation and allowing the reuse of training infrastructure for predicting the response y. The KL-divergence provides this reduction since log splits products into sums. This expected KL-divergence is called the additional mutual information (AMI) (Ranganath and Perotte, 2018).
Computing the expected value of δj to get the p-value in Equation (1) requires estimation from a finite sample. Recall that DN is a collection of N datapoints sampled iid from q(x,y). These samples from the population come from the provided data. Similarly, D̃j,N is sampled iid from q(x−j ,y)q(x̃j | x−j). The samples x−j ,y come from the population distribution, however sampling q(x̃j | x−j) requires learning a model for this distribution also known as the complete conditional. The other two pieces of δj require building a model from x to y and (x̃j ,x−j) to y. To learn these models, the data DN is split into training and test sets. Let θ, β be parameters, qθ(x̃j | x−j) and qβ(y | x) are fit to the training set. The distribution qγ(y | x̃j ,x−j) is fit to the training part of D̃j,N . To compute the expectations in δj , these models are evaluated on their corresponding test sets. Finally, a Monte Carlo estimate of the expectation for the p-value in Equation (1) requires replicating this procedure K times. We call this procedure AMI-CRT (Algorithm 1).
The models qθ(x̃j | x−j) and qβ(y | x) can be shared across the replications, however qγ(y | x̃j ,x−j) must be recomputed for each replication. While this is an embarrassingly parallel problem, it involves the computation of a new model for each draw from the null distribution. To speed up the computation of AMI-CRT, we use the average of the original model qβ(y | x) and a single model trained for qγ(y | x̃j ,x−j). This averaged model is used to estimate the two terms in the AMI-difference δj . Under the null, the estimates are identically distributed because each estimate evaluates the same function on identically distributed samples. Thus, averaging produces p-values that are uniform under the null, but the average may not result in a proper test-statistic. However,
the averaged model performs almost as well as AMI-CRT empirically. We call this procedure the FAST-AMI-CRT, which is summarized in Algorithm 2.
Averaging is needed because models trained on data drawn from the same distribution have variance (Friedman et al., 2001). The averaged model provides FAST-AMI-CRT with several advantages. First, it is more conservative than using just the original model as in HRTs (Tansey et al., 2018) since the averaged model both predicts better on the null data and worse on the real data. We show empirically that this guards against errors in the estimation of the complete conditional distribution. Second, it requires only a single null model per feature instead of one per replication.
To estimate each of qθ and qβ , standard regression models like logistic regression, neural networks, and random forests can be used at no more computational cost than training. Nonparametric regression can be used as well. The choice between these estimators should be made by using the best fitting regression on validation data. The estimation procedure is straightforward, yet effective as we demonstrate in Section 4. In the next section, we show how the building blocks for FAST-AMI-CRT can be used to provide feature importances on an instance-wise level.

heading: 3 TESTS FOR INSTANCE-WISE FEATURE SELECTION
text: So far, we show how to recover features important across the whole population. We have not yet addressed the issue that different samples could have different important features. We call this problem of recovering important features for each sample, instance-wise feature selection (IWFS). To identify important features instance-wise, we can use the probability of observing a particular label y(i) given a set of features x(i). This suggests a candidate definition for important features:
Definition 2. (Candidate) Feature importance in IWFS: Let q be the true population distribution. Let {(x(i),y(i))}Ni=1 be a dataset where each (x(i),y(i)) ∼ q . The jth feature for the ith sample, x (i) j , is an important feature if
q(y = y(i) | x = x(i)) > q(y = y(i) | x−j = x(i)−j).
Definition 2 says that a feature x(i)j is important if observing it increases the probability of y (i). This formulation is exploited in (Yoon et al., 2019; Chen et al., 2018) to obtain instance-wise important features. However, Definition 2 can sometimes fail to identify relevant features, even with access to the true conditional distribution q(y | x). While important features may satisfy this condition, so will a few unimportant features. As a demonstrative example, consider the data generating process where y = zx1+(1−z)x2+ , z ∼ Bernoulli(0.5), x1,x2 ∼ N (0, σ2x), and ∼ N (0, σ2 ). Assume we have the true q(y | x1,x2), and let z be unobserved. Pick any sample (x(i)1 ,x (i) 2 ,y
(i)) where the corresponding z(i) = 1, meaning that x(i)1 is important for this instance. We can expand:
q(y(i)|x(i)1 ,x (i) 2 )− q(y(i)|x (i) 1 ) =
1 2 N ( y(i);x (i) 2 , σ 2 ) − 1 2 N ( y(i); 0, σ2x + σ 2 ) For all i such that y(i) lies in a non-0 interval around x(i)2 , we see that q(y (i)|x(i)1 ,x (i) 2 ) − q(y(i)|x(i)1 ) > 0. For example if σ2 = σ2x = 1, and x (i) 2 = 5, then y
(i) ∈ [3, 7] will violate this inequality. In all of those cases, the wrong feature will be selected as important as per the candidate Definition 2. We show the full derivation of this example in Appendix C.1. The fundamental issue with the formulation in Definition 2 is that noise can act as a “selection” mechanism, but cannot be estimated because it is unobserved. While predictive models qβ for q(y | x) suffice for understanding population distributions, they might not be sufficient to perform IWFS.

heading: 3.1 SUFFICIENT CONDITIONS FOR INSTANCE-WISE FEATURE SELECTION
text: We develop the following condition under which q(y | x) is sufficient to perform IWFS:
Proposition 1. Sufficient conditions for instance-wise feature selection: For each sample (x(i),y(i)), let S(i) be a set of features that contribute to the prediction of y(i) defined as:
S(i) := { x (i) j : q(y = y (i) | x = x(i)) > q(y = y(i) | x−j = x(i)−j) } . (4)
If y is discrete, and q(y = y(i) | x = x(i)) = 1 for each sample (x(i),y(i)), i.e. we have perfect predictions on our dataset, then it is possible to recover such a set S(i) for all i.
The set in Equation (4) consists of only features x(i)j that help increase the likelihood of observing y(i) given the remaining features x(i)−j . If the perfect predictions property of q(y | x) is true, then q(y = y(i) | x−j = x(i)−j) can only be less than or equal to q(y = y(i) | x = x(i)), with equality when x(i)j is not important to y
(i). Assuming the sufficient conditions in Proposition 1, we can now construct an IWFS procedure using the same estimators from AMI-CRT or FAST-AMI-CRT.

heading: 3.2 AMI INSTANCE-WISE FEATURE SELECTION
text: Instance-wise feature selection can be performed using the building blocks of the FAST-AMI-CRT. Starting from Definition 2, we begin by manipulating q(y = y(i) | x = x(i)) and marginalizing out x (i) j . We then use Jensen’s inequality to upper-bound the log of this expectation as follows:
E x̃
(i) j ∼q(xj |x (i) −j)
[ − log q̃j(y = y(i) | xj = x̃j ,x−j = x(i)−j) ] ≥ − log q(y = y(i) | x−j = x(i)−j)
(5)
This suggests the following instance-wise test. If the inequality in Equation (5) is strict, the feature is considered important. If equality holds in Equation (5), the feature is considered unimportant. Notice that Jensen’s inequality could introduce slack in this bound that could make a feature seem relevant when it is not. We use Proposition 1 to show that this is not an issue.
Recall that given a model qβ for q(y | x) that satisfies the instance-wise sufficient conditions in Proposition 1, qβ(y = y(i) | x = x(i)) ≥ qβ(y = y(i) | x−j = x(i)−j). In the case where x (i) j does not help predict y (i), q̃j(y = y(i) | xj = x̃(i)j ,x−j = x (i) −j) cannot be greater than q(y = y(i) | x−j = x(i)−j), as the former does not depend on x̃ (i) j . Then the left-hand side of Equation (5) becomes − log q(y = y(i) | x−j = x(i)−j), a constant with respect to x̃ (i) j , implying an equality. Therefore, checking for equality in Equation (5) is a valid test to see if a feature is either important or unimportant. In Appendix C.2, we detail an example that shows how scores computed using Equation (5) can help rank features from most to least helpful for prediction. We term this procedure the additional mutual information instance-wise feature selection (AMI-IW).
If we computed an expectation over x(i),y(i) of Equation (5), this procedure resembles FASTAMI-CRT. We can reuse the estimators from FAST-AMI-CRT to compute these instance-wise logprobability differences for AMI-IW. Therefore, we only use one null estimator for q̃j(y | x̃j ,x−j). Like FAST-AMI-CRT, we can potentially use a mixture of estimators in AMI-IW, but at the cost of power to select important features.

heading: 4 EXPERIMENTS
text: We compare our methods, the AMI-CRT [ami-crt] and fast additional mutual information conditional randomization test (FAST-AMI-CRT) [fast-ami-crt] to widely-used approaches on various performance metrics. The baselines are:
• Correlation [corr-crt]: Difference between Correlation(xj ,y) and Correlation(x̃j ,y) as a test statistic for a CRT
• Local interpretable model-agnostic explanations (LIME) [lime] (Ribeiro et al., 2016)
• Shapley additive explanations (SHAP) [shap] (Lundberg and Lee, 2017)
• Random forest [rf] feature importance scores • Zero-one loss [loss-hrt]: Binary classifi-
cation loss of qβ(y | x) where qβ is a model for y | x is used as a test statistic in a HRT (Tansey et al., 2018)
HRTs construct a test by comparing a loss function evaluated on x and x̃j ,x−j . The choice of loss, however, is left to the practitioner. We study the 0-1 loss, which is a proper scoring rule, in all of our experiments. In specific settings, we equip an HRT with our AMI test-statistic. This method is called the ami-hrt. We show that ami-hrt is better calibrated than 0-1 HRT. Table 1 presents a summary comparison of the properties of each selection method. We use the regression approach using conditional categorical distribution parameterized with neural networks highlighted in (Miscouridou et al., 2018) to model q(xj |x−j) for all experiments unless specified otherwise.
Simulations: We simulate data for evaluating each selection method. These tests are designed to highlight the differences between each method.
[xor]: To test the case where features on their own are not informative, but together provide information, we use the xor dataset. We first sample x ∼ N (0,ΣD) N times, where ΣD is a Ddimensional covariance matrix. We translate the first two dimensions of each sample x(i) away from the origin in 4 different directions : {(s, s), (−s, s), (s,−s), (−s,−s)} with uniform probability. If the resulting translation has first two coordinates with the same sign, the label is one. Otherwise, it is zero. All but the first two features are independent of y. We set N = 2000, D = 20.
[orange](Chen et al., 2018): To test the case where y is some nonlinear function of x, we use the orange dataset. In this dataset, x ∼ N (0,ΣD), y = 1 if exp (∑` j=1 x 2 j − ` ) > 0.5 and 0 otherwise, where ` < D is the number of important features. We chooseN = 3000, D = 20, and ` = 4.
[selector, noisy-selector]: These experiments test instance-wise feature selection methods. We first sample x ∼ N (0,ΣD) N times, where ΣD is a D-dimensional covariance matrix, and D ≥ 11. The first feature x1, called the “selector” feature, determines the feature selection mechanism. We generate y ∈ {0, 1} as Equation (6). We also investigate the effectiveness of each feature selection method in the presence of noise, and generate y ∈ {0, 1} as Equation (7). We set the parameter N = 2000, D = 20.
q(y = 1 | x) = { 〈β1,x2:6〉 > 0 if x1 > 0 〈β2,x7:11〉 > 0 if x1 ≤ 0 (6) q(y = 1 | x) = { σ(〈β1,x2:6〉) if x1 > 0 σ(〈β2,x7:11〉) if x1 ≤ 0 (7)
Results. For methods based on CRTs or HRTs, we select features using p-values. For the baselines that do not produce p-values, we select features using the importance scores provided by each method. We threshold p-values or importance scores respectively, and compute an ROC curve for each method. We present the mean area under each curve over 100 simulations for the xor and orange datasets in Table 2. We notice that this task is easily solved by most methods apart from corr-crt. This test fails to account for dependencies between features. The ami-crt achieves a higher AUROC than baselines, while fast-ami-crt achieves similar performance.
To identify important features in practice, a threshold for importance scores must be chosen. If a method produces p-values, we can control the false discovery rate (FDR) (Benjamini and Hochberg, 1995). This is the expected proportion of falsely identified features. An assumption for standard FDR-controlling procedures is independent p-values (Benjamini and Hochberg, 1995). Therefore, we investigate the calibration of p-values across ami-crt, fast-ami-crt, corr-crt, and loss-hrt. We omit other baselines in this comparison as they do not produce p-values and therefore have no direct FDR-control.
To evaluate each method, we use the generating process for the orange dataset, and set N = 3000, D = 104, ` = 4. If the p-values are independent, null p-values should resemble iid draws from a Uniform(0,1) distribution. Figure 3 (Appendix D) shows a quantile-quantile plot of null p-values. We also perform a Kolmogorov-Smirnov (KS) (Massey Jr, 1951) test where the null distribution is Uniform(0,1). All CRT-based methods produce independent p-values, while loss-hrt produces deflated and significantly non-uniform p-values (p = 0.0006), implying dependence. As a result loss-hrt incorrectly identifies many null features as important. Models for each replication of null features ami-crt potentially decrease the correlation of the p-values. The fast-ami-crt achieves a middle ground between the hrt and ami-crt by yielding well-calibrated p-values while requiring only one null model qγ(y | x̃j ,x−j) per feature. We also investigate the use of ami-hrt, which uses the AMI as a test statistic in an HRT in Figure 3.
To better understand the difference between refitting (CRTs) and not refitting (HRTs), we inspect the robustness of each method to poor simulations from the null (poor estimation of q(xj | x−j)). We use the orange dataset and set all off-diagonal values of ΣD to 0.5. We sample x̃j from N (0, 1). We see in Figure 4 (Appendix D) that refitting protects against poor approximations of q(xj | x−j); the p-value distribution for fast-ami-crt is uniform, while the p-value distribution for both HRT methods is significantly non-uniform (p = 5× 10−6 for ami-hrt and p = 10−7 for loss-hrt). The robustness of fast-ami-crt comes from averaging which makes predictions for the original data worse and predictions on the null better regardless of the quality of the null simulations.
For instance-wise feature selection, we perform two tests for each selection method. To test HRTs in this setting, we use the procedure prescribed by Burns et al. (2019) but with different test statistics. For precision, we identify the 7 (the true # of relevant features per instance) most important features as dictated by each selection method and report average precision scores across a held-out test set in Table 3. For selector identification, we count the number of instances where the selector variable x1 was identified. For the selector task, we notice that the ami-iw achieves the highest precision, followed by the loss-hrt. In the noisy-selector task, we notice a decrease in scores across all methods with the largest decrease for ami-iw and loss-hrt. The noisy-selector case violates the sufficient conditions for instance-wise feature selection (Equation (4)) meaning that the noise in sampling the response can obscure which features are important. This explains the reduction in performance. Even with the decrease in performance, ami-iw performs best.
The ami-iw and loss-hrt identify the selector variable x1 in nearly every sample in our test set. Linear methods like lime fail because the selection mechanism is highly non-linear. Further, rf and corr-crt are not designed to assign importance at the level of an individual sample and therefore do not provide meaningful scores per instance.
Wellcome Trust Celiac disease: We study data from a genomic analysis on Celiac disease (Dubois et al., 2010). For each individual in this dataset, we have a set of single nucleotide polymorphisms (SNPs). SNPs represents genetic variance in the individual with respect to some reference population. This dataset consists of two classes of individuals: cases (n = 3796) and controls (n = 8154), where the cases are those with Celiac disease. After standard preprocessing steps as prescribed by Bush
and Moore (2012), 1759 SNPs remain. To model q(xj |x−j), we use the same procedures as (Candes et al., 2018) where we estimate q(xj |xSj ) where Sj is only the set of SNPs (not including xj) known to be correlated with xj . To model q(y|x), we use an L1-logistic regression model.
Results. In Table 4, we show the results for all methods with FDR-control. We identify the SNPs that most likely contribute to distinguishing between those with Celiac disease and those without it. Since these methods produce pvalues, we can select features at a theoretical FDR of 20% using the Benjamini and Hochberg (1995) procedure. We report the percentage of selected SNPs that have been previously shown to be associated with Celiac disease in a biological context as reported by one of (Dubois
et al., 2010; Sollid, 2002; Adamovic et al., 2008; Hunt et al., 2008). There are 40 SNPs in total that are both in our dataset and in these papers. Ami-crt outperforms all other methods tested; fast-ami-crt performs similarly. We also list the SNPs returned by ami-crt in Appendix E. As expected, corr-crt selects a large set of features, but achieves fairly low precision. This is because many SNPs are correlated with each other, and all of these seem relevant marginally. The AMI-based methods have better precision and recall compared to loss-hrt potentially both due to aforementioned deviation from uniform and that the zero-one loss may not change when only one out of more than a thousand features gets perturbed.
Hospital readmission: We use a dataset consisting of ten years of medical logs from over 130 hospitals (Strack et al., 2014). Features in the dataset include time spent in the hospital, medical specialty of attending doctor, age, and various other diagnostic information. Labels for each sample represent one of three events: readmitted within the next 30 days (n = 35, 545), readmitted after 30 days (n = 11, 357), or not readmitted (n = 54, 864). Due to class imbalance, we grouped all readmitted patients into one category (n = 46, 902). We detail further preprocessing steps in Appendix F. To model q(y|x), we use a random forest classifier with 100 estimators.
0.0 0.2 0.4 0.6 0.8 1.0 False Positive Rate
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
Tr ue
P os
iti ve
R at
e
ROC curve for selected features
ami-crt loss-hrt corr-crt fast-ami-crt
lime shap rf
Rather than consider each pixel as an individual feature xj , we consider a patch of pixels xS as a single feature, such that no two patches contain overlapping pixels. To model the distribution
q(xS |x−S), we make use of a generative inpainting model πg (Yu et al., 2018). We split the image into an 8 × 8 grid so that there are 64 non-overlapping xS patches. To model q(y|x), we use a VGG-16 network (Simonyan and Zisserman, 2015). To perform our instance-wise test, we compute log-probability differences using fifty generated samples from q(xS |x−S) for each patch.
Results. In Figure 2, we show a subset of results of AMI-IW. The first and third columns show the original images for each class: ambulance and policevan respectively. The second and fourth columns mask out the original image in patches where the patch is not found to be relevant to the prediction. The model used to estimate q(y|x) is able to achieve roughly 90% accuracy on a heldout test set. We see that our predictive model uses relevant details like the words “ambulance” or “police” printed on the vehicle to distinguish between each class. The model also tends to ignore objects like windscreens and other features shared across classes, as is expected. These results indicate that the difference in log probabilities between a model using the true data, and one using x̃S sampled from q(xS | x−S) works well in determining a relevant set of features even on an instancewise level. We show several additional images in Figure 5, in Appendix G. We also compare our method to local interpretable model-agnostic explanations (LIME) and shapley additive explanations (SHAP) (Figures 6 and 7). Both methods perform reasonably well on this task, but identify objects that are known to be common to both classes like wheels and headlamps. Neither method identifies writing on the vehicles in the images. This is likely because of the simplifying assumptions made by these locally-linear methods. They assume that the set of relevant features is independent of the set of irrelevant features, which may not be the case in images. For example, the location of the word “ambulance” may depend on the window position.

heading: 5 DISCUSSION
text: We develop AMI-CRT for testing for conditional independence of each feature xj ⊥ y | x−j from a finite sample from the population distribution. AMI-CRT uses the KL-divergence to cast independence testing as regression and allows for the reuse of code from building the original model from the features to the response. We develop FAST-AMI-CRT which requires less computation than AMI-CRT and is robust to poor estimation of the null conditional. We define sufficient conditions under which to perform instance-wise feature selection and develop the AMI-IW, an instance-wise feature selection method built from the pieces of FAST-AMI-CRT. AMI-CRT, FAST-AMI-CRT, and AMI-IW all outperform several popular methods. in various simulated tasks, in identifying biologically significant genes, selecting the most indicative features to predict hospital readmissions, and in identifying distinguishing features in an image classification task.

heading: A ALGORITHMS
text: A.1 AMI-CRT ALGORITHM
We show the AMI-CRT procedure in Algorithm 1, and the FAST-AMI-CRT procedure in Algorithm 2.
Algorithm 1: k-fold AMI feature selection Input: x ∈ RN×D, feature matrix; y ∈ RN , labels Output: p, the p-values for xj ⊥ y |x−j ∀j Split x and y into M folds, G1, G2, . . . , GM for j ∈ [1, 2, . . . , D] do
Estimate qθ := qθ(xj |x−j) h← 0 for m ∈ [1, 2, . . . ,M ] do
Estimate q(m)β := qβ(y (G−m) |x(G−m)) h← h+ 1MLβ(y (Gm) |x(Gm)), where Lβ is an log-likelihood estimate using q(m)β
end Let h̃ be a K-dimensional vector of 0s for k ∈ [1, 2, . . . ,K] do
Sample x̃j ∼ qθ for m ∈ [1, 2, . . . ,M ] do
Estimate q(k,m) β̃ := qβ̃(y (G−m) | x̃(G−m)j ,x (G−m) −j ) Let h̃ (k) ← h̃(k) + 1MLβ̃(y (Gm) | x̃(Gm)j ,x (Gm) −j ), where Lβ̃ is an log-likelihood
estimate using q(k,m) β̃
end end
Let pj = 1K+1 [ 1 + ∑K k=1 1 ( h ≤ h̃(k) )] end
A.2 FAST-AMI-CRT ALGORITHM
Algorithm 2: FAST-AMI-CRT Input: x ∈ RN×D, feature matrix; y ∈ RN , labels Output: p, the p-values for xj ⊥ y |x−j ∀j Fit qβ := qβ(y | x) for j ∈ [1, 2, . . . , D] do
Fit qθ := qθ(xj | x−j) Let x̃ be a dataset such that x̃−j = x−j , and x̃j is randomly sampled from qθ(xj | x−j) Fit qγ := qβ(y | x̃j ,x−j) Let hj ← E [ log ( 1 2qβ(y | x) + 1 2qγ(y | x) )] with respect to dataset (x,y) for k ∈ [1, 2, . . . ,K] do Let x̃(k) be a dataset such that x̃−j = x−j , and x̃j is randomly sampled from qθ(xj |x−j)
Let h̃(k)j ← E [ log ( 1 2qβ(y | x̃) + 1 2qγ(y | x̃) )] with respect to dataset (x̃,y)
end Let pj = 1K+1 [ 1 + ∑K k=1 1 ( h ≤ h̃(k)j )] end

heading: B PROOFS AND DERIVATIONS
text: B.1 DIVERGENCES ARE PROPER TEST STATISTICS In this section, we prove that divergences are proper test statistics. Let DN := {x(i)j ,y(i),x (i) −j}Ni=1 and D̃j,N := {x̃(i)j ,y(i),x (i) −j}Ni=1. We first list the assumptions here:
1. Let t(DN ) be a consistent estimator of Eq(x−j)[K(q(xj ,y | x−j), q(xj | x−j)q(y | x−j))].
2. Let t(D̃j,N ) be a consistent estimator of Eq̃j(x−j)[K(q̃j(x̃j ,y | x−j), q̃j(x̃j | x−j)q̃j(y | x−j))].
3. The cumulative distribution functions of t(DN ) and t(D̃j,N ) are both continuous everywhere.
4. We have access to complete conditionals q(x̃j | x−j)
Proof. We prove that t is a proper test statistic if and only if t(En) is a consistent estimator of Er(x−j)[K(r(xj ,y | x−j), r(xj | x−j)r(y | x−j))]. We do this by showing t yields p-values that are zero under the alternate hypothesis and uniformly distributed under the null.
Recall that the p-value for our test is:
pj (DN ) = E x̃
(i) j ∼q(xj |x (i) −j)
[ 1 ( t (DN )− t(D̃j,N ) ≤ 0 )] .
Under the alternate hypothesis Consider the case where xj 6⊥ y | x−j . As N →∞,
t(D̃j,N ) q−→ Eq̃j(x−j)K (q̃j(x̃j ,y |x−j) ‖ q̃j(x̃j |x−j)q̃j(y |x−j)) = 0,
where q̃j−→ indicates a convergence in probability.
Since xj 6⊥ y |x−j , notice also that
t (DN ) q−→ Eq(x−j)K (q(xj ,y |x−j) ‖ q(xj |x−j)q(y |x−j)) > 0,
Therefore, the term inside the expectation in the pj(DN ) above is always 0, yielding a p-value of 0 in the limit of N . Since these p-values converge in probability to a single point, the p-values converge in distribution to a delta mass at 0.
Under the null hypothesis In the case where xj ⊥ y | x−j , the samples in qN (y,x) and q̃j,N (y, x̃j ,x−j) are both sampled from the same distribution q = q̃j . Therefore, the distribution of t(DN ) as a function of qN , is the same as that of t(D̃j,N ) as a function of q̃j,N .
Let FN be the cumulative distribution function of t(D̃j,N ) which in this case is the same as that of t(DN ).
We rewrite the p-value expression as pNj := pj(DN ) = P ( t(D̃j,N ) ≤ t(DN ) ) = FN (t(DN )). Now let F−1N (·) be the generalized inverse cumulative distribution function which exists because FN is a continuous everywhere function. With this, we derive the distribution of the p-value:
P (pNj ≤ ρ) = P (F−1N (p N j ) ≤ F−1N (ρ)) = P (t(DN ) ≤ F −1 N (ρ)) = FN (F −1 N (ρ)) = ρ.
This means pNj is uniformly distributed under the null. This result holds for all values ofN . Thus p N j forms a sequence of random variables, indexed by N , that are identically distributed as a uniform random variable over [0, 1]. This means that the sequence converges in distribution to a uniform distribution over [0, 1].
This shows that t is a proper test statistic.
Continuity of FN Earlier, we assumed the cumulative distribution function (CDF) FN (t(DN )) = P (t(D̃j,N ) ≤ t(DN )) was continuous everywhere. This is required for the generalized inverse distribution function F−1N to be well-defined on the full range [0, 1].
Discontinuities could occur when the event t(DN ) = t(D̃j,N ) occurs with some non-zero probability c. This means that the p-value does not take all the values in [0, 1]. To see this, note that FN (·) 6∈ [x, x+ c) for some x ∈ [0, 1− c).
To remedy this, we can replace the indicator function in our test-statistic with the following function:
1(a, b) =  1 if a < b 0 if a > b Uniform([0, 1]) if a = b ,
where Uniform([0, 1]) is a continuous uniform random variable.
With this new function, consider the events, t(DN ) = t(D̃j,N ) and t(DN ) 6= t(D̃j,N ). Conditioning on t(DN ) = t(D̃j,N ), the distribution of the p-value is the same as the uniform random variable : Uniform([0, 1]). Also note that the distribution F (· | t(DN ) 6= t(D̃j,N )) is continuous everywhere in its support because t(DN ) = t(D̃j,N ) occurs with zero probability.
Thus, this modification ensures that the p-values are distributed uniformly.
C INSTANCE-WISE FEATURE SELECTION EXAMPLES C.1 CONSISTENT PREDICTIONS ALONE ARE INSUFFICIENT FOR INSTANCE-WISE FEATURE
SELECTION
Recall our sufficiency condition for instance-wise feature selection as mentioned in Proposition 1. In this example, we see what happens when this condition is not met. We notice that this definition does not suffice to reject an unimportant feature. Consider a simple data generating process:
y = zx1 + (1− z)x2 + z ∼ Bernoulli(0.5)
x1,x2 ∼ N (0, σ2x) ∼ N (0, σ2 )
where z is not observed. We can now write out the probability distributions we care about. Note that taking an expectation like Ex̃1∼q(x1|x2)q(y|x̃1x2) yields q(y|x2). For simplicity, we leave out the use of complete conditions and work directly with the latter probability distributions:
p(y|x1,x2, z) = N ( zx1 + (1− z)x2, σ2 ) p(y|x1,x2) = ∫ p(y|x1,x2, z)p(z|x1,x2)dz = ∫ p(y|x1,x2, z)p(z)dz
= 1 2 N ( x1, σ 2 ) + 1 2 N ( x2, σ 2 ) p(y|x2) = 1 2 N ( 0, σ2x + σ 2 ) + 1 2 N ( x2, σ 2
) p(y|x1) = 1 2 N ( x1, σ 2 ) + 1 2 N ( 0, σ2x + σ 2
) Now consider an instance (x(i)1 ,x (i) 2 ,y
(i), z(i)) where z(i) = 1 which means that y(i) depends only on feature x(i)1 . Now we check if p(y (i)|x(i)1 ,x (i) 2 ) > p(y (i)|x(i)1 ). Using our definitions from
earlier, we can expand this inequality:
p(y(i)|x(i)1 ,x (i) 2 )− p(y(i)|x (i) 1 ) =
1 2 N ( y(i);x (i) 2 , σ 2 ) − 1 2 N ( y(i); 0, σ2x + σ 2 )
For all i such that y(i) lies in a non-0 interval around x(i)2 , we have that p(y (i)|x(i)1 ,x (i) 2 ) − p(y(i)|x(i)1 ) > 0. For example let σ = σx = 1, then x (i) 2 = 5, we have that y ∈ [3, 17] satisfies this. This means that x2 will be deemed important as per the candidate Proposition 1.
C.2 INSTANCE-WISE SCORE EXAMPLE In this example, we see how scores computed using Equation (5) can help identify important features for a given instance, under the assumptions stated in Proposition 1. Consider a simple data generating process:
y = zx1 + (1− z)x2 z,x1,x2 ∼ Bernoulli(0.5)
where all random variables are observed. Let us now consider the following observed instance: (y(i),x
(i) 1 ,x (i) 2 , z (i)) = (1, 0, 1, 0). We can now devise a test for each of x(i)1 , x (i) 2 , and z (i). For x (i) 1 , we want to check:
p(y(i)|x(i)1 ,x (i) 2 , z (i)) > p(y(i)|x(i)2 , z(i)) (8)
We can create similar tests for the other two variables as well:
p(y(i)|x(i)1 ,x (i) 2 , z (i)) > p(y(i)|x(i)1 , z(i)) (9)
p(y(i)|x(i)1 ,x (i) 2 , z (i)) > p(y(i)|x(i)1 ,x (i) 2 ) (10)
We can use Table 5 to help evaluate Equations (8) to (10):
p(y(i) = 1|x(i)1 = 0,x (i) 2 = 1, z (i) = 0) = 1
p(y(i) = 1|x(i)2 = 1, z(i) = 0) = 1
p(y(i) = 1|x(i)1 = 0, z(i) = 0) = 0.5
p(y(i) = 1|x(i)1 = 0,x (i) 2 = 1) = 0.5
meaning x(i)1 is not important to y (i), but x(i)2 and z (i) are important.

heading: D SIMULATED DATA FEATURE SELECTION - ADDITIONAL RESULTS
text: In this section, we present additional results that use AMI as a test statistic in a HRT framework. This offers a significant speedup as the HRT framework avoids having to refit estimators using xj ∼ q(xj | x−j). Figure 3 shows a quantile-quantile plot of the null p-values for each FDR-controlling feature selection method. We notice that both HRT-based methods tend to deflate p-values. This often results in features being mistakenly selected as important. Using the same test statistic, AMI, in a CRT helps mitigate this issue.

heading: E CELIAC DISEASE GENOMIC FEATURE SELECTION
text: Table 6 shows the set of SNPs deemed significant by AMI-CRT. We annotate each SNP with its position in the human genome, and whether it was previously reported as significant in a biological study.

heading: F HOSPITAL READMISSION FEATURE SELECTION
text: For the hospital readmission dataset, we applied several standard pre-processing techniques. We filtered each sample the data in a manner similar to (Strack et al., 2014):
• It is an inpatient encounter (a hospital admission).
• It is a diabetic encounter, that is, one during which any kind of diabetes was entered to the system as a diagnosis.
• The length of stay was at least 1 day and at most 14 days.
• Laboratory tests were performed during the encounter.
• Medications were administered during the encounter.
Further, we binarized the labels so that a label of 1 indicates a readmission event, and a label of 0 indicates no readmission event. We then encoded each categorical feature as a one-hot encoding. We then imputed missing values using the median across the dataset, and dropped the “weight” feature as it was found to be 97% missing.
To sample from the complete conditional distributions q(xj | x−j), we used a neural network to fit the complete conditional regression detailed in (Miscouridou et al., 2018). For continuous values of xj , we first discretized the data into bins, then used our neural network to predict the bins. To map the bins back to values in the domain of xj , we used the mean of the range of values in each bin.
G IMAGENET INSTANCE-WISE FEATURE SELECTION
Figure 5 shows some of the results of instance-wise feature selection on ImageNet data using AMICRT. Figures 6 and 7 show results on the same task, using LIME and SHAP respectively. We notice
that AMI-CRT identifies patches that seem more likely to help differentiate between ambulances and policevans. AMI-CRT identifies relevant text like the words “ambulance” or “police” that are very likely to help distinguish between the two classes. LIME identifies some relevant features of the image like wheels and lights, but fails to identify relevant words. SHAP does a good job at identifying distinguishing symbols like the caduceus and the FBI logo, but occasionally misses out on relevant text.

Review summaries:
This paper addresses supervised feature selection: given a D-dimensional input variable x = (x_1, ..., x_D), and a response variable y, the goal is to find a subset of "useful" features in x. Here, a feature x_j is useful if it is dependent on y even when conditioning on all other input variables (denoted by x_{-j}, which is a set). A generic procedure that can produce a p-value for each feature (allowing on to test each feature whether it is useful) is the conditional randomization test (CRT) proposed in Candes et al., 2018.  For the CRT to produce a valid p-value for each feature (input dimension) x_j, one needs to specify a test statistic that measures conditional dependence between x_j and y given the rest of the features.

This paper presents a method to provide some level of interpretation on the influence of input features on the response of a machine level model all the way down to the instance level. The proposed method is model agnostic. Quoting the authors, they advocate for methods that look at interpretability “as understanding the population distribution through the lens of the model” without restriction on the models fit. The problem is posed as a hypothesis testing problem. The paper proposes “proper test statistics” for model agnostic feature selection. It is argued that f-divergence tests are proper statistic tests, with the KL being particularly interesting as it provides computational advantages.