Evaluating and Improving Domain Invariance in Contrastive Self-Supervised Learning by Extrapolating the Loss Function

Despite the remarkable progress of self-supervised learning (SSL), how self-supervised representations generalize to out-of-distribution data remains little understood. In this paper, we study the effects of distribution shifts on self-supervised representations. Our findings indicate that self-supervised representation learning is more robust than traditional supervised learning (52.8% versus 17.1% on the CMNIST dataset, 63.6% versus 60.6% on the Waterbirds dataset). However, self-supervised representations still suffer significantly from domain shifts, especially when spurious correlations are present. Motivated by this limitation, we propose a risk-extrapolated information NCE (ReinformNCE) to facilitate self-supervised learning algorithms to learn more stable representations. Our approach integrates the infoNCE loss function and a robust optimization approach that extrapolates the risks of training domains. Extensive experiments show that ReinformNCE helps to extract domain-invariant self-supervised representations and it substantially improves the robustness of the self-supervised representations (68.2% versus 52.8% on the CMNIST dataset, 77.9% versus 63.6% on the Waterbirds dataset). To the best of our knowledge, this is the first work demonstrating the feasibility of learning domain-invariant representations based on robust optimization theory and without supervised information.