# Towards a Unified Framework for Uncertainty-aware Nonlinear Variable Selection with Theoretical Guarantees

arXiv: https://arxiv.org/abs/2204.07293 · HTML: https://mlsanity.com/paper/2204.07293/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 98 references that are in the map.
- First posted on arXiv 15 Apr 2022.

## Builds on (most cited first, 30 of 36)

- [Deep Neural Networks as Gaussian Processes](https://mlsanity.com/paper/1711.00165/) (2017)
- [Deep Kernel Learning](https://mlsanity.com/paper/1511.02222/) (2015)
- [Nonparametric regression using deep neural networks with ReLU activation function](https://mlsanity.com/paper/1708.06633/) (2017)
- [Lipschitz regularity of deep neural networks: analysis and efficient estimation](https://mlsanity.com/paper/1805.10965/) (2018)
- [The generalization error of random features regression: Precise asymptotics and double descent curve](https://mlsanity.com/paper/1908.05355/) (2019)
- [Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness](https://mlsanity.com/paper/2006.10108/) (2020)
- [Breaking the Curse of Dimensionality with Convex Neural Networks](https://mlsanity.com/paper/1412.8690/) (2014)
- [Gaussian Process Behaviour in Wide Deep Neural Networks](https://mlsanity.com/paper/1804.11271/) (2018)
- [Scalable Bayesian Optimization Using Deep Neural Networks](https://mlsanity.com/paper/1502.05700/) (2015)
- [Distilling a Neural Network Into a Soft Decision Tree](https://mlsanity.com/paper/1711.09784/) (2017)
- [A Survey on Causal Inference](https://mlsanity.com/paper/2002.02770/) (2020)
- [Generalization Properties of Learning with Random Features](https://mlsanity.com/paper/1602.04474/) (2016)
- [Panning for Gold: Model-X Knockoffs for High-dimensional Controlled Variable Selection](https://mlsanity.com/paper/1610.02351/) (2016)
- [Random Features for Kernel Approximation: A Survey on Algorithms, Theory, and Beyond](https://mlsanity.com/paper/2004.11154/) (2020)
- [Stochastic Variational Deep Kernel Learning](https://mlsanity.com/paper/1611.00336/) (2016)
- [Manifold Gaussian Processes for Regression](https://mlsanity.com/paper/1402.5876/) (2014)
- [Deep Neural Decision Trees](https://mlsanity.com/paper/1806.06988/) (2018)
- [Triple descent and the two kinds of overfitting: Where & why do they appear?](https://mlsanity.com/paper/2006.03509/) (2020)
- [Adaptive Neural Trees](https://mlsanity.com/paper/1807.06699/) (2018)
- [Implicit Regularization of Random Feature Models](https://mlsanity.com/paper/2002.08404/) (2020)
- [Posterior Concentration for Sparse Deep Learning](https://mlsanity.com/paper/1803.09138/) (2018)
- [DeepPINK: reproducible feature selection in deep neural networks](https://mlsanity.com/paper/1809.01185/) (2018)
- [Sparse-Input Neural Networks for High-dimensional Nonparametric Regression and Classification](https://mlsanity.com/paper/1711.07592/) (2017)
- [Benchmarking the Neural Linear Model for Regression](https://mlsanity.com/paper/1912.08416/) (2019)
- [Gaussian Process Boosting](https://mlsanity.com/paper/2004.02653/) (2020)
- [Neural Random Forests](https://mlsanity.com/paper/1604.07143/) (2016)
- [The Random Forest Kernel and other kernels for big data from random partitions](https://mlsanity.com/paper/1402.4293/) (2014)
- [Learning Accurate Decision Trees with Bandit Feedback via Quantized Gradient Descent](https://mlsanity.com/paper/2102.07567/) (2021)
- [Uncertainty Quantification for Sparse Deep Learning](https://mlsanity.com/paper/2002.11815/) (2020)
- [Trajectories, bifurcations and pseudotime in large clinical datasets: applications to myocardial infarction and diabetes data](https://mlsanity.com/paper/2007.03788/) (2020)

## Built on by (most cited first, 4 of 4)

- [Statistical Advantages of Oblique Randomized Decision Trees and Forests](https://mlsanity.com/paper/2407.02458/) (2024)
- [TrIM: Transformed Iterative Mondrian Forests for Gradient-based Dimension Reduction and High-Dimensional Regression](https://mlsanity.com/paper/2407.09964/) (2024)
- [FALCON: Fine-grained Activation Manipulation by Contrastive Orthogonal Unalignment for Large Language Model](https://mlsanity.com/paper/2502.01472/) (2025)
- [Calibrated Derivative-Process Sensitivity for Gaussian-Process Variable Selection](https://mlsanity.com/paper/2609.33549/) (2026)

## Abstract (arXiv)

We develop a simple and unified 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). In particular, for a learned nonlinear model f(x), we consider quantifying the importance of an input variable xj using the integrated partial derivative Psij = Vert fracpartialpartial xj f(x)Vert2PX. We then (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 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 classic and recent variable selection methods.