Optimal Representations for Covariate Shifts

Yann Dubois, Yangjun Ruan, Chris J. Maddison

Keywords: distribution shift, domain generalization, representation learning, self-supervised learning, invariance, robustness

TL;DR: We give a simple variational objective whose optima are exactly the set of all representations that are robust under covariate shift.

Abstract: Machine learning often experiences distribution shifts between training and testing. We introduce a simple objective whose optima are \textit{exactly all} representations on which risk minimizers are guaranteed to be robust to Bayes preserving shifts, e.g., covariate shifts. Our objective has two components. First, a representation must remain discriminative, i.e., some predictor must be able to minimize the source and target risk. Second, the representation's support should be invariant across source and target. We make this practical by designing self-supervised methods that only use unlabelled data and augmentations. Our objectives achieve SOTA on DomainBed, and give insights into the robustness of recent methods, e.g., CLIP.