Towards a Unified Framework for Uncertainty-aware Nonlinear Variable Selection with Theoretical Guarantees
Authors: Wenying Deng
Beau Coker
Rajarshi Mukherjee
Jeremiah Zhe Liu
Brent A. Coull
Venue: NeurIPS 2022 Accept (2022)
Keywords: variable selection
posterior distribution
Bayesian nonparametric

Abstract: We develop a simple and unified framework for nonlinear variable importance estimation that incorporates uncertainty in the prediction function and is compatible with a wide range of machine learning models (e.g., tree ensembles, kernel methods, neural networks, etc). In particular, for a learned nonlinear model $f(\mathbf{x})$, we consider quantifying the importance of an input variable $\mathbf{x}^j$ using the integrated partial derivative $\Psi_j = \Vert \frac{\partial}{\partial \mathbf{x}^j} f(\mathbf{x})\Vert^2_{P_\mathcal{X}}$. We then (1) provide a principled approach for quantifying uncertainty in variable importance by deriving its posterior distribution, and (2) show that the approach is generalizable even to non-differentiable models such as tree ensembles. Rigorous Bayesian nonparametric theorems are derived to guarantee the posterior consistency and asymptotic uncertainty of the proposed approach. Extensive  simulations and experiments on healthcare benchmark datasets confirm that the proposed algorithm outperforms existing classical and recent variable selection methods. 

Towards a Unified Framework for Uncertainty-aware Nonlinear Variable Selection with Theoretical Guarantees
Wenying Deng Harvard University wdeng0g. harvard.
Beau Coker Harvard University beaucokergg harvard. edu
Rajarshi Mukherjee
Harvard University
ram521 Cmail. hhhaarrddd... edu
Jeremiah Zhe Liu"
Brent A. Coull"
Harvard University Google Research jereliu@google.com
Harvard University
bcoull@hsph. harvard. edu
Abstract
We develop simple and unified framework for nonlinear variable importance estimation that incorporates uncertainty in the prediction function and is compatible with wide range of machine learning models (c.g.. tree ensembles, kernel methods, neural networks, etc). In particular, for learned nonlinear model f fxx), we consider quantifying the importance of input variable x² using the integrated partial derivative v quantifying uncertainty variable importance by deriving its posterior distribution, and (2) show that the approach is generalizable even to non-differentiable models such as tree ensembles. Rigorous Bayesian nonparametric theorems are derived to guarantee the posterior consistency and asymptotic uncertainty of the proposed approach. Extensive simulations and experiments on healthcare benchmark datasets confirm that the proposed algorithm outperforms existing classical and recent variable selection methods.
क [()) We then (1) provide a principled approach for
#
Introduction
Variable selection is often of fundamental interest in many data science applications, providing benefits in prediction error, interpretability, and computation by excluding unnecessary variables. As datasets grow in complexity and size, crucial that variable importance estimation methods can account for complex dependencies among variables while remaining computationally feasible. Furthermore, as the number of approaches to model such datasets has increased, it is crucial that the importance of each variable can be compared across model classes and extended to new ones as they are developed.
While there established approaches for quantifying variable importance in linear models (e.g.. LASSO regression Hastie et al. (2015)). there is little consensus as to the preferred methodology theory for variable importance in nonlinear models. Generalized additive models Hastie and Tibshirani [990] use similar methods as their linear counterparts Waan et al. (214]). but the additivity assumption for nonlinear functions of the variables is too restrictive in many applications. Random Forests (RF) Breiman [200T] measure variable importance using an impurity measure, which based on the average reduction of the loss function were given variable removed from the model.
Friedman (200T] extended this method to boosting, where the definition of variable importance is generalized by considering the average over all of the decision trees. Deep neural networks (DNNs) are widely-used for many artificial intelligence applications, and a substantial effort has been invested
"Co-senior author. Work done at Harvard University.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

into developing DNNs with variable selection capabilities. Typically, this class of models involves manipulating the input layer, for example by imposing an L1 penalty Castellano and Fanelli [2000]. Feng and Simon| 2019], using backward selection Castellano and Fanelli [2000], or knockoffs Lu et al.] 2018]. Unfortunately, each model class based on DNNs requires a tailored procedure, which limits comparability across different model formulations.
Bayesian variable selection methods provide prineipled uncertainty quantification in variable im- portance estimates as well as complete characterization of their dependency structure. These methods allow the variable importance estimation procedure to tailor its decision rule with respect to the correlation structure Liu [021]. Yet, as in frequentist models, each method has different definition of a variable's 's importance, For example, in Bayesian additive regression trees (BART), a variable's 's importance can be measured by the proportion of trees that use it Chipman et al. (2010). while in Gaussian process (GP) models, a variable's importance can be measured by the frequency of the fluctuations of the estimated outcome- prredicctrr function (e.g., the length-scale parameter as controlled by the automatic relevance determination) in the direction of the variable Neal 1996),
Wipl and Nagarajan [2007]. Recently, a closely-related line of work uses the norm of the kernel gradient to quantify variable importance under classical GP models [He et al. 202T or deep Bayesian neural networks [Luu 2021]. However, these work either do not incorporate uncertainty, or are restricted to a particular model class (see Appendix D}. Furthermore, the traditional Bayesian model- ing procedures tend to be computationally burdensome, making them less feasible for large-scale
applications [[Anrriuu et al.] 2003).
Our work starts with the observation that many machine learning models can be written as kernel methods by constructing a corresponding feature map. For example, random forests can be written as kernel methods by partitions Davies and Ghahramani (2014), and deep neural networks can be
written as kernel methods by using the last hidden layer as the feature map Snoek et al. 2015), Hinton
and Salakhutdinov [007]] Calandra et al. (016). Each of these feature maps can be constructed before Bayesian learning of the GP (e.g.. g.... by pre-training on the same or a separate dataset), providing additional modeling expressiveness and representational capacity. Then, the GP learning is equivalent to performing Bayesian inference with respect to the (linear) weighting parameters of the feature-map map basis functions and the posterior inference proceeds analogously to that of a Bayesian linear regression (see Section 2. for details). The ability of a GP model to incorporate these adaptive feature maps becomes especially important in high-dimensional applications, where effective dimension reduction is necessary to circumvent the curse of dimensionality and ensure good finite-sample performance Bach 2016).
Contributions. We propose a unified variable importance estimation framework that is compatible with a wide range of machine learning models and can be defined by, or be closely approximated by, a differentiable feature map. Notable members include neural networks and random forests (Appendix B Our approach defines variable importance
the norm of the function's 's partial derivative, as was previously studied in the context of frequentist nonparametric regression Rosasco et al. 2013). We extend it to much wider class of models than previously considered (Section 2). propose a principled Bayesian approach to quantify the variable importance uncertainty in finite data (Section 3.1, and derive rigorous Bayesian nonparametric theorems guarantee the method's consistency and asymptotic optimality (Section B.2). To incorporate powerful non-differentiable models into our framework, we also show how to apply this approach partition-based methods (e.g.. decision trees) by leveraging their (soft) feature representation (Appendix F.I.. This leads to the first serivative-based Bayesian variable importance estimation approach for tree-type models that is both theoretically grounded and empirically powerful. This method strongly outperforms other variable importance estimation approaches tailor-designed for random forests (e.g.. impurity or candom-foreee knockoff [Breiman et al. 1984, Candes et al. 2017)). We conduct extensive empirical validation of our approach and compare its performance to that of many existing methods across a wide range of data generation scenarios. The results show a clear advantage of the proposed approach, especially in complex scenarios or when the input a mixture of discrete and continuous features (Section
2 Preliminaries
Problem Setup. We consider the classical nonparametric regression setting with d-dimensional features x (x', ,...,."'' xx) E X = Rd and a continuous response y E R. The features x are allowed to have a flexible nonlinear effect on y, such that:
y fx( + e;, where e; ~ NN(0,
(1)
2

---

with homoscedastic noise level o². The data dimension d is allowed to be large but assumed to be constant and does not grow with the sample size n. Here the date-gengg function fo is flexible nonlinear function that resides in an reproducing kernel Hilbert space (RKHS) Ho induced by certain positive definite kernel function ko, and the input space Xo of the true function spans only a small subset of the input features (x* xx!....,x)) xx),, i.e., Xo C X.
To this end, the goal of global variable importance estimation is produce a variable importance score , for each of the input features (x' (x!.... x") such that it can be used as a classification signal for whether x' E Xg. As a result, the variable selection decision can be made by threhsolding di > 8 with pre-defined threshold The quality of a variable selection signal t/ can be evaluated comprehensively 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 1 [ > s) over range of thresholds 8.
2.1 Quantifying Model Uncertainty via Featurized GP
In the nonlinear regression scenario given by Equation ① a classical approach to uncertainty-aware model learning is the Gaussian process (GP). Specifically, assuming that Jo be described by flexible RKHS Hy governed by the kernel function k; the GP model imposes a Gaussian process prior f GP k), such that the function evaluated at any collection of examples follows multivariate normal (MVN) distribution
f (f(x).... f(x,)) ~ MVN (mmx1- Knxn),
with mean m; = m(x) and covariance matrix K,, k(x,, xy). The choice of the prior mean m and kernel i enables prior specification directly in the function space. For example, the Matérn kernel with parameter y places prior over [v] - 1 times differentiable functions, with length-scale 1 and amplitude variance o². As -> Do, this reduces to the common radial basis function (RBF) kernel k(x;, x,) = o² exp(IIx; x, 12)
Under the above construction, the posterior predictive distribution of f evaluated new observations xi, x is also multivariate normal,
f f 3 MVN(E Cov(ff)) where
B[f*] = m + K" (K o (y - m)); Cov[f*] = K K"KK + +0 K
with m m(x;) Ki xx),,,,), and Ki₃ k(x,,x x;) Equation 2 is known the kernel-
based representation (or dual representation) of a GP Rasmussen and Willimms (005). Although
mathematically elegant, the posterior 2 is expensive compute due to the need to invert the matrix (K + IT).
Feature-based Representation of A GP. Alternatively, Mercer's 's theorem Cristianini and Shawe Taylor [2000] states that as long as the kernel function -)::-) can be written as the inner product of a set of basis functions o(x) {dv (x) 2-1- such that i(x, x') <(x) Φ(x'), then elements of the RKHS| f E Ht can be written in terms of a linear expansion of basis functions Rasmussen and
Williams [2005]:
D
f(x) B TB, where B ~ MVN MMM33MMMMMM ID).
k=1
This is known as the feature- reature-based representation (or primal representation) of a GP. Notice that is not an approximation method but an exact reparametrization of the GP model whose kernel function is induced by feature representation @xx). Also note that under this featurized representation (). the predictive model f is linear in terms of the model parameters 3 = 1 However, this "Inearity" in the model parameters does not restrict the expressiveness of f. since the GP model is flexibly combine the nonlinear basis functions {ou) to best fit the outcome. Furthermore, the basis functions (x)) can be updated as part of the learning process, which we discuss in the sequel.
essentially learning to use the weights (Bx 2 ]
Scalable Posterior Computation via Minibatch Updates. The above feature- based representation
is powerful in that it reduces the GP posterior inference into a Bayesian linear regression problem for B. This brings two concrete benefits. First, the posterior of (3 in Equation adopts closed form:
e ~ (E Cov(3))
where
(4)
B
++++E の (y /p)/0²; Cov[ß] = Eg (8 0/02+1+
3

---

where &b = (Ø(x,) ......
x(x,) yT E RnXD is the feature matrix evaluated on the training data
Rasmussen and Williams 2005) For large-scale applications, Equation H enables us to compute
the exact posterior of B3 in a mini-batch fashion. For example, the posterior matrix Cov[3] = Eg can be updated using the Woodbury identity:
EB. = Σs. Σz, (o 2m 2m) 4 Dm
(5)
where Lm is the D-dimension batch-specific feature matrix evaluated on the mini-batch. Similarly, the posterior mean E[33 can be computed by accumulating the D x 1 vector ①b (y 4/p) Em (ym - 1i) and computing the posterior mean according to Equation
at the end.
The posterior distribution of (3 induces a GP posterior for the prediction function f Φb where क' T
the feature map evaluated on the test data, with mean E[[*] = $ ş + $ " (y - $p///²² and covariance Cov [f*] 4 "== This distribution is equivalent to the kernel-basee representation ( but reduces the computational complexity from cubic time ((n²) to linear time (n)) and is minibatch compatible (i.e.. Equation (5) Algorithm | and B provides summary of the learning algorithm. Finally, we note that the basis functions @ {ok]k- can also be updated as part of the learning procedure (e.g., via maximum a posteriori (MAP) inference), which we discuss in Appendix A.4
Incorporating Modern ML Model Classes. The second key advantage of the feature-bass
representation (3) is its generality: wide range of machine learning models can be written in the ceature-based form (x) = <(() 3 Rahimi and Recht 2007 Davies and Ghahramani [2014), Lee al. 2017]. making the GP unified framework for quantifying model uncertainty for wide array of modern MI models. Appendix
B summarizes important examples including GAMs, decision trees,
manoum-fcature models. deep neural networks and their ensembles. Appendix B.I summarizes a list of general conditions the model should satisfy for it to be compatible with the proposed framework (i.e., weak differeetiabilityy. Lipschitz condition, and growth rate of model complexity). Furthermore, when a deterministically-rraiine
33 is available (e.g., via sophisticated adaptive shrinkage procedure that is not available in a Bayesian context), we can incorporate this as prior knowledge into GP modeling by setting μ /3 (Equation (3).
2.2 Bayesian Nonparametric Guarantees for Probabilistic Learning
The quality of a Bayesian learning procedure commonly measured by the learning rate of its posterior distribution II, II(- {x,, y1 )=z). Intuitively, the rate of this convergence is measured by the size of the smallest shrinking balls around fo that contains most of the posterior probability.
Specifically, we consider the size of the set A., = {g ||g - foll?: < EM,) } such that 11, (A,) + 1 Ghhssl and Vaart 2007 Polson and Rockova 2018). The concentration rate e,, here indicates how fast the small ball A., concentrates towards fu the sample size increases. Below state the formal
definition of posterior convergence Ghosal and Vaart (207)).
Definition 1 (Posterior Convergence). For fo : X' -
R where X' Rdd, let Ho denote the true RKHS induced by a kernel function ko, and let H. denote the RKHS induced by the feature function & X' + RDP Let fo E Ho be the true function, and let Eg denote the expectation with respect to the true data-generation distribution. Assuming H4 is dense in Ho. then, the posterior distribution n1 (f) concentrates around fo at the rate en if there exists an 6n - O such that, for any Mn - cc,
Eoll, (f: IIf foll: 2 Mnnx) + O.
Notice that we allow the model space HiS and the true function space Ho to be different, but His must be dense in Hu for the convergence to happen. Fortunately, this condition is shown to hold for a wide variety of ML models, including random features, random forests, and neural networks [Biaul 2012 Hornik et al. [989] Rahimi and Rechi 2008 [Shmidt-- Hieber 20200] Ročková and van der Pas The notion of posterior convergence can also be used to discuss the learning quality of other probabilistic estimates (e.g., variable importance v/.). In that case, we can simply replace (J, fo) in
2020
1 by their variable importance counterparts. This is the focus of Section 3.2
3
Methods
3.1
Quantifying Variable Importance under Uncertainty
In this work, we consider quantifying the global importance of variable based on the norm of the corresponding partial derivative. This is motivated by the observation that, if function f is

---

differentiable, the relative importance of a variable x at a point x can be captured by the magnitude of the partial derivative function, lax [x)] Rosasco et al. 2013]. This quantity requires the consideration of two issues. First, instead of quantilyıng the relevance of variable on a single input point, we need to define a proper global notion of variable importance. Therefore, it is natural to integrate this partial derivative over the input space x E X (f( = ox fll =
SxEx | (x (x). Second, since Px(x) is not known from the training observations, vf( can be approximated by its empirical counterpart,
0 vソ( = xx flR
n
Σ
Ox f(x)))
Notice that (f) is an estimator that is derived from the prediction function f estimated using finite data. Consequently, to make proper decision regarding the importance of an input variable x² it is important to take into account uncertainty in f. To this end, by leveraging the featurized GP representation introduced in Section 3.T we show that this can be done easily for a wide range of ML models (x) = (xx) T ,3 by studying the posterior distribution of 4//
Posterior Distribution of Variable Importance. After we obtain the posterior distribution of 3 the posterior distribution of variable importance can be derived according to Equation 7
1 a - Ox
2 Ox f(x)
1 BT
H
0$ T
(f)
T
dx B,
(8)
Oxi
where ax
8+h E Ig is the derivative of the feature map with respect to x², across training samples.
The posterior distribution of (f) adopts closed form as generalized chi-squared distribution (see Appendix A.. for derivation). In practice, we can sample dy conveniently from its posterior
distribution by computing a axe f(X) (bxc "g("), where g(*) are Monte Carlo samples from the
closed-form posterior
There are two ways in which uncertainty aids the variable importance estimation process. First, the posterior survival function (f) > 8) of the variable importance utilizes the full posterior distribution of (f) identify the probability that the variable x exceeds given threshold s. By increasing E (0, oo), P(eb > s) provides an intuitive sense of how model's 's belief about the importance of variable x changes as the criteria s becomes more stringent, similar to the regularization path used by LASSO methods [Friedman al. 2010] but with the incorporation of posterior uncertainty about the variable importance. See Appendix [ for an application to a Bangladesh birth cohort study. Second, by integrating the survival function over the threshold, i.e.. Sa>o P(ay(f) > s) ds, we obtain the posterior mean of (f). and this too incorporates uncertainty in f. To see this, notice that by using the "trace trick" we can write
T3 O 0Φ T
E|STT OD 04
H OΦ Oxi Ox Cov|
E E(y =
Oxi 0xi 3)
E[S] +
(9)
Ox
Oxi
where all expectations are taken with respect to the posterior. Therefore, the posterior mean of (f) depends on the covariance structure of B, and how it interacts with the eigenspace of the partial 0s 20"
derivative functions (encoded by oxe axe In Section 4 we provide extensive investigation of AUROC scores using the posterior mean of (f) for quantifying variable importance.
In Appendix A.3 we summarize the algorithms for computing the posterior distributions of the featurized Gaussian process (Equation 4 and for the posterior distributions of variable importance (Equation and discuss their space and time complexity.
3.2 Theoretical Guarantees
From a theoretical perspective, the variable importance measure vi; introduced in can be under- stood as a quadratic functional of the GP model f Efrrmovvihh and Low 1996). To this end, rigorous Bayesian nonparametric guarantees can be obtained for Jy 's ability in learning the true variable importance in finite samples (i.e., posterior convergence, Theorem | and its statistical optimality from frequentist perspective, in providing low-variance estimator that attains the Cramér-Rao bound (i.e., Bernstein von-Mises phenomenon, Theorem 2 Note that for a given general model T /3, only need to satisfy three mild regularity conditions to be fully compatible with the proposed framework (i.c., weak differentiability, Lipschitz condition, and growth rate of model
f(x) = =(x)TA
5

---

complexity). We summarize these conditions in Appendix B.1 and explain them in detail in the sequel.
Posterior Convergence. We first show that, for an MI model f that can learn the true function fo with rate Cn (in the sense of Definition D.. the entire posterior distribution of the variable importance measure by (f) converges consistently to a point mass at the true v (fo) at speed that is equal or faster than Ent.
Theorem 1 (Posterior Convergence of Variable Importance 1bz). Suppose yi = [o(x;) + es, ei did. NO and denote as Es the expectation with respect to the true data-generation distribution centered around fo. For the RKHS Hi induced by the feature function & : X R.D and f E HH. if:
(1) The posterior distribution IT.(F) converges toward fo at rate cn;
(2) The difserentiation operator Dj
Jxt S is bounded: |D inf(C 0
|D, I Cll for all f E Hol:
Then the posterior distribution for v; (f) = 8xx flk contracts toward V (So)) = - sollix at a rate not slower than ex. That is, for any M,, -
Eg TT
sup JE(I
(0)(s)) - v,(so) 2 Mn - O.
The proof is in Appendix C Theorem is generalization of the classical result of quadratic functional convergence under linear models and sparse neural networks to a much wider range of ML models in the context of Bayesian variable importance estimation [Efromovich and Low 1996 Liu 2021 Wang and Rocková 2020 It confirms the important fact that, for an ML model that can accurately learn the true function So under finite data, we can consistently recover the true variable importance at a fast rate by using the proposed variable importance estimate ر5 (/), despite the potential lack of identifiablity in the model parameters (e.g., weights in neural network). Importantly, although our main setting assumes fixed data dimension d (see Problem Setup), the posterior concentration result Theorem I does not rely on this assumption in its proof, and is in fact compatible with the high-dimensional setting where d is allowed to grow with sample size at a rate of o(n). See Appendix C (in particular, Remark 4) for further discussion.
From a practical point of view, Theorem reveals that the finite-sample performance of variable importance 1) (f) depends on two factors: (1) the inite- generalization performance of the prediction 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 take care to select a model class f that has a theoretical guarantee in capturing the target function fo. empirically delivers strong generalization performance under finite data, and is well-conditioned in terms of the behavior of its partial derivatives. To this end, we note that, under the featurized Gaussian process f == o(x) /3 discussed in this work, users are free choose a performant model class (c.g.. random forest, mandom-feature or DNN) whose feature representation spans an RKHS He that is function space (therefore f enjoys a convergence guanrantee, see Remark 3 in Appendix C for further discussion) [Biau 2012 Hornik et al. 1989 Rahimi and Recht 2008 (Schmidt-Hieber 2020 Ročková and van der Pas, 2020 and is empirically more effective than the GP methods based on classical kernels such as RBF. We discuss the Lipschitz condition of these models in Appendix E. Indeed, as we will verify in experiments (Section 4. there does not exist an "ooptimal" model class that performs universally well across all data settings (i.e., no free lunch theorem WWolprr and Macready 1997]). ). This highlights the importance of having general-purposs purpose framework for variable importance estimation that can flexibly incorporate the most effective model for the task at hand. Finally, we notice that although Theorem T is stated as an asymptotic result, when finite-sample error bound e, for the model class is available (i.e., Condition (1) in Theorem [ it is trivial to obtain a finite- saaple error bound for variable importance (f) by extending the proof of Theorem | Appendix C. provides an example of such bound based on the Bernstein inequality.
dense in the
Statistical Efficiency & Uncertainty Quantification. Next, we verify the uncertainty quantification ability of the variable importance measure vbs()) under a featurized GP by showing that it exhibits the Bernstein-von Mises (BvM) phenomenon. That is, its posterior measure II (#)( converges towards Gaussian distribution that is centered around the truth v (fo), so that its (1 a)% level

---

credible intervals achieve the nominal coverage probability for the true variable importance. More importantly, the BvM theorem verifies that the posterior distribution of b is statistically optimal. in the sense that its asymptotic variance attains the Cramér-Rao bound (CRB) that cannot be improved upon [Bickel and Kleijn 2012].
Theorem 2 (Bernstein-von Mises Theorem for Variable Importance 4yy). Suppose y, =xx +
e,, e, hid. o²(0,²2,, = I,...,n. Denote Dj : +
0 8x exe f the differentiation operator and
11, = Dj D2 the inner product of Dj, such that:
= =
(10)
Assuming conditions (1)-2) in Theorem 1 holld and additionally:
(3) fo is square over the support X and 2| 1;
(4) rank( op(vT);
Then
(f(v by(f)) 4 (0, foll
The proof is in Appendix D Theorem provides rigorous theoretical justification for (f)sss ability to quantify its uncertainty about the variable importance. More importantly, it verifies that (f) has the good frequentist property that it quickly converges to minimum- at a fast speed, which is important for obtaining good variable importance estimation performance in practice. Compared to the previous BvM results that tend to focus a specific Bayesian ML model, Theorem is considerably more general (i.c.. applicable to a much wider range of models) and comes with simpler set of conditions [Rockova 2020, Wang and Rocková 2020 Liu 2021]
estimator
Specifically, (3) is a standard assumption in nonparametric analysis. It ensures the true function fo does not diverge towards infinity and makes learning possible (Castillo and Rousseau 2013]. The unit norm assumption fl = 1 is only needed to simplify the exposition of the proof, and the theorem can be trivially extended to ||.olla C for any C > O. The most interesting condition is (4). Let us denote Hl as the space of partial derivatives functions 2 f of the model functions f E Hog Then, intuitively, (4) says that to attain the BvM phenomenon, the effective dimension of the derivative function space H, (as measured by rank H,) rank ())) cannot be too large. Since the effective dimension of the derivative space is bounded above by that of the original RKHS f H.s (4) essentially states that the effective dimension of the model space H.g. cannot grow too fast with data size (i.e., op( (V))). Fortunately, this condition is satisfied by a wide range of ML models
including trees and deep networks (Rockova 2020 Wang and Rocková 2020 See Appendix for further discussion.
4 Experiment Analysis
In this section, we investigate the finite-sample performance of the derivative norm metric v for variable importance estimation 7 under wide variety of ML methods. We illustrate the breath of our framework by applying it to tree ensembles (Appendix F.I. where a principled and gradient- based mucertainty-awar variable importance estimation approach has been previously unavailable. We also apply it to linear models and (approximate) kernel machines, which are standard approaches to variable selection in data science practice [Tibshirani 1996 Bobb al. 2015). Over wide range of complex and realistic data scenarios (e.g., discrete features, interactions, between-featur correlations) derived from socioeconomic and healthcare datasets, we investigate the method's statistical performance in accurately recovering the ground-truth features (in terms of the Type and Type II errors), and compare it other well-established approaches in each of the model classes (Table D.. Our main observations are:
01: Importance of generality. There does not exist a model class that performs universally well across all data scenarios (i.e., no free lunch theorem (Wolpert and Macready 1997]. Figures I 15 This highlights the importance of an unified 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.
02: Good prediction translates to effective variable importance estimation. Comparing between different model classes, the ranking of models' predictive accuracy is generally consistent with the
7

---

ranking of their variable importance estimation performance under 3 (i.e., better prediction translates to better variable importance estimation, as suggested in Theorem I
03: Statistical efficiency of 1... Comparing within each model class, the derivative norm metric e generally outperforms other measures of variable importance. The advantage is especially pronounced in small samples and for correlated features. This empirically verifies that 1/ has good finite-sample
statistical efficiency even under complex data scenarios (as suggested in Theorem 2
Model Class
(Ours)
Baselines
Tree Ensembles
RF-FDT
RF- RF-Impurity, RF-Knookoff, BART
Kernel Methods & NNs
RFF, NN
BKMR, BAKR
Linear Models
GAM
BRR, BL
Table I: Summary of methods considered in the experiments.
Models & Methods. We consider three main classes of models (Table D: (1) Random Forests (RF). Given a trained forest, we quantify variable importance using 3/ by translating it to an ensemble of featurized decision trees (FDT) (Appendix F.I. and compare it three baselines: impurity (RF-impurity) [Breiman et al. 19841 RF-based kernel knockoff (RF-knockoff) [Candes et al. 2017]. and Bayesian Additive Regression Trees (BART). (II) (Approximate) Kernel Methods & Neural Networks. We apply ey to a random-feature model that approximates a GP with an RBF kernel Rahimi and Recht [200]] and set the number of features to (oglgg() to ensure proper approximation of the exact RBF-GP Rudi and Rosasco (018)) which is termed Random Fourier Feature model (RFF). We also apply v to Neural Networks (NN) based on wide ReLU neural network with 512 hidden units and LASSO regularization in the hidden layer weights [Lembadri 2021] We compare them to Bayesian Kernel Machine Regression (BKMR) Bobb et al. 20155
et al. based on a GP with an exact RBF kernel and spike-and-slab prior, and Bayesian Approximule Kernel Regression (BAKR) based on meatureeature model with a projection-based feature importance measure and an adaptive shrinkage prior (Crawford et al.) 2018). (III) Linear Models. We apply bh to a featurized GP representation of the Generalized Additive Model (GAM), with the prior center u set at the frequentist estimate of the original GAM model obtained from a sophisticated REML procedure [Wood] 2006). We compare it to two baselines: Bayesian Ridge Regression (BRR) Hoerl and Kennard [1970] and Bayesian LASSO (BL) Park and Casella (2088] Appendix G provides further detail. Previously. (Liu 2021] studied the specialization of our framework to the deep neural networks (DNNs).
do not repeat that work here as DNN is not yet standard data science
model for tabular data.
To quantify variable importance while accounting for posterior uncertainty of the variable importance Wy (f), we examine its posterior survival function Ss>o ((bby()) > 8) ds (i.e., the posterior likelihood of sb;(f) being 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 and BKMR. See Appendix G
Datasets and Tasks. We consider two synthetic benchmark datasets and three real-world socio- economic and healthcare datasets, encapsulating challenging phenomena such as between-featur correlations and interaction effects. For the synthetic benchmark datasets, we generate data under the outcome-generation Gaussian noise model ~ (foo 0.01) for four types of outcome- generation functions fo (1 (linear, rbf, matern32 and complex, see Appendix G.2 for a full description) with the number of causal variables set at d' = 5. We consider two types of feature distribution: (1) synthetic-continuous: all features follow x" Unif(-2, 2, 2): (2) synthetic-mixturg two of the causal features and two of the non-causal features are distributed as Bern(0.5) and the rest are distributed as Unif(-2) 2). Features both distributions are independent. We vary sample size n E (100, 200, 500, 1000) and data dimension d E (25, 50, 100, 200), leading to 128 total scenarios.
For real-world data, we consider (1) adult: 1994 U.S. census data of 48842 adults with eight categorical and six continuous features Kohavi, (2) heart: a coronary artery disease dataset of 303 patients from Cleveland clinic database with seven categorical and six continuous features Detrano et al. 1989;; ; and (3) mi: disease records of myocardial infarction (MI) of 1700 patients from Krasnoyarsk interdistrict clinical hospital during 19922-195, with 113 categorical and 11 continuous features Golovenkin et al. (2020). All datasets exhibit non-trival correlation structure among features (Appendix Figures 3 Since the ground-truth causal features on these datasets are not known, in order to rigorously evaluate variable importance estimation performance, we follow the standard practice in causal MI to simulate the outcome based on causal features selected from data (Yao
et al. 2021
We use the four outcome-generating functions described previously and evaluate
8

---

over the same data size x dimension combinations, leading to 192 total scenarios We repeat the simulation 20 times for each scenario, and use AUROC to measure the variable importance estimation performance (in terms of Type I and Type II errors) of each method.
In Appendix ! we further evaluate the method on well-studied environmental health dataset (Bangladesh birth cohort study [Kile et al. 2014)) with respect to the real outcome (infant development scores). We visualize the "Bayesian" regularization path as introduced earlier. The selected variables correspond well with the established toxicology pathways in the literature [Gleason et al. 2014].
RF-FDT (Ours)
RF-ummpurty RFF (Ours)
RF-KKnckOr
BART
method
NN (Ours)
BKMA BL
BAKR
GAM (Ouns)
BRR
syrthalic-mixttuee
synthelic-comiinuooa continuuus
adult
heart
35.
AUROC
37-
250
300
753
1000
250
300
750
1000
150
30
753
1000
30
100
100
200
250
250
200
750
1000
Teasig
Figure 1: Method performance in variable importance estimation (measured by AUROC, row 1) and prediction (measured by test MSE, row 2) under matern32 data-generation function and with input dimension 100 (five causal features). Therefore, the variable importance scores of the five causal features are expected to be higher than the other 95 variable importance estimations. The x-axis represents the training sample sizes v E {100, (00, 500, 1000). GAM does not produce valid results for case S so the results from this model in these cases are not shown. The ranking of FDT (solid purple) outperforms other methods in most of the data settings, and GAM outperforms in the setting of large data size and high percentage of categorical features (adult and mi). The rankings of performance are roughly consistent between prediction and variable importance estimation.
4.1 Results
Figure I shows the methods' performance in variable importance estimation (Row 1) and prediction (Row 2 in an exemplary setting, where the true function fo is matern32 with an input dimension == 100. It represents the tabular data setting that we are the most interested in: nonlinear feature- response relationship with interaction effects and high input dimension. This is because fo is sampled from an RKHS induced by Mátern 3 kernel, which contains a large space of continuous and at least once differentiable functions [Rasmussen and Williams 2005). We delay complete visualizations for all 320 scenarios to Appendix Recalling the three observations introduced earlier:
O1 ('No free lunch"): No method performs universally well. For example, BAKR performs robustly in correlated datasets (heart and mi), but poorly otherwise. Kernel approaches (RFF and BKMR) perform competitively in low dimension, but their performances deteriorate quickly as This is likely due to the classical kernel method's well-known inability to learn an adaptive feature representation, which consequently leads to suffering from the curse of dimensionality and unstable and suboptimal variable importance performance in high dimensions [Bach] 2017). FDT is generally the strongest method in small samples and high dimensions, but be outperformed by GAM in large samples and data with high percentage of categorical features (adult and mi). Notably, they often outperform NN, which is traditionally regarded as the go-to model for nigh-dimensional nonlinear settings. This highlights the importance of unified framework that allows users to select the most appropriate model for variable importance estimation depending on the data setting.
dimension d increases (Figure and Figure 6 7
PI the setting where required data dimension is higher than that of the real data, generate additional synthetic features from Unif(-2, 2). We use n E (50, 100, 150, 257} heart due to data size restrictions. For the prediction plots. method will not be visualized if they share the model fit with another method (RF-impurity and RF- qfooknokofff,, or if it does not produce valid results due to small sample size (GAM).

---

02 ("Good prediction implies effective variable importance estimation"), "): Fixing the variable importance ; and comparing the variable importance estimation performance of each model class (i.e., FDT, GAM, NN and RFF, which are solid lines in Figure D. we see that their rankings in prediction (row 2) are largely consistent with the corresponding rankings in variable importance estimation. It is worth noting that this pattern is occasionally violated (e.g., GAM in adult, 500 and heart, 7 = 250), but that does not contradict our conclusion (Theorem D since the convergence rate of the prediction function only forms an upper bound for the convergence rate of eyg. Finally, when models have comparable generalization performance, we observe that the Lipschitz condition plays a role in variance importance performance (which is consistent with our theoretical observations in Theorem I For example, in Figure | NN and FDT are largely comparable in predictive performance among the real datasets (adult, heart and mi). However, tree-based FDT are known to have well- conditioned Lipschitz behavior when compared to NN (Appendix F. I. which is consequently translated to improved sinite-sample performance in variable importance estimation.
03 (Statistical efficiency of +) When comparing among variable importance estimation methods from the same class (especially for tree models, i... FDT v.s. RF-impurity / RF-knockoff' / BART), we see that FDT is competitive or strongly outperforms its baselines in variable importance estimation, despite being based on exactly the same fitted model (RF (RF-impurity / RF-knockoff), or not accounting for the uncertainty in the tree growing process (BART). This pattern is consistent in most data settings, and the advantage is especially pronounced in high dimensions, small data sizes, and correlated datasets (Appendix H Figure 6 To]. This provides strong empirical evidence for the fact that by is statistically efficient estimator for variable importance with good finite-same sample behavior (as suggested in Theorem 2 and can deliver strong performance for tabular data when combined with a performant ML model like random forests. Appendix H contains further discussion.
5 Discussion and Future Directions
The modern data analysis pipeline typically involves fitting multiple models, comparing their perfor- mance, and iterating as necessary. When variable selection is involved, the practitioner may ask are the variable importance scores across models measuring the same behavior? And, what if the most suitable model does not have a satisfactory variable importance estimation procedure? ? By framing model choice as the specification of kernel which includes kernels corresponding machine learning methods like neural networks and random forests in addition to the long list of traditional kernels propose unified variable importance estimation procedure that is compatible across models and prove strong guarantees for this procedure.
Limitations. We do not consider uncertainty in the feature map itself. For example, the kernel induced by the featurized decision tree studied here does not consider uncertainty in the tree's s partitioning process. Meanwhile, the fact that the full posterior inference is performed only with respect to 3 indeed places a limitation on the model's ability in uncertainty quantification, as the uncertainty in the model hyperparameters is not accounted for. Yet, this does not seem to be significant limitation in the method's empirical performance (e.g., FDT outperforms BART in our experiments), although this point still merits further investigation in the future. On the other hand, in the future, it would be worth expanding this framework to other model classes (c.g.. MARS Friedman| [199]) and estimating the importance of interaction effects and higher-order terms (see Appendix K for details).
Societal Impacts. We expect the method proposed to provide a set of powerful tools for practitioners to understand the importance of input variables in their ML models with limited data, which is especially important for scientific investigations in the fields of epidemiology and computational biology. However, we recognize that this approach can potentially be utilized by bad actors to probe the input-variable uncertainty of existing ML system, and use it to engineer more targeted white- box adversarial attacks. To this end, we recommend system developers to incorporate this approach into the formal verification procedure of an ML system, so to monitor and understand the model uncertainty with respect to input variables, and devise proper improvement and prevention strategies (e.g., data augmentation randomized smoothing targeted at specific variables) accordingly.
A Acknowledgement
This research was supported by NIH grants ES000002, ES028800, ES028811, 11, ES028688, and ES030990. We would also like to thank Dr. Andrew Beam for the helpful discussions.
10

---

References
C. Andricu, N. De Freitas, A. Doucet, and M. I. Jordan. An introduction to meme for machine learning, Machine learning, 50(1):5-43, 2003.
F. Bach. Breaking the Curse of Dimensionality with Convex Neural Networks. arXiv:1412.86990 [cs, math, stat), Oct. 2016. URL http://arxiv... org/ abs/1412 8690
Bach. Breaking the Curse of Dimensionality with Convex Neural Networks. Journal of Machine Learning Research, 18(19):1-533 2017. ISSN 1533-7928.
G. Biau. Analysis of random forests model. The Journal of Machine Learning Research, 13(1): 1063-1095, 2012.
G. Biau, Scornet, and J. Welbl. Neural Random Forests. Sankhya A, 81(2)347-386, Dec. 2019. ISSN 0976-8378. doi: 10.10077/13177-001-01133. htLpttss: 1tt.:/...r1111111...1... 513171-018-0133-
J. Bickel and B. J. Kleijn. The semiparametric bernstein von mises theorem. The Annals of Statistics, 40(1):206-23) 237, 2012.
F. Bobb, L. Valeri, B. Claus Henn, D. C. Christiani, R. O. Wright, M. Mazumdar, J. Godleski, and B. A. Coull. Bayesian kernel machine regression for estimating the health effects multi multi-pollus pollutant mixtures. Biostaristics, 16(3):493-500 July 2015. ISSN 1465-4644. doi: 10.1093..1/// 1093//iostattisiess/ kxu058.
D. Bontemps. Bernstein-von mises theorems for gaussian regression with increasing number of regressors. The Annals of Statistics, 39(5):2557-2588 2011.
L. Breiman. Random Forests. Machine Learning, 45(1) Oct 2001. ISSN 1573-0565. doi: 1010023/A:10100334444444 htll:ttps:///do
L. Breiman, J. Friedman, C. J. Stone, and R. A. Olshen. Classification and Regression Trees. Taylor & Francis, Jan. 1984. ISBN 978-0-412-04844-88
E. Burnaev, A. Zaytsev, and V. Spokoiny. The bernstein-von mises theorem for regression based on gaussian processes. Russ. Math. Surv, 68(5):954-956, 2013.
R. Calandra, J. Peters, C. F.. Rasmussen, and M. P. Deisenroth. Manifold gaussian processes for regression. 2016 International Joint Conference on Neural Networks (IJCNV), pages 3338-3345. IEEE, 2016.
E. Candes, Y. Fan, L. Janson, and J. Lv. Panning for Gold: Model-X Knockoffs for High-dimensional Controlled Variable Selection. arXiv:1610.0235/ [math, stat]. Dec. 2017. URL http:://rxiiv org/abs/16t 11600. 02351 arXiv: 1610.02351.
G. Castellano and A. M. Fanelli. Variable Selection Using Newral-Network Models. 2000.
I. Castillo and J. Rousseau. A bernstein-von mises theorem for smooth functionals in semiparametric models. The Annals of Statistics, 43(6):2333-2383 2015,
H. A. Chipman, E. George, and McCulloch. BART: Bayesian additive regression trees. The Annals of Applied Statistics, 4(1):266-29) Mar. 2010. ISSN 1932-6157, 1941-7330. doi: 2214/09-AAAS200 Publisher: Institute of Mathematical Statisties.
K. Choromanski, M. Rowland, Sarlos, V. Sindhwani, Turner, and A. Weller. The geometry of random features. In A. Storkey and F. Perez-Cruz, editors, Proceedings of the Twenty-Fing First International Conference on Artificial Intelligence and Statistics, volume 84 of Proceedings of Machine Learning Research, pages 1-9, PMLR, 09-11 Apr 2018. httlhttts:///roceeelli mlr. wt..res//vv]] prress//eA/chorroanskk... html
Crawford, K. C. Wood, X. Zhou, and S. Mukherjee. Bayesian Approximate Kernel Regression With Variable Selection. Journal of the American Statistical Association, 113662)):7110-1721 Oct. 2018. ISSN 0162-1455, doi: 10. (098000662149992111.111 htt https://t::///o...
10. 00l10080016215555 2017 1361830

---

N. Cristianini and J. Shale-Taylor. An Introduction to Support Vector Machines and Other Kernel- based Learning Methods. Cambridge University Press, Cambridge, 2000. ISBN 978-0-521-78019
3. doi: 10.1017/CB097803 1100051100111 1801389.
S. d' Ascoli, L.. Sagun, and G. Biroli. Triple descent and the kinds of overfitting: Where & why do they appear? Advances in Neural Information Processing Systems, 33:3058-3069, 2020.
A. Davies and Z. Ghahramani. The random forest kernel and other kernels for big data from random partitions. arXiv preprint arXiv: 4293, 2014.
R. Detrano, A. Janosi, W. Steinbrunn, M. Pfisterer, J. J. Schmid, S. Sandhu, K. H. Guppy, S. Lee, and V. Froelicher. International application of a new probability algorithm for the diagnosis of coronary artery disease. Am J Cardiol, (15))304-3100 Aug. 1989. ISSN 0002-9149, doi: 10.1016/0002-99 (00066/0002-1144889)99554444 49989)90524--9
Efromovich and M. Low. On optimal adaptive estimation of a quadratic functional. The Annals of Statistics, 24(3): (113):1106-1100 102-1 1996.
Feng and N. Simon. Sparse-Input Neural Networks for nigh-dimensional Nonparametric Regres- sion and Classification. 1711.07592
[stat], June 2019. URL http:://arxi..org///4
J. Friedman, T. Hastie, and R. Tibshirani. Regularization paths for generalized linear models via coordinate descent. Journal of statistical software, 33(1):1, 2010
J. H. Friedman. Multivariate Adaptive Regression Splines. The Annals of Statistics, 19(1): 67, 1991. doi: 1...214//o//11764496666 URL https:s////doi.oo org/10 00t.../1///////17777777771
J. H. Friedman. Greedy function approximation: A gradient boosting machine. Annals of Statistics,
189-121189 Oct. 2001. ISSN 0090-5364, 2168-866 doi: 10.1214/ass/1013304444 URL
https://projecteuucce http:ppoojettteccc.d.orrrrreeaaa http://..r...eeee.ddaa///001133.
N. Frosst and G. Hinton. Distilling a Neural Network Into a Soft Decision Tree. arXiv: [cs, stat), Nov. 2017.
arXiv: 1711.09784.
Geurts, Ernst, and Wehenkel. Extremely randoizze trees. Mach Learn, 63(1):3-42, Apr. doi: 10.10077/1199944000--6222 URL https://odoo www.oii..rrg/// wtt//0..//// 1007/
2006. ISSN 1573-0565. 810994-006- 002-000066622222
S. Ghosal and A. v. d. Vaart, Convergence rates of posterior distributions for noniid observations. The Annals of Statistics, 35(1):192-223, Feb. 2007. ISSN 0090-5364, 2168-8966 doi: 10.1214/ 009053606000001111
K. Gleason, J. P. Shine, N. Shobnam, L. B. Rokoff, H. S. Suchanda, M. O. S. Ibne Hasan, G. Mostofa, C. Amarasiriwardena, Q Quamruzzaman, M. Rahman, et al. Contaminated turmeric is a potential source of lead exposure for children in rural bangladesh. Journal of Environmental and Public Health, 2014, 2014.
S. E. Golovenkin, J. Bac, A. Chervov, E. M. Mirkes, Y. V. Orlova, F. Barillot, A. N. Gorban, and A. Zinovyev. Trajectories, bifurcations, and pseudo-time in large clinical datasets: applications to myocardial infarction and diabetes data. GigaScience, 911) ((iaa)laaal22)) 28, Nov. 2020. ISSN 2047-217X. doi: 2o.10099gigasciencc//aaaa
URL BtRp:htss:://at
org 10. 1093/ www..///gggiiinnnggiiiaa..
J. D. Hamadani, F. Tofail, B. Nermell, R. Gardner, S. Shiraji. M. Bottai, S. Arifeen, S. N. Huda, and M. Vahter. Critical windows exposure for arsenic-associated impairment of cognitive function in pre-school girls and boys: population-based cohort study. International journal of epidemiology. 40(6):1593-160 2011.
D. A. Harville. On the distribution of linear combinations of non-central chi-squares. The Annals of Mathematical Statistics, 42(2):809-83 -11, 1971.
Hastie, R. Tibshirani, J. Friedman, and H. Friedman. The elements of statistical learning: data mining, inference, and prediction, volume 2. Springer, 2009.
12

---

T. Hastie, R. Tibshirani, and M. Wainwright. Statistical Learning with Sparsity: The Lasso and Generalizations. Chapman Hall/CRC, 2015. ISBN 978- 778-1-4987--216-31
T. J. Hastie and R. J. Tibshirani. Generalized Additive Models. Chapman and Hall/CRC, Boca Raton, Fla, 1st edition edition, June 1990. ISBN 978-0-412-34390-2.
X. He, J. Wang, and S. Lv. Efficient kernel-base -aase variable selection with sparsistency. Statistica Sinica, 31(4):2123-2155 2021.
G. E. Hinton and R. R. Salakhutdinov. Using deep belief nets to learn covariance kernels for gaussian processes, Advances in neural information processing systems, 20, 2007.
A. E. Hoerl and R. W. Kennard. Ridge Regression: Biased Estimation for Nonorthogonal Problems. Technometrics, 12(1):55 67, Feb. 1970. ISSN 0040-1706. doi:
K. Hornik, M. Stincheombe, and H. White. Multilayer feedforward networks are universal approxi- mators. Neural networks, 2(5):359-366, 1989.
O. Irsoy, O. T. Yildiz, and E. Alpaydin. Soft decision trees. Proceedings of the 21 International Conference on Pattern Recognition (ICPR2012), 2012.
A. Jacot, B. Simsek, F. Spadaro, C. Hongler, and F. Gabriel. Implicit regularization of random feature models. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 4631 16400. PMLR, 13-18 18 Jul 2020. URL
mmlr. (ttsss//v11/jjccottoea html
A. Karthikeyan, N. Jain, N. Natarajan, and P. Jain. Learning accurate decision trees with bandit feedback via quantized gradient descent. arXiv preprint arXiv:2102.07567, 2021.
M. L. Kile, H. G. Rodrigues, M. Mazumdar, C. B. Dobson, N. Diao, M. Golam, Q. Quamruzzaman, M. Rahman, and D. C. Christiani. A prospective cohort study of the association between drinking water arsenic exposure and self-reported maternal health symptoms during pregnancy in bangladesh. Environmental Health, 13(1):1-13, 2014.
R. Kohavi. Scaling Up thea aarprccciiiiccyy-ooooreeaatlllao
Classi ers: page 6.
J. Lee, Y. Bahri, R. Novak, S. S. Schoenholz, J. Pennington, and J. Sohl-Dickstein. Deep neural networks as gaussian processes. arXiv preprint arXiv:/711.00165,, 2017.
I. Lemhadri, F. Ruan, and R. Tibshirani. Lassonet: Neural networks with feature sparsity. In International Conference on Artificial Intelligence and Statistics, pages 10-18. PMLR, 2021.
F. Liu, X. Huang, Y. Chen, and A. Suykens. Random features for kernel approximation: A survey on algorithms, theory, and beyond. IEEE Transactions on Pattern Analysis Machine Intelligence, (01):1-1, 2021.
J. Liu. Variable selection with rigorous uncertainty quantification using deep bayesian neural networks: Posterior concentration and bernstein-von mises phenomenon. In International Conference on Artificial Intelligence and Statistics, pages 3122-3132. PMLR, 2021.
J. Liu, Z. Lin, S. Padhy, D. Tran, T. Bedrax Weiss, and B. Lakshminarayanan. Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. Advances in Neural Information Processing Systems, 33:7498-7512, 2020.
J. 7.. Liu. Gaussian Process Regression and Classification under Mathematical Constraints with Learning Guarantees. /rXiv:/904.0000 09632000033022000 [cs, math, stat], Apr. 2019. URL http:://aaxxv...rrg abs abs/1904.09632
Y. Lu. On the bernstein von mises theorem for high dimensional nonlinear bayesian inverse problems. arXiv preprint /706. 2017.
Y. Lu, Y. Fan, J. Lv, and W. S. Noble. DeepPINK: reproducible feature selection in deep neural networks, arXiv: 01 185 [cs, stat]. Sept. 2018. URL http:http///aet hwr///b//1009/// 01185
13

---

J. Mairal and B. Yu. Complexity analysis of the lasso regularization path. arXiv preprint arXiv: 0079, 2012.
S. Mei and A. Montanari. The generalization enror of random features regression: Precise asymptotics and the double descent curve. arXiv: Statistics Theory, 2019.
C. A. Micchelli, Y. Xu, and H. Zhang. Universal kernels. Journal of Machine Learning Research, 7 (12), 2006.
R. M. Neal. Bayesian Learning for Neural Networks. Springer-Verlag, Berlin, Heidelberg, 1996. ISBN 978-0-387-94724--
S. W. Ober and C. E. Rasmussen. Benchmarking the neural linear model for regression. In 2nd Symposium on Advances in Approximate Bayesian Inference, 2019.
T. Park and G. Casella. The Bayesian Lasso. Journal of the American Statistical Association, 103(482):681-668 June 2008. ISSN 0162-1459. doi: 10.119800122 URL httts:://doi.org/10 http:///.....///00/1100000000000
N. Polson and V. Rockova. Posterior Concentration for Spars Deep Learning. arXiv: 1rX.vv/803.091.33 [cs, star), Mar. 2018.
A. Rahimi and B. Recht. Random features for large-scale kernel machines. In Proceedings of the 20th International Conference on Neural Information Processing Systems, NIPS'07, pages 1177-1184, Red Hook, NY, USA, Dec. 2007. Curran Associates Inc. ISBN 978-1-60550-352-0.
A. Rahimi and B. Recht. Uniform approximation of functions with random bases. In 2008 46th Annual Allerton Conference on Communication, Control, and Computing, pages 555-561. IEEE, 2008
A. Rahimi and B. Recht. Weighted sums of random kitchen sinks: Replacing minimiza- tion with randomization in learning. In D. Koller, D. Schuurmans, Y. Bengio, and L. Bottou, editors, Advances in Neural Information Processing Systems, volume 21. Cur- ran Associates, Inc., 2009, URL htttps://proceeddnggg net..ips..c//ppee////////
C. E. Rasmussen and C. K. I. Williams. Gaussian Processes for Machine Learning (Adaptive Computation and Machine Learning series). The MIT Press, hardcover edition, 11 2005. ISBN 078-0262182539, URL https:///eae: to/ wwo/mmazoo///ooooooooooooa/ooooa horpa///cccoopppptaaaaaeeeeaaaaa 026218253X
V. Rockova. On semi parametric inference for bart. In International Conference on Machine Learning. pages 8137-8146. PMLR, 2020.
V. Ročková and S. van der Pas. Posterior concentration for bayesian regression trees and forests. The Annals of Statistics, 48(4):2108-2133 2131, 2020.
L. Rosasco, S. Villa, S. Mosci, M. Santoro, and A. Verri. Nonparametric sparsity and regularization. The Journal of Machine Learning Research, 14(1):16651771, Jan. 2013. ISSN 1532-4435.
A. Rudi and L. Rosasco. Generalization Properties of Learning with Random Features. SEKY::662..04113 [cs, stat], Jan. 2018 URL http::/aaxiv..rrg/////// http///xxivv..r//////////// 04474
J. Shhmidt-HHebee Nonparametric regression using deep neural networks with ReLU activation function. The Annals of Statistics, 48(4):187511877, Aug. 2020. ISSN 0090-5364, 2168-8966. doi: 10.1214/19-A0S1877
Sigrist. Gaussian process boosting. arXiv preprint arXiv:2004.02655 2020.
J. Snoek, O. Rippel, K. Swersky, R. Kiros, N. Satish, N. Sundaram, M. M. A. Patwary, Prabhat, and R. Adams. Scalable bayesian optimization using deep neural networks. In Proceedings of the 32nd International Conference Machine Learning, 2015.
14

---

R. Tanno, K. Arulkumaran, D. Alexander, A. Criminisi, and A. Nori. Adaptive Neural Trees. In Proceedings of the 36th International Conference on Machine Learning, pages 6166- 6175. PMLR, May 2019. URL https://proceeddngs mlr. wtw.....ss///e.... v97/tanno19a html ISSN: 2640-3498.
Thakur, C. Lorsung, Y. Yacoby, F. Doshi-Velez, and W. Pan. Uncertainty- (una) bases for deep bayesian regression using multi-ha headed auxiliary networks. arXiv preprint wrXt::2006.116930 2021.
R. Tibshirani. Regression Shrinkage and Selection Via the Lasso. Journal of the Royal Statistical Society: Series B (Methodological), 58(1):267-288, 1996. ISSN 2517-6161. doi: https://doi.org. 1111661166.999966600088080
L. Valeri, M. M. Mazumdar, J. F. Bobb, B. Claus Henn, E. Rodrigues, O. I. Sharif, M. L. Kile,
Q. Quamruzzaman, S. Afroz, M. Golam, al. The joint effect of prenatal exposure to metal mixtures on neurodevelopmental outcomes at 20-40 months of age: evidence from rural bangladesh. Environmental health perspectives, 125(6))067015, 2017.
A. van der Vaart and H. Zanten. Information Rates of Nonparametric Gaussian Process Methods. Journal of Machine Learning Research, 12:2095-2119, June 2011.
A. Virmaux and K. Scaman. Lipschitz regularity of deep neural networks: analysis and efficient estimation. Advances in Neural Information Processing Systems, 31, 2018.
L. Wang, L.. Xue, A. Qu, and H. Liang. Estimation and model selection in generalized additive partial linear models for correlated data with diverging number of covariates. Annals of Statistics, 42(2): 692-624, 2014.
Y. Wang and V. Rocková. Uncertainty quantification for sparse deep learning. In International Conference on Artificial Intelligence and Statistics, pages 298-308. PMIR, 2020.
A. G. Wilson, Z. Hu, R. Salakhutdinov, and E. P. Xing. Deep kernel learning. Artificial intelligence and statistics, pages 370-378. PMLR, 2016a.
A. G. Wilson, Z. Hu, R. R. Salakhutdinov, and E. P. Xing. Stochastic variational deep kernel learning. Advances in Neural Information Processing Systems, 29, 2016b.
D. Wipf and S. Nagarajan. A new view of automatic relevance determination. Advances in neural information processing systems, 20, 2007.
D. H. Wolpert and W. Macready. No free lunch theorems for optimization. IEEE transactions on evolutionary computation, 1(1):67-82, 1997.
N. Wood. Generalized additive models: introduction with R. chapman and hall/CRC, 2006.
Y. Yang, G. Cheng, and D. B. Dunson. Semiparametric bernstein-von mises theorem: Second order studies. arXiv preprint
2015.
Y. Yang, I. G. Morillo, and T. M. Hospedales. Deep Neural Decision Trees. arXiv: 88880006888 [cs, stat], June 2018. URL
06988 arXiv: 1806.06988.
L. Yao. Z. Chu, S. Li. Y. Li, Gao, and A. Zhang. A survey on causal inference. ACM Transactions on Knowledge Discovery from Data (TKDD), 15(5):1-46, 2021.
15

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes] See Contributions paragraph Section
|
(b) Did you describe the limitations of your work? [Yes] See Limitations paragraph in
Section 5
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See Societal Impacts paragraph in Section 5
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See Section
2
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix C and
Appendix D
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as URL)? [Yes] Instructions for downloading and processing data in the Appendix G.4 Code repository is in
preparation at GitHub and will be released in camera-ready version.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Appendix G
(c) Did you report error bars (c.g., with respect to the random seed after running experi- ments multiple times)? [Yes] See Appendix G
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Appendix G.4
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you eitet the creators? [Yes] See Appendix G.2
(b) Did you mention the license of the assets? [Yes] See Appendix G.2
(c) Did you include any new assets either the supplemental material or as URL? Yes]
See Appendix G.2
(d) Did you discuss whether and how consent was obtained from people whose data you're re using/curating? [N/A] All data involved in this work are already previously open-sourced. They all can be downloaded via external URL linked in Appendix G.2
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [Yes] See Appendix 2
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
16


Reviews:
[{"review_id": "nzEdsXrM2YW", "reviewer": "Reviewer_smu5", "summary": "This paper considers a method of evaluating the importance of variables of a (possibly) non-convex function f under a Bayesian perspective. Given a Gaussian process over a function set (including true function f) as a prior distribution, the evaluation measure for variable j is the expectation of magnitude of the partial gradient of f w.r.t. j, where the expectation is taken w.r.t. some distribution from which instances are drawn. The paper proves convergence results that the estimate of the measure converges to the true value when the sample size increases up to infinity.", "questions": "- Do you have any assumption on the distribution from which instances are drawn?\n\n-Do you assume that the function is fully differentiable?\n\n- I don't understand how to compute AUCROC. Please explain.", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n- Consistent approach from a Bayesian perspective\n\n\nWeakness:\n-  Theoretical results would be more convincing if finite sample analyses are possible.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "S91lBUuQRU7", "reviewer": "Reviewer_AEuT", "summary": "The importance of a variable in this paper is derived from the norm the corresponding partial derivative.\nThe authors use a Bayesian linear regression to capture the uncertainty of the model parameters.\nBecause of the form of feature represented method, the proposed method is independent to the model and thus can be applied to tree-based structures. ", "questions": "- Even though the title uses the word \"nonlinear\", it seems that the variable selection part is still based on linear structure. If so, how could the model be extended to more complicated cases? For example, a multilayer neural network. \n- The above question also reflected in the current experiments. Except for the synthetic experiments, the results on the real dataset is not very consistent. Why does the performance of RFNN is not good? \n- Why does the points of GAM not show in the Figure 1? \n\n", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n- The paper is well structured and written. \n- The uncertainty is counted into the calculation of variable selection. \n- Use the feature-based representation of Gaussian process to accelerate training and generalization on different machine learning models.\n- The authors also added theoretical proofs to show the convergence and efficiency. \n\nWeakness:\n- The experimental results does not fully reflect the profoundness of the theory. \n- The variable selection still relies on linear structure and could be hard to extend to more complex scenarios. \n- It will be better if the legends of Figure 1 are enlarged.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "O5eCvRo9Jw", "reviewer": "Reviewer_adBr", "summary": "The paper develops a unified framework for uncertainty-aware nonlinear variable selection for machine learning models that can be represented by RKHS including linear models, neural networks, random forests, and tree ensembles. Two major theoretical results about the approach are derived, namely, posterior convergence of variable Importance and a Bernstein-von Mises Theorem verifying that posterior distribution of variance importance is statistically optimal.\n\nThe authors then apply and validate the approach on several synthetic and real datasets, and conclude that:\n\n- there does not exist a model class that performs universally well across all data scenarios, highlighting the necessity of a unified framework\n- good prediction translates to effective variable selection", "questions": "n/a", "limitations": "None noted", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n- Significance:\n    - The paper considers an important problem in machine learning: uncertainty-aware nonlinear variable selections across multiple classes of models. The multiple-model-class aspect is particularly of great interest to the machine learning community since in many practical applications, it is often unknown a priori which specific class of model is more appropriate for the data.\n- Quality and Clarity:\n    - In general, the paper is well-written and well-organized. Discussions about pre-existing literature and related works are sufficient (with some caveats detailed below).\n    - The theoretical analyses of the work are rigorous and of high quality.\n    - One of the stronger points of the paper is that it provides a very careful treatment of the non-trivial gradient function assumption and studies in detail the complications arising when the features are discrete or the model is non-differentiable. These issues are often neglected in the field but are unavoidable in studying a unified framework, and the paper’s efforts in this direction should be appreciated.\n\n### Weaknesses\n\n- Literature: One concern I have with the paper is its lack of direct discussions/comparisons with the work\n\n     > J. Liu (2019). Variable selection with rigorous uncertainty quantification using deep Bayesian neural networks: Posterior concentration and Bernstein-von Mises phenomenon.\n\n   cited as [45] in the manuscript. [45] is mentioned in the paper with minor references (correlated features, simpler conditions), but I believe it’s hard to ignore the similarities between the two works (the same measure of feature importance, similar approaches, emphasis and questions, same theoretical results, even the titles are somewhat similar). \n\n  More importantly, the two papers lead to two contrasting conclusions:\n    - the paper observes that good prediction translates to effective variable selection\n    - [45] hinted that a powerful model alone is not sufficient to guarantee effective variable selection\n\n    Given the connections (and the fact that the authors are aware of the work by Liu), a thorough comparison in both theoretical analyses and (perhaps) experiments would be more appropriate.\n- Quality:\n  - The paper conflates the two concepts of variable importance and variable selection, and this may lead to an incorrect interpretation of the results.\n      - Mathematically, posterior concentration of variance important and Bernstein-von Mises phenomenon do not translate to efficient/consistent variable selection. For example, for standard settings of linear models, ordinary least square satisfies both of the above properties but is not selection consistent/efficient, thus necessitating the use of lasso-based methods. Thus Theorem 1 and 2 do not support the claim that “good prediction translates to effective variable selection” as the authors stated in Line 280.\n      - In practice, to quantify the efficiency of variable selection procedures, a good measure of variable importance and a rigorous decision rule is equally important, and AUROC (which integrates over threshold) is not an appropriate measure in these comparisons. This might be the reason why the conclusions of the paper and [45] (which uses F1 scores on three different modes of thresholding) are pointing in two different directions.\n  - Since the paper attempted to provide a framework that works across models, they seem to avoid directly stating the assumptions of the analyses. For example, the feature map could be differentiable, non-differentiable, or even discontinuous and there seem to be separate treatments for each case. However, this leads to a very difficult reading of the work, and it seems impossible to verify from the main text whether a general model would fit into the framework. I suggest adding a list of assumptions on the feature map, and the parameter spaces of the models in the mathematical descriptions.\n  - The decisions of evaluating the methods on tabular data only, and thus remove neural network models from the comparisons are contradictory to the goal of providing a unified framework of variable selection. Some additional evaluations, even on a synthetic dataset, perhaps in comparison with Liu (2019) would strengthen the manuscript.\n- Significance\n    - The paper assumes that the number of significant features (d) is fixed while data (n) increases. This setting is significantly more restrictive than typical results about posterior concentrations in the field for both generalized linear models and neural network models.\n    - In essence, the models considered in this work need to be compatible with an RHKS and have to be linear in terms of the parameters. As pointed out by the authors, this places a significant restriction on the application of the framework. For example, for neural network models, this means all layers of the network except the last (linear) one have to be fixed, and the same limit applies for tree-based models. I believe this should be acknowledged more directly in the title, the abstract, and the introduction.\n\n\n#####Post-rebuttal comment####\n\nThe authors have addressed all of my comments with their extensive efforts with the revision.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "A2-PtGC3F1l", "reviewer": "Reviewer_2frY", "summary": "This paper proposes an uniﬁed framework for nonlinear variable selection that incorporates uncertainty in the prediction function and is compatible with a wide range of machine learning models (e.g., tree ensembles, kernel methods, neural networks, etc).  The authors (1) provide a principled approach for quantifying variable selection uncertainty by deriving its posterior distribution, and (2) show that the approach is generalizable even to non-differentiable models. Rigorous Bayesian nonparametric theorems are derived to guarantee the posterior consistency and asymptotic uncertainty of the proposed approach. Extensive simulations and experiments verify the effectiveness of the proposed method.", "questions": "+  In line 76, the data-generating function f_0 is assumed as a ﬂexible nonlinear function that resides in a reproducing kernel Hilbert space (RKHS).  How restrictive is this assumption? in other words, does this assumption cover most real data generation functions.\n+  In Theorem 1, what is the specific convergence rate of $\\epsilon$ mentioned in condition (1) (line 203).\n+  He et al. [1] also employ the partial derivative-based kernel method to realize variable selection. It would be appreciated if the authors could explain the differences between theirs and the advantages of this paper.\n\nRef:\n[1] He et al. SCALABLE KERNEL-BASED VARIABLE SELECTION WITH SPARSISTENCY.", "limitations": "All my concerns are presented in “Weakness” and “Questions”. Moreover, I do not have any concern on negative social impact.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n+  The paper is well-written.\n\n+ The proposed method is supported by theory and experiments.\n\nWeakness:\n\nThe numerical experiments can be improved, e.g., considering larger input dimensions in the experiment.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept