Environment Inference for Invariant Learning
This addresses domain generalization for robust ML, but is incremental as it builds on existing invariant learning frameworks.
The paper tackles the problem of domain-invariant learning when environment partitions are not provided, proposing EIIL to infer partitions for invariant learning, resulting in outperforming methods on CMNIST without labels and improving worst-group performance on Waterbirds and CivilComments datasets.
Learning models that gracefully handle distribution shifts is central to research on domain generalization, robust optimization, and fairness. A promising formulation is domain-invariant learning, which identifies the key issue of learning which features are domain-specific versus domain-invariant. An important assumption in this area is that the training examples are partitioned into "domains" or "environments". Our focus is on the more common setting where such partitions are not provided. We propose EIIL, a general framework for domain-invariant learning that incorporates Environment Inference to directly infer partitions that are maximally informative for downstream Invariant Learning. We show that EIIL outperforms invariant learning methods on the CMNIST benchmark without using environment labels, and significantly outperforms ERM on worst-group performance in the Waterbirds and CivilComments datasets. Finally, we establish connections between EIIL and algorithmic fairness, which enables EIIL to improve accuracy and calibration in a fair prediction problem.