[Paper Review] Contrastive learning of global and local features for medical image segmentation with limited annotations

Krishna Chaitanya, Ertunç Erdil+2 more|arXiv (Cornell University)|Jun 18, 2020

Radiomics and Machine Learning in Medical ImagingMedicine241 citations

TL;DR

The paper proposes domain- and problem-specific contrastive pre-training for segmentation in medical images, adding a local contrastive loss and leveraging volume structure to improve performance with few annotations.

ABSTRACT

A key requirement for the success of supervised deep learning is a large labeled dataset - a condition that is difficult to meet in medical image analysis. Self-supervised learning (SSL) can help in this regard by providing a strategy to pre-train a neural network with unlabeled data, followed by fine-tuning for a downstream task with limited annotations. Contrastive learning, a particular variant of SSL, is a powerful technique for learning image-level representations. In this work, we propose strategies for extending the contrastive learning framework for segmentation of volumetric medical images in the semi-supervised setting with limited annotations, by leveraging domain-specific and problem-specific cues. Specifically, we propose (1) novel contrasting strategies that leverage structural similarity across volumetric medical images (domain-specific cue) and (2) a local version of the contrastive loss to learn distinctive representations of local regions that are useful for per-pixel segmentation (problem-specific cue). We carry out an extensive evaluation on three Magnetic Resonance Imaging (MRI) datasets. In the limited annotation setting, the proposed method yields substantial improvements compared to other self-supervision and semi-supervised learning techniques. When combined with a simple data augmentation technique, the proposed method reaches within 8% of benchmark performance using only two labeled MRI volumes for training, corresponding to only 4% (for ACDC) of the training data used to train the benchmark. The code is made public at https://github.com/krishnabits001/domain_specific_cl.

Motivation & Objective

• Motivate reducing labeled data needs for medical image segmentation by leveraging self-supervised learning.
• Introduce global and local contrastive learning extensions tailored for volumetric medical images.
• Incorporate domain-specific cues from MRI/CT and local-region distinctions to improve pixel-wise predictions.
• Evaluate the approach on three MRI datasets and compare with pretext tasks, data augmentation, and semi-supervised methods.

Proposed method

• Extend global contrastive loss to exploit cross-volume slice alignment as domain-specific similarity cues.
• Introduce a local contrastive loss that operates on decoder feature maps to encourage region-wise distinctions for segmentation.
• Define two domain-informed global contrastive strategies (G^D- and G^D) and two local strategies (L^R and L^D) for constructing positive/negative pairs.
• Pre-train an encoder with global loss, then freeze it and pre-train decoder blocks with local loss before full fine-tuning for segmentation.
• Combine global and local pre-training in a staged pipeline without needing joint loss weighting.

Experimental results

Research questions

• RQ1 Can domain-specific global cues from aligned volumes improve contrastive pre-training for segmentation?
• RQ2 Does adding a local contrastive loss on decoder features yield better per-pixel representations for segmentation than global alone?
• RQ3 How do domain-informed contrastive strategies compare to random augmentations or pretext tasks in limited-annotation regimes?
• RQ4 Is the proposed pre-training complementary to data augmentation and semi-supervised approaches?

Key findings

• Global strategies leveraging cross-volume slice similarity (G^D) outperform random-augmentation baselines.
• Local decoder pre-training with L^R or L^D provides added gains over encoder-only pre-training.
• The proposed combination (G^D plus L^R) yields substantial improvements across ACDC, Prostate, and MMWHS datasets with 1–2 training volumes.
• The method approaches benchmark performance with only two labeled volumes, and gains persist when combined with data augmentation or semi-supervised methods.