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

원문 초록

A common approach to out-of-distribution prediction restricts models to causal or invariant covariates to avoid spurious associations that may change across environments. Despite its theoretical appeal, this strategy can underperform empirical risk minimization when only a subset of the causal parents of the outcome is observed. In such settings, non-causal covariates can serve as proxies for unobserved causal parents and improve prediction when the proxy relationship is stable, but they can hurt when shifts disrupt that relationship. Thus, the optimal covariate set can depend on the specific shift encountered. Because different shifts can leave distinct signatures in the covariate distribution, we propose an environment-adaptive covariate selection algorithm that chooses which covariates to use in each environment based on summaries of its unlabeled covariates. The summaries may be hand-crafted or learned from multi-environment data, and prior causal knowledge can be incorporated as constraints. We give sufficient conditions under which such summaries are informative in linear causal models, and explain when adapting only the covariate set is preferable to letting every coefficient vary with the environment. In simulations and two applications, adapting the covariate set improves prediction over static baselines, including causal and invariant covariate selection.