Towards a Uniﬁed Framework for Uncertainty-aware Nonlinear Variable Selection with Theoretical Guarantees
AfﬁliationAddress
Abstract
We develop a simple and uniﬁed framework for nonlinear variable selection thatincorporates uncertainty in the prediction function and is compatible with a widerange of machine learning models (e.g., tree ensembles, kernel methods, neuralnetworks, etc).
In particular, for a learned nonlinear model f ( x ) , we considerquantifying the importance of an input variable x
j using the integrated partialderivative j = k @@ x j f ( x ) k 2 P X .
We then (1) provide a principled approach forquantifying variable selection uncertainty by deriving its posterior distribution,and (2) show that the approach is generalizable even to non-differentiable modelssuch as tree ensembles.
Rigorous Bayesian nonparametric theorems are derivedto guarantee the posterior consistency and asymptotic uncertainty of the proposedapproach.
Extensive simulations and experiments on healthcare benchmark datasetsconﬁrm that the proposed algorithm outperforms existing classic and recent variableselection methods.
Introduction
Variable selection is often of fundamental interest in many data science applications, providingbeneﬁts in prediction error, interpretability and computation by excluding unnecessary variables.
Asdatasets grow in complexity and size, it is crucial that variable selection methods can account forcomplex dependencies among variables while remaining computationally feasible.
Furthermore, asthe number of approaches to model such datasets have increased, it is crucial that the importance ofeach variable can be compared across model classes and extended to new ones as they are developed.
While there are established approaches for variable selection in linear models (e g., LASSO regression[29]), there is little consensus in methodology or theory for variable selection in nonlinear models.
Generalized additive models [30] use similar variable selection methods as their linear counterparts[66] but the additivity assumption for nonlinear functions of the variables is too restrictive in manyapplications.
Random Forests (RF) [7] measure variable importance using an impurity measure,which is based on the average reduction of the loss function were a given variable be removed fromthe model.
[22] extended this method to boosting, where the variable importance is generalized byconsidering the average over all of the decision trees.
Deep neural networks (DNNs) are widely-usedfor many artiﬁcial intelligence applications, and a substantial effort has been invested into developing
DNNs with variable selection capabilities.
Typically, this class of models involves manipulatingthe input layer, for example by imposing an L 1 penalty [11, 19], using backward selection [11], orknockoffs [44].
Unfortunately, each model class based on DNNs requires a tailored variable selectionprocedure, which limits comparability across different model formulations.
Bayesian variable selection methods provide principled uncertainty quantiﬁcation in variable importance estimates as well as a complete characterization of their dependency structure.
These methods allow the variable selection procedure to tailor its decision rule with respect to the correlation structure[42].
Yet, as in frequentist models, each method has a different deﬁnition of a variable’s importance.
For example, in Bayesian additive regression trees (BART), a variable’s importance can be measuredby the proportion of trees that use it [13], while in Gaussian process (GP) models, a variable’simportance can be measured by the frequency of the ﬂuctuations of the resulting function (e.g.,the length-scale parameter as controlled by the automatic relevance determination) in the directionof the variable [48, 68].
Furthermore, the traditional Bayesian modeling procedures tend to becomputationally burdensome, making them less feasible for large-scale applications [1].
Our work starts with the observation that many machine learning models can be written as kernelmethods by constructing a corresponding feature map.
For example, random forests can be writtenas kernel methods by partitions [16], and deep neural networks can be written as kernel methodsby using the last hidden layer as the feature map [62, 31, 9].
Each of these feature maps can beconstructed before the Bayesian learning of the Gaussian process (e.g., by pre-training on the sameor a separate dataset), providing additional modeling expressiveness and representational capacity.
The ability of a GP model to incorporate these adaptive feature maps becomes especially importantin high-dimension applications, where effective dimension reduction is necessary to circumvent thecurse of dimensionality and to ensure good ﬁnite-sample performance [2].
Contributions .
We propose a uniﬁed variable selection framework that is compatible with a widerange of machine learning models that can be deﬁned by, or be closely approximated by, a differentiable feature map.
Notable members include neural networks and random forests (Appendix B).
Our approach deﬁnes variable importance as the norm of the function’s partial derivative, as waspreviously studied in the context of frequentist nonparametric regression [58].
We extend it to apply toa much wider class of models than previously considered (Section 2), propose a principled Bayesianapproach to quantify the variable selection uncertainty in ﬁnite data (Section 3.1), and derive rigorous Bayesian nonparametric theorems to guarantee the method’s consistency and asymptotic optimality(Section 3.2).
To incorporate powerful non-differentiable models into our framework, we also showhow to apply this approach to partition-based methods (e.g., decision trees) by leveraging its (soft) feature representation (Appendix F).
This leads to the ﬁrst derivative-based Bayesian variable selectionapproach for tree-type models that is both theoretically grounded and empirically powerful, stronglyoutperforming other variable selection approaches tailor-designed for random forest (e.g., impurityor random-forest knockoff [8, 10]).
We conduct extensive empirical validation of our approach andcompare its performance to that of many existing methods across a wide range of data generationscenarios.
Results show a clear advantage of the proposed approach, especially in complex scenariosor when the input is a mixture of discrete and continuous features (Section 4).
Preliminaries
Problem Setup .
We consider the classic nonparametric regression setting with d -dimensional featuresx = ( x 1 , . . .
, x d ) 2 X
=
R d and a continuous response y 2 R .
The features x are allowed to have aﬂexible nonlinear effect on y , such that:
[Equation]
with homoscedastic noise level  2 .
The data dimension d is allowed to be large but assumed to beconstant and does not grow with the sample size n .
Here the data-generating function f 0 is a ﬂexiblenonlinear function that resides in a reproducing kernel Hilbert space (RKHS)
H 0 : X 0 !
R induced 76 by a certain positive deﬁnite kernel function k 0 , and the input space of the true function
X 0 spans 77 only a small subset of the input features ( x 1 , . . .
, x d ) , i.e., X 0 ⇢ X .
To this end, the goal of global variable selection is to produce a variable importance score j foreach of the input features ( x 1 , . . .
, x d ) such that it can be used as a classiﬁcation signal for whetherx j 2 X 0 .
As a result, the variable selection decision can be made by threhsolding j > s with a predeﬁned threshold s .
The quality of a variable selection signal j can be evaluated comprehensivelyusing a standard metric such as the area under the receiver operating characteristic (AUROC), whichmeasures the Type-I and Type-II errors of variable selection decision I ( j > s ) over a range ofthresholds s .
2.1 Quantifying Model Uncertainty via Featurized GP
In the nonlinear regression scenario given by Equation (1), a classic approach to uncertainty-awaremodel learning is the Gaussian process (GP).
Speciﬁcally, assuming that f 0 can be described by a ﬂexible RKHS H k governed by the kernel function k , the GP model imposes a Gaussian process priorf ⇠ GP (0 , k ) , such that the function evaluated at any collection of examples follows a multivariatenormal ( MVN ) distribution
[Equation]
with mean m i = m ( x i ) and covariance matrix K i,j = k
( x i , x j ) .
The choice of the prior meanm and kernel k enable prior speciﬁcation directly in function space.
For example, the Matérnkernel places a prior over d ⌫ e  1 times differentiable functions, with length-scale l 2 and amplitudevariance  2 .
As ⌫ ! 1 , this reduces to the common radial basis function (RBF) kernel k ( x i , x j ) =
2 exp( k x i  x j k 22 /l 2 ) .
Under the above construction, the posterior predictive distribution of f evaluated at new observationsx ⇤ 1 , . . .
, x ⇤ n ⇤ is also a multivariate normal,
[Equation]
where
[Equation]
with m ⇤ i
= m ( x ⇤ i ) , K ⇤ ij = k
( x ⇤ i , x j ) , and K ⇤⇤ ij = k
( x ⇤ i , x ⇤ j ) .
Equation (2) is known as the kernelbased representation (or dual representation) of f [55].
Although mathematically elegant, the posterior(2) is expensive to compute due to the need to invert the n ⇥ n matrix ( K +  2 I )  1 .
Feature-based Representation of GP .
Alternatively, Mercer’s theorem [15] states that as long asthe kernel function k ( · , · ) can be written as the inner product of a set of basis functions
( x ) =
{  k ( x ) } Dk =1 , such that k ( x , x 0 ) =
( x ) >  ( x 0 ) , then elements of the RKHS
f 2 H k can be writtenin terms of a linear expansion of basis functions [55]:
[Equation]
where  ⇠
[Equation]
This is known as the feature-based representation (or primal representation) of a Gaussian process. Notice that (3) is not an approximation method but an exact reparametrization of the Gaussian processmodel whose kernel function is induced by feature representation  ( x ) .
Scalable Posterior Computation via Minibatch Updates .
The above feature-based representationis powerful in that it reduces the GP posterior inference into a Bayesian linear regression problem for .
This brings two concrete beneﬁts.
First, the posterior of  in Equation (3) adopts a closed-form:
[Equation]
where
[Equation]
[Equation]
where  =
(  ( x 1 ) > , . . .
,  ( x n ) > ) >
2 R n ⇥ D is the feature matrix evaluated on the training data[55].
For large-scale applications, Equation (4) enables us to compute the exact posterior of  ina mini-batch fashion.
For example, the posterior matrix Cov [  ] =
⌃  can be updated using the
Woodbury identity:
[Equation]
where  m is the D -dimension batch-speciﬁc feature matrix evaluated on the mini-batch.
Similarly,the posterior mean E [  ] can be computed by accumulating the D ⇥ 1 vector  > ( y   µ ) =
P m  > m ( y m   m µ ) , and compute the posterior mean according to Equation (4) at the end.
The posterior distribution of  induces a Gaussian process posterior for the prediction functionf
⇤ =
⇤  , where  ⇤ is the feature map evaluated on the test data, with mean E [ f ⇤ ] =
⇤ µ + ⇤ ⌃
> ( y   µ ) /  2 and covariance Cov
[ f ⇤ ] =
⇤ ⌃
⇤> .
This distribution is equivalent to thekernel-based representation (2) but reduces the computational complexity from cubic time O ( n 3 ) toa linear time O ( n ) and is minibatch compatible (i.e., Equation (5)).
Incorporating Modern ML Model Classes .
The second key advantage of the feature-basedrepresentation (3) is its generality: a wide range of machine learning models can be written in termof the feature-based form f ( x )
=
( x ) >  [52, 16, 39], making the Gaussian process a uniﬁedframework for quantifying model uncertainty for a wide array of modern ML models.
Appendix Bsummarizes important examples including GAMs, decision trees, random-feature models, deep neuralnetworks and their ensembles.
Furthermore, when a deterministically-trained ˆ  is available (e.g.,via a sophisticated adaptive shrinkage procedure that is not available in Bayesian context), we canincorporate this as prior knowledge into GP modeling by setting µ = ˆ  (Equation (3)).
2.2 Bayesian Nonparametric Guarantees for Probabilistic Learning
The quality of a Bayesian learning procedure is commonly measured by the learning rate of itsposterior distribution ⇧ n = ⇧ ( · | { x i , y i } ni =1 ) .
Intuitively, the rate of this convergence is measured 134 by the size of the smallest shrinking balls around f 0 that contain most of the posterior probability.
Speciﬁcally, we consider the size of the set A n = { g | k g  f 0 k 2 n  M ✏ n } such that ⇧ n ( A n ) !
The concentration rate ✏ n here indicates how fast the small ball A n concentrates towards f 0as the sample size increases.
Below we state the formal deﬁnition of posterior convergence [25].
Deﬁnition 1 (Posterior Convergence) .
For f 0 : X !
R where
X = R d , denote H
the true RKHS 139 induced by a kernel function k 0 , and denote
H  induced by the feature function  : X !
R D .
Let 140 f 0 2 H be the true function, and let E 0 denote the expectation with respect to the true data-generationdistribution.
Assuming H  is dense in H , then, the posterior distribution ⇧ n ( f ) concentrates around 142 f 0 at the rate ✏ n if there exists an ✏ n !
such that, for any M n ! 1 ,
[Equation]
[Equation]
Notice that we allow the model space H  and the true function space H to be different, but the priormust be dense in the later for the convergence to happen.
Fortunately, this condition is shown to holdfor a wide variety of ML models, including random features, random forests, and neural networks[3, 33, 53, 61, 57].
The notion of posterior convergence can also be used to discuss the learningquality of other probabilistic estimates (e.g., variable importance j ).
In that case, we can simplyreplace ( f, f 0 ) in (6) by their variable importance counterparts.
This is the focus of Section 3.2.
Methods
[Figure]
3.1 Quantifying Variable Importance under Uncertainty
In this work, we consider quantifying the global importance of a variable based on the norm ofthe corresponding partial derivative.
This is motivated by the observation that, if a function f isdifferentiable, the relative importance of a variable x
j at a point x can be captured by the magnitudeof the partial derivative function, | @@ x j f ( x ) | [58].
This proposed quantity requires the considerationof two issues.
First, instead of quantifying the relevance of a variable on a single input point, weneed to deﬁne a proper global notion of variable importance.
Therefore, it is natural to integratethe magnitude of the partial derivative over the input space x 2 X :
[Equation]
[Equation]
X ( x ) .
Second, since P X ( x ) is not known from the training observations, j [Equation] f )can be approximated by its empirical counterpart,
[Equation]
Notice that j ( f ) is an estimator that is derived from the prediction function
f estimated using ﬁnitedata.
Consequently, to make a proper decision regarding the importance of an input variable x j ,it is important to take into account uncertainty in f .
To this end, by leveraging the featurized GPrepresentation introduced in Section 3.1, we show that this can be done easily for a wide range of MLmodels f ( x ) =
( x ) >  by studying the posterior distribution of j .
Posterior Distribution of Variable Importance .
After we obtain the posterior distribution of  (4),the posterior distribution of variable importance can be derived according to Equation (7):
[Equation]
where @
@ x j 2 R D ⇥ n is the derivative of the feature map with respect to x j , across n training samples.
The posterior distribution of j ( f ) adopts a closed form as a generalized chi-squared distribution(see Appendix A.2 for derivation).
In practice, we can sample j conveniently from its posteriordistribution by computing @@ x j f ( X ) =
@  @ x
j  >  ( s ) , where  ( s ) are Monte Carlo samples from theclosed-form posterior (4).
There are two ways in which uncertainty aids the model selection process.
First, the posterior survivalfunction P ( j ( f ) > s ) of the variable importance utilizes the full posterior distribution of j ( f ) toidentify the probability that the variable x j exceeds a given threshold s .
By increasing s 2 (0 , 1 ) , P ( j > s ) provides a intuitive sense of how model’s belief about the importance of variable x j changes as the criteria s becomes more stringent, similar to the regularization path in LASSO methods[20] but with the incorporation of posterior uncertainty about the variable importance.
See Appendix I for an application to a Bangladesh birth cohort study.
Second, by integrating the survival functionover the threshold, i.e.,
R s> 0 P ( j ( f ) > s ) ds , we obtain the posterior mean of j ( f ) , and this tooincorporates uncertainty in f .
To see this, notice by using the “trace trick” we can write
[Equation]
where all expectations are taken with respect to the posterior.
Therefore, the posterior mean of j ( f )depends on the covariance structure of  , and how it interacts with the eigenspace of the partialderivative functions (encoded by @  @ x j @  > @ x j ).
In Section 4 we provide an extensive investigation of
AUROC scores using the posterior mean of j ( f ) for variable selection.
In Appendix A.3, we summarize the algorithms for computing the posterior distributions of thefeaturized Gaussian process (Equation (4)) and for the posterior distributions of variable importance(Equation (8)), and discuss their space and time complexity.
3.2 Theoretical Guarantees
From a theoretical perspective, the variable importance measure j introduced in (7) can be understood as a quadratic functional of the Gaussian process model f [18].
To this end, rigorous Bayesiannonparametric guarantees can be obtained for j ’s ability in learning the true variable importance inﬁnite samples (i.e., posterior convergence, Theorem 1) and its statistical optimality from a frequentistperspective, in providing a low-variance estimator that attains the Cramér-Rao bound (i.e., Bernsteinvon-Mises phenomenon, Theorem 2).
Posterior Convergence .
We ﬁrst show that, for a ML model f that can learn the true function f 0with rate ✏ n (in the sense of Deﬁnition 1), the entire posterior distribution of the variable importancemeasure j ( f ) converges consistently to a point mass at the true j ( f 0 ) at a speed that is equal orfaster than ✏ n .
Theorem 1 (Posterior Convergence of Variable Importance j ) .
Suppose y i =
f
( x i )
+
e i , e i i.i.d.
⇠
N (0 ,  2 ) , and denote as E 0 the expectation with respect to the true data-generation distributioncentered around f 0 .
For the RKHS H  induced by the feature function  : X !
R D and f 2 H  , if:
(1) The posterior distribution ⇧ n ( f ) converges toward f at a rate of
✏ n ; (2) The differentiation operator D j : f ! [Equation]
[Equation]
for all f [Equation] H  [Equation]
Then the posterior distribution for j ( f ) =
k @@ x j f k 2 n contracts toward j ( f 0 ) =
k @@ x j f 0 k 2 P X at a 206 rate not slower than ✏ n .
That is, for any M n ! 1 ,
[Equation]
Proof is in Appendix C. Theorem 1 is a generalization of the classic quadratic functional convergenceresult to a much wider range of ML models in the context of Bayesian variable selection [18, 42, 67].
It conﬁrms the important fact that, for a ML model f that can accurately learn the true function f 0under ﬁnite data, despite the potential lack of identiﬁablity in the model parameters (e.g., weights ina neural network), we can consistently recover the true variable importance at a fast rate by using theproposed variable importance estimate j ( f ) .
From a practical point of view, Theorem 1 reveals that the ﬁnite-sample performance of variableimportance j ( f ) depends on two factors: (1) the ﬁnite-sample generalization performance of theprediction function f , and (2) the mathematical property of f in terms of its Lipschitz condition.
Therefore, to ensure effective variable selection in practice, the practitioner should take care to selecta model class f that has a theoretical guarantee in capturing the target function f 0 , empiricallydelivers strong generalization performance under ﬁnite data, and is well-conditioned in terms of thebehavior of its partial derivatives.
To this end, we note that, under the featurized Gaussian process f =  ( x ) >  discussed in this work, users are free to choose a performant model class (e.g., randomforest, random-feature or DNN) whose feature representation spans an RKHS H  that is dense in theinﬁnite-dimensional function space (therefore f enjoys a convergence guanrantee) [3, 33, 53, 61, 57],and is empirically more effective than the classic GP methods.
(We discuss the Lipschitz conditionof these models in Appendix E.1)
Indeed, as we will verify in experiments (Section 4), there doesnot exist an “optimal" model class that performs universally well across all data settings (i.e., nofree lunch theorem [69]).
This highlights the importance of having a general-purpose framework forvariable selection that can ﬂexibly incorporate the most effective model for the task at hand.
Statistical Efﬁciency & Uncertainty Quantiﬁcation.
Next, we verify the uncertainty quantiﬁcationability of the variable importance measure j ( f ) under featurized GP, by showing that it exhibitsthe Bernstein-von Mises (BvM) phenomenon.
That is, its posterior measure ⇧ n ( j ( f )) convergestowards a Gaussian distribution that is centered around the truth j ( f 0 ) , so that its (1  ↵ )% levelcredible intervals achieve the nominal coverage probability for the true variable importance.
Moreimportantly, the BvM theorem veriﬁes that the posterior distribution of j ( f ) is statistically optimal ,in the sense that its asymptotic variance attains the Cramér-Rao bound (CRB) that cannot be improvedupon [5].
Theorem 2 (Bernstein-von Mises Theorem for Variable Importance j ) .
Suppose y i =
f
( x i )
+
e i , e i i.i.d.
⇠ N (0 ,  2 ) , i = 1 , . . . , n .
Denote D j : f ! @ @
x j f the differentiation operator and
H j = D > j D j the inner product of D j , such that:
[Equation]
Assuming conditions (1)-(2) in Theorem 1 hold, and additionally:
(3) f 0 is square-integrable over the support X and k f 0 k 2 = 1 ;
(4) rank ( H j ) = o p ( p n ) ;
Then
[Equation]
Proof is in Appendix D. Theorem 2 provides a rigorous theoretical justiﬁcation for j ( f ) ’s ability toquantify its uncertainty about the variable importance.
More importantly, it veriﬁes that j ( f ) has thegood frequentist property that it quickly converging to a minimum-variance estimator at a fast speed,which is important in practice.
Compared to the previous BvM results that tend to focus on a speciﬁc
Bayesian ML model, Theorem 2 is considerably more general (i.e., applicable to a much wider rangeof models) and comes with a simpler set of conditions [56, 67, 42].
Speciﬁcally, (3) is a standardassumption in nonparametric analysis.
It ensures the true function f 0 does not diverge towards inﬁnityand makes learning impossible [12].
The unit norm assumption k f 0 k 2 = 1 is only needed to simplifythe exposition of the proof, and the theorem can be trivially extended to k f 0 k 2 =
C for any C > 0 .
The most interesting condition is (4).
Let’s denote H
j the space of partial derivatives functions@@ x j f, f 2 H .
Then intuitively, (4) says to attain the BvM phenomenon, the effective dimensionalityof the derivative function space H j (as measured by rank ( H j ) =
rank ( D j ) ) cannot be too large.
Since effective dimensionality of the derivative space is bounded above by that of the original RKHSf 2 H  , (4) essentially states that the effective dimensionality of the model space H  cannot grow 257 too fast with data ( o p ( p n ) ).
Fortunately, this condition is satisﬁed by a wide range of ML modelsincluding trees and deep networks [56, 67].
See Appendix E.2 for further discussion.
Experiment Analysis
In this section, we investigate the ﬁnite-sample performance of the derivative norm metric j forvariable selection (7) under a wide variety of ML methods.
We illustrate the breath of our frameworkby applying it to tree ensembles (Appendix F), where a principled and gradient-based uncertaintyaware variable selection approach has been previously unavailable.
We also apply it to linear modelsand (approximation) kernel machines, which are standard approaches to variable selection in datascience practice [65, 6].
Over a wide range of complex and realistic data scenarios (e.g., discretefeatures, interactions, between-feature correlations) derived from socioeconomic and healthcare datasets, we investigate the method’s statistical performance in accurately recovering the groundtruth features (in terms of the Type I and Type II errors), and compare it to other well-establishedapproaches in each of the model classes (Table 1).
Our main observations are:
O1 : Importance of generality .
There does not exist a model class that performs universally wellacross all data scenarios (i.e., no free lunch theorem [69], Figure 1, 6-15).
This highlights theimportance of an uniﬁed variable selection framework that incorporates a wide range of models, sothat practitioners have the freedom of choosing the most suitable model class for the task at hand.
O2 : Good prediction translates to effective variable selection .
Comparing between differentmodel classes, the ranking of models’ predictive accuracy is generally consistent with the rankingof their variable selection performance under j (i.e., better prediction translates to better variableselection, as suggested in Theorem 1).
O3 : Statistical efﬁciency of j .
Comparing within each model class, the derivative norm metric jgenerally outperforms other measures of variable importance.
The advantage is especially pronouncedin small samples and for correlated features.
This empirically veriﬁes that j has good ﬁnite-samplestatistical efﬁciency even under complex data scenarios (as suggested in Theorem 2).
[Table]
Models & Methods .
We consider three main classes of models (Table 1): (I) Random Forests (RF) .
Given a trained forest, we quantify variable importance using j by translating it to an ensembleof featurized decision trees ( FDT s) (Appendix F), and compare it to three baselines: impurity ( RFimpurity ) [8], RF-based kernel knockoff ( RF-knockoff )
[10], and Bayesian additive regressiontrees ( BART ).
(II) (Approximate) Kernel Methods .
We apply j to a random-feature model thatapproximates a Gaussian process with a RBF kernel [52], and set the number of features to p n log ( n )to ensure proper approximation of the exact RBF-GP [59].
We term this approach Random-feature
Neural Networks ( RFNN ), and compared it to both Bayesian Kernel Machine Regression ( BKMR )[6] based on a GP with exact RBF kernel and spike-and-slab prior and Bayesian Approximate Kernel
Regression ( BAKR ) based on random-feature model with a projection-based feature importancemeasure and an adaptive shrinkage prior [14].
(III) Linear Models .
We apply j to a featurized
GP representation of the Generalized Additive Model ( GAM ), with the prior center  0 set at thefrequentist estimate of the original GAM model obtained from a sophiscated REML procedure [70].
We compare it to two baselines: Bayesian Ridge Regression ( BRR ) [32] and Bayesian LASSO ( BL )[50].
Appendix G provides further detail.
Previously, [42] studied the specialization of our frameworkto the deep neural networks (DNNs), we don’t repeat that work here as DNN is not yet a standarddata science model for tabular data.
To quantify variable importance while accounting for posterior uncertainty the variable importancej ( f ) , we examine its posterior survival function
R s> 0 P ( j ( f ) > s ) ds (i.e., the posterior likelihoodof j ( f ) greater than the threshold s ) integrated over the full range of thresholds s .
For other methods,we use their default metrics to quantify variable importance (e.g., variable inclusion probabilities in
BART , BKMR .
Appendix G).
Datasets and Tasks.
We consider two synthetic benchmarks and three real-world socio-economicand healthcare datasets, encapsulating challenging phenomena such as between-feature correlationsand interaction effects.
For the synthetic benchmark, we generate data under the Gaussian noisemodel y ⇠ N ( f 0 , 1) for four types of outcome-generation functions f 0 ( linear , rbf , matern32and complex , see Appendix G.2 for a full description) with number of causal variables d ?
= 5 .
Two types of feature distribution are considered: (1) synthetic-continuous : all features followx j ⇠ Unif (  2 , 2) ; (2) synthetic-mixture : two of the causal features and two of the non-causalfeatures are distributed as Bern (0 . 5) and the rest are distributed as Unif (  2 , 2) ; We vary samplesize n 2 { 100 , 200 , 500 , 1000 } and data dimension d 2 { 25 , 50 , 100 } , leading to 96 total scenarios.
For real-world data, we consider (1) adult : 1994 U.S. census data of 48842 adults with 8 categoricaland 6 continuous features [38]; (2) heart : a coronary artery disease dataset of 303 patients from
Cleveland clinic database with 7 categorical and 6 continuous features [17]; and (3) mi : disease records of myocardial infarction (MI) of 1700 patients from Krasnoyarsk interdistrict clinical hospitalduring 1992-1995, with 113 categorical and 11 continuous features [27].
All datasets exhibit non-trivalcorrelation structure among features (Appendix Figure 3-5).
Since the ground-truth causal featureson these datasets are not known, in order to rigorously evaluate variable selection performance, wefollow the standard practice in causal ML to simulate the outcome based on causal features selectedfrom data [72].
We use the four outcome-generating functions as described previously and evaluateover the same data size ⇥ dimension combinations, leading to 144 total scenarios 1 .
We repeat thesimulation 20 times for each scenario, and use AUROC to measure the variable selection performance(in terms of Type I and II errors) of each method.
In the Appendix I, we further evaluate the method on a well-studied environmental health dataset(Bangladesh birth cohort study [37]) with respect to the real outcome (infant development scores).
Wevisualize the "Bayesian" regularization path as introduced earlier.
The selected variables correspondwell with the established toxicology pathways in the literature [26]. [Figure]
4.1 Results
Figure 1 shows the methods’ performance in variable selection (Row 1) and prediction (Row 2) 2in an exemplary setting, where the true function f 0 is matern32 with an input dimension of 100.
It represents the tabular data setting that we are the most interested in: nonlinear feature-responserelationship with interaction effects and high input dimension.
This is because f 0 is sampled froma RKHS induced by Mátern 32 kernel, which contains a large space of continuous and [-at
 In the setting where required data dimension is higher than that of the real data, we generate additional synthetic features from Unif (  2 , 2) . We use n 2 { 50 , 100 , 150 , 257 } for heart due to data size restriction.-]at least oncedifferentiable functions [55].
We delay complete visualizations for all 240 scenarios to Appendix H.
Recalling the three observations introduced earlier:
O1 ("No free lunch") : No method performs universally well.
For example, BAKR performs robustlyin correlated datasets ( heart and mi ), but poorly otherwise.
Kernel approaches ( RFNN and BKMR )performs competitively in low dimension, but their performance deteriorates quickly as dimensionincreases (Figure 6-7).
FDT s generally is the strongest method in small samples and high dimensions,but can be outperformed by GAM in large samples and data with high percentage of categoricalfeatures ( adult and mi ).
This highlights the importance of an uniﬁed framework that allows users toselect the most appropriate model for variable selection depending on the data setting.
n due data O2 ("Good prediction implies effective variable selection") : Comparing among the gradient-basedmethods under each model class (i.e., FDT s, GAM and RFNN .
Solid lines in Figure 1), we see thattheir rankings in prediction (row 2) is largely consistent with rankings in variable selection.
It’s worthnoting that this pattern is occasionally violated (e.g., GAM in adult , n=500 and heart , n=250), butthat does not contradict our conclusion (Theorem 1) since the convergence rate of the predictionfunction only forms an upper bound for the convergence rate of j .
O3 (Statistical efﬁciency of j ) :
Comparing between variable selection methods from the same class(especially for tree models.
i.e., FDT s v.s. RF-impurity / RF-knockoff / BART ), we see that FDT sis competitive or strongly outperforms its baselines in variable selection, despite being based onexactly the same ﬁtted model ( RF-impurity / RF-knockoff ), or not accounting for the uncertainty inthe tree growing process ( BART ).
This pattern is consistent in most data settings, and the advantageis especially pronounced in high dimensions, small data sizes, and correlated datasets (Appendix H, Figure 6-10).
This provides strong empirical evidence for the fact that j is a statistically efﬁcientestimator for variable selection with good ﬁnite-sample behavior (as suggested in Theorem 2), andcan deliver strong variable selection performance for tabular data when combined with a performant ML model like random forest.
Appendix H contains further discussion.
Discussion and Future Directions
The modern data analysis pipeline typically involves ﬁtting multiple models, comparing their performance, and iterating as necessary.
When variable selection is involved, the practitioner may ask arethe variable importances across models measuring the same behavior?
And, what if the most suitablemodel does not have a satisfactory variable selection procedure?
By framing model choice as kernelchoice — which we emphasize includes kernels corresponding machine learning methods like neuralnetworks and random forests in addition to the long list of traditional kernels — we propose a uniﬁedvariable selection procedure that is compatible across models and we prove strong guarantees for thisprocedure.
Limitations .
Our proposed framework provides principled uncertainty quantiﬁcation by performingexact Bayesian inference on the weights  of a feature map  ( x ) .
We do not consider uncertainty inthe feature map itself.
This means, for example, that if the feature map is given by the last hiddenlayer of a neural network trained by maximizing the posterior, then our model class correspondsto the neural linear model .
This model is different from a fully Bayesian neural network, whichperforms posterior inference also on the kernel hyperparameters (i.e., the hidden weights) [49, 62, 64].
Likewise, the kernel induced by the featurized decision tree studied here does not consider uncertaintyin the tree’s partitioning process.
Yet, this does not seem to be a signiﬁcant limitation in ourexperiments (e.g., FDT s outperforms BART ), although this point still merits further investigation inthe future.
In our experiments, we focused on kernels based on tree ensembles, kernel methods and linear models.
In the future, it would be worth expanding this framework to other model classes (e.g., MARS [21]or neural network) and estimating the importance of interaction effects and higher-order terms.
Wewould also like to apply this method to large-scale scientiﬁc studies (e.g., epidemiology study basedon extremely large EHR datasets) where an uncertainty-aware nonlinear variable selection method istypically impossible due to challenges with scalability.
Societal Impacts .
The method proposed in this paper provides a theoretically-grounded approach forquantifying variable importance that is applicable to a wide range of ML models.
We expect it toprovide a set of powerful tools for practitioners to understand the importance of input variables intheir ML models with limited data, which is especially important for scientiﬁc investigations such asepidemiology and computational biology.
However, we recognize that this approach can potentiallybe utilized by bad actors to probe the input-variable uncertainty of an existing ML system, anduse it to engineer more targeted white-box adversarial attacks.
To this end, we recommend systemdevelopers to incorporate this approach into the formal veriﬁcation procedure of a ML system, so asto monitor and understand the model uncertainty with respect to input variables, and devise properimprovement and prevention strategies (e.g., data augmentation or randomized smoothing targeted atspeciﬁc variables) accordingly.