Improving Domain Generalization in Contrastive Learning using Adaptive Temperature Control

AI summary

Overview

Research area: Self-supervised representation learning, specifically contrastive learning (SimCLR-style) for domain generalization under covariate shift.

Technical level: Intermediate. The paper uses formal notation for the InfoNCE loss and provides a gradient analysis, but the core idea and experimental design are accessible to readers who know what contrastive learning and embeddings are.

Scope in one sentence: The paper proposes a domain-aware, adaptive temperature in the InfoNCE loss that upweights negative samples likely to come from the same domain as the anchor, and tests it on a colored variant of MNIST for 3-versus-5 digit classification across unseen color domains.

What This Paper Is About

Contrastive pre-training learns representations by pulling augmented views of the same instance together and pushing different instances apart, but its performance drops sharply when the test data's distribution differs from training. The authors study a setting where training data come from multiple domains (mostly unlabeled) and test data come from an unseen domain subject to covariate shift, meaning the distribution of inputs X changes across domains while the conditional distribution of labels given inputs stays the same. Their goal is to use the available domain labels to make the learned embeddings more domain-invariant, and therefore more transferable to the unseen domain, without sacrificing accuracy on the domains seen during training.

Key Contributions

1. A domain-aware adaptive temperature for InfoNCE. Instead of a single fixed temperature, the method assigns each anchor-negative pair its own temperature, computed from the estimated probability that the negative comes from the same domain as the anchor, so that negatives likely from the same domain receive a larger penalty.

2. Two ways to estimate same-domain probability. The paper introduces Domain-Weighted Negatives, which treats the anchor's domain as known and uses a soft label for the negative, and Domain-Weighted Pairs, which uses soft labels for both the anchor and the negative.

3. A domain discriminator trained on the current embedding space. A linear network is trained each pre-training epoch on the encoder's embeddings to output domain probabilities, so the weighting adapts as the representation changes and the loss smoothly reduces to standard contrastive learning when domains become indistinguishable.

4. A controlled evaluation dataset and comparative study. The authors build a colored MNIST variant that lets them vary the strength of the domain shift (the within-domain color variance), and compare against standard contrastive learning, same-domain negatives, an MMD penalty, DANN, and an oracle color-jitter model.

Main Findings

• Both proposed variants beat all non-oracle baselines on out-of-distribution accuracy. At color variance σ = 50 with 1% of digit labels (69 examples), Domain-Weighted Pairs reached 0.819 Test-OOD and Domain-Weighted Negatives reached 0.807, versus 0.784 for Standard CL, 0.728 for Same Domain Negatives, 0.688 for CL with MMD Penalty, and 0.700 for CL with DANN.

• The proposed methods also lead on in-distribution accuracy. On Test-ID at the same setting, Domain-Weighted Pairs scored 0.945 and Domain-Weighted Negatives 0.929, compared to 0.880 for Standard CL, 0.896 for Same Domain Negatives, 0.865 for CL with MMD Penalty, and 0.846 for CL with DANN.

• Domain-invariance comes at a measurable cost in domain-classifier accuracy. The domain classifier accuracy (D-Test-ID) was 0.688 for Standard CL, but 0.530 for Domain-Weighted Negatives and 0.520 for Domain-Weighted Pairs. Only Same Domain Negatives produced a lower value (0.496), yet that method scored worse on both digit tasks, which the authors attribute to its inability to learn from negatives drawn from other domains.

• Enforcing domain invariance with MMD or DANN hurt rather than helped. Both methods scored below Standard CL on Test-OOD, which the authors connect to prior work finding that empirical risk minimization can outperform objectives with a domain-invariance penalty.

• The gap to an oracle remains large. A Standard CL model using color jitter as an augmentation, which the authors treat as an upper bound because it exploits knowledge of the domain difference, reached 0.978 Test-OOD, 0.976 Test-ID, and 0.490 D-Test-ID.

• Advantages are largest with few labels. Across downstream label fractions of 1%, 10%, and 100% at σ = 50, the proposed methods were most clearly ahead at 1%, and they always achieved the highest Test-ID accuracy regardless of label fraction. At 100% labels, Domain-Weighted Negatives reached 0.908 Test-OOD and Domain-Weighted Pairs 0.871, compared to 0.716 for Standard CL.

• Robustness improves as the color shift grows. As σ increased within each domain (tested at σ = 0, 25, 50, 75, and 150), the domain-weighted methods were frequently among the best on Test-OOD accuracy at σ > 25, and showed more stable Test-ID accuracy than baselines. The authors conjecture that higher color variance means more domain information competing for the finite embedding space, so methods that suppress it leave more capacity for digit information.

• Model selection is a real bottleneck. Across all hyperparameter combinations, many models achieved Test-OOD accuracies above 90% that were not the ones selected by validation accuracy on the held-out training domain, and this affected all methods attempting to correct for domain heterogeneity.

• The gradient analysis explains why temperature is the right lever. The relative weight on a negative pair scales with exp(s_ij / τ_ij) × (1 / τ_ij), so a small τ_ij for a likely same-domain negative increases its penalty; the authors note that a weighting term placed outside the exponent would have a negligible effect compared with rescaling the pairwise cosine similarities.

• Global versus batch-level domain discriminators showed no clear winner. In the ablation across σ = 0, 25, 50, 75, and 150, results varied by method and setting without a consistent trend, which the authors suggest may reflect insufficient substructure in the domain heterogeneity of this dataset.

Methodology in Plain English

The authors assume each training example carries a domain label, but that most of the data lacks category labels. They follow the SimCLR recipe: each image is augmented twice (using random resized cropping and Gaussian blur), and the encoder is trained to make the two views of the same image similar while pushing apart embeddings of different images.

Their twist is the temperature, the number that controls how harshly the loss penalizes negatives that are already similar to the anchor. In the standard loss, one temperature applies to every pair. Here, every anchor-negative pair gets its own temperature computed from a formula with three parts: a baseline temperature τ_α, an adjustment term τ_β times (1/N_D minus w_ij), and a floor τ_min that prevents numerical instability. The term w_ij is the estimated probability that the pair comes from the same domain, computed by a small linear domain classifier trained on the encoder's own 16-dimensional embeddings at the start of each epoch, without backpropagating into the encoder. The value 1/N_D is the probability the classifier would output if it were maximally uncertain. When w_ij is large, the temperature shrinks and the pair is penalized harder; when domains become indistinguishable, w_ij approaches 1/N_D and the temperature returns to τ_α, recovering the standard loss.

For evaluation, the authors pre-train encoders for 400 epochs, freeze them, and fit a linear classifier on top using only a small fraction of digit labels, then measure digit accuracy on a held-out set from the training domains (Test-ID) and on the unseen domain (Test-OOD). They also train a classifier to predict domain from the embeddings, using that accuracy as a proxy for how much domain information leaked into the representation.

The dataset is deliberately simple so that the type and strength of the shift can be dialed up and down. Digits 3 and 5 are colored by sampling from a normal distribution around a domain-specific mean color, with a variance σ. Two training domains are red and blue, the validation domain is purple, and the test domain is green. The paper reports two different totals for the dataset: Section 3 states 23,100 samples split 60% Train, 10% Validation, 10% Test-ID, and 20% Test-OOD, while Appendix E states 11,548 digits total, split into 6,930 training, 1,154 validation, 1,154 Test-ID, and 2,310 Test-OOD. These two figures are inconsistent with each other and the paper does not reconcile them. Default results use σ = 50, 1% of digit labels, and a mean over 5 seeds.

Hyperparameters were swept separately at each value of σ. Standard CL, Same Domain Negatives, and the color-jitter oracle swept temperature over {0.05, 0.075, 0.1, 0.125, 0.15, 0.175, 0.2, 0.25, 0.5, 1.0}. The two proposed methods swept τ_α over the same set and τ_β over {0.1, 0.25, 0.5, 1.0}. MMD and DANN swept temperature over the same set plus a penalty weight λ over {0.01, 0.1, 1.0, 5.0}. Models use PyTorch, the Adam optimizer with an initial learning rate of 0.001, a scheduler with step size 20 and γ = 0.9, and a three-block CNN encoder. The authors report that adding a projector head made Standard CL results substantially worse, so they trained directly on the 16-dimensional embeddings.

Why This Matters

The paper targets a practical gap: contrastive pre-training is attractive precisely because labeled data is expensive, but its benefits erode when test data comes from a domain you did not anticipate. The proposed method requires only domain labels, which are often far cheaper to obtain than task labels, and it degrades gracefully back to the standard loss when domain information disappears from the representation.

For research, the work sits at the intersection of two active lines: domain generalization from sparsely labeled data via contrastive pre-training, and temperature scheduling or adaptation in contrastive losses. The authors position their contribution against prior domain-dependent weighting of negatives, which they characterize as fixed rather than adaptive, applied outside the exponent, and reliant on per-domain negative queues that become cumbersome as domains multiply. It also contrasts with existing temperature-adaptation work, which they say targets label imbalance rather than covariate shift.

Real-world applications implied by the paper's framing:

• Medical imaging models deployed at a new hospital, where scanners, protocols, and patient populations shift the input distribution and the paper notes it is unclear what augmentations would simulate those shifts.
• Any multi-site deployment where data is pooled from several sources with distinct acquisition properties and only a fraction carries task labels.
• Sensor or wearable data collected across devices or users, where the device is effectively a domain label.
• Settings where the shift is a nuisance nuisance characteristic (such as a color channel or capture condition) that the model should learn to ignore rather than exploit.

Industry relevance: domain labels are frequently already recorded as metadata (site ID, device model, batch, acquisition date), so the extra input the method requires is often available at no additional annotation cost. The finding that MMD and DANN penalties underperformed plain contrastive learning is also a caution for practitioners who might otherwise reach for adversarial domain-invariance objectives by default.

Future Directions

• Better model selection for domain generalization. The authors report that validation accuracy on a held-out training domain often fails to identify the best-performing hyperparameters, that many models exceeded 90% Test-OOD accuracy without being selected, and that this is a known open challenge they leave to future work.
• Increasing the capacity and calibration of the domain discriminator. The current experiments use a linear classifier; the authors propose exploring more representational capacity, and separately investigating how well-calibrated the discriminator's output distribution is, since miscalibrated domain probabilities would feed directly into the temperature.
• Global versus batch-level discriminator design. The ablation showed no clear trend favoring either level, which the authors suspect stems from limited substructure in their dataset's domain heterogeneity, so the question remains open for data with richer domain structure.
• Theoretical analysis. The paper presents an empirical gradient analysis and explicitly frames the controlled MNIST variant as a way to build intuition that could guide future theoretical work on how adaptation strategies affect which attributes of the data get encoded.

Additional items the paper raises but does not resolve include the large remaining gap to the color-jitter oracle (0.819 versus 0.978 Test-OOD at the default setting), the contradiction between the two dataset-size figures reported, and, in the portion of the paper that is truncated at Appendix G, projector-head ablation results whose full discussion is not included in the available content.

Target Audience

This paper is most useful for machine learning researchers and graduate students working on self-supervised learning, representation learning, or domain generalization, particularly those interested in loss-function design rather than architecture changes. Practitioners who pre-train models on pooled multi-source data and hold metadata about the source will find the method directly actionable. Readers without a background in contrastive learning will need to read up on InfoNCE and the role of temperature first, but the experimental sections and the intuition behind the weighting scheme are reachable without deep mathematical preparation.

Authors’ 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.