Environment-Robust Representation Learning with Empirical Bayes

2026-06-03

AI Summary

This work addresses the degradation in generalization performance in multi-environment prediction tasks caused by distributional shifts in latent variables. The authors propose a Bayesian modeling–based approach for learning environment-robust representations, grounded in the assumption that while the data-generating mechanism remains invariant across environments, the distributions of latent variables may vary. Their method introduces a variational objective augmented with a cross-environment balancing term, employs empirical Bayes to automatically determine the prior, and leverages amortized inference for efficient computation. Evaluated on diverse real-world tasks—including astronomical source identification, microbiome-based disease detection, and sepsis prediction in intensive care units—the proposed approach consistently outperforms existing models, demonstrating superior cross-environment predictive accuracy and transferability.

Abstract

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.

Problem

• multi-environment prediction
• environment-robust representation
• latent variable
• distribution shift
• empirical Bayes

Innovation

• Empirical Bayes
• Environment-Robust Representation
• Variational Inference
• Multi-Environment Prediction
• Latent Variable Modeling