Model-Agnostic Feature Selection with Additional Mutual Information

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 the null to assess significance . AMI-CRT works with any regression with probabilistic outputs like cross-entropy-trained neural networks . Training many regressions for each sample from the null can require substantial computation . While this is an embarrassingly parallel computation , we develop FAST-AMI-CRT that only requires a single model trained from null data . FAST-AMI-CRT uses an average of the model from the null with the model from the original data . We show this mixture guards against both variance in model training and poor estimation of the null . Though simple , AMI-CRT outperforms popular procedures for feature importance on a wide variety of simulated data , hospital records , and biological data . Working with data sometimes requires interpreting individual datapoints . For example , a doctor may benefit from knowing which features for a particular patient relate to their health . The process of identifying features at a datapoint-level is called instance-wise feature selection ( Ribeiro et al. , 2016 ; Lundberg and Lee , 2017 ; Gimenez and Zou , 2019 ) . We identify an issue in instance-wise feature selection , where even features selected using the true population distribution do not yield the features that were used to generate the response of an instance . The crux of this disparity is that the response generation process , conditional on the features , may use randomness to select features . We provide an example to demonstrate where instance-wise feature selection can go awry . We develop sufficient conditions for instance-wise feature selection to avoid this issue . The same regression estimates from AMI-CRT can be used to estimate feature-importances with minimal computational overhead , resulting in a method we term additional mutual information instance-wise feature selection ( AMIIW ) . We demonstrate AMI-IW on multiple simulations and image data . Across all of these tasks AMI-IW outperforms popular baselines . Related Work . Permutation tests ( Fisher , 1937 ) provide a test for marginal independence between each feature and the outcome . However , they fail to test conditional independence , which is required when covariates are dependent on each other . To address this , solutions like Sure Independence Screening ( Barut et al. , 2016 ; Fan and Lv , 2008 ) and Conditional Randomization Tests ( Barber et al. , 2015 ; Candes et al. , 2018 ) have been proposed . These outline frameworks for conditional independence testing . However , they often make linearity or additive noise assumptions about the data generating distribution . Furthermore , they require the choice of a test statistic to capture some notion of conditional independence . The user of such frameworks is often burdened with the task of choosing this test statistic , which may require strong assumptions about the data generating distribution . Extending this approach to neural networks , Lu et al . ( 2018 ) propose a fully connected network whose weights are used as a test statistic . Though moving beyond linear models , their method is specific to fully connected networks . Tansey et al . ( 2018 ) propose holdout randomization tests ( HRTs ) that use empirical loss as a test-statistic . The loss they use for continuous-valued distributions of response is the mean-squared-error ( MSE ) , which may ignore higher order dependencies between the response and features . Using AMI inside an HRT would capture these higher order dependencies . In our results , we adapt our test-statistic AMI to HRTs and show that this produces better calibration than other choices of loss . HRTs provide computational speed-ups over CRTs . However this speedup comes at the cost of robustness to poor estimations of the null feature distribution . We demonstrate empirically that FAST-AMI-CRT is robust to such poor estimations . Beyond understanding the population distribution , some tasks require interpreting a population distribution on the level of an individual datapoint . Methods that test for conditional independence work under distributional notions of feature selection , but are not designed to identify the relevant features for a particular sample . To address this issue of “ instance-wise feature selection , ” several methods have been proposed , including local perturbations ( Simonyan et al. , 2013 ; Sundararajan et al. , 2017 ; Ribeiro et al. , 2016 ) and fitting simpler auxiliary models to explain the predictions of a large model ( Chen et al. , 2018 ; Lundberg and Lee , 2017 ; Yoon et al. , 2019 ; Turner , 2016 ; Štrumbelj and Kononenko , 2014 ; Shrikumar et al. , 2017 ) . Our instance-wise work is most similar to that of Burns et al . ( 2019 ) , who repurpose the HRT framework to perform instance-wise feature selection , or Gimenez and Zou ( 2019 ) , who define a conditional randomization test ( CRT ) procedure for subsets of the feature space . In general , however , the conditions under which instance-wise feature selection with predictive models may be possible are not well developed . We address this issue by first identifying a set of sufficient conditions under which instance-wise feature selection is always possible . We then show how estimators used in AMI-CRT can be repurposed for use in an instance-wise setting , yielding a procedure called the AMI-IW . 2 PROPER TESTS FOR FEATURE SELECTION . Practitioners of machine learning use feature selection to identify important features for their predictive task . One way to filter out important features is to find those that improve predictions given the rest of the features . This can be formalized through conditional independence . Let xj be the jth feature of x and let x−j be all features but the jth one . The goal is to discover a set S such that ∀xj 6∈ S , xj ⊥ y | x−j , where independence is with respect to the true population distribution q . The only knowledge about q comes from a finite set of samples DN : = { ( x ( i ) , y ( i ) ) } Ni=1 sampled from the population . This means that it is impossible to assess exact conditional independence . Therefore , in the finite sample setting , we must formulate a statistical hypothesis test . A conditional randomization test ( CRT ) ( Candes et al. , 2018 ) defines a hypothesis test for conditional independence . For the jth feature , CRTs first compute a test statistic t using the N samples of data DN . CRTs place this statistic in a null distribution where samples of the jth feature xj are replaced by samples of x̃j |x−j which by construction satisfy x̃j ⊥ y |x−j . Letting D̃j , N be a dataset where { x ( i ) j } Ni=1 has been replaced by { x̃ ( i ) j } Ni=1 , the p-value for a CRT is pj ( DN ) = E ∀i=1 ... N x̃ ( i ) j ∼q ( xj |x−j=x ( i ) −j ) [ 1 ( t ( DN ) ≤ t ( D̃j , N ) ) ] = E ∀i=1 ... N x̃ ( i ) j ∼q ( xj |x−j=x ( i ) −j ) [ 1 ( t ( DN ) − t ( D̃j , N ) ≤ 0 ) ] , ( 1 ) Under smoothness constraints , the p-value is uniform under the null because it computes the cumulative distribution function of the test statistic under the null . While CRTs provide a general method for conditional independence testing , they leave several components including the choice of test statistic unspecified . 2.1 CHOOSING THE RIGHT TEST STATISTIC . Imagine a test statistic t ( · ) = t ( { x ( i ) j , y ( i ) } Ni=1 ) that uses only a feature xj and the outcome y . Any p-values computed using this test statistic would be meaningless when testing for conditional independence , as t never considers the remaining features x−j . Therefore , particular choices for test statistics limit what can be tested . To address this , we introduce the concept of a proper test statistic . Definition 1 . Proper Test Statistic : A test statistic t ( DN ) is proper if p-values produced by the statistic converge to 0 when the null must be rejected , and are uniformly distributed otherwise . Using t in Equation ( 1 ) , this is : pj ( DN ) d−−−−→ N→∞ { Uniform ( 0 , 1 ) if xj ⊥ y | x−j 0 with probability 1 if xj 6⊥ y | x−j , ( 2 ) where d−→ indicates a convergence in distribution . Under the alternate hypothesis , which in the case of feature selection is xj 6⊥ y | x−j , the power to reject the null hypothesis must be 1 , implying pj → 0 . A proper test statistic requires that Equation ( 2 ) must hold for all distributions of y , x . Proper tests statistics in a CRT select the features in S as the data grows . Definition 1 mirrors the concept of a scoring rule ( Gneiting and Raftery , 2007 ) , which measures the calibration of a probabilistic prediction by a model . A proper scoring rule is one such that the highest expected score is obtained by a model that uses the true probability distribution to make predictions . Divergences are proper test statistics . Conditional independence means the conditional distribution r factorizes : r ( xj , y | x−j ) = r ( xj | x−j ) r ( y | x−j ) . ( 3 ) Divergences measure the closeness between two distributions . A divergence is zero when the two distributions are the same and positive otherwise . Computing any divergence K , like an integral probability metric ( Müller , 1997 ) or an f -divergence ( Csiszár , 1964 ) , between the left hand side and right hand side of Equation ( 3 ) would be a proper test statistic . LetK ( a , b ) ≥ 0 with equality holding only when a is equal in distribution to b , then a proper test statistic Kj ( r ) : = Er ( x−j ) [ K ( r ( xj , y | x−j ) , r ( xj | x−j ) r ( y | x−j ) ) ] . A consistent estimator of this quantity is a proper test statistic ( see Appendix B.1 ) . Casting conditional independence testing as divergence estimation reduces this test to fitting univariate regressions that can reuse pre-developed model code from the features to the response . Define the resampling distribution q̃j = q̃j ( x̃j |x−j ) q ( y , x−j ) . Using a divergence in a CRT requires estimates of the following conditional distributions : q ( xj | x−j ) , q ( y | x ) , q̃j ( y | x̃j , x−j ) , and q ( y | x−j ) . The first distribution q ( xj | x−j ) is required for any CRT . The next distribution q ( y | x ) corresponds to the standard task of building a good regression model . The third distribution q̃j ( y | x̃j , x−j ) requires a regression model with corrupted inputs . This regression can reuse the model structure and code from the standard regression task q ( y | x ) . However , the last distribution q ( y | x−j ) could require development of new model structures . For example , if x is an image , a good model for q ( y | x ) could be a convolutional neural network . If the conditioning set x−j is a subregion of that image , the convolutional neural network used for q ( y | x ) would need to be modified for different padding and filter sizes . This means new models could be needed for each x−j . In the next section , we show that the KL-divergence removes the need for estimating this distribution , and therefore only requires the piece needed for all CRTs , q ( xj | x−j ) , and model code to fit the response from the features .

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