Robust Representation Learning through Explicit Environment Modeling

TL;DR

This paper introduces a method for learning robust representations across multiple environments by explicitly modeling and marginalizing environment variation, outperforming invariant methods in challenging settings.

Contribution

It proposes a generalized random-intercept model that explicitly accounts for environment variation, providing a new approach to robust representation learning.

Findings

Models outperform invariant-learning methods in various challenging scenarios.

Explicit environment modeling improves generalization to unseen environments.

Theoretical analysis characterizes when these models are preferable.

Abstract

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.