ReSSL: Relational Self-Supervised Learning with Weak Augmentation

1 Introduction . Recently , self-supervised learning ( SSL ) has shown its superiority and achieved promising results for unsupervised visual representation learning in computer vision tasks [ 40 , 27 , 32 , 6 , 9 , 47 , 23 , 24 ] . The purpose of a typical self-supervised learning algorithm is to learn general visual representations from a large amount of data without human annotations , which can be transferred or leveraged in downstream tasks ( e.g. , classification , detection , and segmentation ) . Some previous works [ 5 , 23 ] even have proven that a good unsupervised pretraining can lead to a better downstream performance than supervised pretraining . Among various SSL algorithms , contrastive learning [ 47 , 45 , 6 ] serves as a state-of-the-art framework , which mainly focuses on learning an invariant feature from different views . For example , instance discrimination is a widely adopted pre-text task as in [ 6 , 24 , 47 ] , which utilizes the noisy contrastive estimation ( NCE ) to encourage two augmented views of the same image to be pulled closer on the embedding space but pushes apart all the other images away . Deep Clustering [ 4 , 48 , 5 ] is an alternative pre-text task that forces different augmented views of the same instance to be clustered into the same class . However , instance discrimination based methods will inevitably induce a class ∗Corresponding author youshan @ sensetime.com 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . collision problem [ 1 , 36 , 10 ] , where similar images should be pulled closer instead of being pushed away . Deep clustering based methods cooperated with traditional clustering algorithms to assign a label for each instance , which relaxed the constraint of instance discrimination , but most of these algorithms adopt a strong assumption , i.e. , the labels must induce an equipartition of the data , which might introduce some noise and hurt the learned representations . In this paper , we introduce a novel Relational Self-Supervised Learning framework ( ReSSL ) , which does not encourage explicitly to push away different instances , but uses relation as a manner to investigate the inter-instance relationships and highlight the intra-instance invariance . Concretely , we aim to maintain the consistency of pairwise similarities among different instances for two different augmentations . For example , if we have three instances x1 , x2 , y and z where x1 , x2 are two different augmentations of x , y and z are different samples . Then , if x1 is similar to y but different to z , we wish x2 can maintain such relationship and vice versa . In this way , the relation can be modelled as a similarity distribution between a set of augmented images , and then use it as a metric to align the same images with different augmentations , so that the relationship between different instances could be maintained across different views . However , this simple manner induces unexpectedly horrible performance if we follow the same training recipe as other contrastive learning methods [ 6 , 24 ] . We argue that construction of a proper relation matters for ReSSL ; aggressive data augmentations as in [ 6 , 7 , 41 ] are usually leveraged by default to generate diverse positive pairs that increase the difficulty of the pre-text task . However , this hurts the reliability of the target relation . Views generated by aggressive augmentations might cause the loss of semantic information , so the target relation might be noisy and not that reliable . In this way , we propose to leverage weaker augmentations to represent the relation , since much lesser disturbances provide more stable and meaningful relationships between different instances . Besides , we also sharpen the target distribution to emphasize the most important relationship and utilize the memory buffer with a momentum-updated network to reduce the demand of large batch size for more efficiency . Experimental results on multiple benchmark datasets show the superiority of ReSSL in terms of both performance and efficiency . For example , with 200 epochs of pre-training , our ReSSL achieved 69.9 % on ImageNet [ 14 ] linear evaluation protocol , which is 2.4 % higher than our baseline method ( MoCoV2 [ 8 ] ) . When working with the Multi-Crop strategy ( 200 epochs ) , ReSSL achieved new state-of-the-art 74.7 % Top-1 accuracy , which is 1.4 % higher than CLSA-Multi [ 46 ] . Our contributions can be summarized as follows . • We proposed a novel SSL paradigm , which we term it as relational self-supervised learning ( ReSSL ) . ReSSL maintains the relational consistency between the instances under different augmentations instead of explicitly pushing different instances away . • Our proposed weak augmentation and sharpening distribution strategy provide a stable and high quality target similarity distribution , which makes the framework works well . • ReSSL is a simple and effective SSL framework since it replaces the widely adopted contrastive loss with our proposed relational consistency loss . It achieved state-of-the-art performance under the same training cost . 2 Related Work . Self-Supervised Learning . Early works in self-supervised learning methods rely on all sorts of pretext to learn visual representations . For example , colorizing gray-scale images [ 50 ] , image jigsaw puzzle [ 39 ] , image super-resolution [ 34 ] , image inpainting [ 19 ] , predicting a relative offset for a pair of patches [ 16 ] , predicting the rotation angle [ 35 ] , and image reconstruction [ 2 , 22 , 3 , 17 ] . Although these methods have shown their effectiveness , they lack the generality of the learned representations . Instance Discrimination . The recent contrastive learning methods [ 32 , 40 , 6 , 24 , 41 , 38 , 29 , 27 , 30 ] have made a lot of progress in the field of self-supervised learning . Most of the previous contrastive learning methods are based on the instance discrimination [ 47 ] task in which positive pairs are defined as different views of the same image , while negative pairs are formed by sampling views from different images . SimCLR [ 6 , 7 ] shows that image augmentation ( e.g.Grayscale , Random Resized Cropping , Color Jittering , and Gaussian Blur ) , nonlinear projection head and large batch size plays a critical role in contrastive learning . Since large batch size usually requires a lot of GPU memory , which is not very friendly to most of researchers . MoCo [ 24 , 8 ] proposed a momentum contrast mechanism that forces the query encoder to learn the representation from a slowly progressing key encoder and maintain a memory buffer to store a large number of negative samples . InfoMin [ 41 ] proposed a set of stronger augmentation that reduces the mutual information between views while keeping task-relevant information intact . AlignUniform [ 45 ] shows that alignment and uniformity are two critical properties of contrastive learning . Deep Clustering . In contrast to instance discrimination which treats every instance as a distinct class , deep clustering [ 4 ] adopts the traditional clustering method ( e.g.KMeans ) to label each image iteratively . Eventually , similar samples will be clustered into the same class . Simply apply the KMeans algorithm might lead to a degenerate solution where all data points are mapped to the same cluster ; SeLa [ 48 ] solved this issue by adding the constraint that the labels must induce equipartition of the data and proposed a fast version of the Sinkhorn-Knopp to achieve this . SwAV [ 5 ] further extended this idea and proposed a scalable online clustering framework . PCL [ 36 ] reveals the class collision problem and simply performed instance discrimination and unsupervised clustering simultaneously ; although it gets the same linear classification accuracy with MoCoV2 , it has better performance on downstream tasks . Contrastive Learning Without Negatives . Most previous contrastive learning methods prevent the model collapse in an explicit manner ( e.g . push different instances away from each other or force different instances to be clustered into different groups . ) BYOL [ 23 ] can learn a high-quality representation without negatives . Specifically , it trains an online network to predict the target network representation of the same image under a different augmented view and using an additional predictor network on top of the online encoder to avoiding the model collapse . SimSiam [ 9 ] shows that simple Siamese networks can learn meaningful representations even without the use of negative pairs , large batch size , and momentum encoders . 3 Methodology . In this section , we will first revisit the preliminary work on contrastive learning ; then , we will introduce our proposed relational self-supervised learning framework . After that , the algorithm and the implementation details will also be explained . 3.1 Preliminaries on Self-supervised Learning . Given N unlabeled samples x , we randomly apply a composition of augmentation functions T ( · ) to obtain two different views x1 and x2 through T ( x , θ1 ) and T ( x , θ2 ) where θ is the random seed for T . Then , a convolutional neural network based encoder F ( · ) is employed to extract the information from these samples , i.e. , h = F ( T ( x , θ ) ) . Finally , a two-layer non-linear projection head g ( · ) is utilized to map h into embedding space , which can be written as : z = g ( h ) . SimCLR [ 6 ] and MoCo [ 24 ] style framework adopt the noise contrastive estimation ( NCE ) objective for discriminating different instances in the dataset . Suppose z1i and z 2 i are the representations of two augmented views of xi and zk is a different instance . The NCE objective can be expressed by Eq . ( 1 ) , where the similarity function sim ( · ) represents the dot product between L2 normalized vectors sim ( u , v ) = uTv/‖u‖‖v‖ and τ is the temperature parameter . LNCE = − log exp ( sim ( z1 , z2 ) /τ ) exp ( sim ( z1i , z 2 i ) /τ ) + ∑N k=1 exp ( sim ( z 1 i , zk ) /τ ) . ( 1 ) BYOL [ 23 ] and SimSiam [ 9 ] style framework add an additional non-linear predictor head q ( · ) which further maps z to p. The model will minimize the negative cosine similarity ( equivalent to minimize the L2 distance ) between z to p. Lcos = − p1 ‖p1‖ · z 2 ‖z2‖ , Lmse = ‖p1 − z2‖22 . ( 2 ) Tricks like stop-gradient and momentum teacher are often applied to avoid model collapsing . 3.2 Relational Self-Supervised Learning . In classical self-supervised learning , different instances are to be pushed away from each other , and augmented views of the same instance is expected to be of exactly the same features . However , both constrains are too restricted because of the existence of similar samples and the distorted semantic information if aggressive augmentation is adopted . In this way , we do not encourage explicit negative instances ( those to be pushed away ) for each instance ; instead , we leverage the pairwise similarities as a manner to explore their relationships . And we pull the features of two different augmentations in this sense of relation metric . As a result , our method relaxes both ( 1 ) and ( 2 ) , where different instances do not always need to be pushed away from each other ; and augmented views of the same instance only need to share the similar but not exactly the same features . Concretely , given a image x in a batch of samples , two different augmented views can be obtained by x1 = T ( x , θ1 ) , x2 = T ( x , θ2 ) and calculate the corresponds embedding z1 = g ( F ( x1 ) ) , z2 = g ( F ( x2 ) ) . Then , we calculate the similarities between the instances of the first augmented images . Which can be measured by sim ( z1 , zi ) . A softmax layer can be adopted to process the calculated similarities , which then produces a relationship distribution : p1i = exp ( sim ( z1 , zi ) /τt ) ∑K k=1 exp ( sim ( z 1 , zk ) /τt , ) . ( 3 ) where τt is the temperature parameter . At the same time , we can calculate the relationship between x2 and the i-th instance as sim ( z2 , zi ) . The resulting relationship distribution can be written as : p2i = exp ( sim ( z2 , zi ) /τs ) ∑K k=1 exp ( sim ( z 2 , zk ) /τs , ) . ( 4 ) where τs is a different temperature parameter . We propose to push the relational consistency between p1i and p 2 i by minimizing the Kullback–Leibler divergence , which can be formulated as : Lrelation = DKL ( p1||p2 ) = H ( p1 , p2 ) −H ( p1 ) . ( 5 ) Since the p1 will only be used as a target , we only minimize H ( p1 , p2 ) in our implementation . More efficiency with Momentum targets . However , the quality of the target similarity distribution p1 is crucial , to make the similarity distribution reliable and stable , we usually require a large batch size which is very unfriendly to GPU memories . To resolve this issue , we utilize a “ momentum update '' network as in [ 24 , 8 ] , and maintain a large memory buffer Q of K past samples { zk|k = 1 , ... , K } ( following the FIFO principle ) for storing the feature embeddings from the past batches , which can then be used for simulating the large batch size relationship and providing a stable similarity distribution . Ft ← mFt + ( 1−m ) Fs , gt ← mgt + ( 1−m ) gs , ( 6 ) where Fs and gs denote the most latest encoder and head , respectively , so we name them as the student model with a subscript s. On the other hand , Ft and gt stand for ensembles of the past encoder and head , respectively , so we name them as the teacher model with a subscript t. m represents the momentum coefficient which controls how fast the teacher Ft will be updated . Sharper Distribution as Target . Note , the value of τt has to be smaller than τs since τt will be used to generate the target distribution . A smaller τ will result in a “ sharper '' distribution which can be interpreted as highlight the most similar feature for z1 . Align p2 with p1 can be regarded as pulling z2 towards the features that are similar with z1 . Weak Augmentation Strategy for Teacher . To further improve the quality and stability of the target distribution , we adopt a weak augmentation strategy for the teacher model since the standard contrastive augmentation is too aggressive , which introduced too many disturbances and will mislead the student network . Please refer to more details in our empirical study . Compare with SEED and CLSA . SEED [ 21 ] follows the standard Knowledge Distillation ( KD ) paradigm [ 26 , 49 , 18 ] where it aims to distill the knowledge from a larger network into a smaller architecture . The knowledge transfer happens in the same view but between different models . In our framework , we are trying to maintain the relational consistency between different augmentations ; the knowledge transfer happens between different views but in the same network . CLSA [ 46 ] also introduced the concept of using weak augmentation to guide a stronger augmentation . However , the “ weak '' augmentation in CLSA is equivalent to the “ strong '' augmentation in our method ( We do not use any stronger augmentations such as [ 12 , 13 ] ) . On the other hand , CLSA still adopts the InfoNCE loss ( 1 ) for instance discrimination , where our proposed method only utilized the relational consistency loss ( 5 ) . Finally , CLSA requires at least one additional sample during training , which will slow down the training speed . Algorithm 1 : Relational Self-supervised Learning with Weak Augmentation ( ReSSL ) Input : x : a batch of samples . Tw ( · ) : Weak augmentation function . Tc ( · ) : Contrastive augmentation function . Ft and Fs : the teacher and student backbone network . gt and gs : the non-linear projection head for teacher and student . Q : the memory buffer while network not converge do for i=1 to step do Fetch x from current batch B z1 = gt ( Ft ( Tw ( x , θ1 ) ) ) ; z2 = gs ( Fs ( Tc ( x , θ2 ) ) ) ; p1 = SoftMax ( z1QT / τt ) ; p2 = SoftMax ( z2QT / τs ) ; // Eq . ( 3 ) ( 4 ) Calculate Lrelation loss by CrossEntropy ( p1 , p2 ) ; // Eq . ( 5 ) Update Fs and gs with loss Lrelation Update Ft and gt by Ft ← mFt + ( 1−m ) Fs , gt ← mgt + ( 1−m ) gs ; // Eq . ( 6 ) Update the memory buffer Q by z1 end end Output : The well trained model Fs

Summary: This work focuses on self-supervised learning for image classification. Previous work assumes that images and the corresponding augmentation should have similar features (or cluster into the same class) as the objective function to train the model. In contrast, this paper proposes to make the contrastive loss (positive and negative) of two different weak augmentations of a sample to be consistency. Moreover, knowledge distillation is also adopted to sharpen the target distribution for better capturing the most important relationship. Experimental results show that the proposed approach outperforms state-of-the-art approaches, e.g, SimCLR, BYOL, SimSiam, MoCoV2, with less computation requirement.