Robust Representation Learning through Explicit Environment Modeling

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Slavutsky, Yuli, Blei, David M.
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866908999662698496
author Slavutsky, Yuli
Blei, David M.
author_facet Slavutsky, Yuli
Blei, David M.
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2604_26128
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Robust Representation Learning through Explicit Environment Modeling
Slavutsky, Yuli
Blei, David M.
Machine Learning
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.
title Robust Representation Learning through Explicit Environment Modeling
topic Machine Learning
url https://arxiv.org/abs/2604.26128