Towards a Uniﬁed Framework for Uncertainty-aware Nonlinear Variable Selection with Theoretical Guarantees
Abstract
We develop a simple and uniﬁed framework for nonlinear variable importanceestimation that incorporates uncertainty in the prediction function and is compatiblewith a wide range of machine learning models (e.g., tree ensembles, kernel methods,neural networks, 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 uncertainty in variable importance 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.
As datasets grow in complexity and size, it is crucial that variable importance estimation methodscan account for complex dependencies among variables while remaining computationally feasible.
Furthermore, as the number of approaches to model such datasets have increased, it is crucial that theimportance of each variable can be compared across model classes and extended to new ones as theyare developed.
While there are established approaches for variable importance in linear models (e.g., LASSOregression [36]) there is little consensus in methodology or theory for variable importance innonlinear models.
Generalized additive models [37] use similar methods as their linear counterparts[79] but the additivity assumption for nonlinear functions of the variables is too restrictive in manyapplications.
Random Forests (RF) [9] 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.
[27] 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 [14, 24], using backward selection [14], orknockoffs [54].
Unfortunately, each model class based on DNNs requires a tailored procedure, whichlimits 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 importance estimation procedure to tailor its decision rule with respect to thecorrelation structure [50].
Yet, as in frequentist models, each method has a different deﬁnition ofa variable’s importance.
For example, in Bayesian additive regression trees (BART), a variable’simportance can be measured by the proportion of trees that use it [16], while in Gaussian process (GP)models, a variable’s importance can be measured by the frequency of the ﬂuctuations of the resultingfunction (e.g., the length-scale parameter as controlled by the automatic relevance determination)in the direction of the variable [58, 83]. Recently, a closely-related line of work uses the norm ofkernel gradient to quantify variable importance under classic GP models [38] or deep Bayesian neuralnetworks [50].
However, these work either does not incorporate uncertainty, or is restricted to aparticular model class (see Appendix J).
Furthermore, the traditional Bayesian modeling procedurestend to be computationally 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 [21], and deep neural networks can be written as kernel methodsby using the last hidden layer as the feature map [72, 39, 12].
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 importance estimation framework that is compatiblewith a wide range of machine learning models that can be deﬁned by, or be closely approximatedby, 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 partialderivative, as was previously studied in the context of frequentist nonparametric regression [68].
We extend it to apply to a much wider class of models than previously considered (Section 2),propose a principled Bayesian approach to quantify the variable importance uncertainty in ﬁnitedata (Section 3.1), and derive rigorous Bayesian nonparametric theorems to guarantee the method’sconsistency and asymptotic optimality (Section 3.2).
To incorporate powerful non-differentiablemodels into our framework, we also show how to apply this approach to partition-based methods (e.g.,decision trees) by leveraging its (soft) feature representation (Appendix F.1).
This leads to the ﬁrstderivative-based Bayesian variable importance estimation approach for tree-type models that is boththeoretically grounded and empirically powerful, strongly outperforming other variable importanceestimation approaches tailor-designed for random forest (e.g., impurity or random-forest knockoff[10, 13]).
We conduct extensive empirical validation of our approach and compare its performance tothat of many existing methods across a wide range of data generation scenarios.
Results show a clearadvantage of the proposed approach, especially in complex scenarios or when the input is a mixtureof 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 81 by a certain positive deﬁnite kernel function k 0 , and the input space of the true function
X 0 spans 82 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 importance estimation is to produce a variable importancescore j for each of the input features ( x 1 , . . .
, x d ) such that it can be used as a classiﬁcation signalfor whether x j 2 X 0 .
As a result, the variable selection decision can be made by threhsoldingj > s with a pre-deﬁned threshold s .
The quality of a variable selection signal j can be evaluatedcomprehensively using a standard metric such as the area under the receiver operating characteristic(AUROC), which measures the Type-I and Type-II errors of variable selection decision I ( j > s )over a range of thresholds 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 (2)
[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 [65].
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 [19] 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 [65]:
[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]
where  =
(  ( x 1 ) > , . . .
,  ( x n ) > ) >
2 R n ⇥ D is the feature matrix evaluated on the training data[65].
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 term of the feature-based form f ( x )
=
( x ) >  [62, 21, 47], 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, deepneural networks and their ensembles.
Appendix B.1 summarizes a list of general conditions themodel should satisfy for it to be compatible with the proposed framework (i.e., weak differentiability,
Lipschitz condition, and growth rate of model complexity). Furthermore, when a deterministicallytrained ˆ  is available (e.g., via a sophisticated adaptive shrinkage procedure that is not available in a Bayesian context), we can incorporate 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 142 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 [30].
Deﬁnition 1 (Posterior Convergence) .
For f 0 : X !
R where
X = R d , denote H
the true RKHS 147 induced by a kernel function k 0 , and denote
H  induced by the feature function  : X !
R D .
Let 148 f 0 2 H 0 be the true function, and let E 0 denote the expectation with respect to the true data-generationdistribution.
Assuming H  is dense in H 0 , then, the posterior distribution ⇧ n ( f ) concentrates around 150 f 0 at the rate ✏ n if there exists an ✏ n !
such that, for any M n ! 1 ,
[Equation]
Notice that we allow the model space H  and the true function space H 0 to be different, but the H  152 must be dense in the H 0 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[4, 41, 63, 70, 67].
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
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 ) | [68].
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 : j ( f ) = k @
[Equation]
x j f
[Equation]
P X
[Equation]
R x 2 X | @ @ x j f ( x ) | 2 dP X ( x ) . Second, since P X ( x ) is not known from the training observations, j ( f )can be approximated by its empirical counterpart,
[Equation]
169 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]
[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 variable importance estimation process.
First,the posterior survival function P ( j ( f ) > s ) of the variable importance utilizes the full posteriordistribution of j ( f ) to identify 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 aboutthe importance of variable x j changes as the criteria s becomes more stringent, similar to theregularization path in LASSO methods [25] but with the incorporation of posterior uncertainty aboutthe variable importance.
See Appendix I for an application to a Bangladesh birth cohort study.
Second, by integrating the survival function over the threshold, i.e.,
R s> 0 P ( j ( f ) > s ) ds , weobtain the posterior mean of j ( f ) , and this too incorporates uncertainty in f .
To see this, notice byusing 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 quantifying variable importance.
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 [23]. To this end, rigorous
Bayesian nonparametric guarantees can be obtained for j ’s ability in learning the true variableimportance in ﬁnite samples (i.e., posterior convergence, Theorem 1) and its statistical optimalityfrom a frequentist perspective, in providing a low-variance estimator that attains the Cramér-Raobound (i.e., Bernstein von-Mises phenomenon, Theorem 2). Note that for a given general modelf ( x ) =  ( x ) >  , it only need to satisfy three mild regularity conditions to be fully compatible withthe proposed framework (i.e., weak differentiability, Lipschitz condition, and growth rate of modelcomplexity). We summarize these conditions in Appendix B.1 and explain them in detail in thesequel.
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
0 at a rate of ✏ n ;
(2) The differentiation operator D j : f ! @@
x j f is bounded: k D j k 2 op = inf {
[Equation]
[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 220 rate not slower than ✏ n .
That is, for any M n [Equation]
[Equation]
Proof is in Appendix C. Theorem 1 is a generalization of the classic result of quadratic functionalconvergence under linear models and sparse neural networks to a much wider range of ML models in the context of Bayesian variable importance estimation [23, 50, 80].
It conﬁrms the importantfact that, for a ML model f that can accurately learn the true function f 0 under ﬁnite data, wecan consistently recover the true variable importance at a fast rate by using the proposed variableimportance estimate j ( f ) , despite the potential lack of identiﬁablity in the model parameters (e.g.,weights in a neural network). Importantly, although our main setting assumes ﬁxed data dimension d(see Problem Setup), the posterior concentration result Theorem 1 does not rely on this assumption inits proof, and is in fact compatible with the high-dimensional setting where d is allowed to grow withsample size in a rate of o ( n ) . See Appendix C (in particular, Remark 4) for further discussion.
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 importance estimation in practice, the practitioner should takecare to select a model class f that has a theoretical guarantee in capturing the target function f 0 ,empirically delivers strong generalization performance under ﬁnite data, and is well-conditioned interms of the behavior of its partial derivatives.
To this end, we note that, under the featurized Gaussianprocess f =  ( x ) >  discussed in this work, users are free to choose a performant model class (e.g.,random forest, random-feature or DNN) whose feature representation spans a RKHS H  that is densein the inﬁnite-dimensional function space (therefore f enjoys a convergence guanrantee, see Remark3 in Appendix C for further discussion) [4, 41, 63, 70, 67], and is empirically more effective thanthe GP methods based on classic kernels such as RBF.
We discuss the Lipschitz condition of thesemodels in Appendix E.1.
Indeed, as we will verify in experiments (Section 4), there does not existan “optimal" model class that performs universally well across all data settings (i.e., no free lunchtheorem [84]).
This highlights the importance of having a general-purpose framework for variableimportance estimation that can ﬂexibly incorporate the most effective model for the task at hand.
Finally, we notice that although Theorem 1 is stated as an asymptotic result, when a ﬁnite-sampleerror bound ✏ n for the model class is available (i.e., Condition (1) in Theorem 1), it is trivial to obtaina ﬁnite-sample error bound for variable importance
j ( f ) by extending the proof of Theorem 1.
Appendix C.1 provides an example of such a bound based on the Bernstein inequality.
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 [6].
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]
[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 abilityto quantify its uncertainty about the variable importance.
More importantly, it veriﬁes that j ( f )has the good frequentist property that it quickly converges to a minimum-variance estimator at afast speed, which is important for obtaining good variable importance estimation performance inpractice.
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 range of models) andcomes with a simpler set of conditions [66, 80, 50].
Speciﬁcally, (3) is a standard assumption in nonparametric analysis.
It ensures the true function f 0 does not diverge towards inﬁnity and makeslearning impossible [15].
The unit norm assumption k f 0 k 2 = 1 is only needed to simplify theexposition of the proof, and the theorem can be trivially extended to k f 0 k 2 =
C for any C > 0 .
Themost interesting condition is (4).
Let’s denote H
j the space of partial derivatives functions @@ x j f ofthe model functions f 2 H  .
Then intuitively, (4) says to attain the BvM phenomenon, the effectivedimensionality of the derivative function space H j (as measured by rank ( H j )
= rank ( D j ) ) cannotbe too large.
Since effective dimensionality of the derivative space is bounded above by that of theoriginal RKHS f 2 H  , (4) essentially states that the effective dimensionality of the model space H  cannot grow too fast with data size (i.e., o p ( p n ) ).
Fortunately, this condition is satisﬁed by awide range of ML models including trees and deep networks [66, 80].
See Appendix E.2 for furtherdiscussion.
Experiment Analysis
In this section, we investigate the ﬁnite-sample performance of the derivative norm metric j forvariable importance estimation (7) under a wide variety of ML methods.
We illustrate the breath ofour framework by applying it to tree ensembles (Appendix F.1), where a principled and gradient-baseduncertainty-aware variable importance estimation approach has been previously unavailable.
We alsoapply it to linear models and (approximation) kernel machines, which are standard approaches tovariable selection in data science practice [75, 7].
Over a wide range of complex and realistic datascenarios (e.g., discrete features, interactions, between-feature correlations) derived from socioeconomic and healthcare datasets, we investigate the method’s statistical performance in accuratelyrecovering the ground-truth features (in terms of the Type I and Type II errors), and compare it toother well-established approaches 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 [84], Figure 1, 6-15).
This highlights theimportance of an uniﬁed framework for variable importance that incorporates a wide range of models,so that practitioners have the freedom of choosing the most suitable model class for the task at hand.
O2 : Good prediction translates to effective variable importance estimation .
Comparing betweendifferent model classes, the ranking of models’ predictive accuracy is generally consistent with theranking of their variable importance estimation performance under j (i.e., better prediction translatesto better variable importance estimation, 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.1), and compare it to three baselines: impurity( RF-impurity ) [10], RF-based kernel knockoff ( RF-knockoff ) [13], and Bayesian additive regressiontrees ( BART ).
(II) (Approximate) Kernel Methods & Neural Networks .
We apply j to a randomfeature model that approximates a Gaussian process with a RBF kernel [62], and set the numberof features to p n log ( n ) to ensure proper approximation of the exact RBF-GP [69].
We term thisapproach Random Fourier feature model ( RFF ), and compare it to both Bayesian Kernel Machine
Regression ( BKMR ) [7] based on a GP with exact RBF kernel and spike-and-slab prior. Bayesian
Approximate Kernel Regression ( BAKR ) based on random-feature model with a projection-basedfeature importance measure and an adaptive shrinkage prior [18].
Neural Networks ( NN ) based onwide ReLU neural network with 512 hidden units and LASSO regularization in the hidden layerweights [48].
(III) Linear Models .
We apply j to a featurized GP representation of the Generalized
Additive Model ( GAM ), with the prior center µ set at the frequentist estimate of the original GAMmodel obtained from a sophiscated REML procedure [85].
We compare it to two baselines: Bayesian Ridge Regression ( BRR ) [40] and Bayesian LASSO ( BL ) [60].
Appendix G provides further detail.
Previously, [50] studied the specialization of our framework to the deep neural networks (DNNs), wedon’t repeat that work here as DNN is not yet a standard data 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 , 0 . 01) 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 [46]; (2) heart : a coronary artery disease dataset of 303 patients from
Cleveland clinic database with 7 categorical and 6 continuous features [22]; and (3) mi : diseaserecords of myocardial infarction (MI) of 1700 patients from Krasnoyarsk interdistrict clinical hospitalduring 1992-1995, with 113 categorical and 11 continuous features [32].
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 importance estimationperformance, we follow the standard practice in causal ML to simulate the outcome based oncausal features selected from data [88].
We use the four outcome-generating functions as describedpreviously and evaluate over the same data size ⇥ dimension combinations, leading to 144 totalscenarios 1 .
We repeat the simulation 20 times for each scenario, and use AUROC to measure thevariable importance estimation 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 [45]) 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 [31].
4.1 Results
Figure 1 shows the methods’ performance in variable importance estimation (Row 1) and prediction(Row 2) 2 in an exemplary setting, where the true function f 0 is matern32 with an input dimensiond = 100 .
It represents the tabular data setting that we are the most interested in: nonlinear featureresponse relationship with interaction effects and high input dimension.
This is because f 0 is sampledfrom a RKHS induced by Mátern 32 kernel, which contains a large space of continuous and atleast once differentiable functions [65].
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 ( RFF and BKMR )perform competitively in low dimension, but their performances deteriorate quickly as dimensionincreases (Figure 1 and Figure 6-7).
This is likely due to the classical kernel method’s well-knowninability to learn an adaptive feature representation, which consequently leads to suffering fromthe curse of dimensionality and unstable and suboptimal variable importance performance in highdimensions [3].
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 ).
Notably, they often outperforms NN , which is traditionally regarded as thego-to model for high-dimensional nonlinear settings.
This highlights the importance of an uniﬁed [Figure] framework that allows users to select the most appropriate model for variable importance estimationdepending on the data setting
[Figure]
O2 ("Good prediction implies effective variable importance estimation") : Fixing the variableimportance j and comparing the variable importance estimation performance of each model class(i.e., FDT s, GAM , NN and RFF .
Solid lines in Figure 1), we see that their rankings in prediction(row 2) is largely consistent with rankings in variable importance estimation.
It’s worth noting thatthis pattern is occasionally violated (e.g., GAM in adult , n=500 and heart , n=250), but that does notcontradict our conclusion (Theorem 1) since the convergence rate of the prediction function only formsan upper bound for the convergence rate of j . Finally, when models have comparable generalizationperformance, we observe that the Lipschitz condition plays a role in variance importance performance(which is consistent with our theoretical observations in Theorem 1). For example, in Figure 1, NN and FDT s are largely comparable in predictive performance among the real datasets ( adult ,heart and mi ).
However, tree-based FDT s are known to have well-conditioned Lipschitz behaviorwhen compared to NN (Appendix Appendix E.1), which is consequently translated to improvedﬁnite-sample performance in variable importance estimation.
O3 (Statistical efﬁciency of j ) :
Comparing between variable importance estimation methodsfrom the same class (especially for tree models.
i.e., FDT s v.s. RF-impurity / RF-knockoff / BART ), we see that FDT s is competitive or strongly outperforms its baselines in variable importanceestimation, despite being based on exactly the same ﬁtted model ( RF-impurity / RF-knockoff ), ornot accounting for the uncertainty in the tree growing process ( BART ).
This pattern is consistentin most data settings, and the advantage is especially pronounced in high dimensions, small datasizes, and correlated datasets (Appendix H, Figure 6-10).
This provides strong empirical evidence forthe fact that j is a statistically efﬁcient estimator for variable importance with good ﬁnite-samplebehavior (as suggested in Theorem 2), and can deliver strong performance for tabular data whencombined 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 importance estimation procedure?
By framing modelchoice as kernel choice — which we emphasize includes kernels corresponding machine learningmethods like neural networks and random forests in addition to the long list of traditional kernels —we propose a uniﬁed variable importance estimation procedure that is compatible across models andwe prove strong guarantees for this procedure.
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) [59, 72, 74].
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 themethod’s empirical performance (e.g., FDT s outperforms BART in our experiments), althoughthis point still merits further investigation in the future.
The theoretical results of the current paper assume ﬁxed data dimension d = O (1) .
However, thisdoes not restrict the signiﬁcance and practicality of our theoretical results. On the theoretical side,even for ﬁxed d , a general framework to nonparametric Bayesian inference of feature importanceinference is currently missing in the ﬁeld. On the application side, the majority of the machinelearning applications fall into the ﬁxed dimension setting. For example, in vision tasks, we usuallyhandle images with ﬁxed dimensions (i.e. height, width, and number of channels); in language tasks,we handle sentences with ﬁxed vocabulary size and maximum sentence length; and in tabular tasks,we often work with tables with ﬁxed number of columns.
Furthermore, as commented in Section 3.2,the posterior concentration of variable importance (Theorem 1) does not rely on this assumptionin its proof.
Therefore posterior concentration can occur even for high-dimensional settings withd
= o ( n ) .
In contrast, the BvM theorem usually represents a much stronger sense of convergence (i.e.,convergence in the predictive CDF of the entire Bayesian predictive posterior) and usually requires astricter set of conditions.
This is especially true in our setting, where our BvM results consider thequadratic functional of the nonlinear model, rather than the predictive function of the model itself.
Tothis end, we highlight that most of the modern BvM results in ML models are derived by assumingﬁxed dimension.
This includes [80] for Bayesian neural networks, [11, 86] for Gaussian processmodels, [66] for BART (i.e., Bayesian tree models), and ﬁnally [15] for general semi-parametricmodels.
Therefore, our setting is consistent with the modern literature in Bernstein von-Mises resultsfor ML models.
At the same time, it should be acknowledged that in recent years, there have been afew "high-dimensional" Bernstein von Mises results for specialized settings.
The more recent onesinclude [8] for linear regression, and [53] for Bayesian inverse problems with nonlinear forwarddynamics.
However, extending these results to the quadratic functionals of a general nonlinear modelis non-trivial and out-of-scope for our current work, but is a important and interesting direction offuture theoretical investigations.
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 [26])and estimating the importance of interaction effects and higher-order terms.
We would also liketo apply this method to large-scale scientiﬁc studies (e.g., epidemiology study based on extremelylarge EHR datasets) where an uncertainty-aware nonlinear variable importance estimation 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.