Environment-Robust Representation Learning with Empirical Bayes
This work addresses the problem of robust prediction across environments with distribution shifts, offering a principled Bayesian framework that improves generalization in diverse applied domains.
The paper proposes a Bayesian model for multi-environment prediction where environments affect a latent variable but not the conditional mechanisms, and derives a variational objective with a cross-environment balancing term. Using empirical Bayes for prior setting, the method outperforms prior approaches in simulations and real-world tasks including astronomical source identification, microbiome disease detection, and ICU sepsis prediction.
We consider multi-environment prediction problems. We assume the environments change the distribution of a latent variable, while the mechanisms generating observed covariates and targets remain stable conditional on that variable. For example, hospitals or clinical cohorts may differ in the prevalence of latent patient states, even though the relationships between those states, physiological measurements, and outcomes remain unchanged. Given a dataset from multiple environments, we formulate a Bayesian model for such problems and derive the corresponding variational objective. We show that this objective decomposes into per-environment terms and an additional cross-environment balancing term induced by the model's structure. We use an empirical Bayes method to set the prior and incorporate it into the objective. Based on this objective, we develop an amortized variational algorithm for posterior approximation, and use the resulting learned latent variables to form predictions in new environments.We study our approach through simulations and real-world studies of astronomical source identification, microbiome-based disease detection, and ICU sepsis prediction. Across these settings, our method outperforms previous approaches for prediction in new environments.