We thank reviewer 789K for their feedback and suggestions, which we address below.

Questions:

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?

Our 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.

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.

Our 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$.

We 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. The runtime statistics per epoch (averaged across 3 runs) for this experiment are given in the table below:

For $d=15$, we run 50-64 epochs, while for $d=50$, we run 126-150 epochs.

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.

We 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.

We 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.

Additional comments

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.

We 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.

Wrt 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:

• Additional downstream model for evaluating performance: we now include an MLP, along with the existing random forest (RF), as a downstream task model.
• 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).
• Additional data generation settings to explore different covariate-outcome functional relationships.

Experiment

Population Structure:

• A (40%): $Y = 8X_0 + 6X_1 - 4X_2 + 3X_3 + 2X_4 + \epsilon$
• B (35%): $Y = -8X_0 - 6X_1 + 4X_2 - 3X_3 - 2X_4 + 8X_5 + 6X_6 + \epsilon$
• C (25%): $Y = 10X_7 + 8X_8 + 6X_9 - 5X_{10} + \epsilon$ $\text{Noise:} \quad \epsilon \sim \mathcal{N}(0, 0.1^2)$

We 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).

---

Discussion

Although 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.

For 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.

We thank reviewer Guhq for their feedback and suggestions, which we address below.

Questions:

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.

We 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:

Theorem 1 (Population-Level Objective) Under 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:

$\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)$

where $f_i^*(s)=\mathbb{E}[Y|S(\alpha)=s, P_i]$ is the Bayes-optimal predictor, is equivalent to:

$\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)$

where $\mu_{i}(X)=\mathbb{E}[Y|X, P_i]$.

---

Given 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:

$\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)$

where 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)$.

Theorem 2 (Kernel Form Equivalence)

The empirical expectation

$\hat{\mathbb{E}}_{S(\alpha)}\left[\hat{\mathbb{E}}[\hat{\mu_i}(X)|S(\alpha)]^2\right]$

is equivalent to:

$\hat{\mathbb{E}}\_{S(\alpha)}[(\sum_{j=1}^{n_i} w_i^j (S(\alpha),\alpha) \hat{\mu}_i (X_i^j))^2]$

where the weights are:

$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))\})}$

---

The 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.

Q2: Discuss other methods for feature selection and compare their performance to the proposed approach's.

We now include Forward Selection (FS) and Backward Elimination (BE) (+ their DRO variants) as baselines in our experiments. The 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. The 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. We 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.

Q3: Add other data sets AND and other downstream predictors.

We 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.

Q4: Interpretability-report the selected features and compare to other methods of feature selection to demonstrate the advantage of the proposed method.

We 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.

Additional comments

C1: How does runtime scale with d (features) and n (samples)?

Our 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.

C2: After learning continuous α, how is the final subset obtained?

The 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.

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.

Our new experiments include:

• 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.
• New baseline methods: we've added Forward Selection and Backward Elimination, with their DRO variants
• 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.

C4: Typos

We 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.

New Synthetic Experiment: Linear generative process

We set dimension $d=15$.

Population Structure:

• A (40%): $Y = 8X_0 + 6X_1 - 4X_2 + 3X_3 + 2X_4 + \epsilon$
• B (35%): $Y = -8X_0 - 6X_1 + 4X_2 - 3X_3 - 2X_4 + 8X_5 + 6X_6 + \epsilon$
• C (25%): $Y = 10X_7 + 8X_8 + 6X_9 - 5X_{10} + \epsilon$ $\text{Noise:} \quad \epsilon \sim \mathcal{N}(0, 0.1^2)$

The 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.

---

Discussion

Here, 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.

We select those features $i$ whose corresponding $\alpha_{i}$ values are lowest. Our 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.

We thank the reviewer for their response. We recognize the importance of comparing our method with other embedded methods. To this end, we add a new 'dro-embedded-mlp' baseline that uses an MLP with a learnable feature mask trained via DRO. We include the tabulated results at the end of this comment.

With respect to the papers shared by the reviewer, we acknowledge these as new works in the space of feature selection, and will include them in our discussion on related works. However, these methods address different problem settings from our Group-DRO setting, and thus cannot be directly translated to our task: Yang et al. (LSPIN) and Svirsky & Lindenbaum (IDC) select different features for individual samples and clusters respectively—while our method requires selecting a single global feature subset that works universally across all known population groups (e.g., health systems often adopt screeners universally for all patients instead of asking an entirely custom set of questions for each patient or subgroup).

Both methods allow different features for different samples/clusters, which does not align with our constraint of selecting a fixed feature subset for universal collection across populations. These represent complementary approaches to feature selection rather than direct alternatives to our distributionally robust framework. The new baselines implemented by us adapt the approach of the first paper suggested by the reviewer. However, instead of a sample-specific mask, we learn a global mask over features for all samples to better suit our Group DRO setting.

Experiments with embedded baselines

To the results shared in our rebuttal we now add two new baselines (Embedded MLP 1, Embedded MLP 2). Both use an MLP with a learnable feature mask trained via DRO. MLP 1 has a single hidden layer of size 100, while MLP 2 has two hidden layers of sizes 64 and 32. We train both models with the joint objective of MSE minimization (for the regression task), and L1 regularization (weighted by hyperparameter $\lambda=0.01$ ; we found this $\lambda$ value to work best out of [0.1, 0.01, 0.001]). Our method outperforms these baselines as well, with lower variance.

We thank reviewer ggn5 for their feedback.

We have conducted additional synthetic experiments and include three of them in this rebuttal. Our new experiments include:

• 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.
• New baseline methods: we've added Forward Selection and Backward Elimination, and their DRO variants
• Additional data generation settings to explore different covariate-outcome functional relationships.

Experiments

For all experiments, we set dimension $d=50$, and use a 40-60 train-(test+val) split for each population

Expt 1: Linear generative process

We use a dataset of size 36000 with the following composition:

• A (40%): $Y = 8X_0 + 6X_1 - 4X_2 + 3X_3 + 2X_4 + \epsilon$
• B (35%): $Y = -8X_0 - 6X_1 + 4X_2 - 3X_3 - 2X_4 + 8X_5 + 6X_6 + \epsilon$
• C (25%): $Y = 10X_7 + 8X_8 + 6X_9 - 5X_{10} + \epsilon$ $\text{Noise:} \quad \epsilon \sim \mathcal{N}(0,0.1^2)$

---

Discussion

Here, 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.

When 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.

Expt 2: Mixed Linear-Nonlinear with Heterogeneous Noise

We use a dataset of size 44000 with the following population structure:

• A (30%):$Y=4X_0+3X_1+X_2^2+\epsilon_A$
• B (30%):$Y=4X_0+3X_1+X_2^2+\epsilon_B$
• C (25%):$Y=2X_0+3X_5 X_6+4\sin(2X_7)+\epsilon_C$
• D (15%):$Y=3X_0+2X_1+\epsilon_D$

Heterogeneous Noise:

• $\epsilon_A \sim \mathcal{N}(0, (0.05)^2)$ (reduced noise)
• $\epsilon_B=\exp(0.5X_3+0.3X_4) \cdot \eta \cdot 0.1$ where, $\eta \sim \mathcal{N}(0,1)$ (heteroscedastic)
• $\epsilon_C \sim \mathcal{N}(0, 0.1^2)$ (standard)
• $\epsilon_D \sim t_3 \cdot 0.2$ (heavy-tailed, t-distribution with df=3)

Results

Discussion

For 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.

Expt 3: Sparse Linear

We use a dataset of size 36000 with the following composition:

• A (35%):$Y=5X_0+4X_{15}+3X_{30}+\epsilon$
• B (35%):$Y=6X_5+5X_{20}+4X_{35}+\epsilon$
• C (30%):$Y=7X_{10}+6X_{25}+5X_{40}+4X_{45}+\epsilon$

Noise:

• $\epsilon\sim\mathcal{N}(0, 0.1^2)$ (base noise)
• Correlated noise features: $X_{i+1}=0.3X_i+0.7\eta$ where $\eta \sim \mathcal{N}(0,1)$

Results

Discussion

For 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$.

We thank the Area Chair for bringing to our attention the reference [4] (Grandvalet, 2000), and for the broader comment on literature on noise injection for regularization. We originally meant to claim that the use of noise was novel as a relaxation of 0-1 variable selection (as opposed to the broader use of noise as a form of regularization to improve predictive performance). However, we were not aware of [4] (Grandvalet, 2000) and acknowledge this oversight in our literature review. We do wish to emphasize that all of the main technical developments (Sections 3.2-3.4) are still novel.

We will revise the related work section accordingly. Specifically, we will replace the sentence on lines 131ff with the following, more detailed explanation of how our contributions relate to prior work and where they differ:

Previous work has proposed to use the injection of random noise to identify the most important variables for a predictive task in a single population [Grandvalet, 2000], as well as through Bayesian relevance estimation methods [Neal, 1996; Tipping, 2001]. However, tunable-noise-based variable selection has seen limited adoption since it was originally proposed, in contrast to the large body of work on noise as a form of regularization [Bishop, 1995]. In this work, we revisit noise-based relaxations in the distributionally robust setting and show how the optimization problem can be reformulated in ways that create significant, previously-unrecognized advantages. In particular, our approach separates variable selection from fitting the predictive model, so that variable selection is agnostic with respect to the downstream model used for each different subpopulation. This also avoids the need to differentiate through the predictive models, as required in previous work (which may be impossible in the case of frequently-employed models like decision trees or random forests).

Additionally, we wish to highlight that our contributions remain distinct from prior works in the following ways:

1. Distributionally robust setting: previous work does not discuss the DRO setting, the main motivation for our paper
2. Model agnostic framework: Our approach works with arbitrary downstream models (including non-differentiable ones like random forests) by targeting the performance of the Bayes-optimal predictor rather than specific trained models.
3. Training decoupled optimization: We provide a tractable, gradient based optimization framework directly over the task of feature selection, which can make use of a single set of pretrained models as input.