Environment-Adaptive Covariate Selection: Learning When to Use Spurious Correlations for Out-of-Distribution Prediction

TL;DR

This work addresses the gap where invariant or purely causal covariate selection underperforms ERM in out-of-distribution (OOD) prediction when not all causal factors are observed. It shows that proxy covariates can be informative under some shifts but unreliable under others, making the optimal covariate set environment-dependent. The authors propose Environment-Adaptive Covariate Selection (EACS), which learns a mapping from unlabeled target-environment covariate summaries to environment-specific covariate subsets, with discrete and soft-gating implementations and theoretical guarantees. They further integrate prior causal knowledge as constraints to stabilize selection and demonstrate improved predictive performance across simulations and two real-world datasets (bike sharing and ACS Income). Overall, EACS offers a principled, data-driven alternative to fixed causal/invariant rules, enabling robust OOD predictions by leveraging observable signatures of distribution shifts.

Abstract

Out-of-distribution (OOD) prediction is often approached by restricting models to causal or invariant covariates, avoiding non-causal spurious associations that may be unstable across environments. Despite its theoretical appeal, this strategy frequently underperforms empirical risk minimization (ERM) in practice. We investigate the source of this gap and show that such failures naturally arise when only a subset of the true causes of the outcome is observed. In these settings, non-causal spurious covariates can serve as informative proxies for unobserved causes and substantially improve prediction, except under distribution shifts that break these proxy relationships. Consequently, the optimal set of predictive covariates is neither universal nor necessarily exhibits invariant relationships with the outcome across all environments, but instead depends on the specific type of shift encountered. Crucially, we observe that different covariate shifts induce distinct, observable signatures in the covariate distribution itself. Moreover, these signatures can be extracted from unlabeled data in the target OOD environment and used to assess when proxy covariates remain reliable and when they fail. Building on this observation, we propose an environment-adaptive covariate selection (EACS) algorithm that maps environment-level covariate summaries to environment-specific covariate sets, while allowing the incorporation of prior causal knowledge as constraints. Across simulations and applied datasets, EACS consistently outperforms static causal, invariant, and ERM-based predictors under diverse distribution shifts.

Figures (13)

• Figure 1: (a) Causal graph illustrating the relationships among the outcome $Y$, an unobserved cause $C_1$, an observed cause $C_2$, and a covariate $X$ influenced by both causes and acting as a proxy for $C_1$. (b--c) Perturbation scenarios, indicated by the hammer symbol (), applied either to the latent cause $C_1$ or directly to the proxy covariate $X$. Perturbations to $C_1$ and $X$ can induce similar marginal covariate patterns, making it difficult to determine which covariates should be used for prediction.
• Figure 2: MSE of predictive models using only the observed causal covariate $C_2$ versus models using both $C_2$ and the proxy covariate $X$, under perturbations to $C_1$, $C_2$, and $X$. Shaded regions indicate 95% CIs based on 1,000 simulation replications. Causal-only covariate selection outperforms ERM only when perturbations directly disrupt the proxy relationship; when proxy relationships remain stable, including non-causal covariates substantially improves OOD prediction.
• Figure 3: Summary statistics of the covariate distribution across environments corresponding to the perturbation scenarios in Figure \ref{['fig:sim']}. Different types of distribution shift induce distinct and observable signatures in the covariate distribution, such as changes in variance and dependence structure. These signatures can be computed from unlabeled covariates in the target OOD environment and used to assess whether proxy covariates remain reliable or become harmful.
• Figure 4: OOD prediction MSE of the EACS algorithm under varying numbers of training environments per perturbation type and training noise levels $\sigma \in \{1,5,10\}$ (test noise fixed at $\sigma_{\mathrm{test}}=1$). Darker to lighter green indicates fewer training environments. Shaded regions denote 95% CIs over 1,000 replications. With a sufficient number of environments, EACS approaches the performance of the oracle predictor that uses the optimal covariate set for each type of shift.
• Figure 5: OOD prediction MSE of the EACS algorithm under varying sample sizes per environment and training noise levels $\sigma \in \{1,5,10\}$ (test noise fixed at $\sigma_{\mathrm{test}}=1$). Darker to lighter green indicates fewer samples. Shaded regions give 95% CIs over 1,000 replications. Increasing the sample size per environment improves EACS and moves it toward the oracle predictor.

Theorems & Definitions (6)

• Theorem 1: Finite-sample oracle inequality
• Theorem 2: Asymptotic optimality
• Remark 1: Fallback when summaries are uninformative
• Remark 2: Environment-specific predictors
• Theorem 3: Oracle inequality with causal constraints
• Theorem 4: Asymptotic optimality with causal constraints