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

TL;DR

This paper introduces a unified Bayesian framework for nonlinear variable selection that quantifies uncertainty and guarantees theoretical consistency across various machine learning models, including non-differentiable ones.

Contribution

It proposes a generalizable approach for uncertainty-aware variable importance measurement with theoretical guarantees, applicable to diverse models like trees, kernels, and neural networks.

Findings

Outperforms existing variable selection methods in healthcare datasets

Provides posterior distributions for variable importance, enabling uncertainty quantification

Guarantees posterior consistency and asymptotic properties through rigorous theorems

Abstract

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(\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 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.