Learning Rich Nearest Neighbor Representations from Self-supervised Ensembles

1 INTRODUCTION . The widespread application of pretrained convolutional neural networks in computer vision is one of the most important tools in the field . State-of-the-art on many benchmarks ranging from classification , to object detection , to pose estimation has been set using a pretrained model , such as an ImageNet classifier , as a network initialization ( Kornblith et al. , 2019 ; He et al. , 2019 ; Chen et al. , 2020a ; Grill et al. , 2020 ; Kolesnikov et al. , 2019 ) . Transfer learning is an entire field focused on studying and utilizing this phenomenon . While supervised ImageNet classifiers were the dominant feature extractors of choice for many years , recently self-supervised models have begun to take their place . Methods such as MoCo ( v2 ) , SimCLR ( v2 ) , SimSiam , PIRL , Swav , BYOL , and Barlow Twins all claim transferability competitive with or superior to that of ImageNet classifiers ( He et al. , 2019 ; Chen et al. , 2020c ; a ; b ; Chen & He , 2020 ; Misra & Maaten , 2020 ; Caron et al. , 2021 ; Grill et al. , 2020 ; Zbontar et al. , 2021 ) . As such , the question of what initialization to use has arisen ; benchmark studies have sought to compare methods under dozens of different settings ( Goyal et al. , 2019 ; Zhai et al. , 2019 ) . Even when a decision has been made to use a particular feature extractor , the utility and knowledge of other options is then left unutilized . To address this concern , we consider ensembling , a common practice in the supervised setting ( Dietterich , 2000 ; Hinton et al. , 2015 ) . Ensembling models involves combining the predictions obtained by multiple different models in some way , typically by averaging or a similar operation . In the supervised setting , such outputs are aligned and such an operation easily captures the combined knowledge of the models . In the self-supervised setting , however , such alignment is not guaranteed , particularly when dealing with independently trained contrastive learners which many pretrained models of choice are . Averaging the features still is useful , and obtains reasonably strong image representations ( Section 4 ) , but we show that it is possible to build an ensembling strategy that yields richer , stronger representations than the mean feature . We do so without training any new CNN components , allowing for the same frozen backbone to be used across applications . How to approach ensembling in the self-supervised setting ? We contend that the goal of a model ensemble is to capture the useful information provided by the different models in a single representation . We consider the “ capture '' of information from a recoverability perspective : if every network ’ s features can be recovered by some fixed operation on a representation vector , for all data samples , Algorithm 1 PyTorch-like pseudocode of our method # Training Phase # Initialize representations to average feature train_feature_list = [ net ( images ) for net in ensemble ] avg_feat = average ( train_feature_list ) # shape = ( n_points x feature_dimension ) learned_train_reps = Parameter ( avg_feat.detach ( ) ) # initialize params with avg_feat mlps = [ MLP ( ) for net in ensemble ] # 1 net per feature extractor opt = SGD ( mlps.parameters ( ) + learned_train_reps ) # optimize both mapping MLPs and input representations for images_idx , images in trainloader : ensemble_feats = [ net ( images ) for net in ensemble ] # Get ensemble features outputs = [ mlp ( learned_train_reps [ images_idx ] for mlp in mlps ] # Map learned representations through different MLPs loss = cosine_loss ( ensemble_feats , outputs ) loss.backward ( ) # Inference Phase test_feature_list = [ net ( images ) for net in ensemble ] learned_test_reps = average ( test_feature_list ) opt = SGD ( learned_test_reps ) # Freeze MLPs at inference time for images_idx , images in testloader : # Same as training loop MLP ( ) : a multi-layer perceptron model Parameter ( t ) : Pytorch function that takes the argument array t and initializes trainable parameters with that value then such representations are useful . While concatenation of features can trivially achieve this object , we show that such an operation is in fact suboptimal in terms of the behavior of the derived feature space . We instead propose to directly learn via gradient descent a set of representations that contain all of the information necessary to derive the ensemble features . Our architecture is shown in Figure 1 , with example pseudocode in Algorithm 1 . Specifically , we show that extracting features from an ensemble of self-supervised models using this technique improves the nearest neighbor ( NN ) performance when evaluated on downstream supervised tasks . 2 RELATED WORK . Supervised ensembling : It is a ubiquitous technique in machine learning ( Dietterich , 2000 ) . In addition to the large number of online contests won through such approaches ( Andres ) , ensembling has also been demonstrated to achieve state-of-the-art performance on standard computer vision benchmarks ( Huang et al. , 2017 ) . Most approaches employ Bayesian ensembling , where the predictions of separate networks are averaged ( Wu et al. , 2021 ; Hinton et al. , 2015 ; Dietterich , 2000 ) . While this relied on the alignment of objectives between the networks , we show that such an averaging on the intermediate features does indeed generate representations superior to that of individual models . Our method differs from this literature , however , in the learning done on top of the ensemble as well as the no-label setting of our representation learning . Knowledge distillation ( Hinton et al. , 2015 ) : This is a related vein of work is that of where the knowledge of an ensemble is distilled into a single neural network , or a single model is distilled into another network . This second version , known as self-distillation ( Zhang et al. , 2019 ) , is highly relevant to our single-model setting , as we are able to obtain improvements while operating purely in the feature space of a single model . Our goal is not to discard the ensembled models as it is in knowledge distillation , but our method bears similarities to that of Hydra ( Mandt et al . ) , where a network is trained to output representations capable of recovering the ensembled outputs . We note that our resulting accuracies consistently surpass average ensembling , a baseline that Hydra considers as an upper-bound to their method ( but the differing settings do not lend themselves to apples-to-apples comparisons ) . Xu et al . ( 2020 ) propose a self-supervised knowledge distillation algorithm that uses both supervised and self-supervised loss to train a teacher network , and then distill this combined knowledge to a student network . They show that this combination improves the student performance compared to traditional KD . In contrast to their work , our goal is not to learn a separate student network and we do not assume access to labeled data , but instead , our main goal is to extract rich nearest neighbor representations from an ensemble of self-supervised models . k-Nearest neighbor ( k-NN ) classifiers : k-NN is a non-parametric classifier which has been shown to be a consistent approximator ( Devroye et al . ( 1996 ) ) , i.e. , asymptotically its empirical risk goes to zero as k →∞ and k/N → 0 , where N is the number of training samples . While these theoretical conditions may not be true in practice , a main advantage of using k-NN is that it is parameterfree outside of the choice of k ( which typically does not change qualitative rankings ) and thus is a consistent and easily replicable measurement of representation quality . Additionally , k-NN makes the decision process interpretable ( Papernot & McDaniel , 2018 ; Dalitz , 2009 ) , which is important in various applications ( Mehrabi et al. , 2021 ; Vellido , 2019 ; Gilpin et al. , 2018 ) . For these reasons , our paper focuses on extracting rich features from self-supervised ensembles that are condusive to k-NN classification . On a related note , SLADE ( Duan et al. , 2021 ) leverages unlabeled data to improve distance metric learning , which is essentially the setting of our evaluation framework . While the goal is similar , SLADE uses supervised training to initialize learning and generate pseudolabels for the unlabeled points , our method assumes zero label access . Additionally , SLADE is concerned with learning a new network from scratch , as opposed to an ensembling framework . Pretrained Models : Our method relies on the efficacy of pretrained self-supervised models . Specific methods we employ are SimCLR , SwAV , Barlow Twins , RotNet , PIRL , as well as traditional label supervision . In addition to the above works : Goyal et al . ( 2019 ) ; Zhai et al . ( 2019 ) ; Kolesnikov et al . ( 2019 ) benchmark and demonstrate the generalization efficacy of pretrained models in various settings . Gradient descent at inference time : One atypical facet to our method is the use of gradient descent at inference time to directly learn new representations . While this approach is quite uncommon , we are not the first to leverage backpropagation at inference-time . Zadeh et al . ( 2019 ) uses backpropagation at inference time to learn representations from images using a generative “ auto-decoder '' framework and Park et al . ( 2019 ) ; Sitzmann et al . ( 2020 ) employ similar approaches to learn implicit representations of shapes . Sun et al . ( 2020 ) considers new samples as one-image self-supervised learning problems , and perform a brief optimization of a self-supervised learning objective on a new image ( modifying the feature extractor ’ s parameters ) before performing inference , and Shocher et al . ( 2018 ) train a CNN at test-time for the purpose of super-resolution . 3 METHODS . We present a method for directly learning rich ensembled unsupervised representations of images through gradient descent . Consider a training collection of imagesX = { xi } ni=1 and an ensemble of convolutional neural networks feature extractors Θ = { θj } mj=1 . In this work the θj have previously been trained in a self-supervised manner on ImageNet classification and are ResNet-50s ( Deng et al. , 2009 ; He et al. , 2016 ) . Denote the L2-normalized features obtained by removing the linear/MLP heads of these networks and extracting intermediate features post-pooling ( and ReLU ) as Z = { { z ( j ) i } ni=1 } mj=1 , here z ( j ) i denotes the intermediate feature corresponding to θj ( xi ) . We initialize a memory bank of representations of X , with one entry for each xi these entries have the same feature dimensionality as the zji . This memory bank is analagous to the type use in early contrastive learning works such as Wu et al . ( 2018 ) . We denote this memory bank Ψ = { ψk } nk=1 . Each ψk is initialized to the L2-normalized average representation of the ensemble ψk = ∑m j=1 z j k || ∑m j=1 z j k|| , note that the sum operation is equivalent to averaging due to the normalization being performed . To map the memory bank to the ensembled features , we employ a set of multi-layer perceptrons ( MLPs ) , Φ = { φ ` } m ` =1 , each corresponding to a feature extractor θj . Unless noted otherwise in our experiments , these φ ` are 2 layers , both of output dimension the same as their input ( 2048 for ResNet50 features ) . ReLU activations are used after both layers , the first as a traditional activation function , and the second to align the network in mapping to the post-relu set Z . During training , a batch of images { xi } i∈I are sampled with indices I ⊂ { 1 ... n } . The corresponding ensemble features , ZI = { { z ( j ) i } i∈I } mj=1 , are retrieved as are the memory bank representations ΨI = { ψk } k∈I . Note that no image augmentations are included in our framework , meaning that the z ( j ) i are typically cached to lessen computational complexity . Each banked representation is then fed through each of the m MLPs , Φ , resulting in a set of mapped representations Φ ( ΨI ) = { φ ` ( ψi ) } ` ∈ { 1 ... m } , i∈I . The goal of the training is to maximize the alignment of these mapped features Φ ( ΨI ) with the original ensemble features ZI . This is done by training both the networks Φ and the representations Ψ using a cosine loss between Φ ( ΨI ) and ZI , gradients are computed for both the MLPs and representations for each batch . Once training is completed , the φ ` are frozen . During inference , when a new image x′ is given , the above process is repeated with the frozen MLPs . Concretely , the ensemble features φ ` ( x′ ) are computed and averaged to initialize a new representation ψ′ . ψ′ is then optimized via gradient descent to maximize the cosine similarity of each φ ` ( ψ′ ) with θ ` ( x′ ) , ψ′ then serves as our representation of x′ . The described method results in representations superior to either the average or concatenated feature in terms of nearest-neighbor accuracy . We highlight several exciting aspects of our method : • The learning of a representation at test time via gradient descent is an uncommon approach . Existing methods do exist , such as auto-decoders in Zadeh et al . ( 2019 ) or implicit 3D representation literature ( Park et al. , 2019 ; Sitzmann et al. , 2020 ) . There are also techniques , such as Sun et al . ( 2020 ) for generalization , which use gradients to shape the parameters prior to inferenc . • The vast majority of ensembling literature focuses on the supervised setting , where the training objectives of the ensembled networks are identical and thus aligned . Very little work has been performed on improving a group of self-supervised features with outside auxiliary signal . Hydra ( Mandt et al . ) considers a similar setting , but with a focus on knowledge distillation of the ensemble into a single network . • Our method is extremely adaptable to different settings . Networks trained on multiple objectives , with different head architectures , can be usefully ensembled as demonstrated in Sec . 4 . This could trivially be extended to using networks of different architectures as well ( e.g . VGG + ResNet + AlexNet ) ( Simonyan & Zisserman , 2014 ; He et al. , 2016 ; Krizhevsky et al. , 2012 ) . The flexibility of our approach additionally extends to the input data . While we use networks pretrained on ImageNet , the ensembling provides benefit on both ImageNet as well in the self-supervised transfer learning setting . This transfer can either be performed by training new Φ on the target dataset or by training Φ on ImageNet and then using the frozen MLPs to learn representations on the target dataset . • Our method as presented requires only a single forward pass of each image through each ensembled model as we do not use data augmentation . This allows caching of CNN features and fast training .

Summary: This paper propose a new way to learn self-supervised model ensembling. Their novel approach learns representations via gradient descent directly at inference time after having pretrained feature extractors. And the authors conduct a series of experiments to show the efficacy of their method.