Intermediate Layers Matter in Momentum Contrastive Self Supervised Learning

1 Introduction . Self-supervised learning ( SSL ) [ Noroozi and Favaro , 2016 , Doersch et al. , 2015 , Zhang et al. , 2016 , He et al. , 2020 , Chen et al. , 2020 , Caron et al. , 2020 , Zbontar et al. , 2021 , Caron et al. , 2021 ] helps in learning useful representations from large unlabeled dataset that can be used as pre-trained initialization points for different downstream tasks . These pretrained models are particularly useful in settings where we have a small labeled dataset to learn from . Recent advances have shown that such self-supervised methods are capable of learning representations that are competitive with those learned in a fully supervised manner for natural images [ Wu et al. , 2018 , Hjelm et al. , 2018 , Zbontar et al. , 2021 , Caron et al. , 2021 ] . In application domains such as medical imaging , where high quality labeled data is scarce and data modalities shift rapidly , SSL has potential to make a significant impact . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Contrastive self-supervised learning methods [ He et al. , 2020 , Chen et al. , 2020 , Caron et al. , 2020 , Zbontar et al. , 2021 , Caron et al. , 2021 ] are currently the state of the art in SSL . In these methods , final layer ’ s feature representations for positive pairs of images are brought closer and those for negative pairs are pushed away using a contrastive loss functions such as InfoNCE [ Oord et al. , 2018 ] or NT-Xent [ Chen et al. , 2020 ] . A positive pair of image could be two different views of the same image while a negative pair could be views of different images . These contrastive methods have also been applied to various medical imaging tasks such as chest x-ray interpretation [ Sowrirajan et al. , 2021 , Vu et al. , 2021 ] , image segmentation [ Chaitanya et al. , 2020 ] , covid diagnosis [ Chen et al. , 2021 ] and dermatology classification [ Azizi et al. , 2021 ] In this work , we extend the idea of contrastive learning further and hypothesize that ensuring closeness of representations for the intermediate layers , and not just the last one , significantly improves the learning of network parameters . To this end , we hypothesize that the features learned by the method should be more reusable for the downstream task than a standard contrastive learning method . We use Momentum Contrastive ( MoCo ) method [ He et al. , 2020 ] as the base SSL method , and study two additional loss terms based on closeness of feature representation between intermediate layers . The first loss function is based on mean squared error , and the second loss function is based on distance between cross-correlation matrix and identity matrix [ Zbontar et al. , 2021 ] . We test our method on three diverse medical datasets : NIH chest x-ray , diabetic retinopathy and breast cancer histopathology . We show that both proposed loss functions either outperform standard MoCo or have comparable performance for classification tasks on all three datasets . The performance gap between the standard MoCo and our approach is higher in the case where we have a small labeled dataset to fine tune our model . This suggests that the features learned via our self-supervision approach are of higher quality than the standard MoCo . To confirm this finding , we first investigate how much the features are reused by the model , by comparing the similarity of the features before and after fine-tuning on the labeled dataset . Based on prior work [ Neyshabur et al. , 2021 ] , higher feature reuse indicates better feature quality . Additionally , we also perform layer-wise probing of the model to conclude that indeed the features learned earlier in the model by our approach are more informative than a standard MoCo . Finally , to understand the importance of the feature quality for learning from a small labeled dataset , we compare the output probability distribution of models fine tuned on 1 % and 6 % of labeled data versus on 100 % of labeled data . The proposed method of pre-training leads to lower Kolomongrov-Smrinov ( KS ) distance against the best performing model , as compared to a standard MoCo . To the best of our knowledge , this is the first work to show that the proposed approach has high performance on such diverse non-natural imaging datasets . Additionally , this is the first work to investigate the feature quality of the self-supervised method by going beyond the linear probing and focusing on feature reuse and on learning informative features earlier in the model . 2 Relevant work . 2.1 Pretext Task based SSL . Self-supervised learning helps to learn feature representation of data in an unsupervised manner . This is accomplished by performing a formulated task where the labels for the task are extracted from the data itself . Classic examples for handcrafted formulated tasks include solving a jigsaws puzzle [ Noroozi and Favaro , 2016 ] , relative patch prediction [ Doersch et al. , 2015 ] and colorization [ Zhang et al. , 2016 ] . These tasks are based on heuristics and therefore the representations learned using these tasks may not be useful for performing downstream tasks . Consequently , contrastive methods to learn representations in an unsupervised manner emerged . They outperformed the heuristics task based methods and gradually have also reached a fully supervised level of performance [ Caron et al. , 2021 ] . Some of the recent contrastive methods include SimCLR [ Chen et al. , 2020 ] , MoCo [ He et al. , 2020 ] , BYOL [ Grill et al. , 2020 ] , SwAV [ Caron et al. , 2020 ] , Barlow Twins [ Zbontar et al. , 2021 ] and Dino [ Caron et al. , 2021 ] . All the contrastive methods treat the backbone model as a black box and work with the final layer ’ s feature representations only . A recent work [ Chengyue Gong , 2021 ] attempted to incorporate contrastive loss for the intermediate layers of a model trained via MoCo , with a focus on accelerating the model training . They employed partial back-propagation - by randomly choosing a layer to back-propagate from - for reducing model computations and training time . In contrast , our work focuses on investigating the impact of the intermediate loss terms on feature quality , using state of the art analysis methods such as feature re-use , layer-wise probing and model output similarity . We investigate this intermediate loss functions with two distinct formulations ( mean squared error , and a loss based on cross correlation matrix ) and on three diverse medical imaging datasets . We do not employ partial back propagation , and our formulation of intermediate layer loss does not require storing intermediate representations for the negative pairs . 2.2 Contrastive self supervised learning in medical imaging . Out of above mentioned contrastive methods , MoCo has been widely adopted by the research community across many domains including medical imaging . The momentum contrastive method ( MoCo ) differentiates from other contrastive methods by having an additional memory bank of negative examples . The memory bank helps the model to learn even when the batch size is small , and is especially important for medical imaging tasks , because medical images are typically large ( e.g . 1024 × 1024 ) and it becomes computationally infeasible to train with a large batch size . Increasingly , many recent works have used momentum contrastive method with minor modifications as the choice of self-supervised method to perform diverse medical imaging tasks such as histopathology classification [ Ciga et al. , 2020 , Dehaene et al. , 2020 ] , chest x-ray interpretation [ Sowrirajan et al. , 2021 ] and dermatology classification [ Azizi et al. , 2021 ] . Given a widespread adoption of the momentum contrastive method , we focus on it in our work . 3 Methodology . In section 3.1 , first , we briefly discuss the standard MoCo method . Then , we describe our proposed approach in section 3.2 . In section 3.3 , we describe the fine tuning and evaluation methodology for the self-supervised methods . Finally , in section 3.4 , we describe the methods used for feature analysis which helps us quantify the quality of the representations learned by different methods . 3.1 Standard MoCo . With the help of MoCo , we want to learn feature representation for images in an unsupervised manner . We achieve this by minimizing InfoNCE loss [ Oord et al. , 2018 ] LInfoNCE ( x ) = −log exp ( g ( x̃1 ) .h ( x̃2 ) /τ ) exp ( g ( x̃1 ) .h ( x̃2 ) /τ ) + ∑K i=0 exp ( g ( x̃1 ) .h ( z̃i ) /τ ) ( 1 ) where x̃1 and x̃2 are different views/augmentations of the same image x. x̃1 and x̃2 are called a positive pair . There are K negative pairs of x̃1 and z̃is where z̃i is random augmentation of a different image . z̃i either comes from a queue in the case of MoCo or a minibatch in the case of SimCLR [ Chen et al. , 2020 ] . We minimize equation 1 to optimize the parameters of the encoder g. h is an identical copy of g whose parameters are an exponential moving average of g. τ is a temperature-scaling hyper-parameter . 3.2 MoCo with intermediate feature similarity . In our work , we extend the idea of contrastive learning and propose that the encoder should be encouraged to learn augmentation-invariant representations not only at the end of the encoder but also for intermediate layers . We make use of an additional loss function to ensure closeness in the representations for the intermediate layers . Such intermediate loss have shown to improve performance of deep learning models such as InceptionNet [ Szegedy et al. , 2015 ] , PSPnet [ Zhao et al. , 2017 ] and DARTs [ Liu et al. , 2018 ] . Intermediate loss helps to improve performance for various reasons . In Mostajabi et al . [ 2018 ] , it helped as a regularizer . In InceptionNet , PSPnet and DARTs , it improved training of the model by facilitating gradient flow . In our case , we hypothesize that it would facilitate learning of augmentation-invariant features early in the model thereby contributing to learning of high quality features for the downstream tasks . We experiment with two intermediate loss functions that ensure closeness of representations : mean square error , and a loss function inspired by Zbontar et al . [ 2021 ] that minimizes the difference between cross-correlation matrix of representations and an identity matrix . Figure 3.2 includes an overview of our approach . 3.2.1 Mean squared error of intermediate features . We compute the mean squared error of the intermediate features between the two views of the same image . We extract the intermediate features and pass them through an adaptive average pooling layer to reduce the spatial variances in the two views that may be caused due to augmentations such as rotations . For ResNet-50 as the backbone architecture , the spatial resolutions of the features after the average pooling layers are 16× 16 , 16× 16 , 4× 4 and 4× 4 for `` Block 1 '' , `` Block 2 '' , `` Block 3 '' and `` Block 4 '' of ResNet-50 , respectively . The mean squared error is calculated between these spatially down-sampled features of the two views for each block .

Summary: The authors propose to learn augmentation invariant features early on in a network by, in addition to a global contrastive loss, also optimizing for an intermediate loss function to ensure similarity between features in early layers of the network. They test two loss functions for this purpose - MSE and a minimization of the difference between the cross-correlation matrix of the intermediate representations and the identity matrix. They hypothesize that this way the network learns high quality features early on. The authors also test the quality of features learnt with their method in a self-supervised setting, by comparing the similarity of the features learnt under self-supervision to features learnt after fine-tuning on a labeled dataset following the self-supervision. They do so for a variety of different labelled set sizes, as well as end-to-end fine-tuning, and only tuning the last layer.  They test their method on three publicly available medical datasets (chest X-Ray, diabetic retinopathy and histopathological whole slide images), which are of different imaging modalities.