# Contextual Feature Selection with Conditional Stochastic Gates

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 54 references that are in the map.
- First posted on arXiv 21 Dec 2023.

## Builds on (most cited first, 17 of 17)

- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Learning Sparse Neural Networks through L0 Regularization](https://mlsanity.com/paper/1712.01312/) (2017)
- [Learning to Explain: An Information-Theoretic Perspective on Model Interpretation](https://mlsanity.com/paper/1802.07814/) (2018)
- [Implicit Reparameterization Gradients](https://mlsanity.com/paper/1805.08498/) (2018)
- [Learning to Maximize Mutual Information for Dynamic Feature Selection](https://mlsanity.com/paper/2301.00557/) (2023)
- [Locally Sparse Neural Networks for Tabular Biomedical Data](https://mlsanity.com/paper/2106.06468/) (2021)
- [Reducing Reparameterization Gradient Variance](https://mlsanity.com/paper/1705.07880/) (2017)
- [Contextual Explanation Networks](https://mlsanity.com/paper/1705.10301/) (2017)
- [Kernel Feature Selection via Conditional Covariance Minimization](https://mlsanity.com/paper/1707.01164/) (2017)
- [Probabilistic Best Subset Selection via Gradient-Based Optimization](https://mlsanity.com/paper/2006.06448/) (2020)
- [Deep Unsupervised Feature Selection by Discarding Nuisance and Correlated Features](https://mlsanity.com/paper/2110.05306/) (2021)
- [Support Recovery with Stochastic Gates: Theory and Application for Linear Models](https://mlsanity.com/paper/2110.15960/) (2021)
- [The Trimmed Lasso: Sparsity and Robustness](https://mlsanity.com/paper/1708.04527/) (2017)
- [Probabilistic Robust Autoencoders for Outlier Detection](https://mlsanity.com/paper/2110.00494/) (2021)
- [DiSC: Differential Spectral Clustering of Features](https://mlsanity.com/paper/2211.05314/) (2022)
- [Contextualized Machine Learning](https://mlsanity.com/paper/2310.11340/) (2023)

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

- [Hybrid Autoencoders for Tabular Data: Leveraging Model-Based Augmentation in Low-Label Settings](https://mlsanity.com/paper/2511.06961/) (2025)
- [Dynamic Feature Selection from Variable Feature Sets Using Features of Features](https://mlsanity.com/paper/2503.09181/) (2025)
- [Conditional Deep Canonical Time Warping](https://mlsanity.com/paper/2412.18234/) (2024)
- [Hypernetworks for Dynamic Feature Selection](https://mlsanity.com/paper/2605.12278/) (2026)
- [Learning Permutation from Structure Without Supervision](https://mlsanity.com/paper/2605.25551/) (2026)
- [Fourier Geometric Wind Power Forecasting with Numerical Weather Prediction](https://mlsanity.com/paper/2607.17095/) (2026)
- [Measure Less, Know More: Self-Supervised Test-Time Feature Acquisition](https://mlsanity.com/paper/2610.03454/) (2026)

## Abstract (arXiv)

Feature selection is a crucial tool in machine learning and is widely applied across various scientific disciplines. Traditional supervised methods generally identify a universal set of informative features for the entire population. However, feature relevance often varies with context, while the context itself may not directly affect the outcome variable. Here, we propose a novel architecture for contextual feature selection where the subset of selected features is conditioned on the value of context variables. Our new approach, Conditional Stochastic Gates (c-STG), models the importance of features using conditional Bernoulli variables whose parameters are predicted based on contextual variables. We introduce a hypernetwork that maps context variables to feature selection parameters to learn the context-dependent gates along with a prediction model. We further present a theoretical analysis of our model, indicating that it can improve performance and flexibility over population-level methods in complex feature selection settings. Finally, we conduct an extensive benchmark using simulated and real-world datasets across multiple domains demonstrating that c-STG can lead to improved feature selection capabilities while enhancing prediction accuracy and interpretability.