Universal Adaptive Environment Discovery

TL;DR

This work tackles generalization under distribution shift caused by spurious correlations. It introduces Universal Adaptive Environment Discovery (UAED), a framework that learns a distribution over data transformations to instantiate environments and optimizes robust objectives averaged over this distribution, unifying IRM, REx, GroupDRO, and CORAL without predefined groups. The authors provide PAC-Bayes guarantees and show robustness to environment shifts within KL-balls, alongside empirical validation on synthetic benchmarks and the Waterbirds dataset, where adaptive variants improve worst-case accuracy while maintaining competitive mean performance. The results demonstrate that making environments adaptive is a practical route to out-of-distribution generalization and that robustness objectives are effectively complemented by learned environments.

Abstract

An open problem in Machine Learning is how to avoid models to exploit spurious correlations in the data; a famous example is the background-label shortcut in the Waterbirds dataset. A common remedy is to train a model across multiple environments; in the Waterbirds dataset, this corresponds to training by randomizing the background. However, selecting the right environments is a challenging problem, given that these are rarely known a priori. We propose Universal Adaptive Environment Discovery (UAED), a unified framework that learns a distribution over data transformations that instantiate environments, and optimizes any robust objective averaged over this learned distribution. UAED yields adaptive variants of IRM, REx, GroupDRO, and CORAL without predefined groups or manual environment design. We provide a theoretical analysis by providing PAC-Bayes bounds and by showing robustness to test environment distributions under standard conditions. Empirically, UAED discovers interpretable environment distributions and improves worst-case accuracy on standard benchmarks, while remaining competitive on mean accuracy. Our results indicate that making environments adaptive is a practical route to out-of-distribution generalization.

Figures (4)

• Figure 1: waterbirds dataset examples sagawa2020waterbirds: the training correlation between label $y$ and spurious attribute $a$ does not hold at test time.
• Figure 2: a-irm environment discovery on colored-mnist. (Left) Worst-case test accuracy—a-irm achieves 72.3% vs irm's 66.8%, a 5.5% improvement. (Middle) Test accuracy across different color correlations shows a-irm's flatter profile, indicating successful invariant learning. (Right)a-irm discovers intermediate correlations ($e \approx 0.35$) rather than the extremes used by standard irm (blue dashed lines at 0.1, 0.2).
• Figure 3: a-irm environment discovery on rotated-mnist. (Left) Worst-case accuracy comparison shows 28.4% improvement (94.2% vs 65.8%). (Middle) Test accuracy across rotation angles—irm fails catastrophically on unseen rotations (45°, 135°, 180°) while a-irm maintains consistent performance. (Right)a-irm learns approximately uniform distribution over all rotations, unlike irm which only uses 0° and 90° (indicated by dashed lines).
• Figure 4: Evolution of the loss variance of a-irm and irm (Left), and the environment $e$ of a-irm (Right) on colored-mnist.

Theorems & Definitions (10)

• Definition 3.2: Universal Adaptive Objective
• Theorem 4.2: PAC--Bayes for policy-averaged risk
• Lemma 4.3: Change of environment via DV+Hoeffding
• Theorem 4.4: uaed robust generalization
• Proposition 4.5: KL--ball DRO equals entropic risk
• Proposition 4.6: LoA vs. LA gap is variance-controlled
• Remark 4.7: Sub-gamma variant
• Lemma A.1: Hoeffding's lemma
• Lemma A.2: Donsker--Varadhan (DV)
• Theorem A.4: Nonconvex alternating SGD