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: The authors target an interesting and valuable problem: how to train an unsupervised/self-supervised small model by distillation large models. They proposed the OSS framework, which is equipped with a distillation strategy, a multi-view generation strategy, and an adversarial feature generation strategy. The proposed framework is shown to work better than current SOTA self-supervised frameworks and distillation methods. In some specific regimes, the proposed framework works even better than the supervised models. This work makes a valuable exploration for unsupervised/self-supervised learning with light-weight models.