Multi-Domain Self-Supervised Learning

Contrastive self-supervised learning has recently gained significant attention owing to its ability to learn improved feature representations without the use of label information . Current contrastive learning approaches , however , are only effective when trained on a particular dataset , limiting their utility in diverse multi-domain settings . In fact , training these methods on a combination of several domains often degrades the quality of learned representations compared to the models trained on a single domain . In this paper , we propose a Multi-Domain Self-Supervised Learning ( MDSSL ) approach that can effectively perform representation learning on multiple , diverse datasets . In MDSSL , we propose a three-level hierarchical loss for measuring the agreement between augmented views of a given sample , agreement between samples within a dataset and agreement between samples across datasets . We show that MDSSL when trained on a mixture of CIFAR-10 , STL-10 , SVHN and CIFAR-100 produces powerful representations , achieving up to a 25 % increase in top-1 accuracy on a linear classifier compared to singledomain self-supervised encoders . Moreover , MDSSL encoders can generalize more effectively to unseen datasets compared to both single-domain and multidomain baselines . MDSSL is also highly efficient in terms of the resource usage as it stores and trains a single model for multiple datasets leading up to 17 % reduction in training time . Finally , for multi-domain datasets where domain labels are unknown , we propose a modified approach that alternates between clustering and MDSSL . Thus , for diverse multi-domain datasets ( even without domain labels ) , MDSSL provides an efficient and generalizable self-supervised encoder without sacrificing the quality of representations in individual domains . 1 INTRODUCTION . Self-supervised contrastive training ( Chen et al. , 2020 ; He et al. , 2020 ; Misra & van der Maaten , 2020 ; Caron et al. , 2020b ) has become a popular paradigm for unsupervised representation learning as it shows impressive results on linear classification tasks , almost matching the performance of a supervised model trained from scratch . However , we find that current self-supervised models are only effective when trained on a single-domain . This can hinder their deployment in large scale realworld settings where data almost always comes from multiple diverse domains . We illustrate this issue in Table 1 , where we show that a popular self-supervised model , SimCLR ( Chen et al. , 2020 ) , trained on CIFAR-10 ( Krizhevsky et al. , a ) does not generalize to other domains at the test time . We observe that the top-1 accuracy of a linear classifier significantly drops on unseen datasets . This means that a different self-supervised model needs to be trained for every new dataset , which can add significant computational overheads given that training these models often require large batch sizes and a large number of training epochs ( Chen et al. , 2020 ; He et al. , 2020 ; Wu et al. , 2018 ) . One potential solution for self-supervised learning on multi-domain datasets is to train the models on the union of all input domains . Unfortunately , this solution performs poorly and fails to obtain a good performance on every individual dataset and does not generalize well to unseen domains . To illustrate this , we trained SimCLR on the union of multiple datasets including CIFAR-10 ( Krizhevsky et al. , a ) , CIFAR-100 ( Krizhevsky et al. , b ) , SVHN ( Netzer et al. , 2011 ) and STL-10 ( Coates et al. ) . The trained model is unfavorable as it significantly decreases the top-1 accuracy in all training datasets compared to the single-domain baselines ( see Table 1 ) . To tackle these issues , we propose Multi-Domain Self-Supervised Learning ( MDSSL ) , a technique for obtaining a unified embedder that can be trained on multiple domains . In MDSSL , we train the model over the union of multiple datasets using a three-level hierarchical loss involving : • Embedding similarities of two views of a sample : In the first level , we maximize agreement ( i.e . the cosine similarity between l2-normalized vectors ) between embeddings of two augmented views of a given sample . • Embedding similarities of samples from a given dataset : In the second level , we minimize the pairwise agreements between embeddings of all samples within a dataset . • Embedding similarities of samples from different datasets : In the third level , we minimize the pairwise agreement between samples across all training datasets . The first two levels ensure that the model learns high quality representations for each individual domain . The third level of the MDSSL loss encourages the model to learn distinguishable representations between domains . This approach assumes that domain labels are known during training . We also extend MDSSL to more realistic multi-domain setups where domain labels are unknown . In such scenarios , we present an iterative approach that alternates between clustering and MDSSL at fixed intervals . We use clustering to detect pseudo-domain-labels for each training dataset and use these labels in the MDSSL loss . We also propose a robust version of clustering by reducing outlier noise which further improves the performance of MDSSL in an entirely unsupervised setup . In summary , the goal of MDSSL is to compute improved latent representations of samples from multiple diverse datasets using a single self-supervised model ( See Figure 1 ) . We summarize our contributions as follows : Train Dataset Top-1 AccuracyCIFAR-10 STL-10 SVHN CIFAR-100 Average Single-Domain Training Si m C LR CIFAR-10 92.35 56.71 55.97 75.37 70.10 STL-10 71.05 77.58 46.06 63.81 64.62 SVHN 62.83 46.77 92.42 48.27 62.57 CIFAR-100 79.58 55.27 61.16 90.29 71.57 Multi-Domain Training Si m C LR CIFAR-10 , CIFAR-100 , 82.30 61.41 66.65 73.41 70.94SVHN , STL-10 M D SS L CIFAR-10 , CIFAR-100 , 88.45 65.95 75.35 83.05 78.20SVHN , STL-10 ( λ1 = 1 , λ2 = 0.1 ) 2 RELATED WORK . Supervised classification techniques involve minimizing a loss function ( e.g . the cross-entropy loss ) to match model predictions to true labels . Unsupervised classification methods , on the other hand , learn to classify data without the use of training labels , usually with the use of clustering techniques ( Bojanowski & Joulin , 2017 ; Dosovitskiy et al. , 2014 ; YM . et al. , 2020 ; Bautista et al. , 2016 ; Caron et al. , 2018 ; 2019 ; Huang et al. , 2019 ) . More recently , new unsupervised techniques called self-supervised representation learning have been proposed . A self-supervised model learns by observing every instance of the given data and assigns its own labels , and then performs a classification task ( Bojanowski & Joulin , 2017 ; Dosovitskiy et al. , 2014 ; Wu et al. , 2018 ; Dosovitskiy et al. , 2016 ) . To simplify the complexity of instancelevel classification , a memory bank ( Wu et al. , 2018 ; He et al. , 2020 ) can be used with the help of contrastive learning ( Gutmann & Hyvärinen , 2010 ; Hjelm et al. , 2019 ; van den Oord et al. , 2019 ; Grill et al. , 2020 ) . Contrastive learning ( Arora et al. , 2019 ; Tosh et al. , 2021 ; Bachman et al. , 2019 ) is a temperature-controlled cross-entropy loss between positive pairs of similar samples and negative pairs of dissimilar samples . Positive pairs are usually considered as multiple transformations ( views ) ( Tian et al. , 2020 ) of a given sample using stochastic data augmentation . SimCLR ( Chen et al. , 2020 ) shows that contrastive learning can be done without the use of a memory bank , using the samples within a batch , if we have large enough batches . SwAV ( Caron et al. , 2020a ) uses a mixture of contrastive learning and clustering to form a swapped prediction problem that can learn even with very small batch sizes . Finally , contrastive learning can benefit from training labels , if available with a simple modification of contrasting between samples within a class and taking samples of other classes as negatives ( Khosla et al. , 2020 ) . Each of these approaches show remarkable linear classification accuracy on single-domain setups . Extending self-supervised learning to multiple diverse domains , other than ImageNet ( Russakovsky et al. , 2015 ) , is a relatively less explored topic ( Wallace & Hariharan , 2020 ) . When multiple related domains are available during training , a possible approach is to use mutual information to simultaneously encode common invariant information and domain-specific information of each image ( Feng et al. , 2019 ) . In our paper , we focus on a general setup where we combine diverse unrelated domains and evaluate individual domain-specific tasks . 3 MULTI-DOMAIN SELF-SUPERVISED LEARNING WITH DOMAIN LABELS . In this section , we define the Multi-Domain Self-Supervised Learning ( MDSSL ) paradigm for D training datasets where domain labels are known . We define xdi ∈ Rr and x̃i d ∈ Rr as two transformed views of the ith sample from the dth dataset , d ∈ { 1 , ... , D } . Similar to SimCLR , we use a base encoder f ( . ) and a two-layer MLP projection head g ( . ) to map a given sample into the latent space . We define the latent representations of the two views of the ith sample from the dth dataset as zdi = f ( g ( x d i ) ) ∈ Rr ′ and z̃id = f ( g ( x̃id ) ) ∈ Rr ′ where r′ is the size of each latent representation . We represent mini-batches containing 2N samples ( 2 views per sample ) fromD datasets as a matrix X ∈ R2ND×r , whose corresponding latent representation is denoted by Z ∈ R2ND×r′ . We then calculate a similarity matrix S ∈ R ( 2ND ) × ( 2ND ) that contains the exponential cosine similarity scaled by a temperature parameter τ , between all the latent representations in a given batch . The ( i , j ) th element of S is : S ( i , j ) : = exp ( 1 τ zTi zj ‖zi‖‖zj‖ ) ( 1 ) where zi ∈ Rr ′ and zj ∈ Rr ′ are the ith and jth row of Z , respectively . Sdi represents the cosine similarity between z d i and z̃i d. Sd , d ′ ij represents the cross-dataset cosine similarity between zdi and z d′ j where d , d ′ ∈ { 1 , ... , D } . MDSSL aims to solve the following optimization problem : max θ ( 1 ND D∑ d=1 N∑ i=1 logSdi ) ( 2 ) − λ1 ( 1 2ND D∑ d=1 2N∑ i=1 log 2N∑ j=1 1j 6=iS d , d ij ) ( 3 ) − λ2 ( 1 2ND D∑ d=1 2N∑ i=1 log D∑ d′=1 2N∑ j=1 1d′ 6=dS d , d′ ij ) ( 4 ) where θ is the set of model parameters and 0 < λ1 ≤ 1 and λ2 ≥ 0 are tunable regularization parameters . This is a three-level hierarchical loss . ( 2 ) maximizes the similarity between two transformed views ( xdi and x̃i d ) in the latent space . ( 3 ) minimizes the similarity between every pair of samples ( xdi and x d j ) within a dataset . ( 4 ) minimizes the similarity between pairs of samples ( x d i and xd ′ j ) across different datasets . When λ1 = 1 and λ2 = 0 , this optimization is simplified to SimCLR ( Chen et al. , 2020 ) . Therefore , we use λ1 = 1 and λ2 = 0 as the baseline in all our experiments . λ1 helps us control the extent to which we want to minimize the similarity within a domain . We empirically observe that relaxing λ1 from the SimCLR baseline ( λ1 = 0 ) to a value slightly less than 1 , in fact generates better structure in the latent space by clustering samples within a domain relatively closer compared to samples outside a domain ( See Appendix Section A.3 ) . λ2 should always be non-negative as we always want to minimize the agreement between samples of different datasets . We can implement the MDSSL loss efficiently as it only involves calculating the similarity matrix S once and then selecting elements according to each term mask . As the number of datasets increases , the size of S increases , gradually increasing the running time of MDSSL .

Summary: This paper proposes a method for multi-domain self-supervised representation learning (MDSSL).  The contrastive MDSSL training objective consists of three  terms: (1) a loss  term  that  maximizes the similarity between two transformed views of a sample in the latent space (identical to SimCLR), (2) a loss term that maximizes the similarity between every pair of samples within a dataset,  and  (3)  a loss term  that minimizes the similarity between pairs of samples  across different datasets. Loss terms (2) and (3) are weighted using hyperparameters. Authors perform experiments using common datasets including CIFAR-10, CIFAR-100,  STL-10, TinyImageNet and  SVHN. Experiments show that MDSSL outperforms (1) representations learnt by naively combining all datasets and (2) representations learnt on individual domains,  on average classification accuracy on training domains as well as held-out domains. Authors also perform  ablation studies to explore the effect of loss term hyperparameters on  downstream performance. The paper further extends the basic MDSSL method by adding an assumption that dataset labels are not available, and determine dataset labels using k-means clustering during the training process, and also explore a robust k-means clustering method. Experiments using 2 training domains show that the downstream classification performance of this unsupervised MDSSL method is comparable to original labeled MDSSL.