Population-Robust Feature Selection via Generalized Welfare Optimization

2026-08-03

AI Summary

This work addresses the challenge of deploying a unified feature set across diverse, heterogeneous populations, where conventional approaches struggle to balance overall predictive performance with fairness for disadvantaged subgroups. The authors propose PopFS, a method grounded in generalized welfare optimization that enables population-specific modeling under a shared feature representation, thereby reconciling aggregate prediction efficiency with protection for the worst-off groups. PopFS incorporates a tunable welfare function to flexibly trade off fairness and efficiency, and combines multi-task sparse learning, a rank-based add-delete-swap search, and a shortlist full refitting strategy to efficiently navigate large-scale feature spaces. Evaluated across six prediction tasks, eight demographic partitions, and real-time COVID-19 forecasting in 43 U.S. states, PopFS consistently improves both average and worst-group performance while yielding highly interpretable selected features.

Abstract

Choosing which features to collect is a deployment decision: the same limited questionnaire, test panel, or sensor set may need to serve several heterogeneous populations. Standard feature-selection methods typically optimize for one large population, while existing robust approaches tend to learn one shared model for every population. We introduce PopFS, a method for learning one shared, deployable feature set that is robust to population differences while letting each pop- ulation train its own model. PopFS uses a tunable welfare objective that lets practitioners balance overall predictive ben- efit against stronger protection of the populations that benefit least. To make this objective practical at scale, PopFS first uses multitask sparse learning to reduce the candidate pool, then searches directly over hard feature sets by ranking promising additions and swaps and fully refitting only a shortlist. Across eight population splits from six prediction tasks drawn from five tabular and public-health datasets, PopFS consistently achieves strong average and worst-population performance while scaling to thousands of candidate features. A 43-state COVID-19 nowcasting study further shows that changing the welfare objective can improve the least-served states with lit- tle change in average performance and yields an interpretable change in the selected symptom signals. Our code is available at https://github.com/Rachel-Lyu/PopFS.

Problem

• feature selection
• population heterogeneity
• robustness
• welfare optimization
• deployable feature set

Innovation

• feature selection
• population robustness
• welfare optimization
• multitask learning
• interpretable machine learning