Improving Domain Generalization in Contrastive Learning using Adaptive Temperature Control

TL;DR

The paper tackles covariate shift across domains in self-supervised contrastive learning by introducing a domain-aware adaptive temperature for the InfoNCE loss. A domain discriminator estimates $P(D=d|z)$ to compute per-pair weights $w_{ij}$ and pairwise temperatures $\tau_{ij}$, upweighting harder negatives from similar domains and promoting domain-invariant representations. Two variants, Domain-Weighted Negatives and Domain-Weighted Pairs, demonstrate improved OOD and ID performance on a Colored-MNIST benchmark with sparse labels, along with analyses of gradient behavior and ablations on discriminator design and projector heads. The approach yields more robust generalization to unseen domains and provides insights into how adaptive temperature can steer learning away from domain-specific cues while preserving task-relevant information.

Abstract

Self-supervised 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 from training to test time. We study this phenomenon in a setting in which the training data come from multiple domains, and the test data come from a domain not seen at training that is subject to significant covariate shift. 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 probability that a negative sample comes from the same domain as the anchor. This upweights pairs from more similar domains, encouraging the model to discriminate samples based on domain-invariant attributes. Through experiments on a variant of the MNIST dataset, we demonstrate that our method yields better out-of-distribution performance than domain generalization baselines. Furthermore, our method maintains strong in-distribution task performance, substantially outperforming baselines on this measure.

Figures (7)

• Figure 1: Accuracy of contrastive learning methods as color variance ($\sigma$) within each domain is increased. Downstream linear classifiers are trained using 1% of digit labels. We see that the proposed domain-weighted methods are frequently among the best performing for Test-OOD Accuracy at $\sigma > 25$. Furthermore, they display higher robustness on Test-ID Accuracy as $\sigma$ is increased.
• Figure 2: Randomly sampled example digits from the Colored-MNIST dataset variant that we use for the experiments in this paper. Here the colors of the digits are sampled around the mean pixel value for each domain with a variance of $\sigma = 50$. We show examples from each of the dataset splits.
• Figure 3: The left-hand side of each subplot shows the distribution of digit colors in RGB-space; the right hand side shows the marginal distribution of their intensities by each color dimension. As $\sigma$ increases there is more color variance within each of the domains, to the extent that domains start overlapping at high $\sigma$.
• Figure 4: Accuracy of all trained models at color variance $\sigma = 50$ and with downstream linear classifiers trained using 1% of digit labels. Red circles represent the models selected using the strategy of highest validation accuracy on the held-out training domain (where accuracy is averaged over 5 random seeds); grey circles are the other models trained but not selected.
• Figure 5: Projector head ablation of Standard Contrastive Learning method. We consider 3 different types of projector head -- 2-layer MLP, linear, and no projector head; 3 different dimensions for the output of the projector head $\{4,8,16\}$; and, temperature $\tau \in \{0.05,0.075,0.1,0.125,0.15,0.175,0.2,0.25,0.5,1\}$. An encoder embedding dimension of 16-d, color variance of 25, and label fraction of 10% are used for this experiment.