Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Distributionally Robust Feature Selection

Maitreyi Swaroop, Tamar Krishnamurti, Bryan Wilder

TL;DR

This work addresses robust feature selection under distribution shifts by casting feature subset selection as a distributionally robust problem and solving it via a noise-based continuous relaxation. By injecting feature-wise noise and optimizing the Bayes-optimal predictor, the method yields a model-agnostic, differentiable objective that does not require differentiating through downstream model training. Empirical results on synthetic and real datasets show substantial improvements over pooled baselines and existing DRO approaches, with strong cross-population performance and reduced variance. The approach offers practical benefits for cost-constrained deployments where a fixed, robust feature set must support diverse downstream models across subpopulations.

Abstract

We study the problem of selecting limited features to observe such that models trained on them can perform well simultaneously across multiple subpopulations. This problem has applications in settings where collecting each feature is costly, e.g. requiring adding survey questions or physical sensors, and we must be able to use the selected features to create high-quality downstream models for different populations. Our method frames the problem as a continuous relaxation of traditional variable selection using a noising mechanism, without requiring backpropagation through model training processes. By optimizing over the variance of a Bayes-optimal predictor, we develop a model-agnostic framework that balances overall performance of downstream prediction across populations. We validate our approach through experiments on both synthetic datasets and real-world data.

Distributionally Robust Feature Selection

TL;DR

This work addresses robust feature selection under distribution shifts by casting feature subset selection as a distributionally robust problem and solving it via a noise-based continuous relaxation. By injecting feature-wise noise and optimizing the Bayes-optimal predictor, the method yields a model-agnostic, differentiable objective that does not require differentiating through downstream model training. Empirical results on synthetic and real datasets show substantial improvements over pooled baselines and existing DRO approaches, with strong cross-population performance and reduced variance. The approach offers practical benefits for cost-constrained deployments where a fixed, robust feature set must support diverse downstream models across subpopulations.

Abstract

We study the problem of selecting limited features to observe such that models trained on them can perform well simultaneously across multiple subpopulations. This problem has applications in settings where collecting each feature is costly, e.g. requiring adding survey questions or physical sensors, and we must be able to use the selected features to create high-quality downstream models for different populations. Our method frames the problem as a continuous relaxation of traditional variable selection using a noising mechanism, without requiring backpropagation through model training processes. By optimizing over the variance of a Bayes-optimal predictor, we develop a model-agnostic framework that balances overall performance of downstream prediction across populations. We validate our approach through experiments on both synthetic datasets and real-world data.
Paper Structure (51 sections, 2 theorems, 27 equations, 6 figures, 6 tables, 3 algorithms)

This paper contains 51 sections, 2 theorems, 27 equations, 6 figures, 6 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Related work
  3. Feature selection:
  4. Distributionally Robust Optimization (DRO):
  5. Problem formulation
  6. Methods
  7. Continuous relaxation
  8. Solving the relaxation
  9. Population-level objective
  10. Empirical estimation and kernel form of objective
  11. Algorithmic approach
  12. Experiments
  13. Data processing
  14. Evaluation Metrics
  15. Synthetic Data Experiments
  16. ...and 36 more sections

Key Result

Theorem 1

Under the noise-based relaxation $S(\bm{\alpha}) = X + \epsilon(\bm{\alpha})$, $\epsilon(\bm{\alpha}) \sim \mathcal{N}(0, \operatorname{diag}(\bm{\alpha}))$, the distributionally robust feature selection problem is equivalent to where $\mu_i(X) = \mathbb{E}_{P_i}[Y \mid X]$.

Figures (6)

  • Figure 1: Performance comparison across populations on synthetic dataset 1 using different feature selection methods. In each subplot, left: Mean Squared Error of downstream MLP model; right: Mean Squared Error of downstream random forest model. The relative performance of feature selection methods is consistent across the choice of downstream prediction model.
  • Figure 2: Performance comparison across populations on synthetic dataset 2. Left: Mean Squared Error of downstream MLP model; right: Mean Squared Error of downstream random forest (RF) model. Our method consistently achieves low error across populations and budgets, outperforming or matching DRO-based baselines.
  • Figure 3: Performance comparison across populations on real datasets using different feature selection methods. In each subplot, left: Mean Squared Error; right: $R^2$ Score.
  • Figure 4: Performance comparison across populations on real datasets using different feature selection methods.
  • Figure 5: Performance comparison across populations on synthetic dataset 1 using different feature selection methods. In each subplot, left: Mean Squared Error of downstream MLP model; right: Mean Squared Error of downstream random forest model. The relative performance of feature selection methods is consistent across the choice of downstream prediction model, with vanilla XGBoost consistently having the worst performance. Here the total number of features is set to $50$.
  • ...and 1 more figures

Theorems & Definitions (2)

  • Theorem 1: Population-Level Objective
  • Theorem 2: Kernel Form Equivalence