# Probabilistic Best Subset Selection via Gradient-Based Optimization

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

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 11 Jun 2020.

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

- [Bayesian Optimization over Discrete and Mixed Spaces via Probabilistic Reparameterization](https://mlsanity.com/paper/2210.10199/) (2022)
- [Contextual Dropout: An Efficient Sample-Dependent Dropout Module](https://mlsanity.com/paper/2103.04181/) (2021)
- [Gradient Estimation with Discrete Stein Operators](https://mlsanity.com/paper/2202.09497/) (2022)
- [Optimization-based Causal Estimation from Heterogenous Environments](https://mlsanity.com/paper/2109.11990/) (2021)
- [Differentiable Sampling of Categorical Distributions Using the CatLog-Derivative Trick](https://mlsanity.com/paper/2311.12569/) (2023)
- [ARMS: Antithetic-REINFORCE-Multi-Sample Gradient for Binary Variables](https://mlsanity.com/paper/2105.14141/) (2021)
- [Coupled Gradient Estimators for Discrete Latent Variables](https://mlsanity.com/paper/2106.08056/) (2021)
- [Double Control Variates for Gradient Estimation in Discrete Latent Variable Models](https://mlsanity.com/paper/2111.05300/) (2021)
- [Gradient Estimation for Binary Latent Variables via Gradient Variance Clipping](https://mlsanity.com/paper/2208.06124/) (2022)
- [Contextual Feature Selection with Conditional Stochastic Gates](https://mlsanity.com/paper/2312.14254/) (2023)
- [CARMS: Categorical-Antithetic-REINFORCE Multi-Sample Gradient Estimator](https://mlsanity.com/paper/2110.14002/) (2021)
- [Bayesian Invariance Modeling of Multi-Environment Data](https://mlsanity.com/paper/2506.22675/) (2025)
- [Efficient compression of neural networks and datasets](https://mlsanity.com/paper/2505.17469/) (2025)
- [Environment-Adaptive Covariate Selection: Learning When to Use Spurious Correlations for Out-of-Distribution Prediction](https://mlsanity.com/paper/2601.02322/) (2026)
- [Nonparametric Discrete Choice Experiments with Machine Learning Guided Adaptive Design](https://mlsanity.com/paper/2310.12026/) (2023)
- [What to Do When Your Discrete Optimization Is the Size of a Neural Network?](https://mlsanity.com/paper/2402.10339/) (2024)
- [Probabilistic and nonlinear compressive sensing](https://mlsanity.com/paper/2509.15060/) (2025)

## Abstract (arXiv)

In high-dimensional statistics, variable selection recovers the latent sparse patterns from all possible covariate combinations. This paper proposes a novel optimization method to solve the exact L0-regularized regression problem, which is also known as the best subset selection. We reformulate the optimization problem from a discrete space to a continuous one via probabilistic reparameterization. The new objective function is differentiable but its gradient often cannot be computed in a closed form. Then we propose a family of unbiased gradient estimators to optimize the best subset selection objectives by the stochastic gradient descent. Within this family, we identify the estimator with uniformly minimum variance. Theoretically, we study the general conditions under which the method is guaranteed to converge to the ground truth in expectation. The proposed method can find the true regression model from thousands of covariates in seconds. In a wide variety of synthetic and semi-synthetic data, the proposed method outperforms existing variable selection tools based on the relaxed penalties, coordinate descent, and mixed integer optimization in both sparse pattern recovery and out-of-sample prediction.