Unsupervised Representation Transfer for Small Networks: I Believe I Can Distill On-the-Fly

A current remarkable improvement of unsupervised visual representation learning is based on heavy networks with large-batch training . While recent methods have greatly reduced the gap between supervised and unsupervised performance of deep models such as ResNet-50 , this development has been relatively limited for small models . In this work , we propose a novel unsupervised learning framework for small networks that combines deep self-supervised representation learning and knowledge distillation within one-phase training . In particular , a teacher model is trained to produce consistent cluster assignments between different views of the same image . Simultaneously , a student model is encouraged to mimic the prediction of on-the-fly self-supervised teacher . For effective knowledge transfer , we adopt the idea of domain classifier so that student training is guided by discriminative features invariant to the representational space shift between teacher and student . We also introduce a network driven multi-view generation paradigm to capture rich feature information contained in the network itself . Extensive experiments show that our student models surpass state-of-the-art offline distilled networks even from stronger self-supervised teachers as well as top-performing self-supervised models . Notably , our ResNet-18 , trained with ResNet-50 teacher , achieves 68.3 % ImageNet Top-1 accuracy on frozen feature linear evaluation , which is only 1.5 % below the supervised baseline . 1 Introduction . Recently , there has been a growing attention in unsupervised and self-supervised learning where the goal is to effectively learn useful features from a large amount of unlabeled data . Current selfsupervised visual representation learning methods appear to approach and possibly even outperform the fully-supervised counterpart [ 8 , 25 , 29 ] . All top-performing self-supervised learning ( SSL ) algorithms for visual representation involve training deep models on powerful computers . In particular , their smallest architecture is ResNet-50 [ 26 ] , and the networks are trained with large batches ( e.g. , 4096 images ) on multiple specialized hardware devices such as 128 TPU cores [ 6 , 8 , 25 ] . Yet , this heavy implementation is not a viable option in a resource-limited environment , and there is evidence that the SSL methods even do not work well on light models ( Figure 1 and [ 22 ] ) . Also , we need strong small networks that can operate on a system on a chip ( SoC ) for real-world applications . Existing methods on SSL have demonstrated that deeper models learn general visual representation more effectively with unlabeled data [ 8 , 40 ] . Moreover , it has been empirically shown that predictive performance of bigger networks is better transferred into smaller ones [ 9 ] . Inspired by these , we are 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . interested in pretraining a light-weight model by distilling representational knowledge from a deep self-supervised network instead of directly performing self-supervised training on a small one . Most of previous distillation methods in unsupervised representation learning literature are offline in the sense that they leverage an already trained self-supervised teacher to transfer feature information to students [ 22 , 31 , 40 , 57 ] . Often , their sequential training pipelines require extra pre/post-processing steps in order to boost performance . Unlike these approaches , we propose a novel unsupervised representation learning framework for small networks that combines deep self-supervised training and knowledge distillation [ 30 ] within one-phase training . To be specific , our teacher learns clusters and representation , and at the same time the student is trained to align its prediction to the clustering of On-the-fly Self-Supervised teacher . From this , we refer to our method as OSS . The main advantage over offline distillation is that our teacher better distills proper signals into the student at each training stage by solving the same task online . In distillation process , it is difficult for a low-capacity student to perfectly mimic the large teacher ’ s behavior [ 36 ] . To deal with this , we incorporate the idea of domain classifier [ 23 ] . That is , we add a feature classifier trained not to be able to distinguish between teacher and student embeddings , and this leads to the emergence of representational-space-invariant features in the course of joint training . There is evidence that increasing the number of different views improves the quality of resulting features in SSL [ 6 , 8 , 55 ] . Existing approaches to multi-view generation rely on random image transformation techniques . Different from these methods , we introduce a new network driven paradigm in order to utilize rich feature information contained in the network itself . In particular , we apply random dropout and empirically show that this apparently improves label efficiency ( Section 5 ) . Our method is conceptually simple but surprisingly works well with typical 256-batch training on a single 8-GPU machine . Extensive experiments show that our student networks outperform state-of-the-art self-supervised models ( Figure . 1 ) as well as students from top-performing distillation approaches even with stronger self-supervised teachers ( Tables 1 and 5 ) on linear evaluation . We also evaluate learned representations on series of vision tasks with multiple network architectures and various benchmark datasets and demonstrate significant performance gain ( Tables 2 , 3 and 4 and Figure 3 ) . Overall , we make the following contributions . • We propose a novel unsupervised visual representation learning framework for small networks that simultaneously conducts self-supervised teacher training and knowledge distillation . • We introduce a network driven view generation paradigm and an adversarial distillation scheme to better capture representational information from networks . This is the first work to show their effectiveness in the context of unsupervised visual representation learning . • The proposed method significantly improves unsupervised pretraining performance of small models , and its effectiveness is shown under multiple settings of network architectures , datasets and vision tasks . 2 Related Work . Unsupervised visual representation learning . In unsupervised or self-supervised visual representation learning , we learn rich feature information by solving pretext tasks whose labels can be obtained from image data itself . Learning methods can be divided into several categories by types of pretext tasks used therein : construction-based methods [ 43 , 62 ] , prediction-based methods [ 10 , 20 , 24 , 25 , 38 , 39 ] , cluster-based methods [ 3 , 5 , 6 ] , and contrastive methods [ 8 , 11 , 29 ] . Our work is mostly related to [ 6 ] , a recent clustering-based method . Multi-view generation . Contrastive and clustering-based self-supervised visual representation learning methods exploit random image transformation techniques to generate different views of each image . This is because comparing more views during training improves the performance of the resulting model [ 6 , 8 , 37 , 55 ] . Our method uses dropout [ 50 ] for obtaining more views to better utilize useful feature information as well as applying existing data augmentation based strategies . Knowledge distillation from self-supervised models . Most of knowledge distillation [ 30 ] techniques are proposed in the context of supervised learning , where a shallow network is trained via supervision signals from both labels and a deep network [ 1 , 41 , 44 , 47 , 52 , 53 , 58 ] . In unsupervised learning settings , recent studies [ 40 ] and [ 57 ] exploit clustering of teacher ’ s embeddings as distillation signals . More recently , CompRess [ 31 ] maintains queues for storing data sample ’ s embeddings and calculates instance-level similarity scores between all the samples in the queue and the input features extracted from teacher ( and student ) . Its training objective is to minimize the distance between teacher and student score distributions . SEED [ 22 ] is a similar concurrent work of [ 31 ] . Unlike aforementioned methods that rely on already trained static teachers , our framework injects a distillation step into clustering-based self-supervised teacher learning . The only existing online technique is [ 4 ] , which combines contrastive learning and co-teaching . Domain adversarial learning . DANN [ 23 ] learns features that are discriminative for the main classification task and indiscriminate with respect to the domain shift . This is achieved by adding a domain classifier with a gradient reversal layer . In particular , network parameters are updated to minimize the main classifier loss and to maximize the domain classifier loss via a gradient reversal layer that reverses the sign of gradient . We adopt this idea in our framework to transfer discriminative features invariant to the representational space change between teacher and student . Similar idea with a generative adversarial network [ 35 ] is applied in supervised knowledge distillation in [ 12 ] . 3 Proposed Framework . Our goal is to learn visual features of small networks effectively without access to any labeled data . In order to do this , we embed a distillation step into the process of on-the-fly self-supervised teacher learning . Inspired by [ 57 ] , which empirically shows clustering improves generalization power of representation , our online framework is developed upon [ 6 ] , the state-of-the-art clustering-based self-supervised method . We also incorporate the idea of domain classifier to facilitate learning of useful features invariant to the representational space change between teacher and student . In a nutshell , the training objective function of our network is given by L = LT + LS + LD , ( 1 ) where LT , LS , and LD are designed for teacher network training ( Sections 3.1 and 3.2 ) , knowledge distillation for student network ( Sections 3.1 and 3.2 ) , and domain classifier ( Section 3.3 ) , respectively . Overall architecture is presented in Figure 2 , and a pseudo code is available in A.1 . 3.1 Teacher training and knowledge distillation . We now outline teacher network ( e.g . ResNet-50 ) training and knowledge distillation framework . Our architecture takes as input two randomly augmented views xn1 and xn2 of an image xn . The two views are processed by teacher ’ s backbone network fθ and a projection head hT that consists in a multi-layer perceptron ( MLP ) . The processed features are then projected to the unit sphere . We denote teacher ’ s backbone features as fn1 = fθ ( xn1 ) and fn2 = fθ ( xn2 ) , and the projected outputs as zn1 = hT ( fn1 ) ‖hT ( fn1 ) ‖2 and zn2 = hT ( fn2 ) ‖hT ( fn2 ) ‖2 . Also , the two images are feed-forwarded to student ’ s backbone network gη ( e.g. , ResNet-18 ) . The outputs are projected via an MLP head hS and then normalized as well . Similarly , we denote student ’ s backbone features as gn1 = gη ( xn1 ) and gn2 = gη ( xn2 ) , and their projected output vectors as wn1 = hS ( gn1 ) ‖hS ( gn1 ) ‖2 and wn2 = hS ( gn2 ) ‖hS ( gn2 ) ‖2 . We randomly initialize prototype vectorsC = { c1 , . . . , cK } , whereK is the number of clusters . Code vectors , qn ’ s , for cluster assignments are computed from mapping teacher ’ s embedding representation zn ’ s to the trainable prototype vectors C and then solving an optimization problem for online clustering . We then solve a swapped prediction problem [ 6 ] for teacher and simultaneously distill teacher ’ s knowledge into student by enforcing a consensus in prediction . Teacher training . A training objective for the teacher network is to solve a swapped cluster prediction problem [ 6 ] of predicting the code qn1 from the feature zn2 and qn2 from zn1 : Lt ( z1 , z2 ) = ∑ n [ Lc ( zn1 , qn2 ) + Lc ( zn2 , qn1 ) ] . ( 2 ) Here , Lc is the cross entropy loss between the codes and softmax probabilities ( with a temperature parameter τ ) of dot products of zni and all prototypes in C : Lc ( zn1 , qn2 ) = − ∑ k q ( k ) n2 log p ( k ) n1 where p ( k ) n1 = exp ( 1 τ z > n1ck ) ∑ k′ exp ( 1 τ z > n1ck′ ) . ( 3 ) The loss function ( 2 ) is jointly minimized with respect to the prototypes C and the parameters of teacher backbone fθ and projection head hT implicitly contained in zn ’ s . The codes qn ’ s are computed using Sinkhorn-Knopp algorithm [ 18 ] under an equi-partition constraint for each augmented mini-batch with stored embeddings as done in [ 6 ] . Knowledge distillation . The student model is trained to follow teacher ’ s code assignment . The training loss has exactly the same form of ( 2 ) with features from student and targets from teacher ’ s codes : Ls ( w1 , w2 ) = ∑ n [ Lc ( wn1 , qn2 ) + Lc ( wn2 , qn1 ) ] . ( 4 ) Specifically , the cross-entropy loss Lc is similarly calculated using student ’ s embedding vectors wn ’ s and the prototype vectors C. Again , minimization of this loss is performed with respect to the prototypes C and the parameters of student backbone gη and projection head hS used to produce wn ’ s .

Summary: This paper presents a self-supervised distillation method that aims to learn better-performing representations for small models.  In particular, the teacher model is trained to produce consistent cluster assignments between different small views of the image while the student model is encouraged to mimic the prediction of the teacher. Comprehensive experiments on ImageNet and other downstream tasks validate the effectiveness of the self-supervised distillation method.