Title: BLACK BOX FEATURE SELECTION WITH ADDITIONAL MUTUAL INFORMATION

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 Im-ageNet. 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.
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 (AMI-IW). We demonstrate AMI-IW on multiple simulations and image data. Across all of these tasks AMI-IW outperforms popular baselines.

1. INTRODUCTION:

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 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.

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 x j 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 ∀x j ∈ S, x j ⊥ 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 D N := {(x (i) , y (i) )} N i=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 D N . CRTs place this statistic in a null distribution where samples of the jth feature x j 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

EQUATION

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.

2.1. CHOOSING THE RIGHT TEST:

STATISTIC Imagine a test statistic t(•) = t({x (i) j , y (i) } N i=1

) that uses only a feature x j 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(D N ) 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:

EQUATION

where d -→ indicates a convergence in distribution. Under the alternate hypothesis, which in the case of feature selection is x j ⊥ y | x -j , the power to reject the null hypothesis must be 1, implying p j → 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(x j , y | x -j ) = r(x j | 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. Let K(a, b) ≥ 0 with equality holding only when a is equal in distribution to b, then a proper test statistic K j (r Define the resampling distribution qj = qj ( x j | x -j )q(y, x -j ). Using a divergence in a CRT requires estimates of the following conditional distributions: q(x j | x -j ), q(y | x), qj (y | x j , x -j ), and q(y | x -j ). The first distribution q(x j | 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 qj (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(x j | x -j ), and model code to fit the response from the features.

) := E r(x-j ) [K(r(x j , y | x -j ), r(x j | x -j )r(y | x -j ))]. A consistent

2.2. AMI-CRT:

Using the KL-divergence as a test statistic in Equation (1) requires the difference δ j :

δ j =E q(x-j ) [D KL (q(x j , y | x -j ), q(x j | x -j )q(y | x -j ))] -E qj (x-j ) [D KL (q j ( x j , y | x -j ), qj ( x j | x -j )q j (y | x -j ))] =E q(x-j ) [E q(xj ,y|x-j ) [log q(x j , y | x -j ) -log(q(x j | x -j )q(y | x -j ))] -E qj ( xj ,y|x-j ) [log qj ( x j , y | x -j ) -log(q j ( x j | x -j )q j (y | x -j ))]] =E q(x-j ) [E q(xj ,y|x-j ) [log q(y | x) -log q(y | x -j )] -E qj ( xj ,y|x-j ) [log qj (y | x -j ) -log qj (y | x -j )]] =E q(x-j ) [E q(xj ,y|x-j ) [log q(y | x)] -E qj ( xj ,y|x-j ) [log qj (y | x j , x -j )]] =E q(x,y) [log q(y | x)] -E q(x-j ,y) E qj ( xj |x-j ) [log qj (y | x j , x -j )]].

The second equality above follows from q(x -j ) = qj (x -j ), and the fourth from q(y | x -j ) = qj (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 D N 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 D N 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.

3. TESTS FOR INSTANCE-WISE FEATURE SELECTION:

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) )} N i=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

= zx 1 +(1-z)x 2 + , z ∼ Bernoulli(0.5), x 1 , x 2 ∼ N (0, σ 2

x ), and ∼ N (0, σ 2 ). Assume we have the true q(y | x 1 , x 2 ), 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, σ 2 x + σ 2

For all i such that y (i) lies in a non-0 interval around 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.

EQUATION

x (i) 2 = 5, then y (i) ∈ [

3.1. SUFFICIENT CONDITIONS FOR INSTANCE-WISE FEATURE SELECTION:

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:

EQUATION

)

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. (Benjamini and Hochberg, 1995) , allow instance-wise feature selection, and make no distributional assumptions about the data-generating process. This table compares these methods to widely-used feature selection methods.

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.

3.2. AMI INSTANCE-WISE FEATURE SELECTION:

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:

EQUATION

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) , qj (y = y (i) | x j = 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 FAST-AMI-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 qj (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.

4. EXPERIMENTS:

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: (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(x j |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 choose N = 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 x 1 , 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(x j | 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(x j | 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 x 1 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 x 1 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(x j |x -j ), we use the same procedures as (Candes et al., 2018) where we estimate q(x j |x Sj ) where S j is only the set of SNPs (not including x j ) known to be correlated with x j . To model q(y|x), we use an L 1 -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. (Strack et al., 2014) as a clinically validated ground truth, we observe that AMI-CRT is able to achieve the highest area under the receiver operating characteristic curve (AUROC) when compared to state-of-the-art benchmarks.

Results. The ground truth features come from clinical validation done by (Strack et al., 2014) . We use importance scores (or pvalues) estimated by each selection method with these ground truth features to compute an ROC curve. Figure 1 shows these curves for each method. We observe that ami-crt and fast-ami-crt achieve a higher area under the ROC curve than state-of-the-art approaches. The loss-hrt performs well, but achieves low power at false positive rates less than 0.5. These methods, unlike locally-linear methods such as lime and shap, do not assume that relevant features are marginally independent of irrelevant features (Lundberg and Lee, 2017) .

ImageNet experiments: We consider the task of differentiating between ambulances and policevans. This task is interesting as both objects are physically very similar and there are only a few features that can be used to differentiate the two. For example, both objects have windows, wheels, and doors, so other features must be used to distinguish between the two classes. Rather than consider each pixel as an individual feature x j , we consider a patch of pixels x S as a single feature, such that no two patches contain overlapping pixels. To model the distribution q(x S |x -S ), we make use of a generative inpainting model π g (Yu et al., 2018) . We the image into an 8 × 8 grid so that there are 64 non-overlapping x S patches. To model q(y|x), we use VGG-16 network (Simonyan and Zisserman, 2015) . To perform our instance-wise test, we compute log-probability differences using fifty generated samples from q(x S |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(x S | 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.

5. DISCUSSION:

We develop AMI-CRT for testing for conditional independence of each feature x j ⊥ 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.

A ALGORITHMS 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 ∈ R N ×D , feature matrix; y ∈ R N , labels Output: p, the p-values for x j ⊥ y | x -j ∀j Split x and

y into M folds, G 1 , G 2 , . . . , G M for j ∈ [1, 2, . . . , D] do Estimate q θ := q θ (x j | x -j ) h ← 0 for m ∈ [1, 2, . . . , M ] do Estimate q (m) β := q β (y (G-m) | x (G-m) ) h ← h + 1 M L β (y (Gm) | x (Gm) )

, where L β is an log-likelihood estimate using q

EQUATION

-j ), where L β is an log-likelihood estimate using q

(k,m) β end end Let p j = 1 K+1 1 + K k=1 1 h ≤ h(k) end A.2 FAST-AMI-CRT ALGORITHM Algorithm 2: FAST-AMI-CRT Input: x ∈ R N ×D , feature matrix; y ∈ R N , labels Output: p, the p-values for x j ⊥ y | x -j ∀j Fit q β := q β (y | x) for j ∈ [1, 2, . . . , D] do Fit q θ := q θ (x j | x -j )

Let x be a dataset such that x -j = x -j , and x j is randomly sampled from q θ (x k) be a dataset such that x -j = x -j , and x j is randomly sampled from

j | x -j ) Fit q γ := q β (y | x j , x -j ) Let h j ← E log 1 2 q β (y | x) + 1 2 q γ (y | x) with respect to dataset (x, y) for k ∈ [1, 2, . . . , K] do Let x (

q θ (x j | x -j ) Let h(k) j ← E log 1 2 q β (y | x) + 1 2 q γ (y | x) with respect to dataset ( x, y) end Let p j = 1 K+1 1 + K k=1 1 h ≤ h(k) j end

Under review as a conference paper at ICLR 2020 

EQUATION

-j } N i=1 . We first list the assumptions here:

1. Let t(D N ) be a consistent estimator of E q(x-j ) [K(q(x j , y | x -j ), q(x j | x -j )q(y | x -j ))].

2. Let t( D j,N ) be a consistent estimator of E qj (x-j ) [K(q j ( x j , y | x -j ), qj (

x j | x -j )q j (y | x -j ))].

3. The cumulative distribution functions of t(D N ) 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(E n ) is a consistent estimator of E r(x-j ) [K(r(x j , y | x -j ), r(x j | 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:

p j (D N ) = E x (i) j ∼q(xj | x (i) -j ) 1 t (D N ) -t( D j,N ) ≤ 0 .

Under the alternate hypothesis Consider the case where

x j ⊥ y | x -j . As N → ∞, t( D j,N ) q -→ E qj (x-j ) K (q j ( x j , y | x -j ) qj ( x j | x -j )q j (y | x -j )) = 0,

where qj -→ indicates a convergence in probability.

Since x j ⊥ y | x -j , notice also that

t (D N ) q -→ E q(x-j ) K (q(x j , y | x -j ) q(x j | x -j )q(y | x -j )) > 0,

Therefore, the term inside the expectation in the p j (D N ) 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 x j ⊥ y | x -j , the samples in q N (y, x) and qj,N (y, x j , x -j ) are both sampled from the same distribution q = qj . Therefore, the distribution of t(D N ) as a function of q N , is the same as that of t( D j,N ) as a function of qj,N .

Let F N be the cumulative distribution function of t( D j,N ) which in this case is the same as that of t(D N ).

We rewrite the p-value expression as p N j := p j (D N ) = P t( D j,N ) ≤ t(D N ) = F N (t(D N )). Now let F -1 N (•) be the generalized inverse cumulative distribution function which exists because F N is a continuous everywhere function. With this, we derive the distribution of the p-value:

P (p N j ≤ ρ) = P (F -1 N (p N j ) ≤ F -1 N (ρ)) = P (t(D N ) ≤ F -1 N (ρ)) = F N (F -1 N (ρ)) = ρ.

This means p N j is uniformly distributed under the null. This result holds for all values of N . 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]. 

F HOSPITAL READMISSION FEATURE SELECTION:

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(x j | x -j ), we used a neural network to fit the complete conditional regression detailed in (Miscouridou et al., 2018) . For continuous values of x j , 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 x j , 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 AMI-CRT. 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.

Ambulance (original) Ambulance (masked) Policevan (original) Policevan (masked) 

:



Figures and Tables:

Figure1: Using results from(Strack et al., 2014)

Figure 2: Instance-wise feature selection using AMI. The first and third columns show the original image of ambulances or policevans respectively. The second and fourth columns show only the patches which were found to have non-zero AMI with the label, given the rest of the patches.

Figure4: Quantile-Quantile plot showing uniformity of p-values across various FDR-controlling methods in the case of poor approximations of q(x j | x -j ). A Kolmogorov-Smirnov (KS) test is performed to test if each set of p-values is uniformly distributed.

Figure 5: Instance-wise feature selection using AMI-CRT. The first and third columns show the original image of ambulances or policevans respectively. The second and fourth columns show only the patches which were found to have non-zero AMI with the label, given the rest of the patches.

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.

Table content: <table/>

Both AMI-CRT and FAST-AMI-CRT produce well-calibrated p-values, provide false discovery rate (FDR) control

Table content: <table><tr><td>Feature ami-crt fast-ami</td><td>loss-hrt corr-crt lime shap rf</td></tr><tr><td>-crt</td><td/></tr><tr><td>Provides p-values</td><td/></tr><tr><td>Well calibrated p-values</td><td/></tr><tr><td>Instance-wise feature selection</td><td/></tr><tr><td>No distributional assumptions</td><td/></tr><tr><td>Robust to q(x j | x -j ) estimation</td><td/></tr></table>

Dataset ami-crt fast-ami-crt loss-hrt corr-crt lime shap rf Simulated data results: Here we use the scores provided by each method to select features. AMI-CRT's area under the receiver operating characteristic (ROC) curve is better than other of state-of-the-art methods.

Table content: <table><tr><td>orange</td><td>0.97</td><td>0.95</td><td/><td>0.94</td><td>0.22</td><td>0.94</td><td>0.95</td><td>0.94</td></tr><tr><td>xor</td><td>1.00</td><td>0.97</td><td/><td>0.95</td><td>0.45</td><td>1.00</td><td>0.99</td><td>0.95</td></tr><tr><td/><td/><td>Dataset</td><td>ami-iw</td><td colspan="5">loss-hrt corr-crt lime shap rf</td></tr><tr><td/><td>Precision</td><td colspan="2">selector noisy-selector 0.67 0.93</td><td>0.78 0.45</td><td>0.34 0.33</td><td>0.58 0.57</td><td>0.64 0.61</td><td>0.33 0.33</td></tr><tr><td colspan="2">Selector identification</td><td colspan="2">selector noisy-selector 0.97 0.95</td><td>0.88 0.85</td><td>0.34 0.33</td><td>0.35 0.25</td><td>0.39 0.37</td><td>0.33 0.33</td></tr></table>

Instance-wise feature selection results. The precision experiment measures the ability of each selection method to identify relevant features while restricted to 7 features. The selector identification experiment counts the instances where each method identified the selector feature x1 (relevant for all instances).

Table content: <table><tr><td>• Correlation</td><td>[corr-crt]:</td><td>Differ-</td></tr><tr><td colspan="3">ence between Correlation(x j , y) and</td></tr><tr><td colspan="3">Correlation( x j , y) as a test statistic for a CRT</td></tr><tr><td colspan="3">• Local interpretable model-agnostic explana-</td></tr><tr><td colspan="3">tions (LIME) [lime] (Ribeiro et al., 2016)</td></tr><tr><td colspan="3">• Shapley additive explanations (SHAP) [shap]</td></tr><tr><td colspan="2">(Lundberg and Lee, 2017)</td><td/></tr></table>

The number of significant features reported at a 20% FDR level for each test, and the percentage of features previously identified in a biological study.

Table content: <table/>

B PROOFS AND DERIVATIONS B.1 DIVERGENCES ARE PROPER TEST STATISTICSIn this section, we prove that divergences are proper test statistics. Let D

Table content: <table/>

Peer review (iclr2020):
This paper proposes a practical improvement of the conditional randomization test (CRT) of (Candes et al., 2018). In the study of (Candes et al., 2018), the choice of the test statistic as well as how one estimates conditional distributions were kept open. The authors proposed *proper test statistic* as a promising test statistic for CRT, and proved that f-divergence is one possible choice. They further shown that KL-divergence has a nice property among possible f-divergences: KL-divergence cancels out some of the conditional distributions, and thus the users need to estimate only two conditional distributions to compute the test statistic. For estimating those conditional distributions in the test statistic, the authors proposed fitting regression models. Overall, I think the paper is well-written and the idea is stated clearly. The use of KL-divergence for CRT seems to be reasonable. The proposed algorithms look simple and easy to implement. My only concern is on the practical applicability of the proposed algorithms (which, however, may be not a unique problem for this paper, but for all the CRT methods). They require fitting regression models for each feature xj. For high-dimensional data with more than thousands of features, fitting regression models for all the features seem to be impractical. For the imagenet data experiment, the authors successfully avoided this problem by using an inpainting model. However, this approach is apparently limited to image data. I am interested in seeing if there is any promising way to make the algorithms scalable to high-dimensional data. ### Updated after author response ### The authors have partially addressed my concern on the scalability of the proposed algorithm to high-dimensional data. I therefore keep my score unchanged.