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

Abstract

Pre-training a recognition model with contrastive learning on a large dataset of unlabeled data has shown great potential to boost the performance of a downstream task, e.g., image classification. However, in domains such as medical imaging, collecting unlabeled data can be challenging and expensive. In this work, we propose to adapt contrastive learning to work with meta-label annotations, for improving the model's performance in medical image segmentation even when no additional unlabeled data is available. Meta-labels such as the location of a 2D slice in a 3D MRI scan or the type of device used, often come for free during the acquisition process. We use the meta-labels for pre-training the image encoder as well as to regularize a semi-supervised training, in which a reduced set of annotated data is used for training. Finally, to fully exploit the weak annotations, a self-paced learning approach is used to help the learning and discriminate useful labels from noise. Results on three different medical image segmentation datasets show that our approach: i) highly boosts the performance of a model trained on a few scans, ii) outperforms previous contrastive and semi-supervised approaches, and iii) reaches close to the performance of a model trained on the full data.

Boundary-aware Information Maximization for Self-supervised Medical Image Segmentation

“…In this approach, a network is pre-trained to bring closer the feature embeddings of an image under different transformations (positive pairs), while pushing away those from different images (negative pairs). This idea was used to learn a global representation at the end of the network's encoder, using meta-labels on anatomical similarity or subject ID to define the positive pairs (Chaitanya et al, 2020;Peng et al, 2021b;Zeng et al, 2021). To also pre-train the decoder, the method in (Chaitanya et al, 2020) defined positive or negative embedding pairs based on their spatial distance in a feature map, those with a large distance considered as negative while those at the same spatial position but coming from different transformations as positives.…”

“…We consider the 3D-MRI scans as 2D images throughplane due to the high anisotropic acquisition resolution, and re-sample them to a fix space ranging of 1.0 × 1.0 mm. Following (Peng et al, 2021b), we normalize the pixel intensities based on the 1% and 99% percentile of the intensity histogram for each scan. Normalized slices are then cropped to 384 × 384 pixels, coarsely centered based on the foreground delineation of the ground truth.…”

“…In such technique, a model is pre-trained to perform a given pretext task, for example puzzle-solving (Noroozi & Favaro, 2016;Taleb et al, 2021), rotation prediction Gidaris et al, 2018), colorization (Zhang et al, 2016) or contrastive-based instance discrimination (Hjelm et al, 2018;Chen et al, 2020;He et al, 2020), and then fine-tuned with a small set of labeled examples. Among these self-supervised methods, contrastive learning has become a prevailing strategy for pre-training medical image segmentation models (Chaitanya et al, 2020;Zeng et al, 2021;Peng et al, 2021b). The core idea of this strategy is to learn, without pixel-wise annotations, an image representation which can discriminate related images (e.g., two transformations of the same image) from non-related ones.…”

“…Most contrastive learning approaches for segmentation apply a contrastive loss on the global representation of images, which typically corresponds to the features produced by the network's encoder. Experimental results have shown that pre-training the encoder with this loss and then fine-tuning the whole network with few labeled examples can lead to significant improvements (Peng et al, 2021b).…”

“…Recent works have also demonstrated the benefit of using contrastive learning on the decoder's feature maps during pre-training (Chaitanya et al, 2020;Peng et al, 2021b). In this case, the contrastive loss is applied at each position of the feature map, which helps learn a local representation of the image.…”