🤖 machine learning

Population-Robust Feature Selection via Generalized Welfare Optimization

The paper introduces PopFS, a scalable feature selection method that optimizes a tunable welfare objective to identify a single, robust set of features for heterogeneous populations, allowing each group to train its own model while balancing overall predictive performance with protection for the least-served populations.

Original authors: Ruiqi Lyu, Alistair Turcan, Bryan Wilder

Published 2026-08-05

📖 9 min read🧠 Deep dive

CC BY 4.0

Original authors: Ruiqi Lyu, Alistair Turcan, Bryan Wilder

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a mystery, but you have a strict rule: you can only ask three questions to anyone you interview. If you are interviewing a group of teenagers, a group of retirees, and a group of office workers, the three questions that work best for the teenagers might make no sense for the retirees. This is the daily struggle of machine learning, a branch of science where computers learn to make predictions by finding patterns in data. Usually, these computer programs are trained on a giant pile of data from everyone mixed together. But in the real world, people aren't all the same. A model trained on "average" people might work great for the majority but fail miserably for smaller or different groups, like a map that shows a clear road for most drivers but leads a few into a swamp. The big question is: how do we pick a single, limited set of questions (or "features") that works well for everyone, even if everyone is different, without needing a separate, expensive questionnaire for each group?

This is exactly the puzzle tackled in a new paper called "Population-Robust Feature Selection via Generalized Welfare Optimization" by researchers at Carnegie Mellon University. They introduce a clever new method called PopFS. Think of PopFS as a super-smart, fair-minded editor for a massive encyclopedia. Usually, editors pick the most popular topics to include, which makes the book great for the majority but leaves out important details for niche readers. PopFS, however, uses a special "fairness dial" (called a welfare objective) that lets the editor decide how much to prioritize the needs of the least-served groups. If you turn the dial up, the editor ensures that even the groups who usually get the worst predictions get a fair shot, without ruining the book for everyone else.

The researchers found that PopFS is a game-changer. Instead of trying to find one perfect model for everyone (which is often impossible) or picking features that just work best on average, PopFS picks one shared set of "questions" that allows each group to build their own custom answer key. In tests involving eight different groups across six real-world tasks—from predicting income to tracking blood lead levels—PopFS consistently outperformed existing methods. It didn't just help the average case; it boosted the performance for the worst-off groups by up to 22% compared to other methods.

Perhaps the most exciting part is how flexible it is. In a case study tracking COVID-19 symptoms across 43 different US states, the researchers showed that by simply turning their "fairness dial," they could improve predictions for the states that were struggling the most by nearly 40%, with almost no drop in performance for the states that were already doing well. It's like tuning a radio to make sure the static is gone for the listener in the far corner of the room, without making the music worse for the person sitting right next to the speaker. The method is also incredibly fast, sorting through thousands of possible questions in under 15 minutes, making it practical for real-world use.

The Problem: One Size Does Not Fit All

Imagine you are a doctor who needs to diagnose patients, but you only have a tiny, fixed checklist of symptoms you can ask about. You have patients from all over the country: some are young and healthy, some are elderly with chronic conditions, and some live in areas with different environmental risks. If you pick the three symptoms that work best for the "average" patient, you might miss the crucial signs for the elderly or the young.

This is the core problem of feature selection. In machine learning, "features" are just the data points or questions you feed into the computer (like "age," "income," or "does the patient have a fever?"). Usually, scientists pick the features that make the computer smartest for the biggest group of people. But this leaves the smaller or different groups behind.

Some previous methods tried to fix this by building one giant, super-robust model that tries to work for everyone. But that's like trying to build one pair of shoes that fits a toddler, a basketball player, and a grandmother perfectly—it's nearly impossible. Others tried to build a separate model for every single group, but that requires a separate, expensive checklist for each group, which defeats the purpose of having a limited, shared resource.

The Solution: PopFS and the Fairness Dial

The authors of this paper, Ruiqi Lyu, Alistair Turcan, and Bryan Wilder, propose a middle ground. They want to find one shared checklist of features that everyone uses, but then let each group build their own custom model (their own "answer key") based on that checklist.

To do this, they invented PopFS. Here is how it works, using a simple analogy:

The Teacher and the Students: Imagine a master teacher (the "Teacher") who has access to all the information in the world. This teacher knows the perfect answer for every single student. PopFS first asks this teacher to look at all the data and figure out what the "perfect" prediction would be for every group.
The Compression Test: Now, the researchers ask: "If we only give the students a tiny, shared list of k questions (features), can they still guess the teacher's answer?" They test different lists of questions to see which one allows the students to get closest to the teacher's wisdom.
The Fairness Dial (Welfare Objective): This is the magic part. The researchers realized that just picking the list that gets the highest average score isn't fair. So, they added a tunable dial called the welfare parameter (represented by the Greek letter alpha, $\alpha$ $α$ ).
- If you set the dial to 1, the system cares mostly about the total score of everyone combined (the average).
- If you turn the dial down to 0 or even -1, the system starts caring much more about the students who are currently getting the lowest scores. It forces the selection to pick questions that help the struggling groups, even if it means the top-performing groups get slightly less help.

How They Made It Fast

You might think, "Wait, if there are thousands of possible questions, and we have to test them on dozens of different groups, that would take forever!" And you'd be right. Trying every possible combination of questions is like trying every possible combination of ingredients in a recipe book with thousands of pages; it would take a lifetime.

PopFS uses a clever two-step trick to solve this:

The Big Screen: First, it quickly scans all the thousands of features and throws away the ones that are clearly useless, keeping only a "shortlist" of the most promising candidates. It's like a talent show judge who quickly eliminates the people who can't sing at all, so they don't have to listen to every single person in the world.
The Smart Swap: Then, instead of testing every single combination, it starts with a good list and tries swapping one question at a time. It uses a fast "guessing game" (a mathematical shortcut called a surrogate) to predict which swaps will work best. It only does the full, slow, expensive test on the very best candidates. This allows it to find the perfect list in under 15 minutes, even with thousands of options.

What They Found

The researchers tested PopFS on real-world data, including:

Income prediction (using data from the ACS and UCI Adult datasets).
Health risks (blood lead levels in children and adults from NHANES, and credit risk from HELOC).
COVID-19 tracking (predicting hospital admissions across 43 US states based on symptom searches).

In almost every test, PopFS beat the other methods.

Better for Everyone: It improved the average performance for all groups.
Better for the Struggling: Most importantly, it significantly boosted the performance for the groups that usually get the worst predictions. In some cases, it improved the "worst-off" group's performance by up to 22% compared to standard methods.
The COVID-19 Case Study: When they applied this to predicting COVID-19 hospitalizations across 43 states, they found something fascinating. By turning the fairness dial to prioritize the states that were struggling the most, they improved those states' prediction accuracy by about 40%, while the average accuracy for all states barely changed.

They also noticed that the features PopFS chose changed depending on the dial. When they prioritized the struggling states, the system started picking different symptoms (like specific timing of fever or loss of taste) that were better for those specific regions, rather than just the most common symptoms. This suggests the method is actually finding the right tools for the right job, rather than just picking the most popular tools.

Why This Matters

The paper doesn't claim to have solved every problem in the world. The authors are careful to say that their method relies on the quality of the initial "teacher" model and that it uses smart shortcuts (heuristics) rather than finding the mathematically perfect answer every time. However, the results are strong and consistent across many different types of data.

The main takeaway is that we don't have to choose between "good for the majority" and "fair for everyone." With PopFS, we can have a single, deployable set of questions that works well for everyone, and we can tune it to ensure that the people who need help the most get it, without sacrificing the quality of the whole system. It's a step toward making machine learning not just smart, but also fair and adaptable to the messy, diverse reality of the world.

Technical Summary: Population-Robust Feature Selection via Generalized Welfare Optimization

Problem Formulation

The paper addresses a critical deployment challenge in machine learning: selecting a single, shared set of features (e.g., a questionnaire, test panel, or sensor set) that must serve multiple heterogeneous populations. While standard feature selection optimizes for a single pooled distribution, and distributionally robust optimization (DRO) typically fits a single shared predictor across groups, these approaches fail when the deployment constraint requires a fixed feature set but allows each population to train its own specific model.

The core difficulty lies in the discrete nature of feature acquisition ( $k \ll p$ ) and the trade-off between aggregate performance and the protection of populations that benefit least. A feature set optimized for average accuracy may discard predictive signals unique to minority or distinct populations. Existing methods like Distributionally Robust Feature Selection (DRFS) address the shared-acquisition setting but rely on continuous relaxations that may not align with the discrete deployment reality and lack a tunable mechanism to balance overall benefit against the protection of the worst-off populations.

Methodology: PopFS

The authors introduce PopFS, a scalable method for learning a robust shared feature set while allowing each population to train a separate "student" model. The method operates through three main stages:

1. Objective: Generalized Welfare Optimization

PopFS defines a tunable welfare objective to balance aggregate performance and protection for lower-benefit populations.

Utility Definition: For a feature subset $S$ , the utility $u_i(S)$ for population $i$ is defined as the normalized improvement over a constant-prediction baseline (either using a "teacher" model trained on all features or direct label prediction).
Welfare Aggregation: Utilities are aggregated using a weighted power mean $A_{\alpha, \rho}(u)$ $A_{α, ρ} (u)$ . The parameter $\alpha$ $α$ controls the trade-off:
- $\alpha = 1$ : Weighted average (maximizes total benefit).
- $\alpha = 0$ : Geometric mean.
- $\alpha \to -\infty$ : Minimum (maximizes the worst-case population).
Theoretical Justification: The paper establishes a "teacher-compression" objective. By fitting a full-feature "teacher" model for each population, the method approximates the preservation of predictive signal. Theorem 1 and 2 show that minimizing the error of approximating the teacher's conditional mean using a subset of features is a valid surrogate for minimizing the true prediction risk, provided the teacher error is small.

2. Scalable Optimization Strategy

Directly searching over all discrete subsets is computationally intractable. PopFS employs a multi-stage screening and search process:

Screening (Candidate Reduction):
- Round 1 (Marginal): Reduces the feature pool from $p$ to $p_0$ based on the root-mean-square covariance between features and the teacher predictions across populations.
- Round 2 (Multitask Sparse Learning): Fits a joint sparse model with group Lasso penalties across populations to select a refined set $C_{scr}$ of size $d$ . This encourages shared support while allowing population-specific coefficients.
Welfare-Guided Subset Search:
- The method initializes with the top $k$ features from $C_{scr}$ .
- It performs a validated one-swap local search. Instead of fully refitting models for every possible swap (which is expensive), it uses a Gauss–Newton-ranked surrogate (a ridge regression approximation) to rank potential feature additions or swaps.
- Only a shortlist of top-ranked swaps undergoes full retraining and validation.
- The search accepts a swap only if it improves the validation welfare by a margin $\delta_{swap}$ , terminating when no further improvements are found or after a set number of swaps.

Key Contributions

PopFS Framework: A method for learning one shared, deployable feature set robust to population differences under a tunable welfare objective, distinct from methods that fit a single shared predictor.
Scalable Discrete Optimization: A novel optimization strategy that bypasses continuous relaxations. It combines multitask sparse learning for high-recall screening with a Gauss–Newton-ranked refit search to directly optimize discrete feature sets.
Empirical Performance: Demonstrated ability to improve both average and worst-population performance across diverse datasets, with a specific case study showing interpretable changes in feature selection based on welfare tuning.

Results

The authors evaluated PopFS on six tasks (including ACS income, NHANES blood lead, HELOC credit risk, UCI Adult, and COVID-19 nowcasting) with up to 2,110 candidate features and 73 defined populations.

Performance vs. Baselines: PopFS consistently outperformed baselines (Pooled Lasso/XGBoost, DRO variants, and DRFS) in both population-average and worst-population metrics. Improvements reached up to 22% over the strongest external baselines.
Efficiency: PopFS completed selection tasks in 3–12 minutes, even for the largest datasets (2,110 features). In contrast, DRFS often failed to finish within 24 hours or required significantly more time.
Welfare Tuning: Adjusting the parameter $\alpha$ allowed practitioners to shift focus toward the least-served populations. In the COVID-19 case study, tuning $\alpha$ from 0 to -1 increased the worst-state performance by approximately 40% with negligible impact on the population average.
Interpretability: The COVID-19 study revealed that tuning the welfare objective led to interpretable shifts in selected features (e.g., prioritizing specific symptom lags like hypoxemia and ageusia), suggesting the method adapts to state-specific lead-lag relationships.

Significance and Claims

The paper claims that PopFS fills a practical gap in robust machine learning by addressing the discrete nature of feature acquisition while explicitly managing the trade-off between overall system utility and equity across populations.

Practicality: Unlike methods relying on continuous relaxations, PopFS searches directly over deployable hard subsets, ensuring the selected features are actionable.
Tunability: The welfare parameter provides a transparent mechanism for practitioners to prioritize the protection of underperforming populations without sacrificing the entire system's efficacy.
Scalability: The screening and surrogate-ranking approach makes population-robust feature selection feasible for high-dimensional problems where exhaustive search is impossible.

The authors modestly note limitations, including reliance on the quality of the fitted teacher models and the fact that the heuristic search does not guarantee a globally optimal feature set. They suggest future work could extend the framework to evolving populations and incorporate heterogeneous acquisition costs.