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 paper studies how contrastive learning (CL) behaves in the unsupervised domain adaptation (UDA) setting. In particular, it finds that CL brings a different mechanism for UDA compared to traditional adversarial DA methods, where the features learned are far apart between the source and target domains. It further develops a conceptual model for explaining the success of contrastive learning under domain shifts, and empirically demonstrates through experiments on several benchmark datasets.