Contextual Feature Selection with Conditional Stochastic Gates

TL;DR

This paper introduces Conditional Stochastic Gates (c-STG), a novel method for context-dependent feature selection that adapts feature importance based on contextual variables, improving flexibility and interpretability.

Contribution

The paper proposes a new architecture for contextual feature selection using conditional Bernoulli variables and a hypernetwork, with theoretical analysis and extensive benchmarking.

Findings

c-STG improves feature selection accuracy

Enhances prediction performance over population-level methods

Increases interpretability of feature relevance in context

Abstract

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.