Improving Domain Generalization in Contrastive Learning via Domain-Aware Temperature Control

Pre-training with contrastive learning is a powerful method for learning from sparsely labeled data. However, performance can drop considerably when there is a shift in the distribution of data available during training and test time. We study this phenomenon in the domain generalization setting in which the training data come from multiple domains, and the test data come from an unseen domain. We present a new method for contrastive learning that incorporates domain labels to increase the domain invariance of learned representations, leading to improved out-of-distribution generalization. Our method adjusts the temperature parameter in the InfoNCE loss -- which controls the relative weighting of negative pairs -- using the likelihood that a negative sample comes from the same domain as the anchor. This upweights pairs from more similar domains, forcing the model to discriminate samples based on domain-irrelevant features. To assess domain similarity, we train a domain discriminator on the learned embeddings -- critically, this allows us to adapt the weighting as the amount of domain information in the embedding space changes. Through preliminary experiments on a variant of the MNIST dataset, we demonstrate that our method yields better out-of-distribution performance compared to baselines, especially in regimes of high label sparsity (e.g., 1\%). Furthermore, our method concurrently maintains strong in-distribution task performance, greatly outperforming baselines on this measure.