# Robust Representation Learning through Explicit Environment Modeling

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 37 references that are in the map.
- First posted on arXiv 28 Apr 2026.

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

- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [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)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Certifying Some Distributional Robustness with Principled Adversarial Training](https://mlsanity.com/paper/1710.10571/) (2017)
- [Learning Models with Uniform Performance via Distributionally Robust Optimization](https://mlsanity.com/paper/1810.08750/) (2018)
- [Domain Generalization using Causal Matching](https://mlsanity.com/paper/2006.07500/) (2020)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [Statistics of Robust Optimization: A Generalized Empirical Likelihood Approach](https://mlsanity.com/paper/1610.03425/) (2016)
- [Simple data balancing achieves competitive worst-group-accuracy](https://mlsanity.com/paper/2110.14503/) (2021)
- [On Calibration and Out-of-domain Generalization](https://mlsanity.com/paper/2102.10395/) (2021)
- [Invariant Rationalization](https://mlsanity.com/paper/2003.09772/) (2020)
- [Anchor regression: heterogeneous data meets causality](https://mlsanity.com/paper/1801.06229/) (2018)
- [Invariant Causal Prediction for Sequential Data](https://mlsanity.com/paper/1706.08058/) (2017)
- [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)
- [Neural Generalized Mixed-Effects Models](https://mlsanity.com/paper/2604.10976/) (2026)

## Built on by (most cited first, 2 of 2)

- [Environment-Robust Representation Learning with Empirical Bayes](https://mlsanity.com/paper/2606.05365/) (2026)
- [NPBoost: Neural Processes with Gradient-Boosted Fixed Effects](https://mlsanity.com/paper/2609.28122/) (2026)

## Abstract (arXiv)

We consider learning from labeled data collected across multiple environments, where the data distribution may vary across these environments. This problem is commonly approached from a causal perspective, seeking invariant representations that retain causal factors while discarding spurious ones. However, this framework assumes that the environment has no direct effect on the target. In contrast, we consider settings in which this assumption fails, but still aim to learn representations that support robust prediction on average across previously unseen environments. To this end, we study representations learned by explicitly modeling variation across environments and then marginalizing that variation out. We analyze the resulting representations and characterize when they are preferable to those learned by causal invariant-representation methods. We propose a concrete method based on generalized random-intercept models, a class of predictors in which such marginalization is possible, and study their generalization properties. Empirically, we show that these models outperform invariant-learning methods across a range of challenging settings.