B LACK B OX F EATURE S ELECTION WITH A DDITIONAL M UTUAL I NFORMATION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 inﬂuence 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 insufﬁcient for instance-wise feature selection.
We show that the estimators from FAST - AMI - CRT can also be reused to ﬁnd 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, identiﬁes biologically signiﬁcant genes, selects the most important predictors of a hospital readmission event, and identiﬁes distinguishing image regions in an image-classiﬁcation task.
I NTRODUCTION
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 ﬁt 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 coefﬁcients (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 simpliﬁes 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 ).
A MI - 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 signiﬁcance.
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 beneﬁt 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 sufﬁcient 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 speciﬁc to fully connected networks.
Tansey et al.
(2018) propose holdout randomization tests ( HRT s) 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 HRT s and show that this produces better calibration than other choices of loss.
H RT s provide computational speed-ups over CRT s.
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 ﬁtting 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 deﬁne 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 ﬁrst identifying a set of sufﬁcient 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 .
P ROPER TESTS FOR FEATURE SELECTION
Practitioners of machine learning use feature selection to identify important features for their predictive task.
One way to ﬁlter out important features is to ﬁnd those that improve predictions given the rest of the features.
This can be formalized through conditional independence.
Let x j be the j th feature of x and let x − j be all features but the j th one.
The goal is to discover a set S such that ∀ x j (cid:54)∈ 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 ﬁnite set of samples D N := { ( x ( i ) , y ( i ) )
} Ni =1 sampled from the population.
This means that it is impossible to assess exact conditional independence.
Therefore, in the ﬁnite sample setting, we must formulate a statistical hypothesis test.
A conditional randomization test ( CRT ) (Candes et al., 2018) deﬁnes a hypothesis test for conditional independence.
For the j th feature, CRT s ﬁrst compute a test statistic t using the N samples of data
D N .
CRT s place this statistic in a null distribution where samples of the j th feature x j are replaced by samples of (cid:101) x
j | x − j which by construction satisfy (cid:101) x j ⊥ y | x − j .
Letting (cid:101)
D j,N be a dataset where { x ( i ) j } Ni =1 has been replaced by { (cid:101)
x ( i ) j } Ni =1 , the p -value for a CRT is p j ( D N ) =
[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 CRT s provide a general method for conditional independence testing, they leave several components including the choice of test statistic unspeciﬁed.
2. C HOOSING THE RIGHT TEST STATISTIC
Imagine a test statistic t ( · ) =
t ( { x ( i ) j , y
( i ) }
Ni =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.
Deﬁnition 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 (cid:54)⊥ 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.
Deﬁnition 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:
[Equation]
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 ) :
=
E r ( x − j )
[ K ( r ( x j , y | x − j ) , r ( x j | 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 ﬁtting univariate regressions that can reuse pre-developed model code from the features to the response.
Deﬁne the resampling distribution ˜ q j = ˜ q j ( (cid:101) 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 ) , ˜ q j ( y | (cid:101) x j , x − j ) , and q ( y | x − j ) .
The ﬁrst 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 ˜ q j ( y | (cid:101) 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 modiﬁed for different padding and ﬁlter 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 CRT s, q ( x j | x − j ) , and model code to ﬁt the response from the features.
2. AMI - CRT
Using the KL -divergence as a test statistic in Equation (1) requires the difference δ j :
[Equation] [Equation]
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 simpliﬁcation 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 ﬁnite 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, (cid:101) D j,N is sampled iid from q ( x − j , y ) q
( (cid:101)
x j | x − j ) .
The samples x − j , y come from the population distribution, however sampling q ( (cid:101) 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 ( (cid:101) 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 θ ( (cid:101) x
j | x − j ) and q β ( y | x ) are ﬁt to the training set.
The distribution q γ ( y | (cid:101) x j , x − j ) is ﬁt to the training part of (cid:101) 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 θ ( (cid:101) x
j | x − j ) and q β ( y | x ) can be shared across the replications, however q γ ( y | (cid:101) 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 | (cid:101)
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 HRT s (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 ﬁtting 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.
T ESTS 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 deﬁnition for important features:
Deﬁnition 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 j th feature for the i th sample, x ( i ) j , is an important feature if
[Equation]
Deﬁnition 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, Deﬁnition 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 − z ) x 2 + (cid:15) , z ∼ Bernoulli (0 . 5) , x 1 , x 2 ∼ N (0 , σ 2 x ) , and (cid:15) ∼ N (0 , σ 2 (cid:15) ) .
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:
[Equation]
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 (cid:15) =
σ 2 x = 1 , and x ( i )= 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 Deﬁnition 2.
We show the full derivation of this example in Appendix C.1.
The fundamental issue with the formulation in Deﬁnition 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 ) sufﬁce for understanding population distributions, they might not be sufﬁcient to perform IWFS .
3. S UFFICIENT CONDITIONS FOR INSTANCE - WISE FEATURE SELECTION
We develop the following condition under which q ( y | x ) is sufﬁcient to perform IWFS :
Proposition 1.
Sufﬁcient 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 ) deﬁned 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 . [Table]
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 sufﬁcient conditions in Proposition 1, we can now construct an IWFS procedure using the same estimators from AMI - CRT or FAST - AMI - CRT .
A MI I NSTANCE -W ISE F EATURE S ELECTION
Instance-wise feature selection can be performed using the building blocks of the FAST - AMI - CRT .
Starting from Deﬁnition 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 satisﬁes the instance-wise sufﬁcient 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 ) | x j = (cid:101) 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 (cid:101) 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 (cid:101) 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 | (cid:101) 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.
E XPERIMENTS
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:
[Table] [Table]
• Correlation [ corr-crt ]: Difference between Correlation( x j , y ) and Correlation( (cid:101) 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 classiﬁ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)
H RT s construct a test by comparing a loss function evaluated on x and (cid:101) 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 speciﬁc 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 speciﬁed 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 ﬁrst sample x ∼ N (0 , Σ D )
N times, where Σ D is a Ddimensional covariance matrix.
We translate the ﬁrst 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 ﬁrst two coordinates with the same sign, the label is one.
Otherwise, it is zero.
All but the ﬁrst 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 (cid:16)(cid:80) (cid:96)j =1 x 2 j − (cid:96) (cid:17) > 0 . 5 and 0 otherwise, where (cid:96) <
D is the number of important features.
We choose N = 3000 , D = 20 , and (cid:96) = 4 .
[ selector , noisy-selector ]: These experiments test instance-wise feature selection methods.
We ﬁrst sample x ∼ N (0 , Σ D )
N times, where Σ D is a D -dimensional covariance matrix, and D ≥ 11 .
The ﬁrst 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 .
[Equation]
Results.
For methods based on CRT s or HRT s, 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 identiﬁed 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 , (cid:96) = 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 deﬂated and signiﬁcantly non-uniform p -values ( p = 0 . 0006 ), implying dependence.
As a result loss-hrt incorrectly identiﬁes 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 | (cid:101)
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 reﬁtting ( CRT s) and not reﬁtting ( HRT s), 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 .
We sample (cid:101) x
j from N (0 , 1) .
We see in Figure 4 (Appendix D) that reﬁtting 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 signiﬁcantly 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 HRT s in this setting, we use the procedure prescribed by Burns et al.
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 identiﬁcation, we count the number of instances where the selector variable xwas identiﬁed.
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 sufﬁcient 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 ( SNP s).
SNP s 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 SNP s 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 S j ) 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.
[Table]
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 SNP s 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 SNP s 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 overhospitals (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 classiﬁer with 100 estimators. [Figure]
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 split the image into an 8 × 8 grid so that there are 64 non-overlapping x S 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 ﬁfty 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 ﬁrst 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 (cid:101) 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 identiﬁes 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.
D ISCUSSION
We develop AMI - CRT for testing for conditional independence of each feature x j ⊥ y | x − j from a ﬁnite 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 deﬁne sufﬁcient 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 .
A MI - CRT , FAST - AMI - CRT , and AMI - IW all outperform several popular methods. in various simulated tasks, in identifying biologically signiﬁcant genes, selecting the most indicative features to predict hospital readmissions, and in identifying distinguishing features in an image classiﬁcation task.
A A LGORITHMS A. AMI - CRT ALGORITHM
We show the AMI - CRT procedure in Algorithm 1, and the FAST - AMI - CRT procedure in Algorithm 2.
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 ←for m ∈
, M ] do
Estimate q ( m ) β := q β ( y ( G − m ) | x ( G − m ) )
h ← h
+
1 M L β ( y ( G m ) | x ( G m ) ) , where L β is an log-likelihood estimate using q ( m ) β end
Let ˜ h be a K -dimensional vector of 0 s 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 )
+ 1 M L ˜ β
( y ( G m ) | (cid:101) x
( G m ) j , x ( G m ) − j ) , where L ˜ β is an log-likelihood estimate using q ( k,m ) ˜ β end end
Let p j = 1 K +1 (cid:104)
+ (cid:80)
Kk
=1 1 (cid:16) h ≤ ˜ h
( k ) (cid:17)(cid:105) end
A. FAST - AMI - CRT ALGORITHM
: 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 (cid:101) x be a dataset such that (cid:101) x
− j = x − j , and (cid:101) x
j is randomly sampled from q θ ( x j | x − j )
Fit q γ :=
q β ( y | (cid:101) x j , x − j )
Let h j ← E (cid:2) log (cid:0) 12 q β ( y | x ) + 12 q γ ( y | x ) (cid:1)(cid:3) with respect to dataset ( x , y ) for k ∈ [1 , 2 , . .
, K ] do
Let (cid:101) x
( k ) be a dataset such that (cid:101) x − j = x − j , and (cid:101)
x
j is randomly sampled from q θ ( x j | x − j )
Let ˜ h ( k ) j ← E (cid:2) log (cid:0)
1 2 q β ( y | (cid:101) x ) + 12 q γ ( y | (cid:101) x ) (cid:1)(cid:3) with respect to dataset ( (cid:101) x , y ) end
Let p j = 1 K [-+1 (cid:104) 
-]+
[Equation]
(cid:104) [Equation] end
B P ROOFS AND DERIVATIONS B. D IVERGENCES ARE PROPER TEST STATISTICS
In this section, we prove that divergences are proper test statistics.
Let D N :=
{ x ( i ) j , y ( i ) , x ( i ) − j } Ni = and (cid:101)
D j,N := { (cid:101) x ( i ) j , y ( i ) , x ( i ) − j } Ni =1 .
We ﬁrst list the assumptions here:
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 ))] .
Let t ( (cid:101)
D j,N ) be a consistent estimator of E ˜ q j ( x − j )
[ K (˜ q j ( (cid:101) x j , y | x − j ) , ˜ q j ( (cid:101)
x
j | x − j )˜
q j ( y | x − j ))] .
The cumulative distribution functions of t ( D N ) and t ( (cid:101) D j,N ) are both continuous everywhere.
We have access to complete conditionals q ( (cid:101) 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:
[Equation]
Under the alternate hypothesis Consider the case where x j (cid:54)⊥ y | x − j .
As N → ∞ ,
[Equation]
where ˜ q j −→ indicates a convergence in probability.
Since x j (cid:54)⊥ y | x − j , notice also that
[Equation]
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 ˜ q j,N ( y , (cid:101) x j , x − j ) are both sampled from the same distribution q = ˜ q j .
Therefore, the distribution of t ( D N ) as a function of q N , is the same as that of t ( (cid:101) D j,N ) as a function of ˜ q j,N .
Let F N be the cumulative distribution function of t ( (cid:101) D j,N ) which in this case is the same as that of t ( D N ) .
We rewrite the p -value expression as p Nj :=
p j ( D N ) =
P (cid:16) t ( (cid:101) D j,N ) ≤ t ( D N ) (cid:17) =
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:
[Equation]
This means p Nj is uniformly distributed under the null.
This result holds for all values of N .
Thus p Nj 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 F N
Earlier, we assumed the cumulative distribution function ( CDF ) F N ( t ( D N )) =
P ( t ( (cid:101) D j,N ) ≤ t ( D N )) was continuous everywhere.
This is required for the generalized inverse distribution function F − 1 N to be well-deﬁned on the full range [0 , 1] .
Discontinuities could occur when the event t ( D N ) =
t ( (cid:101) 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
F N ( · ) (cid:54)∈ [ 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:
[Equation]
where Uniform ([0 , 1]) is a continuous uniform random variable.
With this new function, consider the events, t ( D N )
=
t ( (cid:101) D j,N ) and t ( D N ) (cid:54) =
t ( (cid:101) D j,N ) .
Conditioning on t ( D N ) =
t ( (cid:101) 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 ( D N ) (cid:54) =
t ( (cid:101) D j,N )) is continuous everywhere in its support because t ( D N ) =
t ( (cid:101) D j,N ) occurs with zero probability.
Thus, this modiﬁcation ensures that the p -values are distributed uniformly.
C I NSTANCE - WISE FEATURE SELECTION EXAMPLES C. C ONSISTENT PREDICTIONS ALONE ARE INSUFFICIENT FOR INSTANCE - WISE FEATURE SELECTION
Recall our sufﬁciency 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 deﬁnition does not sufﬁce to reject an unimportant feature.
Consider a simple data generating process:
[Equation]
where z is not observed.
We can now write out the probability distributions we care about.
Note that taking an expectation like E (cid:101) x 1 ∼ q ( x 1 | x 2 ) q
( y | (cid:101) x 1 x 2 ) yields q ( y | x 2 ) .
For simplicity, we leave out the use of complete conditions and work directly with the latter probability distributions:
[Equation]
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 deﬁnitions from earlier, we can expand this inequality:
[Equation]
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 σ (cid:15) =
σ x = 1 , then x ( i ) 2 = 5 , we have that y ∈
[3 , 17] satisﬁes this.
This means that x 2 will be deemed important as per the candidate Proposition 1.
C. I NSTANCE - 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:
[Equation]
where all random variables are observed.
Let us now consider the following observed instance: [Equation] x [Equation]
i )
[Equation]
z
[Equation]
i ) )
[Equation] [Equation] [Equation] [Equation]
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:
[Equation]
We can create similar tests for the other two variables as well:
[Equation]
We can use Table 5 to help evaluate Equations (8) to (10):
[Equation]
meaning x ( i ) 1 is not important to y ( i ) , but x ( i ) 2 and z ( i ) are important.
[Equation]
D S IMULATED DATA FEATURE SELECTION - ADDITIONAL RESULTS
In this section, we present additional results that use AMI as a test statistic in a HRT framework.
This offers a signiﬁcant speedup as the HRT framework avoids having to reﬁt estimators using x j ∼ q ( x j | 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 deﬂate 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. [Figure]
E C ELIAC DISEASE GENOMIC FEATURE SELECTION
Table 6 shows the set of SNP s deemed signiﬁcant by AMI - CRT .
We annotate each SNP with its position in the human genome, and whether it was previously reported as signiﬁcant in a biological study.
Position
SNP
Featured in previous study chr2: rs yes chr2: rs yes chr3: rs yes chr3: rs yes chr3: rs
yes chr4: rs yes chr6: rs yes chr6: rs yes chr6: rs no chr6: rs
no chr6: rs no chr6: rs
[Table]
[Figure]
F H OSPITAL READMISSION FEATURE SELECTION
For the hospital readmission dataset, we applied several standard pre-processing techniques.
We ﬁltered 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 ofindicates 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 ﬁt the complete conditional regression detailed in (Miscouridou et al., 2018).
For continuous values of x j , we ﬁrst 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 I MAGE N ET 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 identiﬁes patches that seem more likely to help differentiate between ambulances and policevans.
A MI - CRT identiﬁes relevant text like the words “ambulance” or “police” that are very likely to help distinguish between the two classes.
L IME identiﬁes some relevant features of the image like wheels and lights, but fails to identify relevant words.
S HAP does a good job at identifying distinguishing symbols like the caduceus and the FBI logo, but occasionally misses out on relevant text.
[Figure]
[Figure] [Figure]