# Environment-Robust Representation Learning with Empirical Bayes

arXiv: https://arxiv.org/abs/2606.05365 · HTML: https://mlsanity.com/paper/2606.05365/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 16 of its 34 references that are in the map.
- First posted on arXiv 3 Jun 2026.

## Builds on (most cited first, 16 of 16)

- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [Learning Models with Uniform Performance via Distributionally Robust Optimization](https://mlsanity.com/paper/1810.08750/) (2018)
- [Gradient Matching for Domain Generalization](https://mlsanity.com/paper/2104.09937/) (2021)
- [VAE with a VampPrior](https://mlsanity.com/paper/1705.07120/) (2017)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [Fishr: Invariant Gradient Variances for Out-of-Distribution Generalization](https://mlsanity.com/paper/2109.02934/) (2021)
- [Statistics of Robust Optimization: A Generalized Empirical Likelihood Approach](https://mlsanity.com/paper/1610.03425/) (2016)
- [On Calibration and Out-of-domain Generalization](https://mlsanity.com/paper/2102.10395/) (2021)
- [Learning Latent Subspaces in Variational Autoencoders](https://mlsanity.com/paper/1812.06190/) (2018)
- [Capturing Label Characteristics in VAEs](https://mlsanity.com/paper/2006.10102/) (2020)
- [Focus on the Common Good: Group Distributional Robustness Follows](https://mlsanity.com/paper/2110.02619/) (2021)
- [Distributionally Robust Post-hoc Classifiers under Prior Shifts](https://mlsanity.com/paper/2309.08825/) (2023)
- [Robust Representation Learning through Explicit Environment Modeling](https://mlsanity.com/paper/2604.26128/) (2026)
- [Multi-Domain Empirical Bayes for Linearly-Mixed Causal Representations](https://mlsanity.com/paper/2603.18404/) (2026)

## Abstract (arXiv)

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.