How does Contrastive Pre-training Connect Disparate Domains?

1 INTRODUCTION . In applications such as image recognition for self-driving cars ( Yu et al. , 2020 ; Sun et al. , 2020 ) or medical image diagnosis ( AlBadawy et al. , 2018 ; Dai & Gool , 2018 ) , machine learning models often fail on examples drawn from a different distribution than training . Pre-training on large-scale unlabeled data ( Chen et al. , 2020 ; He et al. , 2020 ; Radford et al. , 2021 ) can substantially boost accuracy under distribution shift ( Radford et al. , 2021 ; Hendrycks et al. , 2019 ; Fisch et al. , 2019 ; Yogatama et al. , 2019 ; Devlin et al. , 2019 ) . For example , CLIP ( Radford et al. , 2021 ) , which is trained on internet-scale image and text data , has shown impressive robustness benefits on ImageNet variants such as ImageNetV2 . Similarly , pre-training methods such as BERT ( Devlin et al. , 2019 ) have improved robustness in NLP — for example , multilingual pre-training significantly improves performance on unseen languages ( Liu et al. , 2020 ) . While other robustness methods that use unlabeled data such as self-training ( Prabhu et al. , 2021 ; Sohn et al. , 2020 ; Berthelot et al. , 2021 ) and domain adversarial training ( Shu et al. , 2018 ; Ganin et al. , 2016 ) jointly optimize a labeled and unlabeled objective , pre-training is particularly suitable for large unlabeled datasets since we can generically pre-train once on unlabeled data and then specialize the model to downstream labeled data via fine-tuning . How does pre-training improve robustness to distribution shift when the pre-training method is not tailored to any downstream task ? In this paper , we run controlled experiments to understand the robustness benefits of a popular pre-training method , contrastive pre-training ( Chen et al. , 2020 ; He et al. , 2020 ; Caron et al. , 2020 ) , even without large-scale unlabeled data . We directly apply contrastive pre-training to unsupervised visual domain adaptation ( UDA ) , where unlabeled data from a shifted target domain is used to adapt models trained on labeled source dataset . We compare contrastive pre-training ( SwAV ( Caron et al. , 2020 ) ) to UDA methods such as SENTRY ( Prabhu et al. , 2021 ) ( self-training ) and DIRT-T ( Shu et al. , 2018 ) ( domain adversarial training ) where all methods use the same unlabeled data ( source and target ) . We find that SwAV achieves comparable or better results on the DomainNet , BREEDS Living-17 , and BREEDS Entity-30 benchmarks . Furthermore , we show that pre-training once on unlabeled data from many domains further improves the average target accuracy over the consituent source-target pairs , showing that extra unlabeled data improves contrastive pre-training for domain adaptation . Given the strong performance of contrastive pre-training , we seek to understand why contrastive pre-training works for domain adaptation . Domain adversarial training and self-supervision algorithms are commonly based on the intuition that domains get merged in feature space ( Tzeng et al. , 2014 ; Ganin et al. , 2016 ; Tzeng et al. , 2017 ; Sun et al. , 2019 ; Ben-David et al. , 2010 ; Shu et al. , 2018 ; Wang et al. , 2021 ; Thota & Leontidis , 2021 ) . This is support by theoretical notions such asH∆H-divergence , where domains that can be discriminated well and results in worse generalization bounds for target accuracy ( Ben-David et al. , 2010 ) . A natural hypothesis is that contrastive pre-training brings domains together by bringing the representations of strong augmentations of the same image together ( Sun et al. , 2019 ) . Surprisingly , we find that the source and target domains are still far apart in the feature space learned by contrastive pre-training . A domain classifier trained on the pre-trained representations of strongly augmented images from the sketch and painting domains in DomainNet has 94 % accuracy , which is even higher than the accuracy of a domain classifier on the original ( augmented ) input space ( 75 % ) . We study this phenomenon from the lens of the recently proposed augmentation graph ( HaoChen et al. , 2021 ) , which forms connections between examples if they have similar augmentations ( e.g. , a small crop + grayscale can connect sketch and real photos , Figure 1 ) . These connections are central to the contrastive objective , which seeks to bring neighbors in the graph closer together while pushing apart non-neighbors . From both simulations and real datasets , we define a critical connectivity quantity which seems to govern whether domain adaptation will be successful ( Figure 1 ) . In particular , the connection strength between the same classes across different domains should be relatively greater than the connection between different classes across different domains . Importantly , the absolute connectivity between domains can be small and the domains can be far apart in input space . This suggests that contrastive learning circumvents the separation of domains through the connectivity structure of the ( general-purpose ) augmentations used . Using our understanding , we ablate contrastive learning by removing subsets of unlabeled data that contribute most to connecting the two domains . We do this by training a domain classifier to distinguish between the source and target and removing examples which are most uncertain under the classifier . This consistently worsens the target accuracy in comparison to randomly removing the same number of data points ( by 2–3 % on DomainNet ) , showing the importance of domain connectivity for adaptation . 2 SETUP . We consider a classification problem from an input spaceX ⊆Rd to a label spaceY= [ K ] . Data . Let PS and PT denote the source and target distributions of ( x̄ , y ) pairs , respectively , where x̄∈X and y ∈Y . The labeled source dataset S consists of nS input-output pairs from PS . We also have access to an unlabeled target dataset T with nT unlabeled inputs from PT , and we assume the label sets for the source and target are identical . In some cases , we also consider an additional related datasetRwith nR unlabeled inputs from a related distributionR . Model and metrics . The goal is to learn a classifier h : X →Y with low classification error on the target distributionL0−1 ( h ) =Ex̄ , y∼PT [ 1 [ h ( x̄ ) 6=y ] ] . The pre-training algorithms we consider consist of two steps : first , we train an encoder function f : X →Z into a representation spaceZ⊆Rd′ ( where d′ denotes the representation dimension ) , then we train a classification head g : Z →Y . The final classifier h=g◦f composes the encoder with the classification head . 2.1 CONTRASTIVE LEARNING . Contrastive learning ( He et al. , 2020 ; Chen et al. , 2020 ; Caron et al. , 2020 ) trains representations to be similar between augmented views of the same example and dissimilar between augmented views of random examples . Example image augmentations include random cropping , color jitter , and Gaussian blur ( Chen et al. , 2020 ) . The basic assumption is that the semantic content of the data is invariant under the augmentations . Though there are several variants of contrastive learning , we focus on the following objective from HaoChen et al . ( 2021 ) for our conceptual understanding . LetX denote the “ augmented ” space containing all data-augmented views of natural examples ( e.g. , all random crops of natural ImageNet images ) . Data augmentation is defined as a random functionA : X →X , where P [ A ( x ) =x ] denotes the probability that x∈X augments to x∈X . A pair of augmented views is called a positive pair and denotedx , x+ if a natural pointx is first drawn and thenx , x+ are the results of two calls toA ( x ) . On the other hand , a negative pair x , x− arises from two natural data points x , x′ when x=A ( x ) , x−=A ( x′ ) . Minimizing the following loss function simultaneously brings together the representations of positive pairs x , x+ and repels the representations of negative pairs x , x− : L ( f ) =−2·Ex , x+ [ f ( x ) > f ( x+ ) ] +Ex , x− [ ( f ( x ) > f ( x− ) ) 2 ] . ( 1 ) Here the first expectations are taken over the marginal distribution overX and the randomness ofA . 2.2 CONSTRASTIVE PRE-TRAINING FOR DOMAIN ADAPTATION . Let D denote the unlabeled input data used for contrastive pre-training ( e.g. , some combination of the source unlabeled data S , target unlabeled data T , and/or additional unlabeled data R ) . We use the following simple algorithm to apply contrastive pre-training to domain adaptation : 1 . ( Pre-train ) Learn an encoder function f : X →Z via contrastive learning onD . 2 . ( Fine-tune ) On labeled source data S , train a classifier h , g◦f : X →Y as the composition of the pre-trained encoder f and a classifier head g. In this paper , we take “ fine-tuning ” to mean updating both the parameters of f and g . 3 CONTRASTIVE PRE-TRAINING IS A STRONG DOMAIN ADAPTATION METHOD . We evaluate contrastive learning and relevant baselines on 4 benchmark vision datasets and observe that contrastive learning achieves comparable or better performance in all cases . 3.1 DATASETS . BREEDS ( Santurkar et al. , 2020 ) . BREEDS is a subpopulation shift benchmark derived from ImageNet by constructing a hierarchical tree structure of classes from WordNet . Nodes at a specified depth of the tree become the labels for the classification task , and descendant nodes are treated as subpopulations that can be randomly partitioned into source and target domains . DomainNet ( Peng et al. , 2019 ) . DomainNet is a large unsupervised domain adaptation task , consisting of approximately 600,000 images and 345 classes in 6 domains . For our experiments we utilize the same filtered version of DomainNet from Prabhu et al . ( 2021 ) , which uses 40 of the 345 classes and the sketch , painting , photograph , and clipart domains . STL-10→CIFAR-10 . STL-10 and CIFAR-10 are two classic image classification datasets , each consisting of 10 classes ( Coates et al. , 2011 ; Krizhevsky , 2009 ) . Each dataset has a class label that is not present in the other , and hence we follow the procedure of French et al . ( 2018 ) and filter out the examples in the non-overlapping classes , resulting in a 9-class classification problem . 3.2 METHODS . Baselines . We run empirical risk minimization ( ERM ) on the source dataset for each task , i.e. , a “ source-only ” baseline . In addition , for BREEDS and DomainNet we use these ERM models as weight initialization for SENTRY , a self-training algorithm for domain adaptation that achieved state of the art results on several domains of DomainNet when initialized with ImageNet-pretrained models ( Prabhu et al. , 2021 ) . For STL→ CIFAR , we use Dirt-T ( Shu et al. , 2018 ) , the current state of the art algorithm for this task . Contrastive pre-training . For BREEDS and DomainNet we use SwAV ( Caron et al. , 2020 ) , a contrastive learning algorithm that achieved state-of-the-art linear evaluation accuracy on ImageNet . SwAV contrasts examples in feature space rather than input space and uses a novel multi-crop data augmentation strategy using several crops of different sizes . For the lower-resolution task STL→ CIFAR we use SimCLR ( Chen et al. , 2020 ) , which uses only 2 crops and therefore scales straightforwardly to smaller input dimensions .

Summary: The authors tackle the problem of unsupervised domain adaptation (UDA), where most state-of-the-art methods aim to learn a joint embedding space for samples from the (labeled) source domain and (unlabeled) samples from the target domain.   In contrast to this, the authors propose to use contrastive learning pre-training techniques to first learn an encoder using target/source/additional data: Intuitively, this objective pushes positive samples together, and further away from other samples. The notion of "positive set" is defined as all samples which are obtained from data augmentations of a same sample natural image. A classifier is then learned on top of these features to solve the classification task on the labeled source dataset.  Using standard contrastive learning techniques, the authors show that they can achieve similar accuracies as some SoTA baselines on three UDA datasets. The remaining of the paper analyses the features learned by this techniques, and in particular, how they do not learn to align the two domains as is the common intuition in UDA. This analysis relies on a "connectivity graph" defined for contrastive learning techniques, which, in the current setting, can be used to connect samples from the source and target domain, to determine how likely they are to be in the same positive set.