Research map / Label-efficient and robust learning / Recognition and long-tail
Distributionally Robust Feature Selection
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 a noise-based continuous relaxation of feature selection that optimizes the worst-case (over populations) Bayes-optimal MSE via a kernel-smoothed objective, selecting features with the smallest learned noise variances without differentiating through model training.
Problem
Selecting a limited budget of features so that downstream models trained on them perform well simultaneously across multiple subpopulations, with the downstream model unknown in advance.
Approach
Relax binary feature masks to per-feature Gaussian noise variances alpha, and optimize the worst-population Bayes-optimal MSE (conditional variance of E[Y|X] given noisy S) using a closed-form Gaussian-kernel-weighted estimator, so no model is retrained during optimization.
Key results
- On ACS, the proposed method achieves roughly an order-of-magnitude lower MSE than all baselines (e.g., DRO Lasso MSE around 0.23-0.32 vs. the proposed method's lower values reported in the main table).
- On synthetic dataset 1 with budget 5, the method achieves the lowest MLP MSE on population A (0.2251 vs. 0.5071 for DRO Lasso) and balanced performance across populations, while vanilla Lasso and XGB are strong only on population C.
- On synthetic dataset 2 (nonlinear, budget 8), the method is comparable to the best DRO XGB and XGB baselines, while DRO Lasso performs worst.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use distributionally robust feature selection instead of pooled Lasso or XGBoost | Synthetic dataset 1, budget 5 · Population A | MLP MSE: 1.0087 ± 0.0153 -> 0.2251 ± 0.0029; 1.0136 ± 0.0076 -> 0.2251 ± 0.0029 | Helped |
| Increase the selected-feature budget from 5 to 10 | Our Method · Synthetic dataset 1, population A | MLP MSE: 0.2251 ± 0.0029 -> 0.0342 ± 0.0004 | Helped |
| Use distributionally robust feature selection instead of DRO Lasso | Synthetic dataset 2 · Population A | MLP MSE: 0.0054 ± 0.0014 -> 0.0054 ± 0.0007 | No clear effect |
| Use distributionally robust feature selection instead of DRO Lasso | Synthetic dataset 2 · Population B | MLP MSE: 0.4172 ± 0.0533 -> 0.0062 ± 0.0010 | Helped |
| Use distributionally robust feature selection instead of DRO XGB | Synthetic dataset 2 · Population C | MLP MSE: 0.0226 ± 0.0005 -> 0.0123 ± 0.0080 | Helped |
| Use distributionally robust feature selection instead of DRO XGB | Synthetic dataset 2 · Population C | RF MSE: 0.2460 ± 0.0479 -> 0.2460 ± 0.0483 | No clear effect |
| Use the proposed feature selection instead of baseline feature selection | ACS household-income prediction · California, Florida, and New York | MSE | Helped |
| Use the proposed feature selection instead of XGBoost feature selection | UCI Adult Income classification · Female population | Log Loss | Hurt |
| Use the proposed feature selection instead of baseline feature selection | UCI Adult Income classification · Female and Male populations | accuracy | Helped |
Surprising findings
- Vanilla Lasso and XGBoost are near-best on one population while performing poorly on others, showing that pooled selection favors majority-aligned features.
- The Embedded MLP baseline performs poorly on nonlinear data, imbalanced across populations.
Limitations
- Plug-in estimator of E[Y|X] may be biased, especially for small minority populations.
- Theory is derived only for MSE; extension to other losses such as cross-entropy is not established theoretically.
Experimental setup
Datasets: synthetic linear dataset (3 populations, 15 features) synthetic nonlinear dataset (4 populations, 50 features) UCI Adult Income American Community Survey 2018 PUMS (CA, FL, NY)
Benchmarks: synthetic dataset 1 synthetic dataset 2 UCI Adult Income ACS 2018 PUMS
Metrics: MSE R2 score accuracy log loss
Baselines: Lasso XGBoost DRO Lasso DRO XGBoost Embedded MLP
Builds on: Lasso XGBoost group DRO MAML
Closest papers by meaning
- Estimating Conditional Mutual Information for Dynamic Feature Selection (ICLR 2024 · similarity 0.86)
- Towards a statistical theory of data selection under weak supervision (ICLR 2024 · similarity 0.85)
- Differentially Private High-dimensional Variable Selection via Integer Programming (NeurIPS 2025 · similarity 0.84)
- Most Influential Subset Selection: Challenges, Promises, and Beyond (NeurIPS 2024 · similarity 0.84)
- Learning With Multi-Group Guarantees For Clusterable Subpopulations (ICML 2025 · similarity 0.83)
- The Impact of Coreset Selection on Spurious Correlations and Group Robustness (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.