Rethinking Supervised Pre-Training for Better Downstream Transferring

1 INTRODUCTION . Pre-training neural networks on upstream datasets and fine-tuneing the pre-trained model on downstream tasks has been an important methodology in applications of deep learning ( Tan et al. , 2018 ) . Such a pretrain-finetune paradigm generally works with pre-training on large-scale diverse datasets and fine-tuning on small specific datasets , and has been widely applied in a number of applications ( Devlin et al. , 2019 ; Kolesnikov et al. , 2020 ; Brown et al. , 2020 ) . Specifically , in the area of computer vision , we often apply supervised learning methods ( e.g . cross entropy ) to pre-train a model from labeled dataset ( e.g . ImageNet ( Deng et al. , 2009 ) and Kinetics ( Carreira & Zisserman , 2017 ) ) , and fine-tune it on downstream tasks such as object detection ( Ren et al. , 2017 ) , instance segmentation ( He et al. , 2017 ) , and video understanding ( Gu et al. , 2018 ) . Except supervised pre-training , recent works demonstrate that self-supervised pre-training without label information can also learn effective representation from upstream data and even surpass supervised methods when transferring to downstream tasks ( Chen et al. , 2020a ; He et al. , 2020 ; Grill et al. , 2020 ; Chen & He , 2021 ; Zbontar et al. , 2021 ) . Unlike supervised pre-training that focuses on category-level discrimination , self-supervised pre-training is mainly based on instance discrimination , where models are trained to keep each instance and its augmentations close to each other , and at the same time , separate them from other instances and their augmentations . It effectively captures many important and discriminative features that are useful for downstream tasks . However , without appropriate guidance of supervision information , self-supervised pre-training lacks the ability of mining high-level semantic features and may capture detailed but irrelevant features ( e.g . visual features related to special background ) , resulting in unsatisfying performance on challenging downstream tasks , e.g . fine-grained classification ( Islam et al. , 2021 ) . Therefore , we aim to improve the downstream transferability for supervised pre-training . Figure 2 illustrates two representative supervised pre-training methods , cross entropy ( C.E . ) and supervised contrastive learning ( SupCon , Khosla et al . ( 2020 ) ) , a soft-nearest neighbors loss ( Goldberger et al. , 2005 ; Wu et al. , 2018a ) . To distinguish instances from different classes , they are designed to minimize intra-class variance by pushing all the instances of the same class close to each or to certain centers . For a class with diverse visual appearance ( e.g . the four cat examples from Figure 2 ) , these approaches may “ ruin ” the natural representation of images by bringing images with completely different visual appearances next to each other . As a result , they tend to skip features that capture intra-class difference but are less correlated with classes defined in upstream tasks , leading to the problem of overfitting upstream tasks . In this paper , we propose a new supervised pre-training method based on Leave-One-Out Knearest-neighbor classification , or LOOK for short , that effectively alleviates the problem of neglecting intra-class difference and thus significantly improves transferability for downstream tasks . In particular , a weighted kNN classifier is used to replace the linear or MLP predictor in the last layer of deep neural network , and a leave-one-out classification error is used as the loss function for optimization . Because of the nature of kNN classifier , each instance is only required to share the same class with most of its k nearest neighbors , allowing each class to exhibit multi-mode distribution and consequentially better preserving the features related to intra-class difference , as shown in Figure 2 . We also develop efficient implementation for LOOK that scales well to large datasets . Extensive empirical studies demonstrated that LOOK has better transferrability for downstream tasks than the existing methods for supervised and self-supervised pre-training . 2 RELATED WORKS . Pre-training models have achieved remarkable performance on multiple fields of artificial intelligence . In computer vision , various DNN architectures ( Krizhevsky et al. , 2012 ; Simonyan & Zisserman , 2014 ; He et al. , 2016 ; Dosovitskiy et al. , 2020 ) show significant performance improvements on downstream tasks when pre-trained from large labeled datasets , e.g. , ImageNet ( Deng et al. , 2009 ) , JFT-300M ( Sun et al. , 2017 ) . Pre-trained model provide an essential initialization for complex vision tasks , such as detection , segmentation and etc. , which work as a important part for the convergence of the model training . In natural language processing , pre-trained models ( Devlin et al. , 2019 ; Liu et al. , 2019 ; Brown et al. , 2020 ) also have made great progress . Recently , contrastive learning has made impressive progress in the representation learning task of computer vision . Self-supervised learning methods ( Chen et al. , 2020a ; He et al. , 2020 ) are em- Cat Dog parametric centers range of k-NN Cross Entropy ( C.E . ) training samples training samples Supervised Contrastive Learning ( SupCon ) training samples Leave-One-Out K-Nearest-Neighbor ( LOOK ) optimize Figure 2 : Comparison of LOOK and existing supervised pre-training methods . For C.E . and SupCon , they push all samples from the same class to certain centers or closer to each other , respectively , while LOOK only requires samples next to at most their k nearest neighbors . ployed to pre-train on ImageNet dataset and significantly improve the transferability on multiple downstream tasks over supervised counterparts ( Islam et al. , 2021 ) . Derived from contrastive learning , more self-supervised learning methods ( Chen et al. , 2020b ; Grill et al. , 2020 ; Zbontar et al. , 2021 ; Chen & He , 2021 ) also achieve impressive performance on downstream tasks . Unlike earlier studies of contrastive learning , these methods introduce either momentum encoder or additional projector and predictor to avoid the introduction of negative pairs , significantly improving the performance . Apart from self-supervised learning , contrastive learning is also developed in supervised pre-training ( Khosla et al. , 2020 ; Islam et al. , 2021 ) and further boost the transferability of the pretrained model . In addition to self-supervised methods based on contrastive learning , there are also some unsupervised pre-training methods that focus on learning representation features by designing clever pretext tasks ( Noroozi & Favaro , 2016 ; Gidaris et al. , 2018 ; Trinh et al. , 2019 ; Bao et al. , 2021 ) and structure formed clustering ( Caron et al. , 2018 ; Asano et al. , 2020 ; Li et al. , 2021 ) . Besides designing pre-training method , there are also works ( Ganin & Lempitsky , 2015 ; Yosinski et al. , 2014 ; Li et al. , 2018 ; 2019 ; Chen et al. , 2019a ; Tan et al. , 2018 ) focused on designing downstream training methods to improve transferablity through eliminating domain shifts between the upstream and downstream datasets or utilize an inductive bias to improve transferability . 3 LOOK : A LEAVE-ONE-OUT KNN BASED PRE-TRAINING METHOD . 3.1 MOTIVATION OF LOO-KNN . As already discussed in the introduction , existing supervised pre-training methods suffer from the problem of overfitting upstream tasks , i.e . they tend to keep instances from the same class close to each other , thus neglecting intra-class difference that may be important for downstream tasks . We believe that source of this problem arises from simple prediction models ( e.g . linear classifier with fully-connected layer ) are used in supervised pre-training methods . Figure 1 shows the data distribution resulting from C.E . based pre-training . We see clearly that using a simple linear classifier essentially requires all data points of the same class close to the class center . As a remedy , we propose to replace simple prediction model with kNN classifier . The kNN classifier is much more powerful than a linear prediction model . It can fit almost any decision boundary ( Bishop , 2006 ) . Since it makes prediction without training , kNN has also been used for monitoring model ’ s convergence in the process of self-supervised learning ( Wu et al. , 2018b ; Chen & He , 2021 ) . Compared to the linear prediction model , kNN is advantageous for pre-training because it does not require all the examples of the same class to form one tight cluster . Given one query sample , as long as the number of samples with the same label of query takes the majority of its kNN , the label of query could be correctly predicted . Thus , each sample only need to share the same class with most of its k nearest neighbors , and the resulting distribution can be multi-mode , instead of single-mode . Figure 1 shows the data distributions from cross entropy ( C.E . ) based supervised pre-training and the proposed LOOK method , respectively . We can see clearly that for the distribution from LOOK , data points from the same class can be distributed over multiple clusters while , for C.E. , data points from the same class tend to form one cluster . It is this nice property of kNN classifier that allows us to preserve part of intra-class difference , leading to better performance for downstream tasks . 3.2 ALGORITHM LOOK : MATHEMATICAL DESCRIPTION . As mentioned before , we will use kNN for prediction model , and to facilitate training , a leave-oneout loss is used for optimization . Below we provide detailed description . We start by formulating the pre-training process on visual datasets . Given a set of images X = { xi } Ni=1 , with corresponding labels Y = { yi } Ni=1 when training in supervised scenarios , we aim at training a encoder fθ ( · ) , e.g . CNN , to map each image xi into an embedding vector zi = fθ ( xi ) in the high-dimension space . The trained parameters generally work as the initialization to fine-tune on downstream datasets . Specifically , we use the weighted kNN for class prediction . Given sample xi and its embedding zi = fθ ( xi ) , suppose the kNN set of zi in the representation space is Nk ( zi ) , then we aggregate of its kNN labels based on the weights of pairwise distance , which is selected as the cosine distance : L̃i = Σzj∈Nk ( zi ) wi , j · 1yj , wi , j = 〈zi , zj〉 = zi · zj ‖ zi ‖‖ zj ‖ , ( 1 ) where 1yj is the one-hot vector in size of total classes C with value 1 only at the position yj and value 0 at the remaining positions . Then we normalize the aggregated labels with softmax function and feed it into the negative log-likelihood function as the loss of LOOK for the i-th sample : LLOOK ( xi , yi ) = −log exp ( L̃i , yi/τ ) ΣCc=1 exp ( L̃i , c/τ ) = − log Σzj∈Nk ( zi ) exp ( 〈zi , zj〉/τ ) · 1yj=yi Σzj∈Nk ( zi ) exp ( 〈zi , zj〉 ) /τ ) , ( 2 ) where τ is the temperature hyper-parameter of softmax to control the normalization process , and 1yj=yi equals one when yj = yi and otherwise zero . Despite the simplicity , how to make LOOK work efficiently for large datasets remains a challenge . For large datasets , it is too time consuming to compute kNN for the entire dataset in an online fashion . We will then show our efficient implementation that makes LOOK scale to large datasets .

Summary: The paper presents an interesting approach that expands the recent success of self-supervised pre-training with *instance discrimination* to supervised pre-training. The key insight is to have each class being represented not just by a single weight vector, but in a non-parametric fashion via KNN lookup from a MoCo memory bank. The paper avoids to compare the results for direct, supervised learning, but rather focuses on transfer learning to downstream tasks -- to me this has some empirical novelty. The transfer learning experiments are done extensively, including many settings for quantitative comparisons and qualitative visualizations.