Research map / Probabilistic and scientific modeling / Molecules and proteins
Contextual Feature Selection with Conditional Stochastic Gates
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Conditional Stochastic Gates (c-STG), a hypernetwork-driven embedded feature selection method whose Bernoulli gate probabilities depend on contextual variables, plus a weighted variant, and evaluates on synthetic, MNIST-rotation, heart disease, housing, and neuroscience data.
Problem
Traditional supervised feature selection identifies one global feature subset for the whole population, but feature relevance often varies with context variables that do not themselves determine the outcome.
Approach
Model each feature's selection as a conditional Bernoulli gate whose parameters are produced by a hypernetwork from context variables, relaxed with a truncated Gaussian (STG-style) and trained end-to-end with a prediction network; the weighted variant adds a context-dependent weight vector on selected features.
Key results
- Weighted c-STG reaches 100% accuracy on XOR1 and R^2 0.9956 on XOR2, versus 0.7572 for the best competing method (AFS) on XOR2.
- On MNIST (4 vs 9 with rotations), weighted c-STG achieves 98.69% accuracy with about 217 selected features, versus 98.55% for STG and 97.17% for Contextual LASSO.
- On heart disease, weighted c-STG reaches 89.23% 5-fold accuracy with 6.5 selected features, and on housing it reaches R^2 0.5308 versus 0.4616 for Contextual LASSO.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use c-STG instead of AFS | XOR1 · XOR1 | accuracy: 78.95 (19.91) -> 100 (0) | Helped |
| Use weighted c-STG instead of AFS | XOR1 · XOR1 | accuracy: 78.95 (19.91) -> 100 (0) | Helped |
| Use c-STG instead of AFS | XOR2 · XOR2 | R² score: 0.7572 (0.0187) -> 0.8739 (0.0107) | Helped |
| Use weighted c-STG instead of AFS | XOR2 · XOR2 | R² score: 0.7572 (0.0187) -> 0.9956 (0.0008) | Helped |
| Use c-STG instead of STG | MNIST · MNIST | accuracy: 98.55 (0.08) -> 98.66 (0.05) | No clear effect |
| Use weighted c-STG instead of STG | MNIST · MNIST | accuracy: 98.55 (0.08) -> 98.69 (0.06) | No clear effect |
| Use c-STG instead of STG with context | Heart disease · Heart disease | accuracy: 87.88 (1.95) -> 87.55 (3.41) | No clear effect |
| Use weighted c-STG instead of STG with context | Heart disease · Heart disease | accuracy: 87.88 (1.95) -> 89.23 (1.72) | Helped |
| Use c-STG instead of Contextual LASSO | Housing · Housing | R² score: 0.4616 (0.0033) -> 0.3976 (0.0082) | No clear effect |
| Use weighted c-STG instead of Contextual LASSO | Housing · Housing | R² score: 0.4616 (0.0033) -> 0.5308 (0.0052) | Helped |
Surprising findings
- Plain STG with context concatenation sometimes performs worse than STG without context on several datasets, and the prediction model with context is not consistently better than without it.
- A single c-STG model with fewer parameters (218,834) outperforms training separate STG models per categorical context (876,688 parameters) on MNIST.
Limitations
- Hyperparameter tuning (regularization and learning rate) is challenging and requires cross-validation, unlike LASSO's established tuning theory.
- Theoretical results on linear regression are limited; the risk comparisons rely on relaxed assumptions and the Bernoulli equivalence proof defers to external work.
- Some reported gains have overlapping standard deviations (e.g., heart disease accuracy of c-STG vs STG).
Experimental setup
Datasets: XOR1 XOR2 XOR3 XOR4 MNIST (4 vs 9, rotated) UCI Heart disease Housing Neuroscience motor cortex recordings
Benchmarks: XOR1 XOR2 MNIST Heart disease Housing
Metrics: accuracy R^2 score number of selected features
Baselines: LASSO STG Contextual LASSO CEN INVASE AFS Prediction model
Builds on: STG (stochastic gates) INVASE Contextual LASSO Concrete / Gumbel-softmax relaxation
Closest papers by meaning
- SAND: One-Shot Feature Selection with Additive Noise Distortion (ICML 2025 · similarity 0.87)
- Estimating Conditional Mutual Information for Dynamic Feature Selection (ICLR 2024 · similarity 0.87)
- Local Feature Selection without Label or Feature Leakage for Interpretable Machine Learning Predictions (ICML 2024 · similarity 0.85)
- Hypernetworks for Dynamic Feature Selection (NeurIPS 2026 · similarity 0.84)
- Contextual Active Model Selection (NeurIPS 2024 · similarity 0.83)
- Distributionally Robust Feature Selection (NeurIPS 2025 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.