Distributionally Robust Feature Selection
Authors: Maitreyi Swaroop
Tamar Krishnamurti
Bryan Wilder
Venue: NeurIPS 2025 poster ()
Keywords: DRO
feature-selection
Primary area: optimization
TL;DR: A novel approach to perform distributionally robust feature selection in a model-agnostic manner

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
Maitreyi Swaroop Machine Learning Department Carnegie Mellon University
Tamar Krishnamurti
Division of General Internal Medicine
University of Pittsburgh
mswaroop&andrew. cmu edu
tamark@pitt.edi
Bryan Wilder
Machine Learning Department Carnegie Mellon University bwi bwilder@andrew. cmu. edu
Abstract
We study the problem selecting limited features to observe such that models trained on them can perform well simultaneously across multiple subpopulations. This problem has applications settings where collecting each feature is costly, e.g. requiring adding survey questions or physical sensors, and must be able to use the selected features create high-quality downstream models for differ- ent populations. Our method frames the problem continuous relaxation of traditional variable selection using a noising mechanism, without requiring back- propagation 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.
1 Introduction
Many real-world applications impose significant constraints data collection for machine learning models. These constraints often stem from factors such as limited budgets, privacy concerns, or the operational costs associated with acquiring each data point. For instance, in healthcare, hospital system aiming to implement a predictive screener across diverse patient populations must often contend with limitations the number of questions permissible due to patient burden, time constraints, and regulatory considerations. In such scenarios, the ability to identify a minimal yet highly informative set of features feature selection becomes paramount. Oftentimes, we might collect pilot data on a large number of features in order to make the operational decision of which to collect in deployment. For example, medical screeners questions that are then cut down to a smaller number for general-populatin usage; e.g., the PHQ (Spitzer et al. 1999] is 26-item mental health survey that typically reduced standard 9 or 2-question versions in practice [Krennke et al. 2001 2003].
often developed with larger number of
The challenge intensifies when the selected features must ensure reliable model performance not just on average, but across varied and potentially shifting underlying data distributions. This necessitates distributionally robust approach feature selection, ensuring that particular subpopulation is inadvertently neglected by models trained on the selected features. Furthermore, in many practical settings, the specific downstream inference model that will ultimately use the collected data may be known at the time of feature selection, or multiple different models might be employed. Therefore,
Code for implementing our method is available here (linked)
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

an ideal feature selection methodology should be model-agnostie, providing a core set of features that are broadly useful without being tied the idiosyncrasies of a particular predictive algorithm.
This paper addresses the problem of selecting a limited subset of features from a larger pool, using historical data, facilitate accurate and robust predictions across diverse populations, all while adhering to a collection budget. We aim to identify features that are robust to distributional shifts, ensuring that the model performs well simultancously across variety of specified subpopulations (e.g., institutional settings, clinics, or demographic groups).
To our knowledge, this problem has not been previously studied in the literature: feature selection methods focus on a single distribution of interest, while distributionally robust optimization (DRO) methods attempt to find single model that performs well on all populations, instead of selecting limited number of features that can then be used to train high-performing models for each downstream population. The intersection of the two problems complicates both. On the one hand, we require an approach to feature selection that is model-agnostic and naturally extends across multiple populations. On the other hand, standard DRO methods do not apply because we have a discrete optimization problem of selecting features instead of directly optimizing over a single model. Our approach formulates this task by introducing a continuous relaxation of the discrete feature selection problem which injects synthetic noise into the observation of each covariate. We then develop an analytical simplification of the relaxation which eventually allows us solve it with standard stochastic gradient descent-style methods.
Our main contributions are follows:
1. We propose a novel, model- model-agnostic method for distributionally robust feature selection. This method identifies subset of features that minimizes the maximum expected loss (or error) across potential data distributions, without requiring assumptions about the specific architecture or differentiability of the downstream predictive model.
2. We reframe the combinatorial problem of feature selection into a continuous optimization framework. This is achieved by introducing parameters that govern a noise injection process, effectively creating a differentiable measure of each feature's utility and permitting tractable, gradient-based optimization for selection.
3. We demonstrate the efficacy of our approach through experiments on both synthetic and real-world datasets, comparing its performance against naive selection strategies and existing distributionally robust optimization (DRO) baselines adapted for feature selection.
1.1 Related work
Our work lies in the intersection of distributionally robust optimization and feature selection. While there is rich body of work that focuses each topic individually, to our knowledge no previous work studies the intersection - selecting a set of features that will allow high-performing model to be trained for each subpopulation, as opposed to training single model that works well everywhere.
Feature selection: Classical feature selection methods primarily fall into two categories: combina- torial optimization approaches [Guyon and Elisseeff 2003 Kohavi and John 1997], and embedded methods [Breiman [001 Tibshirani] [996]. In the combinatorial optimization vein, a long line of work develops greedy strategies torward/backward selection heuristics for variable selection, histor- ically focused on linear models. Even for linear models, selecting the best set of at most h features is known to be NP-hard [Natarajan] 1995 Das and Kempel 201T), provided theoretical approximation guarantees for greedy selection in linear models via a connection to submodularity. KKanna et al.) (2017) provide faster algorithms for the same problem. However, there no direct path extend these methods beyond linear models, or to incorporate robustness across multiple distributions as a goal. Lasso regression [Tibshirani] 1996] is a classical method frequently used for feature selection, which leverages an & penalty to induce sparsity. However, the Lasso formulation does not account for robustness to distributional shifts and again bakes in an assumption of linearity. Zhao and Yuj (2006) highlight potential pitfalls of Lasso-bast based feature selection, showing that it may not consistently select the correct variables under certain correlation structures. Recent advances embedded methods have explored heterogeneous feature selection, though these address different problem settings from our work. Yang et al. (2022) propose locally sparse neural networks (L.SPIN) that learn sample-spect feature subsets, allowing different features for different samples. Similarly, Svirsky and Lindenbaum

---

introduce interpretable deep clustering (IDC) that performs cluster-level feature selection. While these methods advance feature selection for heterogeneous populations, they fundamentally differ from our approach: our method selects single global feature subset that must perform well universally across all known population groups, a constraint motivated by practical deployment scenarios where systems require fixed set of features for all users. Previous work has also proposed the use of random noise injection for feature selection single population settings [Grandvalet 2000]. as well as through Bayesian relevance estimation methods [Neal 1996 Tipping 2001]. However, umable-noise-based variable selection has seen limited adoption since it was originally proposed. in contrast to the large body of work noise form of regularization [Bishop 1995]. In this work, we revisit noise- based relaxations in the distributionally robust setting and show how the optimization advantages.
2024]
problem be reformulated in ways that create significant, seeviousl-uunrcognni Crucially, our approach separates variable selection from predictive model fitting, making it agnostic to the downstream model and eliminating the need differentiate through model training which may be infeasible for frequently used models like decision trees or random forests.
Distributionally Robust Optimization (DRO): DRO provides principled framework to account for worst-case model performance under distribution shifts, contrasting with traditional ML methods that optimize only for average performance. Duchi and Namkoong (2020) formalize this approach by providing snite-sample minimax bounds for uniform model performance across test distributions. Group DRO, which focuses worst-case performance across predefined population groups, is particularly relevant to our work. Sagawa et al. (019 introduce group DRO approach for neural networks that explicitly optimizes for the worst-case loss over known population groups, demon- strating improved performance on underrepresented groups compared to standard empirical risk minimization. Similarly, |Hashimoto et al. [2018) demonstrate that minimizing the worst-case risk across groups can prevent representation disparity in sequential learning settings.
Problem formulation
We assume our input (X, Y) ~ P to be covariates E R and outcomes Y. We wish to select a subset of & < variables from X that are most informative for predicting Y, while being robust to shifts in the underlying distribution P. Let I denote a set of indices corresponding to the selected variables, with [ = k. The selected sub-vector is Xi E R Informally, we wish to solve the problem:
min max E[Loss of model trained on (Xz,Y) ~ P.] Z]=k P.EP
where P represents the set of distributions we wish to ensure robust performance on. The above can be equivalently expressed as the problem of choosing binary mask E 10m E, a, = & where the model observes X = a C X. To formalize the problem, let L be a loss function (we will work with the mean squared error throughout), and
M,a arg min fEF
Ex,y~x,
@ X))]
be the risk minimizer for population over some model class F when the covariates are masked by a. Accordingly, we can formalize our problem as
min max (.X.Y=P,I((Y,, (a @ X)) s.t. ||a||| = k. P.E
(1)
In order to solve this problem, we have access to samples {(X), Y drawn iid for each population P. E P.
Methods
There are two main challenges to solving this problem. First, the optimization over is discrete, since a is binary- even for linear models and single population, this is NP-hard [Natargjan 1995]. We address this by introducing a continuous relaxation. Second, the population-level minimizer Mi. depends
in nontrivial manner, as the solution to risk minimization problem on the covariates included by a. In practice, we will only have finite data available to solve an empirical version of the risk minimization problem. Moreover, the model training process may induce complex
3

---

dependencies between the decision variable a and the objective function, which is unlikely to have a closed form for general model families F. To circumvent these issues, our method targets the loss of the Bayes-optimal predictor for each distribution and show that this surprisingly allows us to arrive at a significantly more computationally tractable optimization problem.
3.1 Continuous relaxation
While fn-constrained problems are typically hard, we draw inspiration from the Lasso and relax to an ε constraint that restores convexity of the feasible set while still encouraging sparsity. In continuous relaxation, a E [o, 11th now gives the degree to which a variable included instead of binary decision. However, naively scaling inputs as a @ X allows flexible predictors to trivially undo the scaling. For example, if the model involves a linear transformation, it could internally learn cocfficient w,/ax for the input X, A deterministic scaling does not actually remove any information about the covariate unless a = 0 exactly.
To avoid this issue, we introduce an alternative continuous relaxation that incorporates a stochastic component controlled by a;; effectively, a will control the amount of noise added to the observation of X, Formally, let (o1...., 0m) ) E Rgg be a vector of parameters controlling the degradation level for each covariate. We define the observed (noised) variable S(xx) as a random variable whose i-th component S, (ax) distributed as
((a) ~V),,
with the random variables independent across conditionally on X. Equivalently, S(a) = X + (ax). where e(or) (0, E.) with En diag(o),... am)) Here, a 0 implies S; X; (no degradation), while a; + oo implies S: contains no information about X, Our goal is to find an vector that has small values for relevant features and large values for irrelevant ones, while maintaining predictive performance under distribution shifts. Formally, our objective becomes
min max (tti://rrrr.ICCXMMMM.MM P.EP
+ Reg(a)
(2)
Mia arg min BainnEsteen fEF
(3)
so that we optimize the performance of the lisk-mnnmizzing models that observe the noisy version of the covariates. X is the regularization parameter which controls the sparsity of the solution, and can be varied to obtain solution with the desired cardinality. We note that the regularization term here in contrast to standard l regularization, which typically promotes sparsity by shrinking irrelevant parameters toward zero, eg. we may choose Reg(a) =1 1//aa/11
3.2 Solving the relaxation
Directly solving the relaxation in Equation B3 is nontrivial. Each choice of a leads to a different set of models Mt Morcover, the structure this mapping may be highly complex, mediated as is by the inner minimization problem defining M, One strategy employed throughout the machine learning literature to solve problems with inner optimization loop is to differentiate through the solution to the inner problem in order to optimize the outer objective via gradient descent [Amos and Kolter 2017 Finn et al. 2017 Agrawal al. 2019 Liu et al. 2018 Intuitively, this corresponds to computing a gradient V V. which captures how changes in turn change the fitted model. Perhaps the closest analogy to our setting is model-agnostic meta-learning (MAML) [Finn et al. 2017] which solves meta-learning problems by differentiating through the training loops of models for individual tasks.
While it might be possible to adapt such a strategy to our setting, it would incur three key disadvan- tages.
1. Computational Expense: Retraining potentially complex model M for every adjustment to a and for every considered P. during the optimization process is often prohibitively costly. Backpropagating through the model fitting process in every iteration is similarly costly.
2. Optimization Instability: Backpropagating gradients through the iterative training proce- dure of M with respect to a can be numerically unstable, suffer from vanishing/exploding gradients, or converge poorly, especially for deep non-convex models.

---

3. Specificity to a single model class: Such a formulation would necessarily optimize for the performance the particular class of models used to instantiate the inner loop. In order to maintain differentiability, this would likely have to be neural networks, even if other models (e.g., random forests) might be preferred for the tabular settings common in applications like medical risk prediction.
To circumvent these difficulties, we shift our focus from the performance of a specific, explicitly trained model M the performance of a Bayes-optimal optimal predictor, which represents the best possible performance achievable given the selected features. We believe this to be a good target for optimization for multiple reasons. Firstly, represents a model-family agnostic target which will be more closely approached if practitioner makes good choices for the modeling approach within any specific setting and distribution. Secondly, in many settings, after feature selection, we might be able to collect a substantial amount of new data corresponding to the chosen covariates. E.g., consider a health system that decides on new survey items to gather based on pilot study, and can then observe new data from the routine deployment of the selected questions. In such setting, the Bayes-optimal loss may be a better proxy for the long-term performance of the system. Third, starting from the perspective of the Bayes-optimal predictor will allow us to further simplify the objective and arrive at a significantly more computationally tractable approach. the following sections, we formalize this approach and show how it leads to simple closed-form objective.
Population-Ievel objective We now express the population-Ieve formulation of this problem using the Bayes-optimal predictor for Y given S(a), denoted by (S(()) = [(a)]. With respect to the MSE, the expected loss of this optimal predictor is the conditional variance of Y given S(ax) (since the bias term is O):
E ((s(a))))P [Y E|y {1 " [V[\]((a)]]
Applying the law of total variance and dropping terms constant in a leads to equivalent optimiza- tion problem that depends only on the conditional variance of E.YY X) given S(x). This can be formalized as follows.
Theorem 1 (Population-LLev) Objective). Under the noise-based relaxation S(a) = X + c(x).. <(ax) ~ N(, (i)g the distributionally robust feature selection problem is equivalent to
min max PEP
(Bs(a)-P(BxxxRl (ux() | S(aa)] + +A Reg(a),
where (x(X) = Epp [V]
X].
The proof of Theorem 1 is provided in Appendix
A.1
3.3 Empirical estimation and kernel form of objective
So far, have dealt only with population-level quantities. Next, must develop strategy to estimate these using the sampled data that we observe. Given samples {(X), Y2')} from each
population P., let (xxX) be an estimator of (x()) = E|Y trained on these samples. The empirical
form of the objective from Theorem is
min max. a P.EP -E(a))p (X) S(a) + Reg(a)),
(4)
where the expectation (X), S(a) is taken with respect to the empirical distribution P.(X) = 1 Σ/= 5x;( (X). First, we propose to estimate (xx(') simply by fitting machine learning model to the samples {(X), Y.') from population i. This can be done just once, using an arbitrary model well- suited to the application domain. There will be no need to differentiate through the training process, or refit the model during training. To make the empirical objective in Equation (4) computationally tractable, express the conditional expectation È (X), | S(x)) in closed form using Bayes' theorem under the Gaussian noise model. This yields a seerell-wiighee representation of (X). where each observed sample contributes according to its likelihood under the noisy observation S(a).
Theorem 2 (Kernel Form Equivalence). The empirical expectation
Es(a)
[Bjū(()) |S(a))]]
5

---

is equivalent to
Σ wf ((((S
Es(a)
where the weights
- (x: (S(s)(aa) r diag(a) (x² -s((a))
exp
w/ (S,
Th k= exp
(-2 (x - S(x))) " diag(a)-1 (X -s(())))
The proof for Theorem is provided in Appendix A.2 expectation can be estimated via Gaussian kernel smoothing, where the bandwidth of the kernel is determined by a. In summary, this kernel-based formulation provides explicit, differentiable link between the feature- noise parameters a and the resulting predictive uncertainty. It allows us to efficiently approximate the objective using empirical samples, enabling the end-to-end optimization procedure described in the following section.
This equivalence shows that the conditional
3.4 Algorithmic approach
The final objective function takes expectation of this over the random draw of S itself. Accordingly, the final statement of the optimization problem we wish to solve is
27
Σ w (S ()
min max a P:EP
Esta)
We implement gradient descent algorithm for this problem. We draw samples of S to approximate the outer objective function. Specifically, for Monte Carlo samples S we obtain the objective
2
Σ w (S* ())(((?))
t=1
We take gradient descent step with respect to either the maximum loss over all the populations, or the softmax over the population losses, controlled by the temperature parameter 3. Since the distribution of S(x:) itself depends on we construct each sample via the reparameterization trick: X + Va O e, where ~ (0, I). This ensures that gradients can flow not only through the kernel weights w (S, ax). but also through the sampled values S themselves, enabling end-to- end differentiation of the entire objective. This can now be implemented entirely with standard autodifferentiation software since w a closed-form, differentiable function a. All that is required is to fit model (X) once at the start of the process for each population; it can then be reused at each iteration unchanged and we never need to differentiate through the model training process. In practice, we further improve computational efficiency by taking the inner sum only the k-nearest neighbors of (()) the set (X:] 1 since the conditional distribution of X given S typically places a negligible mass outside this set.
The algorithm concludes by selecting the i features with the smallest optimized noise parameters a", as these represent the most informative variables for maintaining prediction performance across all populations. We summarize our complete proposed method in Algorithm I in Appendix B and provide its computational complexity.
Experiments
We conduct experiments on both synthetic and real-world datasets. evaluate our proposed method against baseline models trained on pooled dataset comprising all populations. Specifically, we use Lasso regression (linear regression for real targets, and logistic regression for categorical) and XGBoost Chen and Guestrin, 2016 regression/dlasiifcattio basclines. For Lasso, feature

---

selection is based on the largest coefficients under the regularization path, while for XGBoost, we rely on internal feature importance scores. We also implement distributionally robust (DRO) variants of both models. Further, we include an embedded baseline (Embedded MI.P), which uses multi- layer-perceptron (MI.P) with learnable feature mask trained via DRO. comprehensive description of all baseline methods can be found in Appendix [ a model on the selected features. The downstream model training is carried out independently for each feature selection method. We implemented our method using the PyTorch [Paszke et al. 2019) library, while for the downstream models, we use the scikit-learn [Pedregosa et al.) [011]] library. All baselines plus our method share the pipeline for downstream models, isolating the impact of the feature selections they output as opposed to predictive performance of models that they use en route.
For the downstream prediction task, we train
Data processing We split each dataset into three parts feeattrreseeeccionn-duuee downstream- We first do a 60 40 split of each
hondel-training-dee -aataiii-nnnee and a wawtrrrreammmoe------------e--ee population to obtain the ceature-selections dataset and the downstream model training and evaluation datasets. The latter is split 80 20 for downstream-moot training and evaluation respectively.
Evaluation Metrics In keeping with our motivation to improve downstream prediction accuracy, we first perform feature selection using our method and the baseline methods. We then train the downstream model independently on the wwwnssreaan-modll-rrainin -datasee,,, using the features selected by each method. We then evaluate performance of the fitted downstream models (random mawesrrremmmoll---------e-eee
forest and MLP) on the the primary metric for regression datasets, and Log Loss on the classification dataset. For the real datasets, we also include R -core for regression and prediction accuracy for classification tasks. The tabulated metrics are provided in Appendix D
using the mean-squared-eroor (MSE) as
4.1 Synthetic Data Experiments
To systematically evaluate the behavior of our method under controlled settings, we conduct series of synthetic experiments that vary in functional form, dimensionality, and population heterogeneity. The first experiment considers purely linear data-generating process with population-specille coefficients to study the effect of structured covariate shifts. A higher- dinhee--dinsii variant of this experiment is provided in Appendix [E.]] The second experiment introduces nonlinear relationships and heterogeneous noise distributions across populations, allowing us to assess robustness to more complex, mixed effects. Finally, an additional experiment using a sparse linear prediction model for targets is included in Appendix F.2 For the downstream prediction models, we use a random forest and an MI.P- both models use the standard scikit-learn implementations for fitting to the output. The random forest is an ensemble of 100 decision trees, while the MLP has single hidden layer with 100 neurons, trained for a maximum of 1000 epochs with early stopping based on validation loss convergence (handled internally by scikit-learn).
Synthetic dataset 1: Linear model This synthetic dataset comprises three populations A, B, C, with 15 features each. The populations have the following proportions in the dataset, and (linear) outcome functions:
A (40%)
Y 8X, + 6X, 4X2 2X4 +
B (35%)
Y = -8X, - 6X, + +4X, - 3X, - 2XX, + 8X, + 6X, +
C (25%)
y = -10X, 88, 6X, 5Xo +
Noise
E~~ c~1000,)
We set the budget for feature selection to 5 and, for comparison, also include the results for an increased budget of 10. Figure [ shows the performance of the downstream prediction models on the task of predicting Y for each population, using the features selected using our method and the baselines. The complete metrics table can be found in Appendix D Table More details on the implementation and hyperparameters used are provided in Appendix D.I We also include an additional experiment with this data set in which the number of features increases to 50, the results of which can be found in Appendix E.
7

---

(a) Budget= 5: Our method consistently achieves low еттог across all populations, performing par with DRO-XGBoost and DRO-Lasso, even oulperforming both methods in population
(b) Budget= 10: Overall performance improves pre- dictably upon increasing the feature budget. Our method maintains lower MSE compared to other meth- ods, including DRO-Lasso.
Figure 1: Performance comparison across populations on synthetic dataset 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 feature selection methods is consistent across the choice of downstream prediction model.
MSE
Budget
Budget
0.0
Mean Squared Error
Mean Squ Error
Figure Performance comparison across populations on synthetic dataset 2. Left: Mean Squared Error of downstream MIP 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.
Synthetic experiment 2: Linear model Here we have four populations, A, B, C, with 50 features each. The populations have the following proportions in the dataset, and outcome functions:
(40%)
Y 4X, + +3x, + X2 +
(35%)
Y 4Xo + 3X, + X3 +
C (25%)
Y 2X, + 3X,X, + 4sin((2X7) + EC
(15%)
Y -3x, + 2X, +ED
with heterogeneous noise across populations: Population A experiences reduced noise with с ~ (0, 0.052), while Population B exhibits heteroscedasticity where cg = oX3, X4) . 0.1, with (0.50. 0.3.X,) and y (N 1). Population C follows standard noise model Eg ~ (0, 0.12). and Population is characterized by heavy-tailed noise drawn from a scaled 0.2.
I-distribution: en
set the budget for feature selection to Figure shows the performance of the downstream prediction models on the task of predicting Y for each population, using the features selected using

---

our method and the baselines. The complete metrics table can be found in Appendix D Table More details on the implementation and hyperparameters used are provided in Appendix D.2
2
Results In syntheric experiment 1, although the generative process is linear and variables Xo to X, have strong effects in both populations A and B, the signs of their coefficients are reversed between populations. This reduces the effectiveness of LASSO, which tends to select features based on average effects across all data. As result, vanilla LASSO achieves its best performance on population C, as features relevant for C have conflicting coefficients, but performs poorly A and B, in spite of the linear setting. Vanilla XGBoost shows a similar trend. Our method outperforms most baselines, and has balanced performance across populations, not performing inordinately well in a certain population at the cost of others, unlike vanilla Lasso and XGBoost. For budget=10, our method is comparable with the best performing baselines, namely DRO XGB and DRO Lasso. The Embedded MLP baseline performs poorly even for the higher budget.
In synthetic experiment 2, the nonlinear data setting puts both vanilla Lasso and its DRO variant at disadvantage, as seen in Figure with DRO Lasso having the worst performance. The Embedded MLP baseline also performs poorly, being imbalanced in favor of other populations at the cost of population C, similar to Lasso. Our method is again comparable with the best performing baselines of XGB (DRO and vanilla), for both budgets.
should also be noted that our method has consistently low variance across both the synthetic experiments, the exact values of which may be found in Tables [ feature selection methods is consistent across the choice of downstream prediction model.
and The relative performance of
4.2
Real-datasets
UCI Adult Income Dataset [Becker and Kohavi, [996] We use the UCI Adult Income dataset to predict income across different demographic groups, where each age group represents distinet population. The dataset contains census income data with 14 features including age, education, occupation, and demographic information. The target variable is binary, indicating whether an individual's income exceeds $50K per year. After combining the original train/test splits (es 48, 000 samples), we subsample 10% of the data for faster computation. Populations are defined by sex (Male/Female). We drop any features used to define the population group. After one-hot-encoding we have 44 features, of which we select 5 (i.e. we set our budget to 5). Since the task is classification, a Random Forest Classifier is trained on selected features. Training is performed on each population separately, using each set of selected features. For further details about hyperparameter choices, please see Appendix D.4
American Community Survey (ACS) Dataset [U.S. Census Bureau 2018
We use person-level ACS Public Use Microdata Sample (PUMS) data for the year 2018 to predict household income across state populations. We focus on three populations, namely California (CA), FLorida (FL) and New York (NY), We subsample 5% of each state's population for faster computation. To further ease computation load, we also subsample features by fitting a random forest regressor to the entire pooled dataset of all the populations, and picking the top 18 most important features. We then perform feature selection using our method, and the baselines. For the downstream evaluation, since the target is real valued (similar to the synthetic datasets), we fit Random Forest Regressor using the selected features. For further details about hyperparameter choices, please see Appendix D.3
Data pre-processing Categorical features are one-hot encoded to create numerical representations. Missing values are imputed using median imputation. Any instances with missing target values are dropped. The numerical features standardized to ensure consistent scale across features. From the ACS data, we drop all non non-numeric or weight/geoid features, and retain only the numeric features.
4.3 Results
Our method consistently outperforms all baseline approaches across both the ACS and UCI datasets. On the ACS regression task, we observe (Figure 3ap an order-of-mann maddn-o--maanitudee reduction in MSE and substantial gains in R2 across all populations. This indicates that our feature selection procedure identifies highly informative variables. Notably, even strong baselines such as DRO-XGBBA suffer from significantly higher error and variance, suggesting sensitivity to spurious or less transferable

---

(a) ACS dataset results: Comparison of mean squared
(b) UCI dataset results: Downstream classification performance (mean accuracy and log loss with stan- dard deviations) for each method on Female Male populations. Our method achieves the highest accuracy in both groups while maintaining comparative log loss and generally lower run-to-run variance.
error (MSE) and R score (with standard deviations)
for each method across California (CA), Florida (FL),
and New York (NY). Our method achieves an order-of- magnitude lower MSE and substantially higher I?² than
all baselines, with consistently low run-to-run variance.
Figure 3: Performance comparison across populations on real datasets using different feature selection methods. In each subplot, left: Mean Squared Error; right: R Score.
features across groups. We also note that on the ACS dataset, the Lasso variants outperform XGBoost. For exact metric values, please see Table in Appendix D.3 D.3
the UCI classification task, our method generally outperforms the baselines across both popula- tions, with the exception of the Log Loss over the Female population, where XGBoost outperforms (FigureBbb) 3b). However, our method has more balanced performance across the two population. Furthermore, the consistently low standard deviation across all reported metrics demonstrates the stability of our method across random initializations and training splits. For exact metric values, please see Table 4 in Appendix D.4
Discussion and future work
We have presented novel noise-based continuous relaxation framework for distributionally robust feature selection in group-DRO setting. In contrast to existing methods that optimize for average- case performance train single robust model, our approach directly targets the selection of features that enable high-quality models across multiple subpopulations. By injecting feature-win wise noise and optimizing the Bayes-optimal predictor's variance, we derive model-agnostic and computationally tractable objective that avoids differentiating through model training. This formulation allows us to identify features with stable predictive utility under distribution shifts.
Empirical results highlight the limitations of standard selection techniques as Lasso in capturing non-linear or eopulation-specitie signal. Our method achieves substantially improved performance, including an order-of- magnitude reduction in MSE on the ACS dataset, underscoring its practical value in
settings. In future work, we aim apply our method broader range of real-world datasets further assess its generalizability and practical utility. Replacing our current plug-in estimators with influence sunction-based approaches could reduce estimation bias, especially when dealing with limited samples from minority populations. This would improve the robustness of our feature selection when population sizes are imbalanced.
Additionally, extending our framework beyond MSE loss presents an interesting theoretical challenge. Our derivation leverages the bias-variance decomposition specific to mean squared error, though generalized variance decompositions have been proposed for broader classes of loss functions, including proper scoring rules such as cross-entropy. Incorporating these generalized forms could allow our method to directly optimize alternative objectives, though this would require nontrivial extensions of current theoretical framework. Empirically, our results with cross-entropy loss in the UCI experiment suggest that optimizing MSE within our framework transfers well to other proper scoring losses, indicating that our bias-variance based reasoning may extend beyond MSE in practice. Developing a theoretical foundation for such extensions is important direction for future work.
10

---

Acknowledgments
Research reported this publication was supported by the National Institute of Mental Health the National Institutes of Health under award number ROIMH139097, and the Al Research Institutes Program funded by the National Science Foundation under Al Institute for Societal Decision Making (AI-SDM), Award No. 2229881.
References
Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd, Steven Diamond, and J Zico Kolter. Differentiable convex optimization layers. Advances in neural information processing systems, 32, 2019.
Brandon Amos and J Zico Kolter. Optnet: Differentiable optimization as layer in neural networks. In International conference on machine learning, pages 136-145. PMLR, 2017.
Barry Becker and Ronny Kohavi. Adult. https://oolor/11022422/////
UCI Machine Learning Repository, 1996. DOI:
Chris M Bishop. Training with noise is equivalent to tikhonov regularization. Neural computation, (1):108-116, 1995.
Breiman. Random forests. Machine Learning, 45:5-32, 2001. URL httpt:://iikkk.eeeee springee
Com/article/100 10110/3331009934433333
Tianqi Chen and Carlos Guestrin. Xgboost: scalable tree boosting system. In Proceedings of the 22nd ACM SIGKDD International Conference Knowledge Discovery and Data Mining, KDD '1, page 785-794, ACM, August 2016. doi: 10.1145/2939977.2999775 URL http:
7/dx.dois wwwooogg///1 1141/29396772.... 711.3997772222978533 2939785
Abhimanyu Das and David Kempe. Submodular meets spectral: Greedy algorithms for subset selection, sparse approximation and dictionary selection, 2011. 1102.3975
/bs/
John Duchi and Hongseok Namkoong, Learning models with uniform performance via distributionally robust optimization, 2020. URL https::taaa.ii..trg////////
08750
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agmostic meta-learning for fast adaptation of deep networks. In International conference on machine learning. pages 1126-1135. PMLR, 2017.
Yves Grandvalet. Anisotropic noise injection for input variables relevance determination. JEEE Iransactions neural nelworks, 11 6:1201-12, 2000. URL.iittpp:///eeeepplorrr. .ieeeorrreeeeeeee...... ieee org/
document/883399
Isabelle M Guyon and André Elisseeff. An introduction variable and feature selection. J. Mach. Learn. Res,, 3:1157-1182 2003. URI http::ttts:::w..jim org.oog///aa.. hrggpapers///eeee ctaperssvollune www.evvllleeeeeeeeeoee guyon03a/
guyon03a.pd pdf
Tatsunori B. Hashimoto, Megha Srivastava, Hongseok Namkoong, and Percy Liang. Fairness
without demographics in repeated loss minimization. ArXiv, (bs//800.00800, 2018. URL
//arxiv. org/ abs/ /1tes///06...0 08010
Rajiv Khanna, Ethan Elenberg, Alex Dimakis, Sahand Negahban, and Joydeep Ghosh. Scalable greedy feature selection via weak submodularity. In Artificial Intelligence and Statistics, pages 1560-1568. PMLR, 2017.
Diederik Kingma. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6999 2014.
Ron Kohavi and George John. Wrappers for feature subset selection. Artificial intelligence, 97 (1- (1-2):273-324, 1997.

---

Kurt Kroenke, Robert L Spitzer, and Janet BW Williams. The phq 9: validity of a brief depression severity measure. Journal of general internal medicine, 16(9):606-613 613, 2001.
Kurt Kroenke, Robert L Spitzer, and Janet BW Williams. The patient health questionnaire-2: validity of a two-item depression screener. Medical care, 41(11):1284-129, 2003.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search, arXiv preprint VrXiv:/806.0900) 2018.
Balas Kausik Natarajan. Sparse approximate solutions to linear systems. SIAM journal on computing, 24(2):27-234, 1995.
Radford M. Neal. Bayesian learning for neural networks. 1996. URL https://link wtw.iiii...iigeerr com/book 10. 1007/9778878 1612--4411-777 0745-0
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library, 2019. URL arxiv. org/ /912. 01703
Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg. J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825-2830 2011.
Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprini arXiv: 100831, 2019.
Robert L Spitzer, Kurt Kroenke, Janet BW Williams, Patient Health Questionnaire Primary Care Study Group, Patient Health Questionnaire Primary Care Study Group, et al. Validation and utility of self-report version of prime-md: mnd the phq primary care study. Jama, 282(18):13771744, 1999.
Jonathan Svirsky and Ofir Lindenbaum. Interpretable deep clustering for tabular data. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 47314-47338 PMLR, 21-27 htmll
Jul 2024. URL httts:://proceeeee ings. ngss mlr.
Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society Series B: Statistical Methodology, 58(1):267-2881 1996.
Michael E Tipping. Sparse bayesian learning and the relevance vector machine. Journal of machine learning research, 1(Jun):211-244, 2001.
U.S. Census Bureau. American community survey [1-year] [1-year/5-year estimates. Dataset, U.S. Census Bureau, 2018. Retrieved from [URL], [Access date].
Junchen Yang, Ofir Lindenbaum, and Yuval Kluger. Locally sparse neural networks for tabular biomedical data. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 25123-25153, PMLR, 17-23 Jul 2022. URL htttt:://proceeediiiii
mlr
.html
Peng Zhao and Bin Yu. On model selection consistency lasso. Journal of Machine learning research, (5NNo))5541-256) 2006.
12

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Section 3, Section 4
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess..
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
13

---

Justification: Section 3
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Section 4 Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
14

---

Answer: [Yes]
Justification: Abstract see footnote. Section
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ /pipp/.cc/pullc//uuui wwwpc//ggue////eee CoodeSubeiiiee ssi asionnolli onnolicy)))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Section Guidelines:
Appendix
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Section Appendix
D
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
15

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httpp:://neurippee http://c.cccpp..ttttttttttteeeea
Answer: [Yes]
Justification: Yes
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Section
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
16

---

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: See citations and Section
4
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, wwppppswpptseeee whworee.wthhhe....c. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
17

---

Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
18

---

Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy htte:s/s//eerrippoccc///o/rreerp for what should or should not be described.
19

---

A Theoretical Results: Derivations and Intuition
In this appendix, we provide detailed derivations and additional intuition for the theoretical results presented in the main text. We begin by expanding the population-level formulation of the objective, clarifying its interpretation in terms of conditional variances. We then derive the equivalent kernel- based form used in the algorithmic implementation.
A.1 Population-level objective
I Recall that the Bayes-optimal With respect to the MSE, the expected loss of this optimal predictor the conditional variance of Y given S(a:) (since bias is 0):
In this section, we derive the population-level objective in Theorem predictor for Y given S(x).. is given by f (S(()) = [[
((s(a)))) [(Y BaaBe|(SSsso [) = Bs (MM]S(())
We can expand this using the law of total variance to rewrite
V[Y[S(a)] Ex |V xX))x)())(()))) VIE)
(5)
Given the generative process S(a) X + vε, where the noise is independently sampled, we have Y IL S|X and thus can drop S(x)) from the inner conditioning terms. We obtain:
V[Y[S(a)]] Ex [V|[XX + V[E[{]X][S(a)).
(6)
This expression has the appealing property that a only enters through the outer conditioning: it changes the distribution over which we take the expectation/variance but not the variable we are taking the sopectation/varianee of. Thus we may estimate V[Y[X] and [E[[X]] just once per test distribution P., instead of having to fit a new model for every value a. However, we can pursue further simplifications to arrive at an even more streamlined objective with closed-form dependence on a.
First, recall that our objective is to minimize the expected variance Es( [V(|)(()))) When wrapped in this outer expectation, the first term in Equation (aabvee becomes constant with respect to a: Estaa (xVe///XXXVV))XX(((())))) objective. We are left with the problem
= Ex [V[[] Accordingly, we can drop it from the optimization
min max Bs(a)~P, (EE]))))))))) P.EP
Intuitively the inner variance term V[E(Y| repressentt represents the variance of the true conditional expectation [[ conditioned on us observing S(a).. Specifically, given S(ax) 8, there is distribution possible true covariate (X) values that could have generated the corresponding S(a)) through the noising process. Each of these potential X values has an associated true conditional mean E[|[X] The term ||E[Y| ((] = s] measures the variability or spread of these true [[Y|X' values, given the observed A high value of this conditional variance indicates that single noisy observation is consistent with a wide range of possible values for X, and correspondingly, E|y|X]] values. In this scenario, the noise introduced by has significantly obscured the relationship between the observed S(a) and the underlying true predictive signal E Y Conversely, a low value for this variance implies that S(a) is highly informative about E (X] Our objective, therefore, seeks to choose noise levels a such that, even under the worst-case test distribution, the average uncertainty (variance) about the true predictive function [[[|X] given the noised observation S(a) is minimized. In order to further simplify the objective we expand the variance term above. Let (X) = Ep, (V|X) denote the conditional mean function on the ith population. We have
Es (X))((()(() = (s(a)~ [xx (X (x)) ())((]]]]
Es [E|n (E)))
where the first term again collapses to dropped from the optimization. We are left with the second term, -Esia)>,, This term requires us average (Xx)) over the conditional distribution of X given S.
which is constant with respect to ar and can be
A.2 Kernel Form Equivalence Proof
In this section, we derive the kernel-form of the objective in Theorem 2 conditional expectation over /4 (X). By applying Bayes theorem, we can arrive at estimate for
We need to estimate the
20

---

this expression with closed form in terms of a. Specifically, we can rewrite
(X) [x( /
where ii is averaged over the distribution of X implied by the noise model, P,(X)S(a)). Via Bayes theorem, rewrite
P.(X(S(a)) = P((
P.(S(a))
We propose to estimate this quantity, and hence the objective function, by setting the "prior" (X in this expression to be the uniform distribution over the observed samples of X from population P.. X "!" The likelihood P.(S(x))XX) under our Gaussian noise model for S is given simply by
P((S(a))X) & exp {
1 (X- S) diag(a) - 5)
and then calculate the denominator
ni ΣØ= P((S(a)) Putting this all together, let
P((((a)(X))
i (S, a) E P((S(a))
denote the estimate the probability of observed data point X². We arrive at the estimator
2
n4 of (S))))
which can be interpreted as kernel smoothing of (x() under the Gaussian kernel implied by the noise model, measuring the amount of information lost due to the smoothing.
B DRO Feature Selection Algorithm and complexity
In this section, we summarize the full procedure for optimizing the noise-based relaxation and obtain- ing the final set of robust features described in section 3 of our method, followed by an analysis of its computational complexity. The proposed procedure, is outlined in algorithm
present the algorithmic implementation
Computational complexity per iteration Our method requires O(P b K d) operations per iteration, where is the number of populations, the number Monte Carlo samples, max, np is the maximum population size, K is the number of nearest neighbors used for kemel weight computation, and feature dimensionality d.
C Baseline Methods
Data Pooling & Standardization: Data points (X(P), y((P)) from all populations = 1,..., P are pooled into single dataset [(X,, Y-) Both features X and outcomes Y are then standardized to have mean and unit variance.
We summarize the baseline methods below:
1. Vanilla Lasso regression
2. Vanilla XGBoost regression
3. DRO Lasso regression
4. DRO Lasso regression
5. Embedded MLP
21

---

Algorithm Distributionally Robust Feature Selection
Input: Dataset {(x]() ) for populations p l...., P; budget k;; regularization parameter
j=
X; learning rate yi; Monte Carlo samples b; number of nearest neighbors K
Output: Feature noise parameters a" E
1: Precompute: For each population p.
2: Fit model (X) R [X] using ((x,))) y,P
3: Build k-NN index for (
Optional: for efficiency
4: Initialize: ( 11 E, where 021
5: for epoch = to Tmax do
6:
C
7: for p to do
8:
Lp < 0
9:
for f= to do
10:
Sample noise { (),,
Generate noisy observation 5(() t X Valepoch- O
11:
12:
if using k-NN then
N + k-nearest neighbors of 5(0) in (x 13=1
else
15:
N + {1. (p.
for j = 1 n, do
17:
Compute weights:
exp -x/(///0 -s(0) diag(a) ' (xfo) -s(0))
18:
(5(0), a EyeNN exp -x ((xx)'--500) -500) (a) -1 ((x((p) -s()))
19:
Cp Lp - - w)
(x)x (3))(
20:
C E max (E.,........ (1.... P) Lp + X Reg(a))
D or use softmax
21:
a (spoch) <- a (spoch-1)
V.a
22:
Project a (spoch) to m
Ensure non-negativity
Feature Selection: Select features with smallest values return Selected feature indices I = (indices of k smallest ox
24:
Vanilla Lasso This method applies Lasso regression to the combined dataset from all populations. The Lasso model trained by solving:
1 2N
min
(X Az
BER
5=1
We set YS value to 0.01. This selection is performed by retraining the model for each A and evaluating its performance the full training set. Once the coefficients 5 are determined, the k features with the largest absolute coefficient values (1371) are selected.
Vanilla XGBoost This method utilizes feature importance scores from XGBoost model trained on the pooled dataset. Hyperparameters such as tree depth, learning rate, etc., are set to default values provided by the XGBoost library [Chnn and Guestrin 2016). Feature importance scores extracted from the trained XGBoost model using the default feature feature_importeee importances. attribute). The lo features with the highest importance scores selected.
Distributionally Robust Optimization (DRO) Lasso This method adapts Lasso to be robust by iteratively re-weighting populations to focus on worst-case performance. Features (P) and outcomes (p) are standardized separately for each population p. We then perform DRO Lasso, summarized in Algorithm
The features with the largest absolute coefficient values from the final Lasso model ((Tmax) are selected.
22

---

Algorithm 2 DRO Lasso
1: Initialize uniform weights (0) for each population p ......P. P.
2: Standardize each population's 's data (Xx()) y(r) independently.
3: Select 1.1 regularization parameter Az.
4: for f. = 1... max do
5: Construct Pooled Dataset with Sample Weights:
Form Xau, Yan by concatenating all ( Y(),, and assign weight wp (t)
to each sample from
population p.
Fit the Lasso model:
N₂ w (y - -x(0)* 8)2 A1118112
min BER
6: Compute Population Losses; Calculate the MSE for the current model (80) on each original, unweighted standardized population p:
loss ; (p(), ( (()))
7: Weight Update: Update population weights for next iteration:
w (1+1)) w exp(1 loss
Normalize the weights:
(++1)'
w [t+1)
up
Ej=1 3j (1+1)'
Distributionally Robust Optimization (DRO) XGBoost This method adapts XGBoost to a distri- butionally robust setting by iteratively re-weighting populations based worst-case performance. Each population's 's data is standardized separately. The procedure aims to select the most important features by accounting for heterogeneity in population-level performance. The method is summarized in Algorithm 3
Algorithm 3 DRO XGBoost
Initialize uniform weights up } for each population p= =1 1,...., P.
Standardize each population's 's data (X(p), y()) independently.
for ( = 1... __max do
Construct Pooled Dataset with Sample Weights:
Form Xan, Yan by concatenating all X()), (r),, and assign weight wge to each sample from population
Train XGBoost Model: Fit an XGBoost regressor or classifier (depending on the task) on the pooled data using sample weights.
Compute Population Losses:
No (X, 5(())3
(Regression)
For each population compute: loss(() =
N,
No E logloss(Y. E(01)
(Classification) w w0+
(8) Update Population Weights: up up exp(1 loss
(t+1)'
(e+1)
up
Σ
The & features with the largest feature importance scores from the final XGBoost model are selected.
Embedded baseline Embedded MLP We include embedded baseline, namely Embedded MI.P which uses an MLP with a learnable feature mask trained via DRO. The MI.P has a single hidden layer of size 100. We train the model with the joint objective of MSE minimization (for the
23

---

regression task), and L-1 regularization (weighted by hyperparameter À = 0.01; we found this value work best out of [0.1, 0.01, 0.001J. The training procedure alternates between neural network training and population weight updates over multiple iterations (with the DRO weights being updated once every 10 epochs of the model training). The learnable mask parameters are constrained to [o. 1] and the top-k features are selected based on the final mask values after training convergence. We train for a maximum of 200 epochs, with early stopping with patience of 10 epochs when the average loss improvement falls below 10
D Results (Continued)
D.1 Synthetic dataset 1: Linear model
Table 1: Results for synthetic experiment (Section 4.1
Performance comparison across methods
and populations for different feature budgets. Total number of features is 15.
Method
Pop.
Budget
Budget 10
MLP MSE
RF MSE
MLP MSE
RF MSE
A
0.2251 = 0.0029
0.2486 0.0018
0.0342 + 0.0004
0.0535 + 0.0010
Our Method
B
0.2888 + 0.0113
0.3194
0.0098
0.0207 + 0.0005
0.0804 + 0.0036
C
0.2787 + 0.0150
0.3039 + 0.0130
0.0039 0.0004
0.0339 + 0.0079
A
0,5071 0.0129
0.5434 = 0.0191
0.0338 0.0009
0.0534 + 0.0009
DRO Lasso
B
0.2886 + 0.0086
0.3194 + 0.0026
0.0211
0.0001
0.0807 + 0.0041
C
0.2755 + 0.01550
0.3031 + 0.0113
0.0039 + 0.0004
0.0340 + 0.0078
A
0.5107 + 0.0113
0.5441 = 0.0188
0.0738 0.0370
0.0887 + 0.0327
DRO XGB
B
0.2888 + 0.0081
0.3187 + 0.0030
0.0408 + 0.0198
0.0928 + 0.0145
C
0.2762 + 0.0151
0.3032 = 0.0112
0.0036 + 0.0003
0.0337 + 0.0076
A
1.0087 + 0.0153
1.0571 = 0.0107
0.2056 + 0.1714
0.2106 + 0.1625
Lasso
B
0.5744 = 0.0051
0.6150 = 0.0052
0.1160 = 0.0921
0.1460 + 0.0763
C
0.1080 0.0107
0.1283 = 0.0102
0.0036 0.0003
0.0340 + 0.0076
A
1.0136 = 0.0076
1.0623 = 0.0039
0.4051 = 0.2802
0.3905 + 0.2579
XGB
B
0.5702 0.0008
0.6170 = 0.0076
0.2279 0.1586
0.2373 = 0.1307
C
0.1314 = 0.0291
0.1538 = 0.0352
0.0037 = 0.0002
0.0338 + 0.0077
A
0.7355 0,2702
0.7709 = 0.2636
0.5336 0.3503
0.5011 = 0.3107
Embedded MLP
B
0,5599 + 0313
0.6043 = 0.1310
0.4170 + 0.2951
0.4009 1 0.2521
C
0.2805 + 0.1630
0.3077 = 0.1728
0.2801 + 0.2117
0.2649 + 0.1923
I Values reported as mean = standard deviation.
indicates lower values are better.
Best results per population metric are highlighted in bold.
Implementation details a is initialized values to near 1 by adding random noise a vector ones. We use Adam optimzer Kingmal (2014) with learning rate of 0.1. We also use Cosine Scheduler for the learning rate, and train the model for 200 epochs. For the kernel estimation, we set the number of nearest neighbours A 1000. We take 10 Monte Carlo samples for estimating the objective. At each epoch, we do full-batch gradient descent. For the objective, we use the hard-max formulation (setting the SoftMax parameter to inf). The penalty term is a reciprocal of the L₁ norm of a. The entire dataset (combining all splits) is of size 36000. We ran the experiment for 3 different seeds and reported the average over all runs as seen in Figure Experiments were conducted on an Apple MacBook Pro equipped with an Apple M3. We set the budget to 5 and, for comparison, also include the results for an increased budget of 10. The metrics are summarized in Table
Consistency across seeds Our method consistently selects similar core ordered set of features
across (3) different seeds (in decreasing order of importance) [o, 7, 1, 5, 8, 6, 9, 2, 10,
3], [0, 7, 1, 8, 5, 9, 6, 2, 10, 3] [0, 7, 5, 1, 8, 6, 9, 2, 10, 3]. In con- trast, baseline methods, especially non-DRO versions, show significant variability, often selecting irrelevant noisy features (e.g., 13, 14) and failing to consistently identify the true signal variables.

---

D.2 Synthetic dataset 2: Nonlinear
Table 2: Results for synthetic experiment (Section 4.1 Per- formance comparison across methods and populations for dif- ferent feature budgets. Total number of features is 50, and budget is set to 8.
Method
Pop.
MLP MSE +
RF MSE
0.0054 + 0.0007
0.0150 + 0.0033
0.0062 0.0010
0.0205 0.0016
Our Method
0.0123 = 0.0080
0.2460
0.0483
D
0.0114 = 0.0015
0.0130 + 0.0008
A
0.4115 +: 0.0554
0.4020 1 0.0560
B
0,4172 = 0.0533
0,4230 + 0,0511
DRO Lasso
C
0.3114 = 0.2644
0.3816 0.1018
0.3569 = 0.0259
0.3375 0.0089
0.0054 = 0.0014
0.0150 + 0.0034
0.0064 = 0.0009
0.0206 0.0017
DRO XGB
0.0226 + 0.0005
0.2460 0.0479
D
0.0114 + 0.0012
0.0132 + 0.0010
A
0.0517 = 0.0417
0.0616 0.0433
0.0562 = 0.0443
0.0657 + 0.0390
Lasso
C
0.4654 = 0.0573
0.4581 + 0.0676
0.0114 0.0014
0.0130 + 0.0012
A
0.0051 = 0.0013
0.0151 = 0.0033
0.0067 = 0.0012
0.0206 = 0.0015
XGB
0.0200 = 0.0039
0.2471 0.0475
D
0.0117 + 0.0017
0.0131 + 0.0007
0.0050 = 0.0008
0.0150 0.0035
0.0062 = 0.0014
0.0207 + 0,0019
Embedded MLP
C
0.3030 + 0.2435
0.4119 + 0.1885
0.0119 + 0.0019
0.0128 0.0008
Values reported as + standard deviation for Budget 8.
J. indicates lower values are better.
Budget refers to the number features selected.
Implementation details a is initialized values to near 2 by adding random noise vector of twos. We use Adam optimzer Kingmal (2014) with learning rate 0.1. We also use CosineAnnealligg Scheduler for the learning rate. and train the model for 150 epochs. For the kernel estimation, set the number of nearest neighbours & = 1000. We take 50 Monte Carlo samples for estimating the objective. At each epoch, do full-batch gradient descent. For the objective, hard-max formulation (setting the SoftMax parameter to inf). The penalty term a reciprocal of the Lı norm of The entire dataset (combining all splits) is of size 44000. We ran the experiment for 3 different seeds and reported the average over all runs as seen in Figure Experiments were Apple MacBook Pro equipped with Apple M3. We set the budget The metrics are summarized in Table
the
conducted
D.3 ACS Dataset
Implementation details a is intialized values to near 5 by adding random noise to vector of 5s. We use Adam optimzer, set initial learning rates of 0.01 with Cosine Annealinnn Scheduler for the learning rate, and run the optimization for 120 epochs. For the kernel estimation, we set the number of nearest neighbours
500. We take 10 Monte Carlo samples for estimating the objective. At each epoch, we do full-batch gradient descent. For the objective, we use the SoftMax formulation (setting the SoftMax parameter 10). The penalty term is a reciprocal of the L1 norm of a,, with
X = le 4. We ran the experiment for 5 different seeds and reported the average over all runs as seen Figure 3a For the ACS dataset, we set the budget to out 18 features. The metries are summarized in Table For larger visualizations for greater readability, see Figure
25

---

Mean Squared Error
R2 Score
CA
FL
NY
0.6
Mean Squared Error
R2 Score
Dro Lasso
Lasso
Our Method
Xgh
Xgb
(a) ACS dataset results: Comparison of mean squared error (MSF) R score (with standard deviations) for each method across California (CA). Florida (FL), and New York (NY). Our method achieves an order-of-
magnitude lower MSE and substantially higher It than all baselines. with consistently low run-to-run variance.
Accuracy
Loss
0.2
0.4
0.8
0.0
0.2
0.4
0.6
1.0
Accuracy
Log Loss
Dro Xgb
Xgb
(b) UCI dataset results: Downstream classification performance (mean accuracy and log loss with standard deviations) for each method on Female and Male populations. Our method achieves the highest accuracy both groups while maintaining comparative log loss and lower run-to-run variance.
Figure 4: Performance comparison across populations on real datasets using different feature selection methods.
D.4 UCT Dataset
a is intialized values to near by adding random noise to vector of 5s. We use an Adam optimzer, and initial learning rates to 0.02 with CosineAnnealing Scheduler for the learning rate, and run the optimization for 120 epochs. For the kemne estimation, we set the number of nearest neighbours 500. We take 10 Monte Carlo samples for estimating the objective. At each epoch, we do full-batch gradient descent. For the objective, we use the SoftMax formulation (setting the SoftMax parameter to 10). The penalty term reciprocal the J, norm of a, with 5e 6. Due to computational instability of the Logistic Regression baseline,
the experiment for 3 different
26

---

Table 3: ACS dataset: Performance comparison across methods and populations
Method
Population
MSE
1²² +
CA
0.260 = 0.076
0.758
0.067
Baseline DRO Lasso
FL
0.318 = 0.054
0.688 = 0.056
NY
0.226 0.132
0.766 = 0.064
CA
1.063 = 0.174
0.020 = 0.079
Baseline DRO XGBoost
FL
1.092 = 0.262
-0.049 = 0.078
NY
1.058 = 0.479
-0.099 = 0.060
CA
0.604 + 0.410
0.433 = 0.381
Baseline Lasso
FL
0.628 = 0.309
0.359 = 0.373
NY
0.668 = 0.654
0.373 = 0.445
CA
1.022 = 0.110
0.050 = 0.104
Baseline XGBoost
FL
1.098 + 0.286
-0.052 100
NY
1.022 + 0.426
-0.069 + 0.057
CA
0.016 = 0.012
0.985 0.010
Our Method
FL
0.017 = 0.010
0.983 = 0.009
NY
0.040 + 0.052
0.967 0.030
Values reported mean + standard deviation, rounded three decimal places.
J. indicates lower values are better; 1 indicates higher values are hetter.
Best results per population metric are highlighted in bold.
Table 4: UCI dataset:
Performance comparison across methods and populations
Method
Population
Accuracy f
Log Loss
Female
0.886 0.010
0.654 = 0.273
Baseline DRO Lasso
Male
0.720 = 0.027
0.626 0.192
Female
0.897 0.012
0.339 = 0.137
Baseline DRO XGBoost
Male
0.711 = 0.039
0.594 = 0.096
Female
0.896 + 0.023
0.657 = 0.101
Baseline Lasso
Male
0.716 0.036
0.847 = 0.321
Female
0.904 + 0.021
0.279 = 0.083
Baseline XGBoost
Male
0.699 = 0.035
0.608 0.090
Female
0.918 + 0.006
0.300
0.086
Our Method
Male
0.786 + 0.010
0.562 0.047
Values reported mean 1 standard deviation, rounded to three decimal places.
indicates higher values better; indicates lower values better.
Best results per population metric are highlighted in bold.
seeds and reported the average over all runs as seen in Figure 3bb We set the budget to 5 out of 44 (real and one-hot-encoded) features. The metrics are summarized in Table 4 For larger visualizations for greater readability, see Figure 4b
E
Additional synthetic experiments
E.1
Synthetic experiment 1: Linear model with increased dimension
In this section we include additional experiments for the data setting from Section 4.1 synthetic experiment (linear data setting), when the total number of features is set to 50. Here, the increase in dimension only contributes in additional noise variables, while the meaningful variables are the same as for dimension= 15. We set the budget for feature selection and, for comparison, also include the results for increased budget of 10. Figure shows the performance of the downstream
27

---

1011A0 Supee1
(a) Budget= 5: Our method consistently among the top- 3 feature selection methods along with DRO-XGBoost DRO-Lasso, even outperforming both methods in population C.
(b) Budget= 10: Overall performance improves pre- dictably upon reducing the budget. Our method main- tains lower MSE compared other methods, including DRO-Lasso, and has a balanced performance across populations.
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 downstream random forest model. The relative performance 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.
Table 5: Additional results for synthetic experiment 1 (Section 4.1
Performance comparison across methods and populations for different feature budgets. Total number of features 50.
Method
Pop.
Budget 5
Budget 10
MLP MSE
RF MSE
MLP MSE
RF MSE +
A
0.2227 + 0.0094
0.2420 = 0.0161
0.0241 0.0178
0.0450 + 0.0129
Our Method
B
0.5805 + 0.0375
0.6158 = 0.0448
0.0150 + 0.0096
0.0809 + 0.0035
C
0.1208 = 0.0034
0.1388 = 0.0011
0.0480 = 0.0761
0.0697 + 0.0605
A
0.4155 = 0.1769
0.4418 = 1816
0.0345 + 0.0004
0.0511 + 0.0029
DRO Lasso
B
0.2377 = 0.0909
0.2677 0.0972
0.0206 = 0.0006
0.0835 = 0.0067
C
0.3642 0.1350
0.3849 = 0.1300
0.0040 0.0007
0.0354 + 0.0012
A
0.6827 + 0.3136
0.7120 = 0.3138
0.1041 = 0.0645
0.1139 + 0.0597
DRO XGB
B
0.3857 0.1461
0.4209 = 0.1495
0.0613 0.0416
0.1098 + 0.0326
C
0.2259 + 0.0928
0,2479 = 0.0926
0.0040 + 0.0005
0.0350 1 0.0010
A
1.0181 0.0526
1.0429 0.0387
0.4126 0.5303
0.3859 1 0.4760
Lasso
B
0,5880 1 0.0323
0.6264 1 0.0284
0.2327 1 0.2912
0,2504 1 0.2401
C
0.1351 + 0.0252
0.1550 + 0.0267
0.0037 + 0.0001
0.0354 + 0.0009
A
1.0211 1 0.0474
1.0421 1 0.0400
1.1132 + 0.0439
1.0216 1 0.0507
XGB
B
0.5867 1 0.0319
0.6258 1 0.0282
0.6607 1 0.0314
0.6135 1 0.0252
C
0.1356 = 0.0261
0.1551 = 0.0267
0.0036 + 0.0004
0.0347 + 0.0012
A
0.9346 + 0.1914
0.9670 + 0.1896
0.5275 + 0.2598
0.4928 + 0.2475
Embedded MLP
B
0.8857 + 0.2055
0.9238 0.2096
0.5952 + 0.3581
0.5597 + 0.3116
C
0.3666 0.1390
0.3772 = 0.1274
0.4200 0.1610
0.3756 + 0.1282
Values reported mean standard deviation.
indicates lower values are better.
3 Best results population metric are highlighted in bold.
prediction models on the task of predicting Y for each population, using the features selected using our method and the baselines. The complete metries can be found in Table
Even in the higher dimensional setting, our method's performance is comparable with the best performing baselines (DRO XGBoost and DRO Lasso), while maintaining balanced performance across populations. Vanilla XGBoost shows the greatest degradation performance when compared to the lower dimensional setting. Vanilla Lasso and Embedded MLP show the highest
Results
28

---

variance across different seeds, while our method, along with DRO Lasso and DROXGBoost has the lowest variance across seeds.
Implementation details a is initialized values to near 1 by adding random noise a vector of ones. We use Adam optimzer Kingma [2014] with learning rate of 0.1. We also use CosineAnnealing Scheduler for the learning rate, and train the model for 150 epochs. For the kernel estimation, we set the number of nearest neighbours k- = 1000. We take 10 Monte Carlo samples for estimating the objective. At each epoch, we do full-batch gradient descent. For the objective, we use the hard-max formulation (setting the SoftMax parameter to inf). The penalty term is a reciprocal of the I,1 norm of a. Our entire dataset (combining all splits) is of size 36000. We ran the experiment for 3 different seeds and reported the average over all runs seen in Figure Experiments were conducted on an Apple MacBook Pro equipped with an Apple M3. We set the budget to 5 and, for comparison, also include the results for an increased budget of 10.
E.2 Synthetic experiment Sparse Linear
MGLIMUS Batpess
Mose
(a) Budget== 5: Our method consistently achieves low error across all populations, performing on par with DRO-XGBoost and DRO-Lasso, even outperforming both methods in population A.
(b) Budget 10: Overall performance improves pre-
dictably upon reducing the budget.
Figure 6: Performance comparison across populations on synthetic dataset 3 using different feature selection methods. 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 consistent across the choice of downstream prediction model.
This synthetic dataset comprises three populations A, B, C, with 50 features each. The populations have the following proportions in the dataset, and outcome functions:
A (35%)
Y = 5X, + 1X + 3. +
(35%)
Y 6.,, + 5X20 + 4Xas +
C (30%)
Y 7.X10 6X25 + 5X + 4X +
The noise term follows base distribution <~ (0,,0.13) Feature correlations are introduced via generative process where X,-1 0.3X;+0.77) with 17 NO 1). inducing structured dependencies across dimensions. The dataset contains 50 features in total, the majority of which are noise. We set the budget for feature selection and, for comparison, also include the results for an increased the performance of the downstream prediction models on the task of
budget of 10. Figure predicting Y for each population, using the features selected using our method and the baselines. The complete metrics table can be found in T
Results We observe that our method performs similarly to the remaining baselines. For budget= 5. all models perform similarly, with the exception of vanilla XGBoost, which has the best performance on populations B and C, but seems to do so at the cost of its performance on population A. In group-DRO setting, we would prefer to have a more balanced performance across all populations, as seen in the other methods. For budget= 10, the Embedded ML. performs the worst on all populations. Our method consistent with the other baselines on populations B and C, however, in this setting performs poorly population A, potentially due to A having the weakest signals.
29

---

Table 6: Results for synthetic experiment Performance comparison across methods and popula- tions for different budgets. Total number of features is 50.
Method
Pop.
Budget
Budget 10
MLP MSE +
RF MSE
MLP MSE
RF MSE
A
0.1186 + 0.0055
0.1299 + 0.0064
0.1274 + 0.0054
0.1277 + 0.0047
Our Method
B
0.2099 + 0.0061
0.2290 + 0.0086
0.0034 + 0.0001
0.0093 + 0.0014
C
0.6009 + 0.0356
0.6453 + 0.0318
0.0036 0.0002
0.0282 + 0.0021
A
0.1179 + 0.0053
0.1296 0.0066
0.0038 0.0003
0.0095 + 0.0003
DRO Lasso
B
0.2115 + 0.0070
0.2295 = 0.0088
0.0033
0.0004
0.0092 + 0.0013
0.6054 + 0.0438 C
0.6456 + 0.0310
0.0034 + 0.0001
0.0281 + 0.0021
A
0.1189 + 0.0057
0.1298 = 0.0066
0.0038 0.0000
0.0095 + 0.0004
DRO XGB
B
0.2118 + 0.0071
0.2286 = 0.0092
0.0035 + 0.0003
0.0092 + 0.0014
0.6073 = 0.0432 C
0.6447 = 0.0311
0.0037 = 0.0004
0.0280 - 0.0018
A
0.4729 = 0.0153
0,5066 = 0.0185
0.0039 + 0.0004
0.0095 + 0.0004
Lasso
B
0.2115 = 0.0057
0.2309 = 0.0090
0.0036 = 0.0004
0.0092 = 0.0013
0.3202 + 0.0029 C
0.3485 0.0106
0.0037 0.0001
0.0281 + 0.0018
A
1.0188 + 0.0235
1.0651 = 0.0303
0.0036
0.0003
0.0095 = 0.0004
XGB
B
0.1414 0.1214
0.1564 0.1285
0.0034
0.0002
0.0091 + 0.0013
C
0.1915 0.1125
0.2121 + 0.1196
0.0037
0.0004
0.0280 + 0.0019
A
0.1188 0.0051
0.1306 0.0054
0.1229 0.0163
0.1222 1 0.0138
Embedded MLP
B
0.2381 + 0.0499
0.2595 = 0.0562
0.2579
0.0562
0,2530 1 0.0534
C
0.6004 + 0.0314
0.6364 = 0.0459
0.4461 + 0.0724
0,4168 + 0.0659
1 Values reported as mcan = standard deviation.
indicates lower values are better.
Budget refers to the number of features selected.
Implementation details a is initialized values to near by adding random noise to a vector of twos. We use Adam optimzer Kingma [2014) with learning rate of 0.1. We also use Cosine Annealing Scheduler for the learning rate, and train the model for 150 epochs. For kernel estimation, we set the number of nearest neighbors & 1000. We take 50 Monte Carlo samples for estimating the objective. At each epoch, we do full-batch gradient descent. For the objective, we use the hard-max formulation (setting the SoftMax parameter to inf). The penalty term reciprocal of the L₁ norm of a. Our entire dataset (combining all splits) is of size 36000. We ran the experiment for 3 different seeds and reported the average over all runs Apple MacBook Pro equipped with an Apple M3. We set the budget to 5 and, for comparison, also include the results for increased budget of 10.
in Figure Experiments
conducted on
30


Reviews:
[{"review_id": "ufs70qSE6N", "reviewer": "Reviewer_ggn5", "summary": "This paper addresses the problem of selecting features such that models trained on these features are robust to distributional shifts. The authors formulate this problem as a continuous relaxation of the traditional variable selection using a noising mechanism. By optimziing over the variance of a Bayes-optimal predictor, they develop a model agnostic feature selectiom framework that achieves robust performance across populations. Experiments on synthetic and real-worlds datasets are conducted.", "questions": "See the weakness.", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strength**\n\n**1.**  This paper introduces an important problem of distributional robust feature selection, which lies at the intersaction of feature selection and domain generalization.\n\n**2.** The continuous relaxation of feature selection is novel. It efficiently create a differetiable measure of feature utility and premit tracable optimization.\n\n**3.** The writting of this paper is very clear and easy to follow.\n\n**4.** The authos conduct extension experiments for validation\n\n**Weakness**\n\n**1.** The proposed method lacks theoretical optimality guarantee.", "was_revised": true, "final_justification": "The author's response has addressed my concern. So I keep my positive score.", "rebuttal": "We thank reviewer ggn5 for their feedback. \n\nWe have conducted additional synthetic experiments and include three of them in this rebuttal.\nOur new experiments include:\n- Additional downstream model for evaluating performance: we now include an MLP (single hidden layer of size 100), along with the existing random forest (RF), as a downstream task model. \n- New baseline methods: we've added Forward Selection and Backward Elimination, and their DRO variants\n- Additional data generation settings to explore different covariate-outcome functional relationships.\n\n# Experiments\nFor all experiments, we set dimension $d=50$, and use a 40-60 train-(test+val) split for each population\n# Expt 1: Linear generative process\nWe use a dataset of size 36000 with the following composition:\n- A (40%): $$ Y = 8X_0 + 6X_1 - 4X_2 + 3X_3 + 2X_4 + \\epsilon $$\n- B (35%): $$ Y = -8X_0 - 6X_1 + 4X_2 - 3X_3 - 2X_4 + 8X_5 + 6X_6 + \\epsilon $$\n- C (25%): $$ Y = 10X_7 + 8X_8 + 6X_9 - 5X_{10} + \\epsilon $$\n$$\\text{Noise:} \\quad \\epsilon \\sim \\mathcal{N}(0,0.1^2)$$\n\n|Method|Population|MSE±Std (MLP) — Budget 5|MSE±Std (RF) — Budget 5|MSE±Std (MLP) — Budget 10|MSE±Std (RF) — Budget 10|\n|---|---|---|---|---|---|\n|Our Method| A|0.2227±0.0094| 0.2420±0.0161|0.0241±0.0178| 0.0450±0.0129|\n||B|0.5805±0.0375| 0.6158±0.0448|0.0150±0.0096| 0.0809±0.0035|\n||C|0.1208±0.0034| 0.1388±0.0011|0.0480±0.0761| 0.0697±0.0605|\n|(BL) DRO Lasso|A|0.4155±0.1769| 0.4418±0.1866|0.0345±0.0004| 0.0511±0.0029|\n||B|0.2377±0.0909| 0.2677±0.0972|0.0206±0.0006| 0.0835±0.0067|\n||C|0.3642±0.1350| 0.3849±0.1300|0.0040±0.0007| 0.0354±0.0012|\n|(BL) DRO XGB|A|0.6827±0.3136| 0.7120±0.3138|0.1041±0.0645| 0.1139±0.0597|\n||B|0.3857±0.1461| 0.4209±0.1495|0.0613±0.0416| 0.1098±0.0326|\n||C|0.2259±0.0928| 0.2479±0.0926|0.0040±0.0005| 0.0350±0.0010|\n|(BL) Lasso| A|1.0181±0.0526| 1.0429±0.0387|0.4126±0.5303| 0.3859±0.4760|\n||B|0.5880±0.0323| 0.6264±0.0284|0.2327±0.2912| 0.2504±0.2401|\n||C|0.1351±0.0252| 0.1550±0.0267|0.0037±0.0001| 0.0354±0.0009|\n|(BL) XGB| A|1.0211±0.0474| 1.0421±0.0400|1.1132±0.0439| 1.0216±0.0507|\n||B|0.5867±0.0319| 0.6258±0.0282|0.6607±0.0314| 0.6135±0.0252|\n||C|0.1356±0.0261| 0.1551±0.0267|0.0036±0.0004| 0.0347±0.0012|\n|(BL) DRO Forward Selection|A|0.3735±0.1506|0.3996±0.1636|0.0341±0.0006|0.0514±0.0029|\n||B|0.3729±0.0553|0.3995±0.0553|0.0209±0.0012|0.0833±0.0066|\n||C|0.4735±0.0805|0.5124±0.0883|0.0038±0.0001|0.0354±0.0012|\n|(BL) DRO Backward Elimination| A|0.0320±0.0012|0.0474±0.0024|0.0346±0.0005|0.0512±0.0030|\n||B|0.1757±0.0093|0.2079±0.0149|0.0208±0.0008|0.0834±0.0068|\n||C|1.0008±0.0669|1.0355±0.0635|0.0038±0.0003|0.0354±0.0012|\n|(BL) Forward Selection|A|1.0132±0.0372|1.0408±0.0435|1.1209±0.0483|1.0224±0.0403|\n||B|0.5879±0.0331|0.6276±0.0384|0.6509±0.0360|0.6109±0.0362|\n||C|0.9808±0.1161|1.0040±0.0989|1.1410±0.1413|0.9882±0.1095|\n|(BL) Backward Elimination|A|1.0095±0.0393|1.0463±0.0447|1.1267±0.0415|1.0244±0.0405|\n||B|0.5897±0.0341|0.6267±0.0300|0.6640±0.0208|0.6111±0.0364|\n||C|0.9883±0.1042|1.0071±0.0912|1.1047±0.1156|0.9895±0.1071|\n---\n\n## Discussion\nHere, although the generative process is linear and variables $X_0$ to $X_4$ have strong effects in both A and B, the signs of their coefficients are reversed between populations. This reduces the effectiveness of LASSO, which tends to select features based on average effects across all data. As a result, vanilla LASSO does not achieve the best performance even in this linear setting. For budget=$5$, our method outperforms most baselines, and has a balanced performance across population. For budget=$10$, our method is comparable with the best performing baselines.\n\nWhen we set dimension=$15$, our method is at par with the best performing baselines for all budgets- pl see the experiment section of our response to reviewer Guhq for the results. \n\n## Expt 2: Mixed Linear-Nonlinear with Heterogeneous Noise\nWe use a dataset of size 44000 with the following **population structure**:\n\n- A (30%):$$Y=4X_0+3X_1+X_2^2+\\epsilon_A$$\n- B (30%):$$Y=4X_0+3X_1+X_2^2+\\epsilon_B$$\n- C (25%):$$Y=2X_0+3X_5 X_6+4\\sin(2X_7)+\\epsilon_C$$\n- D (15%):$$Y=3X_0+2X_1+\\epsilon_D$$\n\n**Heterogeneous Noise**:  \n- $\\epsilon_A \\sim \\mathcal{N}(0, (0.05)^2)$ (reduced noise)\n- $\\epsilon_B=\\exp(0.5X_3+0.3X_4) \\cdot \\eta \\cdot 0.1$ where, $\\eta \\sim \\mathcal{N}(0,1)$ (heteroscedastic)\n- $\\epsilon_C \\sim \\mathcal{N}(0, 0.1^2)$ (standard)\n- $\\epsilon_D \\sim t_3 \\cdot 0.2$ (heavy-tailed, t-distribution with df=3)\n\n### Results\n|Method|Population|MSE±Std (MLP) — Budget 5|MSE±Std (RF) — Budget 5|MSE±Std (MLP) — Budget 8|MSE±Std (RF) — Budget 8|\n|:---|:---|:---|:---|:---|:---|\n|Our Method|A|0.0024±0.0005|0.0132±0.0032|0.0054±0.0007|0.0150±0.0033|\n||B|0.0038±0.0010|0.0190±0.0016|0.0062±0.0010|0.0205±0.0016|\n||C|0.8184±0.0750|0.8548±0.0918|0.0123±0.0080|0.2460±0.0483|\n||D|0.0101±0.0003|0.0129±0.0007|0.0114±0.0015|0.0130±0.0008|\n|(BL) DRO Lasso|A|0.4107±0.0218|0.4333±0.0208|0.4115±0.0554|0.4020±0.0560|\n||B|0.4199±0.0138|0.4499±0.0190|0.4172±0.0533|0.4230±0.0511|\n||C|0.2778±0.2397|0.3782±0.1259|0.3114±0.2644|0.3816±0.1018|\n||D|0.3225±0.0252|0.3408±0.0145|0.3569±0.0259|0.3375±0.0089|\n|(BL) DRO XGB|A|0.0026±0.0002|0.0133±0.0030|0.0054±0.0014|0.0150±0.0034|\n||B|0.0037±0.0009|0.0189±0.0015|0.0064±0.0009|0.0206±0.0017|\n||C|0.4312±0.0635|0.4690±0.0652|0.0226±0.0005|0.2460±0.0479|\n||D|0.0102±0.0004|0.0127±0.0010|0.0114±0.0012|0.0132±0.0010|\n|(BL) Lasso|A|0.0510±0.0417|0.0617±0.0447|0.0517±0.0417|0.0616±0.0433|\n||B|0.0547±0.0434|0.0653±0.0402|0.0562±0.0443|0.0657±0.0390|\n||C|0.4310±0.0636|0.4611±0.0745|0.4654±0.0573|0.4581±0.0676|\n||D|0.0100±0.0003|0.0126±0.0011|0.0114±0.0014|0.0130±0.0012|\n|(BL) XGB|A|0.0027±0.0007|0.0132±0.0031|0.0051±0.0013|0.0151±0.0033|\n||B|0.0041±0.0012|0.0188±0.0015|0.0067±0.0012|0.0206±0.0015|\n||C|0.4285±0.0622|0.4637±0.0613|0.0200±0.0039|0.2471±0.0475|\n||D|0.0100±0.0002|0.0129±0.0010|0.0117±0.0017|0.0131±0.0007|\n|(BL) DRO Forward Selection|A|0.3881±0.0598|0.4151±0.0607|0.4088±0.0707|0.4066±0.0570|\n||B|0.3989±0.0510|0.4239±0.0535|0.4208±0.0555|0.4186±0.0489|\n||C|0.2765±0.2391|0.3717±0.1236|0.2980±0.2530|0.3746±0.1020|\n||D|0.3253±0.0192|0.3467±0.0118|0.3522±0.0265|0.3391±0.0083|\n|(BL) DRO Backward Elimination|A|0.3889±0.0516|0.4152±0.0548|0.4076±0.0713|0.4060±0.0579|\n||B|0.3996±0.0496|0.4309±0.0482|0.4188±0.0563|0.4191±0.0494|\n||C|0.2771±0.2395|0.3685±0.1210|0.2930±0.2507|0.3756±0.1046|\n||D|0.3242±0.0151|0.3460±0.0113|0.3519±0.0309|0.3392±0.0062|\n|(BL) Forward Selection|A|0.0506±0.0421|0.0607±0.0441|0.0531±0.0422|0.0613±0.0431|\n||B|0.0545±0.0440|0.0654±0.0403|0.0569±0.0436|0.0654±0.0389|\n||C|0.8114±0.0705|0.8523±0.0810|0.8811±0.0775|0.8445±0.0760|\n||D|0.0099±0.0007|0.0124±0.0010|0.0115±0.0013|0.0130±0.0012|\n|(BL) Backward Elimination|A|0.0504±0.0417|0.0609±0.0442|0.0537±0.0434|0.0614±0.0431|\n||B|0.0551±0.0438|0.0654±0.0408|0.0577±0.0449|0.0655±0.0388|\n||C|0.8173±0.0702|0.8558±0.0753|0.8585±0.0741|0.8437±0.0742|\n||D|0.0100±0.0003|0.0127±0.0013|0.0113±0.0020|0.0129±0.0012|\n\n## Discussion\nFor budget=$10$, our method is comparable with the best performing baseline (XGB). For the smaller budget, lasso and DRO have comparable performance, followed by our method. Forward selection and Backward Elimination consistently underperform other methods.\n\n## Expt 3: Sparse Linear\nWe use a dataset of size 36000 with the following composition:\n\n- A (35%):$$Y=5X_0+4X_{15}+3X_{30}+\\epsilon$$\n- B (35%):$$Y=6X_5+5X_{20}+4X_{35}+\\epsilon$$\n- C (30%):$$Y=7X_{10}+6X_{25}+5X_{40}+4X_{45}+\\epsilon$$\n\n**Noise**:  \n- $\\epsilon\\sim\\mathcal{N}(0, 0.1^2)$ (base noise)\n- Correlated noise features: $X_{i+1}=0.3X_i+0.7\\eta$ where $\\eta \\sim \\mathcal{N}(0,1)$\n\n### Results\n|Method|Population|MSE±Std (MLP) — Budget 5|MSE±Std (RF) — Budget 5|MSE±Std (MLP) — Budget 10|MSE±Std (RF) — Budget 10|\n|:---|:---|:---|:---|:---|:---|\n|Our Method|A|0.1186±0.0055|0.1299±0.0064|0.1274±0.0054|0.1277±0.0047|\n||B|0.2099±0.0061|0.2290±0.0086|0.0034±0.0001|0.0093±0.0014|\n||C|0.6009±0.0356|0.6453±0.0318|0.0036±0.0002|0.0282±0.0021|\n|(BL) DRO Lasso|A|0.1179±0.0053|0.1296±0.0066|0.0038±0.0003|0.0095±0.0003|\n||B|0.2115±0.0070|0.2295±0.0088|0.0033±0.0004|0.0092±0.0013|\n||C|0.6054±0.0438|0.6456±0.0310|0.0034±0.0001|0.0281±0.0021|\n|(BL) DRO XGB|A|0.1189±0.0057|0.1298±0.0066|0.0038±0.0000|0.0095±0.0004|\n||B|0.2118±0.0071|0.2286±0.0092|0.0035±0.0003|0.0092±0.0014|\n||C|0.6073±0.0432|0.6447±0.0311|0.0037±0.0004|0.0280±0.0018|\n|(BL) Lasso|A|0.4729±0.0153|0.5066±0.0185|0.0039±0.0004|0.0095±0.0004|\n||B|0.2115±0.0057|0.2309±0.0090|0.0036±0.0004|0.0092±0.0013|\n||C|0.3202±0.0029|0.3485±0.0106|0.0037±0.0001|0.0281±0.0018|\n|(BL) XGB|A|1.0188±0.0235|1.0651±0.0303|0.0036±0.0003|0.0095±0.0004|\n||B|0.1414±0.1214|0.1564±0.1285|0.0034±0.0002|0.0091±0.0013|\n||C|0.1915±0.1125|0.2121±0.1196|0.0037±0.0004|0.0280±0.0019|\n|(BL) DRO Forward Selection|A|0.5335±0.1147|0.5619±0.1252|0.0470±0.0746|0.0488±0.0687|\n||B|0.5032±0.0173|0.5417±0.0190|0.0826±0.1371|0.0830±0.1267|\n||C|0.6510±0.1143|0.6876±0.1048|0.2283±0.0848|0.2277±0.0756|\n|(BL) DRO Backward Elimination|A|0.1182±0.0049|0.1297±0.0065|0.0039±0.0005|0.0095±0.0003|\n||B|0.2118±0.0074|0.2293±0.0087|0.0035±0.0004|0.0092±0.0014|\n||C|0.6059±0.0430|0.6457±0.0309|0.0034±0.0003|0.0279±0.0019|\n|(BL) Forward Selection|A|0.4928±0.1395|0.5228±0.1508|0.3298±0.3016|0.3049±0.2713|\n||B|0.0013±0.0001|0.0075±0.0012|0.0036±0.0004|0.0091±0.0013|\n||C|0.7542±0.2417|0.7885±0.2349|0.5878±0.1616|0.5532±0.1395|\n|(BL) Backward Elimination|A|0.4056±0.2588|0.4297±0.2712|0.3312±0.3015|0.3055±0.2717|\n||B|0.3015±0.1603|0.3233±0.1707|0.0035±0.0003|0.0091±0.0012|\n||C|0.7559±0.2519|0.7844±0.2232|0.5951±0.1720|0.5529±0.1390|\n\n## Discussion\nFor budget=$5$, our method is comparable with the best performing baseline (XGB and DRO XGB). Despite the linear setting, lasso is not the best performing method. For budget=$10$, lasso and XGB variants, along with DRO Backward Elimination perform best. Our method succeeds in selecting all relevant variables except $X_{30}$, resulting in relatively worse performance on A with budget=$10$."}, {"review_id": "WWMLJUteRR", "reviewer": "Reviewer_Guhq", "summary": "The paper introduces Distributionally Robust Feature Selection (DRO-FS), a feature selection method aiming to perform well simultaneously across multiple sub-populations of the data. Instead of using masks, the authors present an interesting mechanism for feature selection: injecting white Gaussian noise into the features. An informative feature would be injected with a small variance noise, whereas a non-informative feature would be injected with high variance noise. Instead of differentiating through an inner training loop, the authors derive an objective based on the conditional variance of the Bayes-optimal predictor, which depends on the variance of the noise through a closed-form kernel weight expression. They use SGD to optimize the worst-case loss across groups. The proposed method is demonstrated using synthetic data and two real-world datasets for predicting household income. The proposed approach yields lower worst-group error than Lasso, XGBoost, and the Group-DRO variants, sometimes by an order of magnitude.", "questions": "1. Writing - please revise the writing and the derivation of the mathematical formulation, add a concise algorithm box and/or a block diagram of the proposed method to improve clarity.\n\n2. Discuss other methods for feature selection and compare their performance to the proposed approach's. \n\n3. Add other data sets AND and other downstream predictors. \n\n4. Interpretability - report the selected features and compare to other methods of feature selection to demonstrate the advantage of the proposed method.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n* Previous methods are optimized for a specific model, not for a subpopulation, meaning that they ignore distribution shifts.\n* The noise injection, along with the Bayes-risk objective, is rather cool. The method is model-agnostic, and there's no need for backpropagating through training. \n* The derivation seems correct. \n* The authors provide anonymized code\n* Training details, architecture choices, and optimization schedules are neatly provided. \n\nWeaknesses\n\n* While the code idea of the paper is interesting, it is poorly presented. The formulation is cumbersome, lacks clarity, and is hard to follow. There are no clear building blocks. Just a steady stream of claims. I think the authors should reformulate the way they derive the equations, use definitions, lemmas, and theorems, and provide proofs. Even if that would end up taking more space, the proofs could go to the appendix. Still, I think it would make the reading more fluent and add clarity to the paper.   \n\n* A concise algorithm box and/or a block diagram of the proposed method would improve clarity.\n\n* How does runtime scale with d (features) and n (samples)? \n* After learning continuous α, how is the final subset obtained? A hard top-k cut or magnitude-based sorting? I did not find an explicit description. The authors should provide a full description, as this is a fundamental issue, and discuss stability across random seeds.\n* So many other methods for feature selection since lasso were not mentioned at all! \n* The authors claim the proposed method is model agnostic, yet they demonstrate the performance of their method using only one choice of downstream model (Random Forest). To convince their choice of features is indeed robust, they should use additional downstream models that are not tree-based. \n* The synthetic task is highly nonlinear. Including a linear-target setting would indicate when Lasso-type methods might still be sufficient.\n* Report the learned sparsity pattern overlap across different random seeds to quantify selection stability.\n\n* Line 133 >=0 at the end is a sub-index of R?\n* Line 157 – incurr should be incur\n* Line 205 \"with with\"\n* Line 2015 \"involves requires\"\n* Line 222 \"By applying ..\" has no subject\n* Line 443 – \"initialized\"\n* Line 452 a missing reference -  ??\n* \"Moreoever\"  should be Moreover appears twice", "was_revised": true, "final_justification": "The authors have fully addressed all of my comments. Together with their responses to the other reviewers’ remarks, the authors resolved every concern raised during the review process. In light of these revisions, I am updating my score to 4.", "rebuttal": "We thank reviewer Guhq for their feedback and suggestions, which we address below.\n## Questions:\n> Q1: *Writing-please revise the writing and the derivation of the mathematical formulation, add a concise algorithm box and/or a block diagram of the proposed method to improve clarity.*\n\nWe thank the reviewer for their suggestions on improving Section 3. We have added an algorithm box summarizing our method in the appendix for greater clarity. We also incorporate the suggestion to present our main results as follows:\n\n**Theorem 1 (Population-Level Objective)**\nUnder the noise-based relaxation $S(\\alpha) = X + \\epsilon(\\alpha)$ where $\\epsilon(\\alpha) \\sim \\mathcal{N}(0, \\text{diag}(\\alpha))$, the distributionally robust feature selection problem: \n\n$$\\min_{\\alpha} \\max_{P_i \\in \\mathcal{P}} \\mathbb{E}_{S(\\alpha),Y \\sim P_i}[\\mathcal{L}(Y, f_i^*(S(\\alpha)))]+\\lambda \\text{Reg}(\\alpha)$$\n\nwhere $f_i^*(s)=\\mathbb{E}[Y|S(\\alpha)=s, P_i]$ is the Bayes-optimal predictor, is equivalent to:\n\n$$\\min_{\\alpha} \\max_{P_i \\in \\mathcal{P}} ( -\\mathbb{E}_{S(\\alpha) \\sim P_i}[\\mathbb{E}[\\mu_i(X)|S(\\alpha)]^2] )+\\lambda \\mathrm{Reg}(\\alpha)$$\n\nwhere $\\mu_{i}(X)=\\mathbb{E}[Y|X, P_i]$.\n\n---\nGiven samples $\\{(X_i^j\\, Y_i^j)\\}_{j=1}^{n_i}$ , from each population $P_i$, let $\\hat{\\mu}_i(X)$ be an estimator of $\\mu_i(X) = \\mathbb{E}[Y|X]$ trained on these samples. The empirical form of the population objective from Theorem 1 is:\n\n$$\\min_{\\alpha} \\max_{P_i \\in \\mathcal{P}} \\left( -\\hat{\\mathbb{E}}_{S(\\alpha)}\\left[\\hat{\\mathbb{E}}[\\hat{\\mu}_i(X)|S(\\alpha)]^2\\right] \\right)+\\lambda \\text{Reg}(\\alpha)$$\n\nwhere the expectation $\\hat{\\mathbb{E}}[\\hat{\\mu_i}(X)|S(\\alpha)]$ is taken with respect to the empirical distribution $\\hat{P_i}(X)=\\frac{1}{n_i}\\sum_{j=1}^{n_i} \\delta_{X_i^j}(X)$.\n\n**Theorem 2 (Kernel Form Equivalence)**\n\nThe empirical expectation \n\n$$\\hat{\\mathbb{E}}_{S(\\alpha)}\\left[\\hat{\\mathbb{E}}[\\hat{\\mu_i}(X)|S(\\alpha)]^2\\right]$$\n\nis equivalent to:\n\n$$\\hat{\\mathbb{E}}\\_{S(\\alpha)}[(\\sum_{j=1}^{n_i} w_i^j (S(\\alpha),\\alpha) \\hat{\\mu}_i (X_i^j))^2]$$\n\nwhere the weights are:\n\n$$w_i^j(S(\\alpha), \\alpha)=\\frac{\\exp(\\{-\\frac{1}{2}(X_i^j-S(\\alpha))^T \\text{diag}(\\alpha)^{-1}(X_i^j-S(\\alpha))\\})}{\\sum_{k=1}^{n_i} \\exp(\\{-\\frac{1}{2}(X_i^k-S(\\alpha))^T \\text{diag}(\\alpha)^{-1}(X_i^k-S(\\alpha))\\})}$$\n\n---\n\nThe proofs of these results are exactly the step-by-step derivations currently presented in Section 3. We will create a section in the appendix where we expand the proofs further to improve clarity.\n\n> Q2: *Discuss other methods for feature selection and compare their performance to the proposed approach's.*\n\nWe now include Forward Selection (FS) and Backward Elimination (BE) (+ their DRO variants) as baselines in our experiments. \nThe implementations use k-fold cross-validation to score feature subsets. FS iteratively adds the feature that most improves the cross-validation score. BE starts with all features and iteratively removes the one whose absence hurts performance the least.\nThe DRO variants also follow this greedy logic but optimize for robustness- evaluating feature subsets by finding the worst-case cross-validation score across all individual data populations.\nWe have added experimental results on a new synthetic dataset that includes these baselines to this rebuttal. Our method generally outperforms these baselines across the evaluated metric.\n\n> Q3: *Add other data sets AND and other downstream predictors.*\n\nWe now include an MLP (single hidden layer of size 100) as a downstream model in our evaluations, and will do so for all the experiments in the final paper. We have conducted new synthetic experiments-we include only one of them here due to character limits, however our final paper will include more experiments. Pl. see the last section of our rebuttal for the new experiment.\n\n> Q4: *Interpretability-report the selected features and compare to other methods of feature selection to demonstrate the advantage of the proposed method.*\n\nWe will add the selected features for each run in the appendix. Our method is stable across seeds in the subsets of variables that it selects and the order in which it selects them, differing only in a few values. Pl. see the last section of our rebuttal for an example.\n\n## Additional comments\n> C1: *How does runtime scale with d (features) and n (samples)?*\n\nOur method requires $O(P \\cdot b \\cdot n \\cdot K \\cdot d)$ operations per iteration, where $P$ is the number of populations, $b$ is the number of Monte Carlo samples, $n=\\max_p n_p$ is the maximum population size, $K$ is the number of nearest neighbors used for kernel weight computation, and $d$ is the feature dimensionality. \n\n> C2: *After learning continuous α, how is the final subset obtained?* \n\nThe final subset is obtained by selecting those variables $i$ with the smallest $\\alpha_{i}$ values (corresponding to minimum noise). We will explicitly state this in our paper.\n\n> C3: *Experiments: Use additional downstream models that are not tree-based, Including a linear-target setting would indicate when Lasso-type methods might still be sufficient. Report the learned sparsity pattern overlap across different random seeds to quantify selection stability.*\n\nOur new experiments include:\n- Additional downstream model for evaluating performance: we now include an MLP (single hidden layer of size 100), along with the existing random forest (RF), as a downstream task model. \n- New baseline methods: we've added Forward Selection and Backward Elimination, with their DRO variants\n- Additional data generation settings to explore different covariate-outcome functional relationships. Three of these are included in our response to reviewer ggn5: a purely linear model, a mixed linear-nonlinear model with heterogeneous noise, and a sparse linear model with correlated noise variables. Our method outperforms the baselines or closely follows the best performing ones. We include only the first here due to character limits.\n\n> C4: *Typos*\n\nWe thank the reviewer for bringing these typos to our notice, we have corrected them. Wrt ``Line 133 >=0 at the end is a sub-index of R?``, we use this notation for $m$ dimensional vectors whose entries are non-negative real numbers.\n\n# New Synthetic Experiment: Linear generative process\nWe set dimension $d=15$.\n\n**Population Structure**:\n- A (40%): $$ Y = 8X_0 + 6X_1 - 4X_2 + 3X_3 + 2X_4 + \\epsilon $$\n- B (35%): $$ Y = -8X_0 - 6X_1 + 4X_2 - 3X_3 - 2X_4 + 8X_5 + 6X_6 + \\epsilon $$\n- C (25%): $$ Y = 10X_7 + 8X_8 + 6X_9 - 5X_{10} + \\epsilon $$\n$$\\text{Noise:} \\quad \\epsilon \\sim \\mathcal{N}(0, 0.1^2)$$\n\nThe train-test-val sets for Pops A, B, andChad 14400, 12600, and 9000 samples respectively, and we used a 40-60 train-(test+val) split for each population.\n\n|Method|Population|MSE±Std (MLP) — Budget 5|MSE±Std (RF) — Budget 5|MSE±Std (MLP) — Budget 10|MSE±Std (RF) — Budget 10|\n|---|---|---|---|---|---|\n|Our Method|A|0.2251±0.0029|0.2486±0.0018|0.0342±0.0004|0.0535±0.0010|\n||B|0.2888±0.0113|0.3194±0.0098|0.0207±0.0005|0.0804±0.0036|\n||C|0.2787±0.0150|0.3039±0.0130|0.0039±0.0004|0.0339±0.0079|\n|(BL) DRO Lasso|A|0.5071±0.0129|0.5434±0.0191|0.0338±0.0009|0.0534±0.0009|\n||B|0.2886±0.0086|0.3194±0.0026|0.0211±0.0001|0.0807±0.0041|\n||C|0.2755±0.0150|0.3031±0.0113|0.0039±0.0004|0.0340±0.0078|\n|(BL) DRO XGB|A|0.5107±0.0113|0.5441±0.0188|0.0738±0.0370|0.0887±0.0327|\n||B|0.2888±0.0081|0.3187±0.0030|0.0408±0.0198|0.0928±0.0145|\n||C|0.2762±0.0151|0.3032±0.0112|0.0036±0.0003|0.0337±0.0076|\n|(BL) Lasso|A|1.0087±0.0153|1.0571±0.0107|0.2056±0.1714|0.2106±0.1625|\n||B|0.5744±0.0051|0.6150±0.0052|0.1160±0.0921|0.1460±0.0763|\n||C|0.1080±0.0107|0.1283±0.0102|0.0036±0.0003|0.0340±0.0076|\n|(BL) XGB|A|1.0136±0.0076|1.0623±0.0039|0.4051±0.2802|0.3905±0.2579|\n||B|0.5702±0.0008|0.6170±0.0076|0.2279±0.1586|0.2373±0.1307|\n||C|0.1314±0.0291|0.1538±0.0352|0.0037±0.0002|0.0338±0.0077|\n|(BL) DRO Forward Selection|A|0.2257±0.0030|0.2498±0.0005|0.0344±0.0009|0.0534±0.0012|\n||B|0.3274±0.0747|0.3596±0.0808|0.0208±0.0010|0.0806±0.0041|\n||C|0.3752±0.1545|0.3991±0.1581|0.0037±0.0002|0.0341±0.0079|\n|(BL) DRO Backward Elimination|A|0.0314±0.0008|0.0492±0.0008|0.0341±0.0007|0.0535±0.0010|\n||B|0.1833±0.0047|0.2129±0.0066|0.0209±0.0006|0.0806±0.0037|\n||C|1.0281±0.0344|1.0455±0.0548|0.0038±0.0002|0.0342±0.0079|\n|(BL) Forward Selection|A|1.0183±0.0147|1.0430±0.0118|1.0887±0.0176|1.0132±0.0219|\n||B|0.5745±0.0043|0.6221±0.0061|0.6430±0.0226|0.5956±0.0086|\n||C|1.0298±0.0587|1.0571±0.0500|0.3572±0.3066|0.3407±0.2694|\n|(BL) Backward Elimination|A|0.9972±0.0199|1.0427±0.0126|1.1153±0.0405|1.0136±0.0219|\n||B|0.5672±0.0081|0.6077±0.0068|0.6486±0.0288|0.5954±0.0068|\n||C|0.4222±0.2690|0.4523±0.2803|0.3514±0.3013|0.3399±0.2682|\n---\n\n## Discussion\nHere, although the generative process is linear and variables $X_0$ to $X_4$ have strong effects in both A and B, the signs of their coefficients are reversed between populations. This reduces the effectiveness of LASSO, which tends to select features based on average effects across all data. As a result, vanilla LASSO does not achieve the best performance even in this linear setting.\n\nFor budget=$5$, our method outperforms most baselines, and has a balanced performance across population. For budget=$10$, our method is comparable with the best performing baselines.\n\nWe select those features $i$ whose corresponding $\\alpha_{i}$ values are lowest.\nOur method consistently selects a similar core ordered set of features across (3) different seeds (in decreasing order of importance) ``[0, 7, 1, 5, 8, 6, 9, 2, 10, 3], [0, 7, 1, 8, 5, 9, 6, 2, 10, 3],[0, 7, 5, 1, 8, 6, 9, 2, 10, 3]``. In contrast, baseline methods, especially non-DRO versions, show significant variability, often selecting irrelevant noisy features (e.g., 13, 14) and failing to consistently identify the true signal variables."}, {"review_id": "4xE49x2L29", "reviewer": "Reviewer_789K", "summary": "The paper under review proposes a new method for distributional robust feature selection, using an insightful continuous relaxation of an intractable, non-convex problem. The authors provide helpful intuition for their method and evaluate it via simulations and real data analysis.", "questions": "1. The form of the relaxation of problem (1) as detailed in Section 3 is tailored for a square loss. However, the authors use their method in the experimentation with the cross entropy loss to satisfying effect. What is the authors' understanding of the validity of their proposed procedure under different loss functions or distributions on the outcome $Y$? How does the performance of the proposed method change under different, perhaps non-smooth loss functions?\n\n2. How does the proposed method fare under larger dimensions? The simulations are run with only 14 features, and the real data examples are not much larger. Is there a severe computational burden in higher dimensions? It would be very helpful to see a table of runtime statistics, since this is nominally an optimization paper.\n\n3. How does the method perform under different data generating mechanisms between the train and test data? More diverse simulations would be extremely valuable to this paper.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "The primary strength of the paper is its intuitive derivation of the method. Section 3 (Methods) is well written and clear. Furthermore, the authors present a very early (if not explicitly the first) attempt at a solution for distributionally robust feature selection. Since distributionally robust methods are growing in popularity, this constitutes a valuable contribution to the machine learning field.\n\nHowever, the evaluation of the proposed method leaves much to be desired. The authors only evaluate the methods under one data generating mechanism, with the test set being drawn using the same functional relationship between covariates and outcome as the test data. The advantage of DRO in principle is that good performance is guaranteed for nearly arbitrary distribution shifts - this is not explored at all in the simulations. I would like to see a more expansive suite of simulations exploring different functions forms between the covariates and outcome, as well as varying degrees of distribution shift between the training and test datasets. Since the authors do not provide any theoretical guarantees for their proposed method, I would like to see more empirical evidence for the effectiveness of their method in practice.", "was_revised": true, "final_justification": "In the rebuttal period, the authors provided strong simulation results in response to my questions, as well as the points raised by the other reviewers. I think that these simulation results, which cover a wider range of scenarios than those explored in the original version of the paper, improve the evidence of the efficacy of the proposed method. I still hold minor reservations over the fact that the proposed algorithm is derived using the bias-variance decomposition of the MSE, which is why I keep my score at 4 rather than raise it higher. However, I understand that this would constitute a non-trivial extension to the submitted paper and is not feasible to address within the rebuttal period.", "rebuttal": "We thank reviewer 789K for their feedback and suggestions, which we address below.\n\n## Questions:\n> Q1: *The form of the relaxation of problem (1) as detailed in Section 3 is tailored for a square loss. However, the authors use their method in the experimentation with the cross entropy loss to satisfying effect. What is the authors' understanding of the validity of their proposed procedure under different loss functions or distributions on the outcome? How does the performance of the proposed method change under different, perhaps non-smooth loss functions?*\n\nOur derivation uses the bias-variance decomposition specific to MSE. However, there exist generalized variance decompositions for other loss functions. Using such generalized decompositions could allow our method to directly optimize cross-entropy or other losses, but this is a nontrivial extension we leave for future work. Our empirical results with cross-entropy (UCI experiment) indeed suggest that optimizing MSE in our method may transfer well to other losses, particularly losses that are proper scoring rules akin to MSE.\n\n> Q2: *How does the proposed method fare under larger dimensions? The simulations are run with only 14 features, and the real data examples are not much larger. Is there a severe computational burden in higher dimensions? It would be very helpful to see a table of runtime statistics, since this is nominally an optimization paper.*\n\nOur method requires $O(P\\cdot b\\cdot n\\cdot K\\cdot d)$ operations per iteration, where $P$ is the number of populations, $b$ is the number of Monte Carlo samples, $n =\\max_p n_p$ is the maximum population size, $K$ is the number of nearest neighbors used for kernel weight computation, and feature dimensionality $d$.\n\nWe now include a synthetic experiment where we vary the dimension of the data under the same generative process (increasing the number or irrelevant/noise variables). Pl. see the last section of our rebuttal for the experiment results.\nThe runtime statistics per epoch (averaged across 3 runs) for this experiment are given in the table below:\n|Dimension (d)|Time per epoch (s) (Mean±Std)|\n|---|---|\n|15|1.8988±0.0026|\n|50|1.9217±0.0044|\n\nFor $d=15$, we run 50-64 epochs, while for $d=50$, we run 126-150 epochs.\n\n> Q3: *How does the method perform under different data generating mechanisms between the train and test data? More diverse simulations would be extremely valuable to this paper.*\n\nWe wish to clarify that our method is designed for Group DRO, which we discuss in our related work section (Section 1.1, citing Sagawa et al. [2019]). Group-DRO optimizes for worst-case performance across a priori known subpopulations (e.g., demographic groups, different hospitals). This setting is distinct from DRO over arbitrary distribution shifts. Thus our experimental setup uses predefined groups rather than exploring arbitrary functional form changes.\n\nWe have run new synthetic experiments with different generative processes such as a purely linear model, a mixed linear-nonlinear model with heterogeneous noise, and a sparse linear model with correlated noise variables. Across these, our method outperforms the baselines or closely follows the best performing ones. We include only the first experiment in the last section of this rebuttal due to character limits. Our final paper will have more experiments.\n\n## Additional comments\n\n> C1: *The test set being drawn using the same functional relationship between covariates and outcome as the test data. The advantage of DRO in principle is that good performance is guaranteed for nearly arbitrary distribution shifts - this is not explored at all in the simulations. I would like to see a more expansive suite of simulations exploring different functions forms between the covariates and outcome, as well as varying degrees of distribution shift between the training and test datasets. Since the authors do not provide any theoretical guarantees for their proposed method, I would like to see more empirical evidence for the effectiveness of their method in practice.*\n\nWe request the reviewer to refer to our response to Q3. Additionally, our algorithmic framework can be extended to more general f-divergence DRO. This would involve fitting a series of conditional mean models for the worst-case distribution at each optimization step, rather than one model per predefined group. The core procedure remains the same, highlighting the flexibility of our approach. We agree that exploring these broader distribution shift scenarios is a valuable direction for future work.\n\nWrt more simulations, we have run new experiments on synthetic datasets. We include only one of them here due to character limits. However, you can find the other experiments in our response to reviewer ggn5. They also include:\n- Additional downstream model for evaluating performance: we now include an MLP, along with the existing random forest (RF), as a downstream task model.\n- New baseline methods: we've added Forward Selection and Backward Elimination, with their DRO variants, (along with the existing baseline methods in the original paper).\n- Additional data generation settings to explore different covariate-outcome functional relationships.\n\n# Experiment\n**Population Structure**:\n- A (40%): $$ Y = 8X_0 + 6X_1 - 4X_2 + 3X_3 + 2X_4 + \\epsilon $$\n- B (35%): $$ Y = -8X_0 - 6X_1 + 4X_2 - 3X_3 - 2X_4 + 8X_5 + 6X_6 + \\epsilon $$\n- C (25%): $$ Y = 10X_7 + 8X_8 + 6X_9 - 5X_{10} + \\epsilon $$\n$$\\text{Noise:} \\quad \\epsilon \\sim \\mathcal{N}(0, 0.1^2)$$\n\nWe use a dataset of size 36000, with a 40-60 train-(test/val) split for each population. We report the results for two feature budgets (5 and 10).\n\n|Dimension|Method|Population|MSE±Std (MLP) — Budget 5|MSE±Std (RF) — Budget 5|MSE±Std (MLP) — Budget 10|MSE±Std (RF) — Budget 10|\n|---|---|---|---|---|---|---|\n|15|Our Method|A|0.2251±0.0029|0.2486±0.0018|0.0342±0.0004|0.0535±0.0010|\n|||B|0.2888±0.0113|0.3194±0.0098|0.0207±0.0005|0.0804±0.0036|\n|||C|0.2787±0.0150|0.3039±0.0130|0.0039±0.0004|0.0339±0.0079|\n|50||A|0.2227±0.0094|0.2420±0.0161|0.0241±0.0178|0.0450±0.0129|\n|||B|0.5805±0.0375|0.6158±0.0448|0.0150±0.0096|0.0809±0.0035|\n|||C|0.1208±0.0034|0.1388±0.0011|0.0480±0.0761|0.0697±0.0605|\n|15|(BL) DRO Lasso|A|0.5071±0.0129|0.5434±0.0191|0.0338±0.0009|0.0534±0.0009|\n|||B|0.2886±0.0086|0.3194±0.0026|0.0211±0.0001|0.0807±0.0041|\n|||C|0.2755±0.0150|0.3031±0.0113|0.0039±0.0004|0.0340±0.0078|\n|50||A|0.4155±0.1769|0.4418±0.1866|0.0345±0.0004|0.0511±0.0029|\n|||B|0.2377±0.0909|0.2677±0.0972|0.0206±0.0006|0.0835±0.0067|\n|||C|0.3642±0.1350|0.3849±0.1300|0.0040±0.0007|0.0354±0.0012|\n|15|(BL) DRO XGB|A|0.5107±0.0113|0.5441±0.0188|0.0738±0.0370|0.0887±0.0327|\n|||B|0.2888±0.0081|0.3187±0.0030|0.0408±0.0198|0.0928±0.0145|\n|||C|0.2762±0.0151|0.3032±0.0112|0.0036±0.0003|0.0337±0.0076|\n|50||A|0.6827±0.3136|0.7120±0.3138|0.1041±0.0645|0.1139±0.0597|\n|||B|0.3857±0.1461|0.4209±0.1495|0.0613±0.0416|0.1098±0.0326|\n|||C|0.2259±0.0928|0.2479±0.0926|0.0040±0.0005|0.0350±0.0010|\n|15|(BL) Lasso|A|1.0087±0.0153|1.0571±0.0107|0.2056±0.1714|0.2106±0.1625|\n|||B|0.5744±0.0051|0.6150±0.0052|0.1160±0.0921|0.1460±0.0763|\n|||C|0.1080±0.0107|0.1283±0.0102|0.0036±0.0003|0.0340±0.0076|\n|50||A|1.0181±0.0526|1.0429±0.0387|0.4126±0.5303|0.3859±0.4760|\n|||B|0.5880±0.0323|0.6264±0.0284|0.2327±0.2912|0.2504±0.2401|\n|||C|0.1351±0.0252|0.1550±0.0267|0.0037±0.0001|0.0354±0.0009|\n|15|(BL) XGB|A|1.0136±0.0076|1.0623±0.0039|0.4051±0.2802|0.3905±0.2579|\n|||B|0.5702±0.0008|0.6170±0.0076|0.2279±0.1586|0.2373±0.1307|\n|||C|0.1314±0.0291|0.1538±0.0352|0.0037±0.0002|0.0338±0.0077|\n|50||A|1.0211±0.0474|1.0421±0.0400|1.1132±0.0439|1.0216±0.0507|\n|||B|0.5867±0.0319|0.6258±0.0282|0.6607±0.0314|0.6135±0.0252|\n|||C|0.1356±0.0261|0.1551±0.0267|0.0036±0.0004|0.0347±0.0012|\n|15|(BL) DRO Forward Selection|A|0.2257±0.0030|0.2498±0.0005|0.0344±0.0009|0.0534±0.0012|\n|||B|0.3274±0.0747|0.3596±0.0808|0.0208±0.0010|0.0806±0.0041|\n|||C|0.3752±0.1545|0.3991±0.1581|0.0037±0.0002|0.0341±0.0079|\n|50||A|0.3735±0.1506|0.3996±0.1636|0.0341±0.0006|0.0514±0.0029|\n|||B|0.3729±0.0553|0.3995±0.0553|0.0209±0.0012|0.0833±0.0066|\n|||C|0.4735±0.0805|0.5124±0.0883|0.0038±0.0001|0.0354±0.0012|\n|15|(BL) DRO Backward Elimination|A|0.0314±0.0008|0.0492±0.0008|0.0341±0.0007|0.0535±0.0010|\n|||B|0.1833±0.0047|0.2129±0.0066|0.0209±0.0006|0.0806±0.0037|\n|||C|1.0281±0.0344|1.0455±0.0548|0.0038±0.0002|0.0342±0.0079|\n|50||A|0.0320±0.0012|0.0474±0.0024|0.0346±0.0005|0.0512±0.0030|\n|||B|0.1757±0.0093|0.2079±0.0149|0.0208±0.0008|0.0834±0.0068|\n|||C|1.0008±0.0669|1.0355±0.0635|0.0038±0.0003|0.0354±0.0012|\n|15|(BL) Forward Selection|A|1.0183±0.0147|1.0430±0.0118|1.0887±0.0176|1.0132±0.0219|\n|||B|0.5745±0.0043|0.6221±0.0061|0.6430±0.0226|0.5956±0.0086|\n|||C|1.0298±0.0587|1.0571±0.0500|0.3572±0.3066|0.3407±0.2694|\n|50||A|1.0132±0.0372|1.0408±0.0435|1.1209±0.0483|1.0224±0.0403|\n|||B|0.5879±0.0331|0.6276±0.0384|0.6509±0.0360|0.6109±0.0362|\n|||C|0.9808±0.1161|1.0040±0.0989|1.1410±0.1413|0.9882±0.1095|\n|15|(BL) Backward Elimination|A|0.9972±0.0199|1.0427±0.0126|1.1153±0.0405|1.0136±0.0219|\n|||B|0.5672±0.0081|0.6077±0.0068|0.6486±0.0288|0.5954±0.0068|\n|||C|0.4222±0.2690|0.4523±0.2803|0.3514±0.3013|0.3399±0.2682|\n|50||A|1.0095±0.0393|1.0463±0.0447|1.1267±0.0415|1.0244±0.0405|\n|||B|0.5897±0.0341|0.6267±0.0300|0.6640±0.0208|0.6111±0.0364|\n|||C|0.9883±0.1042|1.0071±0.0912|1.1047±0.1156|0.9895±0.1071|\n---\n\n## Discussion\nAlthough the generative process here is linear and variables $X_0$ to $X_4$ have strong effects in both A and B, the signs of their coefficients are reversed between populations. This reduces the effectiveness of LASSO, which tends to select features based on average effects across all data. As a result, vanilla LASSO does not achieve the best performance even in this linear setting.\n\nFor budget=$5$, our method outperforms most baselines, and has a balanced performance across populations. For budget=$10$, our method is comparable with the best performing baselines."}]

Author rebuttal:

Meta-review: This paper proposes a novel approach for feature selection in a distributionally robust context. The problem is framed as a continuous relaxation of hard variable selection, and relies on noise injection. The method decouples feature selection from downstream model training, and is thus applicable to non-differentiable models such as decision trees and random forests. 

The reviewers concur that this is a relevant problem, at the intersection of feature selection and out-of-distribution generalization. The proposed method is technically sound. Empirical results on synthetic and real-world datasets demonstrate improvements over several baselines, in some cases by large margins. The availability of anonymized code and clear reporting of training details further strengthen the contribution. The experimental evaluation could consider more distribution shifts and more varied downstream models. Additional empirical evidence such as the one provided for the rebuttal will help support the claims. Also, an assessment of the robustness of the approach with respect to the Bayesian optimal predictor surrogate would be welcomed. 

The consensus is that the paper offers novel and impactful contributions. Its formulation of distributionally robust feature selection is original, the proposed formalization is simple and the approach is generic and efficient. While further empirical validation would further strengthen the work, the novelty and potential impact of the paper justify acceptance at NeurIPS.
Decision: Accept (poster)