Self-Paced Contrastive Learning for Semi-supervised Medical Image Segmentation with Meta-labels

1 Introduction . Since the emergence of deep learning [ 26 ] , there has been an active debate on the importance of pre-training neural networks . Precursor works [ 17 ] showed that pre-training a convolutional neural network with an unsupervised task ( e.g. , denoising autoencoders [ 49 ] ) could lead to a better performance in the final supervised task . As the amount of labeled training data increased , thanks to large datasets like ImageNet [ 13 ] , it was however found that pre-training could actually hinder performance [ 38 ] . This makes sense in light of recent studies showing , for instance , that symmetries in large networks induce many equivalent local minima [ 16 , 35 , 47 ] in which a pre-trained model can get stuck . Recently , contrastive learning has renewed the interest in unsupervised pre-training [ 37 ] . Several works [ 8 , 10 , 11 , 20 , 62 ] have found that pre-training a model with a contrastive loss can improve its performance on a subsequent supervised training task , often outperforming a network with supervised pre-training on ImageNet . While this has reopened the debate on the benefit of pre-training , it offers little help for domains where data is scarce such as medical imaging . In medical imaging , not only are labels expensive since they come from highly-trained experts like radiologists , but images are also hard to obtain due to the need for costly equipment ( e.g. , MRI or CT scanner ) and privacy regulations . Over the last years , a breadth of semi-supervised learning approaches have been proposed for medical image segmentation , including methods based on attention [ 32 ] , adversarial learning [ 60 ] , ∗Corresponding author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 2 temporal ensembling [ 12 , 55 ] , co-training [ 40 , 63 ] , data augmentation [ 6 , 61 ] and transformation consistency [ 5 ] . The common principle of these approaches is to add an unsupervised regularization loss using unlabeled images , which is optimized jointly with a standard supervised loss on a limited set of labeled images . Despite reducing significantly the amount of labeled data required for training , current semi-supervised learning methods still suffer from important drawbacks which impede their use in various applications . Thus , a large number of unlabeled images is often necessary to properly learn the regularization prior . As mentioned before , this may be impossible in medical imaging scenarios where data is hard to obtain . Moreover , these methods also need a sufficient amount of labeled data , otherwise the learning may collapse [ 36 ] . In a recent work , Chaitanya et al . [ 7 ] showed that unsupervised pre-training can be useful to learn a segmentation task with very few samples , by leveraging the meta information of medical images ( e.g. , the position of a 2D image in the 3D volume ) . While achieving impressive accuracy with as few as two volumes , this work has significant limitations . First , it relies on the strong assumption that the global or local representations of 2D images are similar if their locations within the volume or feature map are related . This assumption does not always hold in practice since volumes may not be well aligned , or due to the high variability of structures to segment . Second , it requires dividing the 2D images of a 3D volume in an arbitrary number of hard partitions that are contrasted , while the structure to segment typically varies gradually within the volume . Third , they do not exploit the full range of available meta data , for instance the patient ID or cycle phase of cardiac cine MRI , nor evaluate the benefit of combining several types of meta information in pre-training . Last , their approach leverages meta data only in pre-training , however this information could further boost performance if used while learning the final segmentation task , in a semi-supervised setting . Our work addresses the limitations of current semi-supervised and self-supervised approaches for segmentation by proposing a novel self-paced contrastive learning method , which takes into account the noisiness of weak labels from meta data and exploits this data jointly with labeled images in a semi-supervised setting . The detailed contributions of this paper are as follows : – We propose , to our knowledge , the first self-paced strategy for contrastive learning which dynamically adapts the importance of individual samples in the contrastive loss . This helps the model deal with noisy weak labels that arise , for instance , from misaligned images or splitting a 3D volume in arbitrary partitions . – We demonstrate the usefulness of a contrastive loss on meta-data for improving the performance of a final task , not only in pre-training but also as an additional loss in semi-supervised training . – We show that combining multiple meta-labels in our self-paced contrastive learning framework can improve performance on the final task , compared to using them independently . Our results also demonstrate the benefit of combining contrastive learning with temporal ensembling to further boost performance . We empirically validate our contributions on five well-known medical imaging datasets , and show the proposed approach to outperform the contrastive learning method of [ 7 ] as well as several state-ofthe-art semi-supervised learning methods for segmentation [ 41 , 43 , 50 , 58 , 60 ] . In the results , our approach obtains a performance close to fully supervised training with very few training scans . 2 Related work . We focus our presentation of previous works on two machine learning sub-fields that are most related to our current work : self-supervision , which includes contrastive learning , and self-paced learning . Self-supervision and contrastive learning Self-supervision is a form of unsupervised learning where a pretext task is used to pre-train a model so to better perform a downstream task . Examples of pretext tasks are learning to sort a sequence [ 28 , 54 ] , predicting rotations [ 19 , 18 ] , solving a jigsaw puzzle [ 33 ] and many others [ 14 , 15 , 25 , 46 , 59 ] . Most of these methods can improve performance on the downstream task when labeled data for training is scarce . However , when a large and generalenough dataset of labeled images like ImageNet [ 13 ] is available , a simple supervised pre-training may sometimes provide better results [ 22 , 34 ] . Recently , unsupervised contrastive learning [ 8 , 20 , 10 ] was shown to boost performance even when learning a downstream task on a large dataset , and to improve over a model pre-trained in a supervised manner on a large dataset . This approach is based on the simple idea of enforcing similarity in the representation of two examples from the same class ( positive pairs ) , and increase representation dissimilarity on pairs from different classes ( negative pairs ) [ 37 , 48 ] . In image analysis tasks , positive pairs are typically defined as two 3 transformed versions of the same image , for instance using a geometric or color transformation , while negative ones are any other pair of images [ 8 , 10 , 11 , 20 ] . In a recent work , Khosla et al . [ 24 ] showed that , when combined with true semantic labels , a contrastive learning based “ pre-train and fine-tune ” pipeline performed surprisingly well , outperforming conventional training with cross-entropy in some cases . The work in [ 62 ] extended this idea to pixel-level tasks like semantic segmentation , clustering the representations of pixels in an image according to their labels . Applying a similar strategy to medical image segmentation , Chaitanya et al . [ 7 ] leveraged meta-labels from 3D scans in a local and global contrastive learning framework to improve performance when training with limited data . However , positive and negative pairs in their contrastive loss are defined using noisy “ weak ” labels , arising for example from misaligned images or an arbitrary partitioning of 3D volumes , which may lead to learning sub-optimal representations . The work in [ 56 ] mitigated this problem by imposing a maximum distance along the z axis between slices forming positive pairs . Moreover , existing approaches only exploit meta-labels in pre-training , instead of considering them jointly with labeled images in a semi-supervised strategy . Our work extends these approaches with a self-paced learning method that adapts the importance of positive pairs dynamically during training , and focuses the learning on the most reliable ones . In contrast to [ 7 ] , we also exploit contrastive learning as regularization loss in semi-supervised training and show that further improvements can be achieved when combining it with a temporal ensembling strategy like Mean Teacher [ 12 , 55 ] . A recent approach by Chen et al . [ 10 ] also performs contrastive learning for image recognition in a semi-supervised setting . In this approach , a large teacher model is trained with a unsupervised contrastive loss and then fine-tuned with a small fraction of labeled data . A student model is trained afterwards using a knowledge distillation technique . Hence , unlike our method , there is no joint optimization of the supervised and contrastive objectives . In this work , we show that significant improvements can be achieved by jointly optimizing these two objectives . Self-paced learning A sub-category of curriculum learning ( CL ) [ 2 ] , self-paced learning ( SPL ) is inspired by the learning process of humans that gradually incorporates easy to hard samples in training [ 27 ] . The effectiveness of such strategy has been validated in various computer vision tasks [ 52 , 23 , 57 ] . Jiang et al . [ 23 ] proposed an SPL method considering both the difficulty and diversity of training examples , which outperformed conventional SPL methods that ignore diversity . Zhang et al . [ 57 ] incorporated SPL in a DNN fine-tuning process for object detection , to cope with data ambiguity and guide the learning in complex scenarios . The usefulness of SPL when training with a limited budget or when the training data is corrupted by noise was also studied in recent work [ 53 ] . So far , the application of SPL to image segmentation remains limited . Wang et al . [ 52 ] presented an SPL method for lung nodule segmentation , where the uncertainty of each sample prediction in the loss is controlled by the SPL regularizer . Similarly , a self-paced co-training method was proposed in [ 51 ] for the semi-supervised segmentation of medical images . Related to our work , Wang et al . [ 31 ] proposed a margin preserving contrastive learning framework for domain adaptation that uses a self-paced strategy in self-training . Compared to this approach , which uses SPL outside the contrastive loss to select confident pseudo-labels for self-training , our method incorporates it within the loss via importance weights that are learned jointly with network parameters . 3 Proposed method . In this section , we present our self-paced contrastive learning approach for segmentation that leverages intrinsic meta information extracted from medical volumetric images . Our method effectively pretrains a segmentation model with a self-paced variant of contrastive learning that is more robust to noisy annotations . The same approach is also used to further boost the segmentation accuracy in a semi-supervised setting , where only very limited pixel-wised annotations are used . In the following subsections , we detail the formulation of the proposed approach .

Summary: The authors propose using “metal labels” in pertaining image encoders on unlabeled datasets. They use the metal-labels in pertaining the encoder, and also as an auxiliary task in training. They show their method works on multiple medical datasets.