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 proposes a leave-one-out kNN supervised method for pre-training on large-scale datasets.  The paper claims such a method will benefit downstream tasks by NOT enforcing each class, whether visually similar or not, to cluster together. Instead, each sample just needs to be close to the nearest K neighbors.  In this way, the model will not be overfitted to pre-train datasets. To implement this method, the paper uses a loss of weighted soft kNN loss with cosine distance and softmax normalization. To solve the large search space and feature update problem of kNN, the paper uses ideas from MoCo, utilizing a memory queue as a feature bank to store features of previous N samples and First-in-first-out to update the memory queue. To solve the convergence issue, the paper uses ideas from SimCLR and utilizes MLP instead of the linear layer to project features into a more compact space.   The proposed is tested on a standard fine-tuning setting, using ImageNet as the up-stream pre-training dataset and multiple fine-grained datasets as down-stream fine-tuning datasets. The proposed method is compared with both the supervised and self-supervised methods,