Towards Environment-Invariant Representation Learning for Robust Task Transfer

Abstract

To train a classification model that is robust to distribution shifts upon deployment, auxiliary labels indicating the various “environments” of data collection can be leveraged to mitigate reliance on environment-specific features. In this paper we attempt to determine where in the network the environment invariance property can be located for such a model, with the hopes of adapting a single pre-trained invariant model for use in multiple tasks. We discuss how to evaluate whether a model has formed an environment-invariant internal representation - as opposed to an invariant final classifier function - and propose an objective that encourages learning such a representation. We also extend color-biased digit recognition to a transfer setting where the target task requires an invariant model, but lacks the environment labels needed to train an invariant model from scratch, thus motivating the transfer of an invariant representation trained on a source task with environment labels.