Contextual Feature Selection with Conditional Stochastic Gates
Ram Dyuthi Sristi Ofir Lindenbaum Maria Lavzin Jackie Schiller Gal Mishne Hadas Benisty
We study the problem of contextual feature selection, where the goal is to
learn a predictive function while identifying subsets of informative features
conditioned on specific contexts. Towards this goal, we generalize the recently
proposed stochastic gates (STG) Yamada et al. [2020] by modeling the
probabilistic gates as conditional Bernoulli variables whose parameters are
predicted based on the contextual variables. Our new scheme, termed
conditional-STG (c-STG), comprises two networks: a hypernetwork that
establishes the mapping between contextual variables and probabilistic feature
selection parameters and a prediction network that maps the selected feature to
the response variable. Training the two networks simultaneously ensures the
comprehensive incorporation of context and feature selection within a unified
model. We provide a theoretical analysis to examine several properties of the
proposed framework. Importantly, our model leads to improved flexibility and
adaptability of feature selection and, therefore, can better capture the
nuances and variations in the data. We apply c-STG to simulated and real-world
datasets, including healthcare, housing, and neuroscience, and demonstrate that
it effectively selects contextually meaningful features, thereby enhancing
predictive performance and interpretability.