SEED: Self-supervised Distillation For Visual Representation

1 INTRODUCTION The burgeoning studies and success on self-supervised learning ( SSL ) for visual representation are mainly marked by its extraordinary potency of learning from unlabeled data at scale . Accompanying with the SSL is its phenomenal benefit of obtaining task-agnostic representations while allowing the training to dispense with prohibitively expensive data labeling . Major ramifications of visual SSL include pretext tasks ( Noroozi & Favaro , 2016 ; Zhang et al. , 2016 ; Gidaris et al. , 2018 ; Zhang et al. , 2019 ; Feng et al. , 2019 ) , contrastive representation learning ( Wu et al. , 2018 ; He et al. , 2020 ; Chen et al. , 2020a ) , online/offline clustering ( Yang et al. , 2016 ; Caron et al. , 2018 ; Li et al. , 2020 ; Caron et al. , 2020 ; Grill et al. , 2020 ) , etc . Among them , several recent works ( He et al. , 2020 ; Chen et al. , 2020a ; Caron et al. , 2020 ) have achieved comparable or even better accuracy than the supervised pre-training when transferring to downstream tasks , e.g . semi-supervised classification , object detection . The aforementioned top-performing SSL algorithms all involve large networks ( e.g. , ResNet-50 ( He et al. , 2016 ) or larger ) , with , however , little attention on small networks . Empirically , we find that existing techniques like contrastive learning do not work well on small networks . For instance , the linear probe top-1 accuracy on ImageNet using MoCo-V2 ( Chen et al. , 2020c ) is only 36.3 % with MobileNetV3-Large ( see Figure 1 ) , which is much lower compared with its supervised training accuracy 75.2 % ( Howard et al. , 2019 ) . For EfficientNet-B0 , the accuracy is 42.2 % compared with its supervised training accuracy 77.1 % ( Tan & Le , 2019 ) . We conjecture that this is because smaller models with fewer parameters can not effectively learn instance level discriminative representation with large amount of data . To address this challenge , we inject knowledge distillation ( KD ) ( Buciluǎ et al. , 2006 ; Hinton et al. , 2015 ) into self-supervised learning and propose self-supervised distillation ( dubbed as SEED ) as a new learning paradigm . That is , train the larger , and distill to the smaller both in self-supervised manner . Instead of directly conducting self-supervised training on a smaller model , SEED first trains a large model ( as the teacher ) in a self-supervised way , and then distills the knowledge to the smaller model ( as the student ) . Note that the conventional distillation is for supervised learning , while the distillation here is in the self-supervised setting without any labeled data . Supervised distillation can be formulated as training a student to mimic the probability mass function over classes predicted by a teacher model . In unsupervised knowledge distillation setting , however , the distribution over classes is not directly attainable . Therefore , we propose a simple yet effective self-supervised distillation method . Similar to ( He et al. , 2020 ; Wu et al. , 2018 ) , we maintain a queue of data samples . Given an instance , we first use the teacher network to obtain its similarity scores with all the data samples in the queue as well as the instance itself . Then the student encoder is trained to mimic the similarity score distribution inferred by the teacher over these data samples . The simplicity and flexibility that SEED brings are self-evident . 1 ) It does not require any clustering/prototypical computing procedure to retrieve the pseudo-labels or latent classes . 2 ) The teacher model can be pre-trained with any advanced SSL approach , e.g. , MoCo-V2 ( Chen et al. , 2020c ) , SimCLR ( Chen et al. , 2020a ) , SWAV ( Caron et al. , 2020 ) . 3 ) The knowledge can be distilled to any target small networks ( either shallower , thinner , or totally different architectures ) . To demonstrate the effectiveness , we comprehensively evaluate the learned representations on series of downstream tasks , e.g. , fully/semi-supervised classification , object detection , and also assess the transferability to other domains . For example , on ImageNet-1k dataset , SEED improves the linear probe accuracy of EfficientNet-B0 from 42.2 % to 67.6 % ( a gain over 25 % ) , and MobileNet-V3 from 36.3 % to 68.2 % ( a gain over 31 % ) compared to MoCo-V2 baselines , as shown in Figure 1 and Section 4 . Our contributions can be summarized as follows : • We are the first to address the problem of self-supervised visual representation learning for small models . • We propose a self-supervised distillation ( SEED ) technique to transfer knowledge from a large model to a small model without any labeled data . • With the proposed distillation technique ( SEED ) , we significantly improve the state-of-theart SSL performance on small models . • We exhaustively compare a variety of distillation strategies to show the validity of SEED under multiple settings . 2 RELATED WORK . Among the recent literature in self-supervised learning , contrastive based approaches show prominent results on downstream tasks . Majority of the techniques along this direction are stemming from noise-contrastive estimation ( Gutmann & Hyvärinen , 2010 ) where the latent distribution is estimated by contrasting with randomly or artificially generated noises . Oord et al . ( 2018 ) first proposed Info-NCE to learn image representations by predicting the future using an auto-regressive model for unsupervised learning . Follow-up works include improving the efficiency ( Hénaff et al. , 2019 ) , and using multi-view as positive samples ( Tian et al. , 2019b ) . As these approaches can only have the access to limited negative instances , Wu et al . ( 2018 ) designed a memory-bank to store the previously seen random representations as negative samples , and treat each of them as independent categories ( instance discrimination ) . However , this approach also comes with a deficiency that the previously stored vectors are inconsistent with the recently computed representations during the earlier stage of pre-training . Chen et al . ( 2020a ) mitigate this issue by sampling negative samples from a large batch . Concurrently , He et al . ( 2020 ) improve the memory-bank based method and propose to use the momentum updated encoder for the remission of representation inconsistency . Other techniques include Misra & Maaten ( 2020 ) that combines the pretext-invariant objective loss with contrastive learning , and Wang & Isola ( 2020 ) that decomposes contrastive loss into alignment and uniformity objectiveness . Knowledge distillation ( Hinton et al. , 2015 ) aims to transfer knowledge from a cumbersome model to a smaller one without losing too much generalization power , which is also well investigated in model compression ( Buciluǎ et al. , 2006 ) . Instead of mimicking the teacher ’ s output logit , attention transfer ( Zagoruyko & Komodakis , 2016 ) formulates knowledge distillation on attention maps . Similarly , works in ( Ahn et al. , 2019 ; Yim et al. , 2017 ; Koratana et al. , 2019 ; Huang & Wang , 2017 ) have utilized different learning objectives including consistency on feature maps , consistency on probability mass function , and maximizing the mutual information . CRD ( Tian et al. , 2019a ) , which is derived from CMC ( Tian et al. , 2019b ) , optimizes the student network by a similar objective to Oord et al . ( 2018 ) using a derived lower bound on mutual information . However , the aforementioned efforts all focus on task-specific distillation ( e.g. , image classification ) during the fine-tuning phase rather than a task-agnostic distillation in the pre-training phase for the representation learning . Several works on natural language pre-training proposed to leverage knowledge distillation for a smaller yet stronger small models . For instances , DistillBert ( Sanh et al. , 2019 ) , TinyBert ( Jiao et al. , 2019 ) , and MobileBert ( Sun et al. , 2020 ) , have used knowledge distillation for model compression and shown their validity on multiple downstream tasks . Similar works also emphasize the value of smaller and faster models for language representation learning by leveraging knowledge distillation ( Turc et al. , 2019 ; Sun et al. , 2019 ) . These works all demonstrate the effectiveness of knowledge distillation for language representation learning in small models , while are not extended to the pre-training for visual representations . Notably , a recent concurrent work CompRess ( Abbasi Koohpayegani et al. , 2020 ) also point out the importance to develop better SSL method for smaller models . SEED closely relates to the above techniques but aims to facilitate visual representation learning during pre-training phase using distillation technique for small models , which as far as we know has not yet been investigated . 3 METHOD . 3.1 PRELIMINARY ON KNOWLEDGE DISTILLATION . Knowledge distillation ( Hinton et al. , 2015 ; Buciluǎ et al. , 2006 ) is an effective technique to transfer knowledge from a strong teacher network to a target student network . The training task can be generalized as the following formulation : θ̂S = argmin θS N∑ i Lsup ( xi , θS , yi ) + Ldistill ( xi , θS , θT ) , ( 1 ) where xi is an image , yi is the corresponding annotation , θS is the parameter set for the student network , and θT is the set for the teacher network . The loss Lsup is the alignment error between the network prediction and the annotation . For example in image classification task ( Mishra & Marr , 2017 ; Shen & Savvides , 2020 ; Polino et al. , 2018 ; Cho & Hariharan , 2019 ) , it is normally a cross entropy loss . For object detection ( Liu et al. , 2019 ; Chen et al. , 2017 ) , it includes bounding box regression as well . The loss of Ldistill is the mimic error of the student network towards a pre-trained teacher network . For example in ( Hinton et al. , 2015 ) , the teacher signal comes from the softmax prediction of multiple large-scale networks and the loss is measured by the Kullback–Leibler divergence . In Romero et al . ( 2014 ) , the task is to align the intermediate feature map values and to minimize the squared l2 distance . The effectiveness has been well demonstrated in the supervised setting with labeled data , but remains unknown for the unsupervised setting , which is our focus . 3.2 SELF-SUPERVISED DISTILLATION FOR VISUAL REPRESENTATION . Different from supervised distillation , SEED aims to transfer knowledge from a large model to a small model without requiring labeled data , so that the learned representations in small model can be used for downstream tasks . Inspired by contrastive SSL , we formulate a simple approach for the distillation on the basis of instance similarity distribution over a contrastive instance queue . Similar to He et al . ( 2020 ) , we maintain an instance queue for storing data samples ’ encoding output from the teacher . Given a new sample , we compute its similarity scores with all the samples in the queue using both the teacher and the student models . We require that the similarity score distribution computed by the student matches with that computed by the teacher , which is formulated as minimizing the cross entropy between the student and the teacher ’ s similarity score distributions ( as illustrated in Figure 2 ) . Specifically , for a randomly augmented view xi of an image , it is first mapped and normalized into feature vector representations zTi = f T θ ( xi ) /||fTθ ( xi ) ||2 , and zSi = fSθ ( xi ) /||fSθ ( xi ) ||2 , where zTi , z S i ∈ RD , and fTθ and fSθ denote the teacher and student encoders , respectively . Let D = [ d1 ... dK ] denote the instance queue where K is the queue length and dj is the feature vector obtained from the teacher encoder . Similar to the contrastive learning framework , D is progressively updated under the “ first-in first-out ” strategy as distillation proceeds . That is , we en-queue the visual features of the current batch inferred by the teacher and de-queue the earliest seen samples at the end of iteration . Note that the maintained samples in queue D are mostly random and irrelevant to the target instance xi . Minimizing the cross entropy between the similarity score distribution computed by the student and teacher based on D softly contrasts xi with randomly selected samples , without directly aligning with the teacher encoder . To address this problem , we add the teacher ’ s embedding ( zTi ) into the queue and form D + = [ d1 ... dK , dK+1 ] with dK+1 = zTi . Let pT ( xi ; θT ; D+ ) denote the similarity score between the extracted teacher feature zTi and dj ’ s ( j = 1 , ... , K + 1 ) computed by the teacher model . pT ( xi ; θT ; D+ ) is defined as pT ( xi ; θT , D + ) = [ pT1 ... p T K+1 ] , p T j = exp ( zTi · dj/τT ) ∑ d∼D+ exp ( z T i · d/τT ) , ( 2 ) and τT is a temperature parameter for the teacher . Note , we use ( ) T to represent the feature from the teacher network and use ( · ) to represent the inner product between two features . Similarly let pS ( xi ; θS , D+ ) denote the similarity score computed by the student model , which is defined as pS ( xi ; θS , D + ) = [ pS1 ... p S K+1 ] , where p S j = exp ( zSi · dj/τS ) ∑ d∼D+ exp ( z S i · d/τS ) , ( 3 ) and τS is a temperature parameter for the student . Our self-supervised distillation can be formulated as minimizing the cross entropy between the similarity scores of the teacher , pT ( xi ; θT , D+ ) , and the student , pS ( xi ; θS , D+ ) , over all the instances xi , that is , θ̂S = argmin θS N∑ i −pT ( xi ; θT , D+ ) · logpS ( xi ; θS , D+ ) = argmin θS N∑ i K+1∑ j − exp ( z T i · dj/τT ) ∑ d∼D+ exp ( z T i · d/τT ) · log exp ( z S i · dj/τS ) ∑ d∼D+ exp ( z S i · d/τS ) . ( 4 ) Since the teacher network is pre-trained and frozen , the queued features are consistent during training w.r.t . the student network . The higher the value of pTj is , the larger weight will be laid on p S j . Due to the l2 normalization , similarity score between zTi and dK+1 remains constant 1 before softmax normalization , which is the largest among pTj . Thus , the weight for p S K+1 is the largest and can be adjusted solely by tuning the value of τT . By minimizing the loss , the feature of zSi can be aligned with zTi and meanwhile contrasts with other unrelated image features in D. We further discuss the relation of these two goals with our learning objective in Appendix A.5 . Relations with Info-NCE loss . When τT → 0 , the softmax function for pT smoothly approaches to a one-hot vector , where pTK+1 equals 1 and all others 0 . In this extreme case , the loss becomes LNCE = N∑ i − log exp ( z T i · zSi /τ ) ∑ d∼D+ exp ( z S i · d/τ ) , ( 5 ) which is similar to the widely-used Info-NCE loss ( Oord et al. , 2018 ) in contrastive-based SSL ( see discussion in Appendix A.6 .

Summary: This paper proposes a  knowledge distillation (KD) approach for self-supervised learning (SSL) with small neural network models. The authors first observe that the state-of-the-art contrastive learning-based SSL does not obtain good performance on small models, due  to the larger  model capacity required for instance discrimination. To tackle this problem, they propose a SEED, a  KD method where the smaller student model learns to mimic its larger teacher model’s similarity distribution between an instance and its augmented views, using a cross-entropy based objective. The authors perform various experiments to show that -- 1)  SEED obtains substantial improvement in SSL-based imagenet classification performance for small models as compared to SSL training without SEED, 2) the performance gains are also substantial for transfer learning on other classification tasks, 3) the performance gains are smaller for downstream tasks of object detection and instance segmentation, with performance gains reducing for the larger COCO dataset, as compared to VOC, 4) SEED is robust to choice of SSL method, and performs better than other KD approaches.